feat(investigation): preserve cumulative reference graphs
Accumulate investigation paths and preserve failed attempts across rounds. Keep graph annotations separate from execution state, synchronize the API contract, and consolidate documentation.
This commit is contained in:
@@ -0,0 +1,182 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user