company-ai-platform/alembic/versions/202607080001_v2_production_controls.py

"""Add V2 production control tables and fields.

Revision ID: 202607080001
Revises: 202607060003
Create Date: 2026-07-08
"""

from alembic import op
import sqlalchemy as sa
from sqlalchemy import inspect


revision = "202607080001"
down_revision = "202607060003"
branch_labels = None
depends_on = None

WORK_TASKS_TABLE = "work_tasks"
RISK_EVENTS_TABLE = "risk_events"
RISK_EVENT_ACTIONS_TABLE = "risk_event_actions"
REPORT_PUSH_RUNS_TABLE = "report_push_runs"


def _table_exists(table_name: str) -> bool:
    inspector = inspect(op.get_bind())
    return table_name in inspector.get_table_names()


def _column_names(table_name: str) -> set[str]:
    inspector = inspect(op.get_bind())
    if table_name not in inspector.get_table_names():
        return set()
    return {column["name"] for column in inspector.get_columns(table_name)}


def _add_column_if_missing(table_name: str, column: sa.Column) -> None:
    if column.name not in _column_names(table_name):
        op.add_column(table_name, column)


def upgrade() -> None:
    _add_column_if_missing(
        WORK_TASKS_TABLE,
        sa.Column("source_system", sa.String(length=64), nullable=True),
    )
    _add_column_if_missing(
        WORK_TASKS_TABLE,
        sa.Column("external_id", sa.String(length=128), nullable=True),
    )
    op.create_index(
        op.f("ix_work_tasks_external_id"),
        WORK_TASKS_TABLE,
        ["external_id"],
        unique=False,
        if_not_exists=True,
    )

    for column in [
        sa.Column("assigned_to", sa.String(length=128), nullable=True),
        sa.Column("resolved_at", sa.DateTime(), nullable=True),
        sa.Column("closed_at", sa.DateTime(), nullable=True),
        sa.Column("closed_reason", sa.Text(), nullable=True),
        sa.Column("review_summary", sa.Text(), nullable=True),
    ]:
        _add_column_if_missing(RISK_EVENTS_TABLE, column)
    op.create_index(
        op.f("ix_risk_events_assigned_to"),
        RISK_EVENTS_TABLE,
        ["assigned_to"],
        unique=False,
        if_not_exists=True,
    )

    if not _table_exists(RISK_EVENT_ACTIONS_TABLE):
        op.create_table(
            RISK_EVENT_ACTIONS_TABLE,
            sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
            sa.Column("code", sa.String(length=64), nullable=False),
            sa.Column("risk_event_id", sa.Integer(), nullable=False),
            sa.Column("action", sa.String(length=64), nullable=False),
            sa.Column("actor", sa.String(length=128), nullable=False),
            sa.Column("from_status", sa.String(length=32), nullable=True),
            sa.Column("to_status", sa.String(length=32), nullable=True),
            sa.Column("assigned_to", sa.String(length=128), nullable=True),
            sa.Column("comment", sa.Text(), nullable=True),
            sa.Column("payload", sa.JSON(), nullable=True),
            sa.Column("created_at", sa.DateTime(), nullable=False),
            sa.PrimaryKeyConstraint("id"),
        )
        for column_name, unique in [
            ("code", True),
            ("risk_event_id", False),
            ("action", False),
            ("actor", False),
            ("assigned_to", False),
            ("created_at", False),
        ]:
            op.create_index(
                op.f(f"ix_{RISK_EVENT_ACTIONS_TABLE}_{column_name}"),
                RISK_EVENT_ACTIONS_TABLE,
                [column_name],
                unique=unique,
            )

    if not _table_exists(REPORT_PUSH_RUNS_TABLE):
        op.create_table(
            REPORT_PUSH_RUNS_TABLE,
            sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
            sa.Column("code", sa.String(length=64), nullable=False),
            sa.Column("report_type", sa.String(length=64), nullable=False),
            sa.Column("title", sa.String(length=255), nullable=True),
            sa.Column("receive_id", sa.String(length=128), nullable=True),
            sa.Column("receive_id_type", sa.String(length=64), nullable=False),
            sa.Column("status", sa.String(length=32), nullable=False),
            sa.Column("task_id", sa.String(length=128), nullable=True),
            sa.Column("actor", sa.String(length=128), nullable=False),
            sa.Column("provider_response", sa.JSON(), nullable=True),
            sa.Column("error_message", sa.Text(), nullable=True),
            sa.Column("queued_at", sa.DateTime(), nullable=False),
            sa.Column("sent_at", sa.DateTime(), nullable=True),
            sa.Column("created_at", sa.DateTime(), nullable=False),
            sa.Column("updated_at", sa.DateTime(), nullable=False),
            sa.PrimaryKeyConstraint("id"),
        )
        for column_name, unique in [
            ("code", True),
            ("report_type", False),
            ("receive_id", False),
            ("receive_id_type", False),
            ("status", False),
            ("task_id", False),
            ("actor", False),
            ("queued_at", False),
        ]:
            op.create_index(
                op.f(f"ix_{REPORT_PUSH_RUNS_TABLE}_{column_name}"),
                REPORT_PUSH_RUNS_TABLE,
                [column_name],
                unique=unique,
            )


def downgrade() -> None:
    if _table_exists(REPORT_PUSH_RUNS_TABLE):
        op.drop_table(REPORT_PUSH_RUNS_TABLE)
    if _table_exists(RISK_EVENT_ACTIONS_TABLE):
        op.drop_table(RISK_EVENT_ACTIONS_TABLE)

    risk_columns = _column_names(RISK_EVENTS_TABLE)
    if "assigned_to" in risk_columns:
        op.drop_index(op.f("ix_risk_events_assigned_to"), table_name=RISK_EVENTS_TABLE)
    for column_name in [
        "review_summary",
        "closed_reason",
        "closed_at",
        "resolved_at",
        "assigned_to",
    ]:
        if column_name in _column_names(RISK_EVENTS_TABLE):
            op.drop_column(RISK_EVENTS_TABLE, column_name)

    task_columns = _column_names(WORK_TASKS_TABLE)
    if "external_id" in task_columns:
        op.drop_index(op.f("ix_work_tasks_external_id"), table_name=WORK_TASKS_TABLE)
        op.drop_column(WORK_TASKS_TABLE, "external_id")
    if "source_system" in _column_names(WORK_TASKS_TABLE):
        op.drop_column(WORK_TASKS_TABLE, "source_system")