"""Where a run came from: a schedule, an alert rule, or an agent over MCP (#5212)

Revision ID: 0031_ansible_run_origin
Revises: 0030_plugin_metric_latest
Create Date: 2026-10-07
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa

revision: str = "0031_ansible_run_origin"
down_revision: Union[str, None] = "0030_plugin_metric_latest"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    op.add_column("ansible_runs", sa.Column("origin", sa.String(16), nullable=True))
    op.add_column("ansible_runs", sa.Column("origin_id", sa.String(36), nullable=True))
    # Recover what the record already says about past runs. These tables are
    # small (runs, alert events that started one, audit entries), unlike the
    # metrics seed 0030 left to the background.
    op.execute("""
        UPDATE ansible_runs r SET origin = 'alert', origin_id = e.rule_id
        FROM alert_events e WHERE e.ansible_run_id = r.id
    """)
    op.execute("""
        UPDATE ansible_runs r SET origin = 'schedule', origin_id = s.id
        FROM ansible_schedules s WHERE s.last_run_id = r.id AND r.origin IS NULL
    """)
    op.execute("""
        UPDATE ansible_runs r SET origin = 'mcp'
        FROM audit_events a
        WHERE a.action = 'ansible.run.started' AND a.entity_id = r.id
          AND a.username LIKE '% via MCP' AND r.origin IS NULL
    """)


def downgrade() -> None:
    op.drop_column("ansible_runs", "origin_id")
    op.drop_column("ansible_runs", "origin")