4596f09e71
- Migration 0003: enforce techniques NOT NULL via batch_alter_table - Migration 0003: remove unused _sims table proxy and dead column/table imports - mitre.py: add _TACTIC_NAMES dict to fix 'Command And Control' → 'Command and Control'
74 lines
2.4 KiB
Python
74 lines
2.4 KiB
Python
"""replace scalar MITRE columns with techniques JSON array
|
|
|
|
Revision ID: 0003
|
|
Revises: 0002
|
|
Create Date: 2026-05-27 00:00:00.000000
|
|
"""
|
|
import json
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.sql import text
|
|
|
|
|
|
revision = "0003"
|
|
down_revision = "0002"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
|
|
def upgrade():
|
|
bind = op.get_bind()
|
|
|
|
# 1. Add techniques column (nullable while we backfill).
|
|
op.add_column("simulations", sa.Column("techniques", sa.Text(), nullable=True))
|
|
|
|
# 2. Backfill: scalar → JSON array.
|
|
rows = bind.execute(
|
|
text("SELECT id, mitre_technique_id, mitre_technique_name FROM simulations")
|
|
).fetchall()
|
|
for row in rows:
|
|
if row[1]: # mitre_technique_id is not null
|
|
val = json.dumps([{"id": row[1], "name": row[2] or ""}])
|
|
else:
|
|
val = "[]"
|
|
bind.execute(
|
|
text("UPDATE simulations SET techniques = :v WHERE id = :id"),
|
|
{"v": val, "id": row[0]},
|
|
)
|
|
|
|
# 3. Make NOT NULL now that every row has a value.
|
|
with op.batch_alter_table("simulations") as batch_op:
|
|
batch_op.alter_column("techniques", existing_type=sa.Text(), nullable=False)
|
|
|
|
# 4. Drop old scalar columns.
|
|
with op.batch_alter_table("simulations") as batch_op:
|
|
batch_op.drop_column("mitre_technique_id")
|
|
batch_op.drop_column("mitre_technique_name")
|
|
|
|
|
|
def downgrade():
|
|
bind = op.get_bind()
|
|
|
|
# 1. Re-add scalar columns.
|
|
with op.batch_alter_table("simulations") as batch_op:
|
|
batch_op.add_column(sa.Column("mitre_technique_id", sa.String(length=32), nullable=True))
|
|
batch_op.add_column(sa.Column("mitre_technique_name", sa.String(length=255), nullable=True))
|
|
|
|
# 2. Back-fill: take first element of techniques array.
|
|
rows = bind.execute(text("SELECT id, techniques FROM simulations")).fetchall()
|
|
for row in rows:
|
|
techniques = json.loads(row[1] or "[]")
|
|
if techniques:
|
|
first = techniques[0]
|
|
bind.execute(
|
|
text(
|
|
"UPDATE simulations SET mitre_technique_id = :tid, mitre_technique_name = :tname WHERE id = :id"
|
|
),
|
|
{"tid": first.get("id"), "tname": first.get("name"), "id": row[0]},
|
|
)
|
|
|
|
# 3. Drop techniques column.
|
|
with op.batch_alter_table("simulations") as batch_op:
|
|
batch_op.drop_column("techniques")
|