83 lines
2.6 KiB
Python
83 lines
2.6 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 column, table, text
|
||
|
|
|
||
|
|
|
||
|
|
revision = "0003"
|
||
|
|
down_revision = "0002"
|
||
|
|
branch_labels = None
|
||
|
|
depends_on = None
|
||
|
|
|
||
|
|
# Lightweight table proxies for data migration (no ORM import).
|
||
|
|
_sims = table(
|
||
|
|
"simulations",
|
||
|
|
column("id", sa.Integer),
|
||
|
|
column("mitre_technique_id", sa.String),
|
||
|
|
column("mitre_technique_name", sa.String),
|
||
|
|
column("techniques", sa.Text),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
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.
|
||
|
|
# SQLite doesn't support ALTER COLUMN, so we skip the nullable constraint
|
||
|
|
# change at DDL level — the application model enforces it.
|
||
|
|
|
||
|
|
# 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")
|