"""Add v2 LLM validation models (tables, columns, indexes) Adds: - validation_sources table - validation_runs table - v2 columns to validation_policies - v2 columns + FK + indexes to llm_validation_results - index on validation_policies.provider_id - optional data backfill: creates ValidationRun placeholders for existing ValidationRecord rows that have a policy_id but no run_id @ADR [MIGRATION-001] Backfill creates one run per distinct policy_id. @RATIONALE Existing records were created before the ValidationRun model existed. Grouping by policy_id avoids creating excessive runs for what is essentially one historic execution batch per policy. The run status is set to 'completed' to avoid confusing downstream consumers that filter on 'running'. Revision ID: c9d8e7f6a5b4 Revises: 2df63b7ce038 Create Date: 2026-05-31 12:00:00.000000 """ from collections.abc import Sequence from alembic import op import sqlalchemy as sa from sqlalchemy import inspect # revision identifiers, used by Alembic. revision: str = "c9d8e7f6a5b4" down_revision: str | Sequence[str] | None = "2df63b7ce038" branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None def _table_exists(table_name: str) -> bool: conn = op.get_bind() inspector = inspect(conn) return table_name in inspector.get_table_names() def upgrade() -> None: """Upgrade schema to v2 validation models.""" _create_validation_sources() _create_validation_runs() _extend_validation_policies() _extend_validation_results() _backfill_validation_runs() def downgrade() -> None: """Downgrade schema by removing v2 additions in reverse order.""" # Drop FK constraints from llm_validation_results op.drop_constraint( "fk_llm_validation_results_source_id", "llm_validation_results", type_="foreignkey", ) op.drop_constraint( "fk_llm_validation_results_run_id", "llm_validation_results", type_="foreignkey", ) op.drop_index( op.f("ix_llm_validation_results_run_id"), table_name="llm_validation_results", ) # Remove v2 columns from llm_validation_results _v2_result_columns = [ "screenshot_paths", "timings", "token_usage", "logs_sent_to_llm", "tab_screenshots", "chart_data_results", "dataset_health", "execution_path", "source_id", "run_id", ] for col in _v2_result_columns: op.drop_column("llm_validation_results", col) # Remove columns + index from validation_policies op.drop_index( op.f("ix_validation_policies_provider_id"), table_name="validation_policies", ) _v2_policy_columns = [ "source_snapshot", "policy_dashboard_concurrency_limit", "llm_batch_size", "execute_chart_data", "logs_enabled", "screenshot_enabled", "prompt_template", ] for col in _v2_policy_columns: op.drop_column("validation_policies", col) # Drop validation_runs table (index first, then table) op.drop_index( op.f("ix_validation_runs_policy_id"), table_name="validation_runs", ) op.drop_table("validation_runs") # Drop validation_sources table (index first, then table) op.drop_index( op.f("ix_validation_sources_policy_id"), table_name="validation_sources", ) op.drop_table("validation_sources") # --------------------------------------------------------------------------- # Upgrade helpers (grouped for readability) # --------------------------------------------------------------------------- def _create_validation_sources() -> None: """Create the validation_sources table (no FK dependencies).""" if not _table_exists("validation_policies"): return op.create_table( "validation_sources", sa.Column("id", sa.String(), nullable=False), sa.Column("policy_id", sa.String(), nullable=False), sa.Column("type", sa.String(), nullable=False), sa.Column("value", sa.String(), nullable=False), sa.Column("parsed_context", sa.JSON(), nullable=True), sa.Column( "status", sa.String(), nullable=False, server_default="valid" ), sa.Column("resolved_dashboard_id", sa.String(), nullable=True), sa.Column("title", sa.String(), nullable=True), sa.Column("last_error", sa.String(), nullable=True), sa.Column("last_checked_at", sa.DateTime(), nullable=True), sa.Column( "created_at", sa.DateTime(), nullable=False, server_default=sa.func.now(), ), sa.ForeignKeyConstraint( ["policy_id"], ["validation_policies.id"], ), sa.PrimaryKeyConstraint("id"), ) op.create_index( op.f("ix_validation_sources_policy_id"), "validation_sources", ["policy_id"], ) def _create_validation_runs() -> None: """Create the validation_runs table.""" if not _table_exists("validation_policies"): return op.create_table( "validation_runs", sa.Column("id", sa.String(), nullable=False), sa.Column("policy_id", sa.String(), nullable=False), sa.Column("task_id", sa.String(), nullable=True), sa.Column( "started_at", sa.DateTime(), nullable=False, server_default=sa.func.now(), ), sa.Column("finished_at", sa.DateTime(), nullable=True), sa.Column( "trigger", sa.String(), nullable=False, server_default="manual" ), sa.Column( "status", sa.String(), nullable=False, server_default="running" ), sa.Column( "dashboard_count", sa.Integer(), nullable=False, server_default="0", ), sa.Column( "pass_count", sa.Integer(), nullable=False, server_default="0" ), sa.Column( "warn_count", sa.Integer(), nullable=False, server_default="0" ), sa.Column( "fail_count", sa.Integer(), nullable=False, server_default="0" ), sa.Column( "unknown_count", sa.Integer(), nullable=False, server_default="0" ), sa.Column( "created_at", sa.DateTime(), nullable=False, server_default=sa.func.now(), ), sa.ForeignKeyConstraint( ["policy_id"], ["validation_policies.id"], ), sa.PrimaryKeyConstraint("id"), ) op.create_index( op.f("ix_validation_runs_policy_id"), "validation_runs", ["policy_id"], ) def _extend_validation_policies() -> None: """Add v2 columns and index to validation_policies.""" if not _table_exists("validation_policies"): return op.add_column( "validation_policies", sa.Column("prompt_template", sa.Text(), nullable=True), ) op.add_column( "validation_policies", sa.Column( "screenshot_enabled", sa.Boolean(), nullable=False, server_default="true", ), ) op.add_column( "validation_policies", sa.Column( "logs_enabled", sa.Boolean(), nullable=False, server_default="true", ), ) op.add_column( "validation_policies", sa.Column( "execute_chart_data", sa.Boolean(), nullable=False, server_default="false", ), ) op.add_column( "validation_policies", sa.Column( "llm_batch_size", sa.Integer(), nullable=False, server_default="1", ), ) op.add_column( "validation_policies", sa.Column( "policy_dashboard_concurrency_limit", sa.Integer(), nullable=False, server_default="3", ), ) op.add_column( "validation_policies", sa.Column("source_snapshot", sa.JSON(), nullable=True), ) op.create_index( op.f("ix_validation_policies_provider_id"), "validation_policies", ["provider_id"], ) def _extend_validation_results() -> None: """Add v2 columns, FKs, and indexes to llm_validation_results.""" if not _table_exists("llm_validation_results"): return op.add_column( "llm_validation_results", sa.Column( "run_id", sa.String(), sa.ForeignKey("validation_runs.id", name="fk_llm_validation_results_run_id"), nullable=True, ), ) op.add_column( "llm_validation_results", sa.Column( "source_id", sa.String(), sa.ForeignKey("validation_sources.id", name="fk_llm_validation_results_source_id"), nullable=True, ), ) op.add_column( "llm_validation_results", sa.Column("execution_path", sa.String(), nullable=True), ) op.add_column( "llm_validation_results", sa.Column("dataset_health", sa.JSON(), nullable=True), ) op.add_column( "llm_validation_results", sa.Column("chart_data_results", sa.JSON(), nullable=True), ) op.add_column( "llm_validation_results", sa.Column("tab_screenshots", sa.JSON(), nullable=True), ) op.add_column( "llm_validation_results", sa.Column("logs_sent_to_llm", sa.JSON(), nullable=True), ) op.add_column( "llm_validation_results", sa.Column("token_usage", sa.JSON(), nullable=True), ) op.add_column( "llm_validation_results", sa.Column("timings", sa.JSON(), nullable=True), ) op.add_column( "llm_validation_results", sa.Column("screenshot_paths", sa.JSON(), nullable=True), ) op.create_index( op.f("ix_llm_validation_results_run_id"), "llm_validation_results", ["run_id"], ) def _backfill_validation_runs() -> None: """Create placeholder ValidationRun entries for existing records. For each distinct policy_id that has ValidationRecord rows without a run_id, creates one ValidationRun (status='completed') and links all matching records to it. Safe no-op when no such rows exist (e.g. fresh database). """ if not _table_exists("llm_validation_results"): return conn = op.get_bind() # Check for existing records that need backfill result = conn.execute( sa.text( "SELECT DISTINCT v.policy_id " "FROM llm_validation_results v " "WHERE v.policy_id IS NOT NULL " "AND v.run_id IS NULL" ) ) rows = result.fetchall() if not rows: return import uuid for (policy_id,) in rows: run_id = str(uuid.uuid4()) conn.execute( sa.text( "INSERT INTO validation_runs " "(id, policy_id, started_at, trigger, status, created_at) " "VALUES (:run_id, :policy_id, NOW(), 'manual', 'completed', NOW())" ), {"run_id": run_id, "policy_id": policy_id}, ) conn.execute( sa.text( "UPDATE llm_validation_results " "SET run_id = :run_id " "WHERE policy_id = :policy_id AND run_id IS NULL" ), {"run_id": run_id, "policy_id": policy_id}, )