"""restrict persisted sources to local

Revision ID: 0006_local_sources_only
Revises: 0005_restore_dry_run
"""

import sqlalchemy as sa
from alembic import op

revision = "0006_local_sources_only"
down_revision = "0005_restore_dry_run"
branch_labels = None
depends_on = None


def _reject_nonlocal_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 != "local")
    )
    if count is None:
        raise RuntimeError("Cannot inspect persisted source kinds before migration.")
    if count:
        raise RuntimeError(
            "Cannot restrict sources to local: "
            f"found {count} non-local source row(s). Remove or migrate them before upgrading."
        )


def upgrade() -> None:
    _reject_nonlocal_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'")


def downgrade() -> 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','sftp','postgresql','mysql')"
        )