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) 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: def upgrade() -> None:
# ── Enums ──────────────────────────────────────────────────────────────── # ── Enums (idempotent via DO/EXCEPTION) ──────────────────────────────────
if not _enum_exists("queue_priority"): op.execute("""
op.execute("CREATE TYPE queue_priority AS ENUM ('critical', 'high', 'medium', 'low')") DO $$ BEGIN
if not _enum_exists("queue_status"): CREATE TYPE queue_priority AS ENUM ('critical', 'high', 'medium', 'low');
op.execute("CREATE TYPE queue_status AS ENUM ('pending', 'in_progress', 'completed', 'dismissed')") EXCEPTION WHEN duplicate_object THEN NULL;
if not _enum_exists("queue_reason"): END $$;
op.execute( """)
"CREATE TYPE queue_reason AS ENUM (" op.execute("""
"'validation_expired', 'infra_change', 'osint_alert', " DO $$ BEGIN
"'mitre_update', 'rule_modified', 'low_confidence', 'manual')" 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 ───────────────────────────────────────────────── # ── technique_ownerships ─────────────────────────────────────────────────
if not _table_exists("technique_ownerships"): if not _table_exists("technique_ownerships"):