mimic/backend/migrations/versions/0003_simulation_techniques_array.py

"""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")