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)
|
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"):
|
||||||
|
|||||||
Reference in New Issue
Block a user