feat(mcp): queue long-running calculations

app/services/calculation_cleanup.py

@@ -0,0 +1,74 @@
+"""Record only job-owned native databases so abrupt cancellation is recoverable."""
+import asyncio
+import json
+import os
+from pathlib import Path
+import re
+from uuid import UUID
+
+
+def _names(directory: Path) -> list[str]:
+    path = directory / "databases.json"
+    if not path.exists():
+        return []
+    if path.is_symlink() or path.stat().st_size > 65536:
+        raise ValueError("Invalid calculation database ledger")
+    names = json.loads(path.read_text())
+    if not isinstance(names, list) or len(names) > 100 or any(
+        not isinstance(name, str) or not re.fullmatch(r"tjw_tmp_[a-z0-9_]+_[a-f0-9]{32}", name)
+        or len(name) > 63 for name in names
+    ):
+        raise ValueError("Invalid calculation database ledger")
+    return names
+
+
+def _write(directory: Path, names: list[str]) -> None:
+    path = directory / "databases.json"
+    staging = directory / "databases.json.new"
+    with staging.open("w") as output:
+        json.dump(names, output)
+        output.flush()
+        os.fsync(output.fileno())
+    staging.replace(path)
+
+
+def record_calculation_database(name: str, *, removed: bool = False) -> None:
+    """Called before clone creation and after successful normal cleanup."""
+    root = os.environ.get("TJWATER_CALCULATION_JOB_DIR")
+    if not root:
+        return
+    directory = Path(root)
+    names = _names(directory)
+    if removed:
+        names = [value for value in names if value != name]
+    elif name not in names:
+        names.append(name)
+    _write(directory, names)
+
+
+async def cleanup_calculation_databases(project_id: str, job_directory: str) -> None:
+    directory = Path(job_directory)
+    names = _names(directory)
+    if not names:
+        return
+    from app.infra.db.metadb.database import SessionLocal
+    from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository
+    from app.infra.db.project_routing import ActiveProjectRouting, activate_project_routing
+    from app.native.wndb.core.projects import have_project, delete_project
+    async with SessionLocal() as session:
+        metadata = MetadataRepository(session)
+        project = await metadata.get_project_by_id(UUID(str(project_id)))
+        route = await metadata.get_project_db_routing(UUID(str(project_id)), "biz_data")
+        if not project or not route or route.db_type != "postgresql":
+            raise ValueError("Calculation cleanup routing unavailable")
+        routing = ActiveProjectRouting(project.code, route.dsn, project_id=project.id)
+        await session.commit()
+
+    def cleanup():
+        with activate_project_routing(routing):
+            for name in names[:]:
+                if have_project(name):
+                    delete_project(name)
+                names.remove(name)
+                _write(directory, names)
+    await asyncio.to_thread(cleanup)