"""Preview, back up, then adopt saved progress graphs without rewriting reports. Run with ``python -m scripts.migrate_investigation_graphs --help`` from the Server. DSNs are resolved from MetaDB and never written to the plan or console. """ from __future__ import annotations import argparse import asyncio import hashlib import json import os from copy import deepcopy from datetime import datetime, timezone from pathlib import Path from uuid import UUID from psycopg import AsyncConnection from psycopg.rows import dict_row from sqlalchemy import text from app.domain.schemas.operational_task import InvestigationDiagram, InvestigationProgressSnapshot from app.infra.db.metadb.database import SessionLocal, engine from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository, ProjectDbRouting from app.infra.db.postgresql.operational_tasks import OperationalTaskRepository from app.services.investigation_graph import adopt_legacy_investigation_graph MIGRATION = "investigation-graph-reference-v1" def prepare_progress_document(document: dict) -> dict | None: """Return a new revision, preserving publication provenance and Mermaid bytes.""" graph = document["progress"].get("diagram") if graph is None or graph.get("edges") is not None: return None diagram = adopt_legacy_investigation_graph(InvestigationDiagram.model_validate(graph)) result = deepcopy(document) result["revision"] += 1 result["progress"]["diagram"] = diagram.model_dump(mode="json") InvestigationProgressSnapshot.model_validate(result) return result def _write_private(path: Path, content: str, *, append: bool = False) -> None: path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) flags = os.O_WRONLY | os.O_CREAT | os.O_NOFOLLOW | (os.O_APPEND if append else os.O_EXCL) with os.fdopen(os.open(path, flags, 0o600), "w", encoding="utf-8") as handle: handle.write(content) handle.flush() os.fsync(handle.fileno()) async def _routing(project_id: UUID) -> ProjectDbRouting: async with SessionLocal() as session: await session.execute(text("SET TRANSACTION READ ONLY")) route = await MetadataRepository(session).get_project_db_routing(project_id, "biz_data") if route is None or route.db_type != "postgresql": raise ValueError(f"Project {project_id} has no business PostgreSQL routing") return route async def preview(project_id: UUID) -> dict: route = await _routing(project_id) async with await AsyncConnection.connect( route.dsn, row_factory=dict_row, connect_timeout=8, options="-c default_transaction_read_only=on -c statement_timeout=30000", ) as conn: cursor = await conn.execute("SELECT current_database() AS database") database = (await cursor.fetchone())["database"] cursor = await conn.execute( """SELECT task_id,run_id,document FROM workbench.operational_task_progress WHERE project_id=%s AND jsonb_typeof(document#>'{progress,diagram}')='object' AND jsonb_typeof(document#>'{progress,diagram,edges}') IS DISTINCT FROM 'array' ORDER BY task_id,run_id""", (project_id,), ) records = [] for row in await cursor.fetchall(): after = prepare_progress_document(row["document"]) if after is not None: records.append({"task_id": str(row["task_id"]), "run_id": str(row["run_id"]), "before": row["document"], "after": after}) return {"project_id": str(project_id), "database": database, "records": records} async def _protected_digest(conn: AsyncConnection, project_id: UUID) -> dict[str, str]: """Compare complete protected rows inside the migration transaction.""" hashes = {} for table in ("operational_tasks", "operational_task_runs", "operational_task_reports"): # Table names are fixed here, never supplied by CLI input. cursor = await conn.execute( f"SELECT to_jsonb(t) AS row FROM workbench.{table} t WHERE project_id=%s", (project_id,), ) rows = sorted(json.dumps(row["row"], sort_keys=True) for row in await cursor.fetchall()) hashes[table] = hashlib.sha256(json.dumps(rows).encode()).hexdigest() return hashes async def apply_project(project: dict) -> dict: project_id = UUID(project["project_id"]) route = await _routing(project_id) async with await AsyncConnection.connect( route.dsn, row_factory=dict_row, connect_timeout=8, options="-c lock_timeout=5000 -c statement_timeout=30000", ) as conn: repo = OperationalTaskRepository(conn, project_id) updated = already_managed = 0 async with repo.transaction(): cursor = await conn.execute("SELECT current_database() AS database") if (await cursor.fetchone())["database"] != project["database"]: raise ValueError("Project database routing changed since preview") before = await _protected_digest(conn, project_id) for item in project["records"]: await repo.lock(f"task:{item['task_id']}") cursor = await conn.execute( """SELECT document FROM workbench.operational_task_progress WHERE project_id=%s AND task_id=%s AND run_id=%s FOR UPDATE""", (project_id, item["task_id"], item["run_id"]), ) row = await cursor.fetchone() if row and (row["document"]["progress"].get("diagram") or {}).get("edges") is not None: already_managed += 1 continue if row is None or row["document"] != item["before"]: raise ValueError(f"Progress changed since preview: {item['task_id']}; prepare a new plan") expected = prepare_progress_document(item["before"]) if expected is None or expected != item["after"]: raise ValueError("Migration plan does not match the graph contract") await repo.save_progress(item["task_id"], expected) await repo.append_event( topic="operational.progress.changed", aggregate_type="operational_task_progress", aggregate_id=item["task_id"], aggregate_revision=expected["revision"], occurred_at=datetime.now(timezone.utc), metadata={ "task_id": item["task_id"], "run_id": item["run_id"], "execution_generation": expected["execution_generation"], "migration": MIGRATION, }, ) updated += 1 if before != await _protected_digest(conn, project_id): raise ValueError("Protected task/run/report rows changed; migration rolled back") return {"project_id": str(project_id), "database": project["database"], "updated": updated, "already_managed": already_managed, "protected_rows_unchanged": True} async def run(args: argparse.Namespace) -> None: try: if args.plan: if not args.project_id: raise ValueError("--plan requires at least one --project-id") projects = [await preview(project_id) for project_id in dict.fromkeys(args.project_id)] plan = {"migration": MIGRATION, "created_at": datetime.now(timezone.utc).isoformat(), "projects": projects} _write_private(args.plan, json.dumps(plan, ensure_ascii=False, indent=2) + "\n") print(json.dumps({"plan": str(args.plan), "projects": [ {"project_id": item["project_id"], "database": item["database"], "progress_graphs": len(item["records"]), "tasks": len({row["task_id"] for row in item["records"]})} for item in projects ]})) else: plan = json.loads(args.apply.read_text(encoding="utf-8")) if plan["migration"] != MIGRATION: raise ValueError("Unsupported migration plan") for project in plan["projects"]: result = await apply_project(project) _write_private(args.apply.with_suffix(".receipts.jsonl"), json.dumps(result) + "\n", append=True) print(json.dumps(result), flush=True) finally: await engine.dispose() def main() -> None: parser = argparse.ArgumentParser(description=__doc__) mode = parser.add_mutually_exclusive_group(required=True) mode.add_argument("--plan", type=Path, help="Write a preview and full progress backup to a new private JSON file") mode.add_argument("--apply", type=Path, help="Apply an existing preview; changed rows abort that project's transaction") parser.add_argument("--project-id", type=UUID, action="append", help="Explicit project scope for --plan") asyncio.run(run(parser.parse_args())) if __name__ == "__main__": main()