75 lines
2.8 KiB
Python
75 lines
2.8 KiB
Python
"""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)
|