vignette/scripts/check-measurement-provenance-census.py

"""Fail-closed census for every provenance-bearing table in the local database.

The report contains counts and identifiers only. It never emits connection
strings, free-text evidence, assessment payloads, or transcript content.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
from pathlib import Path
from typing import Any

import asyncpg


PROVENANCE_COLUMNS = {
    "source_kind",
    "perspective",
    "instrument_id",
    "model_run_id",
    "evidence_turn_ids",
}
KNOWN_PRODUCER_TABLES = {
    "app.calibration_assessment_snapshot",
    "app.calibration_metacognitive_prescription",
    "app.calibration_performance_observation",
    "app.calibration_prediction_revision",
    "app.calibration_subgroup_drift_report",
    "app.calibration_teacher_review_event",
    "app.calibration_transfer_assessment",
    "app.calibration_transfer_execution_event",
    "app.calibration_transfer_suite",
    "app.calibration_transfer_trial",
    "app.competency_graph_snapshot",
    "app.measurement_event",
    "app.measurement_instrument",
    "app.multimodal_axis_measurement",
    "app.outcome_trajectory_observation",
    "app.practice_attempt_evidence",
    "app.practice_coaching_card",
    "app.practice_episode_submission",
    "app.practice_prescription",
    "app.practice_teacher_correction",
    "app.relationship_memory_event",
    "app.rupture_observation_event",
    "app.rupture_reconciliation_revision",
    "app.self_assessment",
    "app.supervision_evaluation_batch",
    "app.supervision_teacher_ai_disagreement",
    "audit.model_run",
    "ds.benchmark_case",
    "ds.benchmark_observation",
}
SOURCE_KINDS = {
    "learner_reported",
    "agent_reported",
    "model_inferred",
    "human_rated",
    "observed_runtime",
    "simulated_state",
}
PERSPECTIVES = {
    "learner_self_report",
    "client_agent_report",
    "independent_observer",
    "supervisor_human",
    "runtime_observation",
    "client_simulation",
}
SOURCE_PERSPECTIVE = {
    "learner_reported": {"learner_self_report"},
    "agent_reported": {"client_agent_report"},
    "model_inferred": {"independent_observer"},
    "human_rated": {"supervisor_human"},
    "observed_runtime": {"runtime_observation"},
    "simulated_state": {"client_simulation"},
}
TABLE_SOURCE_KIND_OVERRIDES = {
    "app.multimodal_axis_measurement": {
        "model_inferred_text",
        "model_inferred_voice",
    },
}
# G7 records provider-specific instrument/model identifiers under its own
# constrained ledger. They are deliberately not foreign keys into the G0
# instrument/model-run tables, so central orphan checks would be category errors.
TABLE_LOCAL_PROVENANCE = {"app.multimodal_axis_measurement"}


class CensusError(RuntimeError):
    pass


def _load_api_env() -> None:
    env_path = Path("apps/api/.env")
    if not env_path.exists():
        return
    for raw_line in env_path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))


def _identifier(value: str) -> str:
    if not value.replace("_", "").isalnum():
        raise CensusError(f"unsafe SQL identifier discovered: {value!r}")
    return '"' + value + '"'


async def _scalar(conn: asyncpg.Connection[Any], sql: str) -> int:
    return int(await conn.fetchval(sql) or 0)


async def census(dsn: str) -> dict[str, Any]:
    conn = await asyncpg.connect(dsn)
    try:
        # The application role is intentionally fail-closed under RLS. The
        # census is an evaluator-plane integrity job, so establish the same
        # explicit AI view used by runtime provenance workers before counting.
        await conn.execute("SELECT set_config('app.ai_context', '1', false)")
        await conn.execute("SELECT set_config('app.current_ai_view', 'evaluator', false)")
        columns = await conn.fetch(
            """
            SELECT table_schema, table_name, column_name, data_type, udt_name
            FROM information_schema.columns
            WHERE table_schema IN ('app','audit','ds')
            ORDER BY table_schema, table_name, ordinal_position
            """
        )
        discovered: dict[str, dict[str, tuple[str, str]]] = {}
        for row in columns:
            key = f"{row['table_schema']}.{row['table_name']}"
            discovered.setdefault(key, {})[str(row["column_name"])] = (
                str(row["data_type"]),
                str(row["udt_name"]),
            )
        discovered = {
            table: table_columns
            for table, table_columns in discovered.items()
            if PROVENANCE_COLUMNS & set(table_columns)
        }
        unknown_tables = sorted(set(discovered) - KNOWN_PRODUCER_TABLES)
        missing_tables = sorted(KNOWN_PRODUCER_TABLES - set(discovered))
        violations: dict[str, int] = {
            "unknown_producer_table": len(unknown_tables),
            "missing_registered_producer_table": len(missing_tables),
        }
        table_reports: list[dict[str, Any]] = []
        for table in sorted(discovered):
            schema, name = table.split(".", 1)
            qualified = f"{_identifier(schema)}.{_identifier(name)}"
            table_columns = discovered[table]
            checks: dict[str, int] = {"row_count": await _scalar(conn, f"SELECT count(*) FROM {qualified}")}
            for column in ("source_kind", "perspective", "instrument_id", "instrument_version"):
                if column in table_columns:
                    checks[f"null_or_blank_{column}"] = await _scalar(
                        conn,
                        f"SELECT count(*) FROM {qualified} WHERE {_identifier(column)} IS NULL OR btrim({_identifier(column)}::text) = ''",
                    )
            if "source_kind" in table_columns:
                allowed_source_kinds = TABLE_SOURCE_KIND_OVERRIDES.get(
                    table, SOURCE_KINDS
                )
                allowed = ",".join(
                    "'" + item + "'" for item in sorted(allowed_source_kinds)
                )
                checks["unknown_source_kind"] = await _scalar(
                    conn,
                    f"SELECT count(*) FROM {qualified} WHERE source_kind IS NOT NULL AND source_kind NOT IN ({allowed})",
                )
            if "perspective" in table_columns:
                allowed = ",".join("'" + item + "'" for item in sorted(PERSPECTIVES))
                checks["unknown_perspective"] = await _scalar(
                    conn,
                    f"SELECT count(*) FROM {qualified} WHERE perspective IS NOT NULL AND perspective NOT IN ({allowed})",
                )
            if (
                {"source_kind", "perspective"} <= set(table_columns)
                and table not in TABLE_SOURCE_KIND_OVERRIDES
            ):
                valid_pairs = " OR ".join(
                    f"(source_kind='{source}' AND perspective='{perspective}')"
                    for source, perspectives in SOURCE_PERSPECTIVE.items()
                    for perspective in perspectives
                )
                checks["incompatible_source_perspective"] = await _scalar(
                    conn,
                    f"SELECT count(*) FROM {qualified} WHERE source_kind IS NOT NULL AND perspective IS NOT NULL AND NOT ({valid_pairs})",
                )
            if (
                {"instrument_id", "instrument_version"} <= set(table_columns)
                and table != "app.measurement_instrument"
                and table not in TABLE_LOCAL_PROVENANCE
            ):
                checks["orphan_instrument"] = await _scalar(
                    conn,
                    f"SELECT count(*) FROM {qualified} p LEFT JOIN app.measurement_instrument i ON i.instrument_id=p.instrument_id AND i.instrument_version=p.instrument_version WHERE p.instrument_id IS NOT NULL AND i.instrument_id IS NULL",
                )
            if (
                "model_run_id" in table_columns
                and table != "audit.model_run"
                and table not in TABLE_LOCAL_PROVENANCE
            ):
                checks["orphan_model_run"] = await _scalar(
                    conn,
                    f"SELECT count(*) FROM {qualified} p LEFT JOIN audit.model_run m ON m.model_run_id=p.model_run_id WHERE p.model_run_id IS NOT NULL AND m.model_run_id IS NULL",
                )
                if "source_kind" in table_columns:
                    checks["missing_required_model_run"] = await _scalar(
                        conn,
                        f"SELECT count(*) FROM {qualified} WHERE source_kind IN ('model_inferred','agent_reported') AND model_run_id IS NULL",
                    )
            if "evidence_turn_ids" in table_columns and table_columns["evidence_turn_ids"][1] == "_uuid":
                checks["orphan_evidence_turn"] = await _scalar(
                    conn,
                    f"SELECT count(*) FROM {qualified} p CROSS JOIN LATERAL unnest(COALESCE(p.evidence_turn_ids, ARRAY[]::uuid[])) evidence(turn_id) LEFT JOIN app.turns t ON t.id=evidence.turn_id WHERE t.id IS NULL",
                )
            for check, count in checks.items():
                if check != "row_count":
                    violations[check] = violations.get(check, 0) + count
            table_reports.append(
                {
                    "producer_table": table,
                    "provenance_contract": (
                        "table_local" if table in TABLE_LOCAL_PROVENANCE else "g0_central"
                    ),
                    **checks,
                }
            )
    finally:
        await conn.close()
    failed = {name: count for name, count in violations.items() if count}
    return {
        "ok": not failed,
        "policy": "fail_closed_on_unknown_producer_or_null_or_orphan_provenance",
        "registered_producer_table_count": len(KNOWN_PRODUCER_TABLES),
        "discovered_producer_table_count": len(discovered),
        "unknown_producer_tables": unknown_tables,
        "missing_registered_producer_tables": missing_tables,
        "violations": violations,
        "failed_checks": failed,
        "tables": table_reports,
    }


async def _run(args: argparse.Namespace) -> dict[str, Any]:
    _load_api_env()
    dsn = args.database_url or os.environ.get("DATABASE_URL")
    if not dsn:
        raise CensusError("DATABASE_URL is required via --database-url or apps/api/.env")
    return await census(dsn)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--database-url", default="")
    parser.add_argument("--out", default="")
    args = parser.parse_args()
    result = asyncio.run(_run(args))
    text = json.dumps(result, ensure_ascii=False, indent=2, default=str)
    if args.out:
        path = Path(args.out)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(text + "\n", encoding="utf-8")
    print(text)
    if not result["ok"]:
        raise SystemExit(1)


if __name__ == "__main__":
    main()