company-ai-platform/alembic/versions/202607150001_architecture_hardening.py

"""Harden event, workflow, and owned-ledger constraints.

Revision ID: 202607150001
Revises: 202607120004
Create Date: 2026-07-15
"""

from alembic import op
import sqlalchemy as sa


revision = "202607150001"
down_revision = "202607120004"
branch_labels = None
depends_on = None


def _scalar_count(sql: str) -> int:
    return int(op.get_bind().execute(sa.text(sql)).scalar() or 0)


def _assert_owned_relations_are_consistent() -> None:
    orphan_risk_actions = _scalar_count(
        """
        SELECT COUNT(*)
        FROM risk_event_actions action
        LEFT JOIN risk_events event ON event.id = action.risk_event_id
        WHERE event.id IS NULL
        """
    )
    if orphan_risk_actions:
        raise RuntimeError(
            "Cannot add risk-event foreign key: orphan risk_event_actions rows exist"
        )

    orphan_workflow_actions = _scalar_count(
        """
        SELECT COUNT(*)
        FROM workflow_actions action
        LEFT JOIN workflow_instances workflow ON workflow.code = action.workflow_code
        WHERE workflow.code IS NULL
        """
    )
    if orphan_workflow_actions:
        raise RuntimeError(
            "Cannot add workflow foreign key: orphan workflow_actions rows exist"
        )

    duplicate_workflows = _scalar_count(
        """
        SELECT COUNT(*)
        FROM (
            SELECT workflow_type, aggregate_type, aggregate_id
            FROM workflow_instances
            WHERE aggregate_id IS NOT NULL
            GROUP BY workflow_type, aggregate_type, aggregate_id
            HAVING COUNT(*) > 1
        ) duplicates
        """
    )
    if duplicate_workflows:
        raise RuntimeError(
            "Cannot add workflow uniqueness constraint: duplicate aggregate workflows exist"
        )


def upgrade() -> None:
    _assert_owned_relations_are_consistent()
    op.execute("UPDATE domain_events SET max_attempts = 3 WHERE max_attempts IS NULL")

    with op.batch_alter_table("domain_events") as batch_op:
        batch_op.alter_column(
            "max_attempts",
            existing_type=sa.Integer(),
            nullable=False,
            server_default=sa.text("3"),
        )

    with op.batch_alter_table("risk_event_actions") as batch_op:
        batch_op.create_foreign_key(
            "fk_risk_event_actions_risk_event_id",
            "risk_events",
            ["risk_event_id"],
            ["id"],
            ondelete="RESTRICT",
        )

    with op.batch_alter_table("workflow_actions") as batch_op:
        batch_op.create_foreign_key(
            "fk_workflow_actions_workflow_code",
            "workflow_instances",
            ["workflow_code"],
            ["code"],
            ondelete="RESTRICT",
        )

    with op.batch_alter_table("workflow_instances") as batch_op:
        batch_op.create_unique_constraint(
            "uq_workflow_aggregate",
            ["workflow_type", "aggregate_type", "aggregate_id"],
        )


def downgrade() -> None:
    with op.batch_alter_table("workflow_instances") as batch_op:
        batch_op.drop_constraint("uq_workflow_aggregate", type_="unique")

    with op.batch_alter_table("workflow_actions") as batch_op:
        batch_op.drop_constraint(
            "fk_workflow_actions_workflow_code",
            type_="foreignkey",
        )

    with op.batch_alter_table("risk_event_actions") as batch_op:
        batch_op.drop_constraint(
            "fk_risk_event_actions_risk_event_id",
            type_="foreignkey",
        )

    with op.batch_alter_table("domain_events") as batch_op:
        batch_op.alter_column(
            "max_attempts",
            existing_type=sa.Integer(),
            nullable=True,
            server_default=None,
        )