fix(migration): use DO/EXCEPTION for idempotent enum creation in b035
Some checks failed
Aegis CI / lint-and-test (push) Has been cancelled

Replace _enum_exists() helper (which had connection context issues in
Alembic) with PostgreSQL DO $$ BEGIN ... EXCEPTION WHEN duplicate_object
THEN NULL; END $$ blocks, which are truly idempotent regardless of
transaction state.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
kitos
2026-05-19 16:51:22 +02:00
parent a8b4518485
commit 36fe4aa250

View File

@@ -21,33 +21,28 @@ def _table_exists(table_name: str) -> bool:
return conn.dialect.has_table(conn, table_name)
def _column_exists(table_name: str, column_name: str) -> bool:
conn = op.get_bind()
insp = sa.inspect(conn)
cols = [c["name"] for c in insp.get_columns(table_name)]
return column_name in cols
def _enum_exists(enum_name: str) -> bool:
conn = op.get_bind()
result = conn.execute(
sa.text("SELECT 1 FROM pg_type WHERE typname = :name"), {"name": enum_name}
).fetchone()
return result is not None
def upgrade() -> None:
# ── Enums ────────────────────────────────────────────────────────────────
if not _enum_exists("queue_priority"):
op.execute("CREATE TYPE queue_priority AS ENUM ('critical', 'high', 'medium', 'low')")
if not _enum_exists("queue_status"):
op.execute("CREATE TYPE queue_status AS ENUM ('pending', 'in_progress', 'completed', 'dismissed')")
if not _enum_exists("queue_reason"):
op.execute(
"CREATE TYPE queue_reason AS ENUM ("
"'validation_expired', 'infra_change', 'osint_alert', "
"'mitre_update', 'rule_modified', 'low_confidence', 'manual')"
)
# ── Enums (idempotent via DO/EXCEPTION) ──────────────────────────────────
op.execute("""
DO $$ BEGIN
CREATE TYPE queue_priority AS ENUM ('critical', 'high', 'medium', 'low');
EXCEPTION WHEN duplicate_object THEN NULL;
END $$;
""")
op.execute("""
DO $$ BEGIN
CREATE TYPE queue_status AS ENUM ('pending', 'in_progress', 'completed', 'dismissed');
EXCEPTION WHEN duplicate_object THEN NULL;
END $$;
""")
op.execute("""
DO $$ BEGIN
CREATE TYPE queue_reason AS ENUM (
'validation_expired', 'infra_change', 'osint_alert',
'mitre_update', 'rule_modified', 'low_confidence', 'manual');
EXCEPTION WHEN duplicate_object THEN NULL;
END $$;
""")
# ── technique_ownerships ─────────────────────────────────────────────────
if not _table_exists("technique_ownerships"):