fix(migration): use DO/EXCEPTION for idempotent enum creation in b035
Some checks failed
Aegis CI / lint-and-test (push) Has been cancelled
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:
@@ -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"):
|
||||
|
||||
Reference in New Issue
Block a user