vignette/scripts/rehearse-measurement-migration.py

"""Rehearse G0 apply/transactional rollback and legacy lossless projection.

A uniquely named temporary local database is created, initialized through G0's
prerequisites, exercised, and removed in ``finally``. No configured database is
modified and no connection details are printed or written to evidence.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import secrets
import subprocess
import sys
from pathlib import Path
from typing import Any
from uuid import UUID, uuid4

import asyncpg


TEMP_DATABASE_PREFIX = "vignette_g0_rehearsal_"
BASE_MIGRATIONS = (
    "01_extensions.sql",
    "02_schema.sql",
    "03_kb.sql",
    "04_audit_eval_rls.sql",
    "05_runtime_auth.sql",
    "06_session_evaluation.sql",
)
G0_MIGRATION = "07_measurement_foundation.sql"


class RehearsalError(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 _database_identifier(value: str) -> str:
    if not value.startswith(TEMP_DATABASE_PREFIX) or not value.replace("_", "").isalnum():
        raise RehearsalError("refusing unsafe temporary database identifier")
    return '"' + value + '"'


def _role_identifier(value: str) -> str:
    if not value.replace("_", "").isalnum():
        raise RehearsalError("refusing unsafe database owner identifier")
    return '"' + value + '"'


async def _g0_objects(conn: asyncpg.Connection[Any]) -> dict[str, bool]:
    row = await conn.fetchrow(
        """
        SELECT
          to_regclass('app.measurement_instrument') IS NOT NULL AS instrument,
          to_regclass('app.measurement_event') IS NOT NULL AS event,
          to_regclass('app.alliance_pulse') IS NOT NULL AS pulse,
          to_regclass('audit.model_run') IS NOT NULL AS model_run,
          to_regclass('ds.benchmark_case') IS NOT NULL AS benchmark
        """
    )
    assert row is not None
    return {key: bool(row[key]) for key in row.keys()}


async def _seed_and_compare_legacy(conn: asyncpg.Connection[Any]) -> dict[str, Any]:
    learner_id = uuid4()
    persona_id = uuid4()
    case_id = uuid4()
    session_id = uuid4()
    turn_id = uuid4()
    await conn.execute(
        "INSERT INTO app.app_user (user_id, external_id, role) VALUES ($1, $2, 'learner')",
        learner_id,
        f"g0-rehearsal:{learner_id}",
    )
    await conn.execute(
        """
        INSERT INTO app.persona_card (
          persona_id, code, version, status, display_name, difficulty,
          demographics, presenting, history, big5, resistance, speech_style,
          affect_baseline, ccd, dsm5_dimensional, source_provenance
        ) VALUES ($1, 'G0-FIXTURE', 1, 'draft', 'G0 Fixture', 'easy',
                  '{}'::jsonb, '{}'::jsonb, '{}'::jsonb, '{}'::jsonb,
                  '{}'::jsonb, '{}'::jsonb, '{}'::jsonb, '{}'::jsonb,
                  '{}'::jsonb, 'local migration rehearsal')
        """,
        persona_id,
    )
    await conn.execute(
        """
        INSERT INTO app.case_profile (case_id, persona_id, learner_id, alliance_level)
        VALUES ($1, $2, $3, 0.375)
        """,
        case_id,
        persona_id,
        learner_id,
    )
    await conn.execute(
        """
        INSERT INTO app.sessions (
          id, case_id, learner_id, persona_id, persona_version, session_no, ended_at
        ) VALUES ($1, $2, $3, $4, 1, 1, now())
        """,
        session_id,
        case_id,
        learner_id,
        persona_id,
    )
    await conn.execute(
        """
        INSERT INTO app.turns (id, session_id, seq, speaker, text, text_masked)
        VALUES ($1, $2, 1, 'counselor', '[masked fixture]', '[masked fixture]')
        """,
        turn_id,
        session_id,
    )
    await conn.execute(
        "INSERT INTO app.session_state (session_id, rapport_credit) VALUES ($1, 0.625)",
        session_id,
    )
    await conn.execute(
        """
        INSERT INTO app.session_evaluation (session_id, status, source, scope, stage, payload)
        VALUES ($1, 'ready', 'fixture', 'session', '정리',
                '{"distribution":{"total":7}}'::jsonb)
        """,
        session_id,
    )
    old = await conn.fetchrow(
        """
        SELECT ss.rapport_credit::double precision AS rapport_credit,
               cp.alliance_level::double precision AS alliance_level,
               (se.payload #>> '{distribution,total}')::double precision AS distribution_total,
               t.id AS turn_id
        FROM app.sessions s
        JOIN app.session_state ss ON ss.session_id=s.id
        JOIN app.case_profile cp ON cp.case_id=s.case_id
        JOIN app.session_evaluation se ON se.session_id=s.id
        JOIN app.turns t ON t.session_id=s.id AND t.seq=1
        WHERE s.id=$1
        """,
        session_id,
    )
    if old is None:
        raise RehearsalError("legacy persisted fixture could not be read")

    api_path = str((Path("apps/api")).resolve())
    if api_path not in sys.path:
        sys.path.insert(0, api_path)
    from app.services.measurement_legacy import (  # pylint: disable=import-outside-toplevel
        adapt_deep_evaluation,
        adapt_legacy_simulation_signals,
    )

    simulation = adapt_legacy_simulation_signals(
        session_id=session_id,
        rapport_credit=float(old["rapport_credit"]),
        alliance_level=float(old["alliance_level"]),
        turn_id=UUID(str(old["turn_id"])),
    )
    deep = adapt_deep_evaluation(
        session_id=session_id,
        evaluation={"distribution": {"total": int(old["distribution_total"])}},
        model_run_id=uuid4(),
        evidence_turn_ids=(turn_id,),
    )
    projected = {
        str(event.metadata["legacy_signal"]): event.value
        for event in (*simulation, deep)
    }
    expected = {
        "session_state.rapport_credit": float(old["rapport_credit"]),
        "case_profile.alliance_level": float(old["alliance_level"]),
        "SessionEvaluation.distribution": float(old["distribution_total"]),
    }
    persisted_after = await conn.fetchrow(
        """
        SELECT ss.rapport_credit::double precision AS rapport_credit,
               cp.alliance_level::double precision AS alliance_level,
               (se.payload #>> '{distribution,total}')::double precision AS distribution_total
        FROM app.sessions s
        JOIN app.session_state ss ON ss.session_id=s.id
        JOIN app.case_profile cp ON cp.case_id=s.case_id
        JOIN app.session_evaluation se ON se.session_id=s.id
        WHERE s.id=$1
        """,
        session_id,
    )
    unchanged = persisted_after is not None and [
        float(persisted_after["rapport_credit"]),
        float(persisted_after["alliance_level"]),
        float(persisted_after["distribution_total"]),
    ] == [
        expected["session_state.rapport_credit"],
        expected["case_profile.alliance_level"],
        expected["SessionEvaluation.distribution"],
    ]
    return {
        "fixture_kind": "persisted_legacy_rows",
        "signal_count": len(expected),
        "old_read_model": expected,
        "new_read_model": projected,
        "numeric_projection_lossless": projected == expected,
        "legacy_rows_unchanged_after_g0_apply": unchanged,
    }


async def _docker_database_command(
    container: str,
    admin_user: str,
    database: str,
    sql: str,
) -> None:
    process = await asyncio.create_subprocess_exec(
        "docker",
        "exec",
        container,
        "psql",
        "-U",
        admin_user,
        "-d",
        database,
        "-v",
        "ON_ERROR_STOP=1",
        "-c",
        sql,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    stdout, stderr = await process.communicate()
    if process.returncode:
        detail = (stderr or stdout).decode("utf-8", errors="replace")[-500:]
        raise RehearsalError(
            f"local Docker PostgreSQL admin command failed ({process.returncode}): {detail}"
        )


async def rehearse(
    dsn: str,
    *,
    docker_container: str,
    docker_admin_user: str,
    database_owner: str,
) -> dict[str, Any]:
    temp_name = TEMP_DATABASE_PREFIX + secrets.token_hex(6)
    identifier = _database_identifier(temp_name)
    created = False
    try:
        await _docker_database_command(
            docker_container,
            docker_admin_user,
            "postgres",
            f"CREATE DATABASE {identifier} OWNER {_role_identifier(database_owner)}",
        )
        created = True
        await _docker_database_command(
            docker_container,
            docker_admin_user,
            temp_name,
            "CREATE EXTENSION IF NOT EXISTS vector WITH SCHEMA public; "
            "CREATE EXTENSION IF NOT EXISTS pgcrypto WITH SCHEMA public; "
            "CREATE EXTENSION IF NOT EXISTS pg_trgm WITH SCHEMA public;",
        )
        temp = await asyncpg.connect(dsn=dsn, database=temp_name)
        try:
            migration_dir = Path("infra/db/init")
            for name in BASE_MIGRATIONS:
                await temp.execute((migration_dir / name).read_text(encoding="utf-8"))
            before = await _g0_objects(temp)
            if any(before.values()):
                raise RehearsalError(f"G0 objects existed before apply: {before}")
            g0_sql = (migration_dir / G0_MIGRATION).read_text(encoding="utf-8")
            transaction = temp.transaction()
            await transaction.start()
            await temp.execute(g0_sql)
            inside_transaction = await _g0_objects(temp)
            if not all(inside_transaction.values()):
                raise RehearsalError(f"G0 apply omitted expected objects: {inside_transaction}")
            await transaction.rollback()
            after_rollback = await _g0_objects(temp)
            if after_rollback != before:
                raise RehearsalError(
                    f"transactional rollback was not schema-lossless: {after_rollback} != {before}"
                )
            await temp.execute(g0_sql)
            after_committed_apply = await _g0_objects(temp)
            legacy = await _seed_and_compare_legacy(temp)
            if not legacy["numeric_projection_lossless"] or not legacy["legacy_rows_unchanged_after_g0_apply"]:
                raise RehearsalError("legacy old/new read-model comparison was lossy")
        finally:
            await temp.close()
    finally:
        if created:
            await _docker_database_command(
                docker_container,
                docker_admin_user,
                "postgres",
                f"DROP DATABASE {identifier} WITH (FORCE)",
            )
    return {
        "ok": True,
        "temporary_database_removed": True,
        "base_migration_count": len(BASE_MIGRATIONS),
        "g0_migration": G0_MIGRATION,
        "schema_before_apply": before,
        "schema_inside_transaction": inside_transaction,
        "schema_after_rollback": after_rollback,
        "schema_after_committed_apply": after_committed_apply,
        "transactional_ddl_rollback_lossless": after_rollback == before,
        "forward_apply_after_rollback_succeeded": all(after_committed_apply.values()),
        "legacy_read_model": legacy,
    }


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 RehearsalError("DATABASE_URL is required via --database-url or apps/api/.env")
    return await rehearse(
        dsn,
        docker_container=args.docker_container,
        docker_admin_user=args.docker_admin_user,
        database_owner=args.database_owner,
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--database-url", default="")
    parser.add_argument("--docker-container", default="vignette-dev-db")
    parser.add_argument("--docker-admin-user", default="vignette_owner")
    parser.add_argument("--database-owner", default="vignette")
    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 __name__ == "__main__":
    main()