"""allow staged SSH source definitions

Revision ID: 0009_allow_ssh_sources
Revises: 0008_notification_outbox
"""

from __future__ import annotations

import sqlalchemy as sa
from alembic import op

revision = "0009_allow_ssh_sources"
down_revision = "0008_notification_outbox"
branch_labels = None
depends_on = None


def _reject_ssh_sources() -> None:
    connection = op.get_bind()
    sources = sa.table("sources", sa.column("kind"))
    count = connection.scalar(
        sa.select(sa.func.count()).select_from(sources).where(sources.c.kind == "ssh")
    )
    if count is None:
        raise RuntimeError("Cannot inspect persisted SSH sources before migration.")
    if count:
        raise RuntimeError(
            "Cannot restrict sources to local: "
            f"found {count} SSH source row(s). Remove them before downgrading."
        )


def upgrade() -> None:
    with op.batch_alter_table("sources") as batch:
        batch.drop_constraint(op.f("ck_sources_kind"), type_="check")
        batch.create_check_constraint(op.f("ck_sources_kind"), "kind IN ('local','ssh')")


def downgrade() -> None:
    _reject_ssh_sources()
    with op.batch_alter_table("sources") as batch:
        batch.drop_constraint(op.f("ck_sources_kind"), type_="check")
        batch.create_check_constraint(op.f("ck_sources_kind"), "kind = 'local'")