44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
"""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'")
|