390 lines
11 KiB
Python
390 lines
11 KiB
Python
"""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},
|
|
)
|