205 lines
7.4 KiB
Python
205 lines
7.4 KiB
Python
"""Repair software-job outputs written under the duplicated legacy user path.
|
|
|
|
The command is dry-run by default. It reads its database URL only from the
|
|
explicit ``ZCBOT_MIGRATION_DB_URL`` environment variable and never loads .env.
|
|
Run schema migration 0033 before applying this repair.
|
|
"""
|
|
# ruff: noqa: I001
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from uuid import UUID
|
|
|
|
from sqlalchemy import create_engine, select
|
|
from sqlalchemy.orm import Session
|
|
from sqlalchemy.orm.attributes import flag_modified
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from core.software_contracts import DEFAULT_CAPABILITIES, get_contract
|
|
from core.software_jobs import software_job_output_path
|
|
from core.storage.models import Artifact, SoftwareJob, Task
|
|
from core.paths import from_db_path
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RepairPlan:
|
|
job_id: UUID
|
|
user_root: Path
|
|
working_dir: Path
|
|
legacy_dir: Path
|
|
output_dir: Path
|
|
move_legacy_dir: bool
|
|
metadata_moves: tuple[tuple[Path, Path], ...]
|
|
|
|
|
|
def _user_root(working_dir: Path, user_id: UUID) -> Path:
|
|
expected = str(user_id)
|
|
for candidate in (working_dir, *working_dir.parents):
|
|
if candidate.name == expected and candidate.parent.name == "users":
|
|
return candidate.resolve()
|
|
raise RuntimeError(f"task working_dir has no user root for {user_id}")
|
|
|
|
|
|
def _within(root: Path, target: Path) -> Path:
|
|
resolved = target.resolve()
|
|
resolved.relative_to(root.resolve())
|
|
return resolved
|
|
|
|
|
|
def build_plan(job: SoftwareJob, task: Task) -> RepairPlan | None:
|
|
capability = getattr(job, "capability", DEFAULT_CAPABILITIES[0])
|
|
contract = get_contract(capability)
|
|
working_dir = from_db_path(task.working_dir).resolve()
|
|
user_root = _user_root(working_dir, job.user_id)
|
|
output_dir = _within(
|
|
user_root, working_dir / contract.output_namespace / str(job.job_id)
|
|
)
|
|
legacy_dir = _within(
|
|
user_root,
|
|
user_root / Path(task.working_dir) / contract.output_namespace / str(job.job_id),
|
|
)
|
|
legacy_exists = legacy_dir.is_dir()
|
|
output_exists = output_dir.is_dir()
|
|
if legacy_exists and output_exists:
|
|
raise RuntimeError(f"job {job.job_id}: legacy and target output directories both exist")
|
|
if not legacy_exists and not output_exists:
|
|
print(f"[WARN] job {job.job_id}: output directory is missing")
|
|
return None
|
|
source_dir = legacy_dir if legacy_exists else output_dir
|
|
metadata_moves: list[tuple[Path, Path]] = []
|
|
for item in job.artifact_manifest or []:
|
|
output_id = str((item or {}).get("source_artifact_id") or "")
|
|
if not output_id or contract.output_spec(output_id).publish:
|
|
continue
|
|
source = source_dir / str((item or {}).get("filename") or "")
|
|
destination = source_dir / software_job_output_path(capability, output_id)
|
|
if source == destination or not source.exists():
|
|
continue
|
|
if destination.exists():
|
|
raise RuntimeError(f"job {job.job_id}: metadata destination already exists")
|
|
metadata_moves.append((source, destination))
|
|
return RepairPlan(
|
|
job_id=job.job_id,
|
|
user_root=user_root,
|
|
working_dir=working_dir,
|
|
legacy_dir=legacy_dir,
|
|
output_dir=output_dir,
|
|
move_legacy_dir=legacy_exists,
|
|
metadata_moves=tuple(metadata_moves),
|
|
)
|
|
|
|
|
|
def apply_files(plan: RepairPlan) -> None:
|
|
for _, destination in plan.metadata_moves:
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
for source, destination in plan.metadata_moves:
|
|
os.replace(source, destination)
|
|
if plan.move_legacy_dir:
|
|
plan.output_dir.parent.mkdir(parents=True, exist_ok=True)
|
|
os.replace(plan.legacy_dir, plan.output_dir)
|
|
|
|
|
|
def update_rows(session: Session, job: SoftwareJob, task: Task, plan: RepairPlan) -> int:
|
|
capability = getattr(job, "capability", DEFAULT_CAPABILITIES[0])
|
|
contract = get_contract(capability)
|
|
artifact_count = 0
|
|
manifest = [dict(item) for item in (job.artifact_manifest or [])]
|
|
task_prefix = plan.working_dir.relative_to(plan.user_root).as_posix()
|
|
for item in manifest:
|
|
output_id = str(item.get("source_artifact_id") or "")
|
|
if not output_id:
|
|
continue
|
|
path = (
|
|
f"{contract.output_namespace}/{job.job_id}/"
|
|
f"{software_job_output_path(capability, output_id)}"
|
|
)
|
|
item["path"] = path
|
|
raw_artifact_id = item.get("artifact_id")
|
|
if not raw_artifact_id:
|
|
continue
|
|
try:
|
|
artifact_id = UUID(str(raw_artifact_id))
|
|
except ValueError:
|
|
continue
|
|
artifact = session.get(Artifact, artifact_id)
|
|
if artifact is None or artifact.user_id != job.user_id:
|
|
raise RuntimeError(f"job {job.job_id}: artifact {artifact_id} is missing")
|
|
artifact.current_path = f"{task_prefix}/{path}"
|
|
if contract.output_spec(output_id).publish:
|
|
artifact.software_job_id = job.job_id
|
|
artifact_count += 1
|
|
job.artifact_manifest = manifest
|
|
flag_modified(job, "artifact_manifest")
|
|
return artifact_count
|
|
|
|
|
|
def repair(session: Session, *, apply: bool, job_id: UUID | None = None) -> tuple[int, int]:
|
|
statement = (
|
|
select(SoftwareJob, Task)
|
|
.join(Task, Task.task_id == SoftwareJob.task_id)
|
|
.where(SoftwareJob.status == "succeeded")
|
|
.order_by(SoftwareJob.created_at, SoftwareJob.job_id)
|
|
)
|
|
if job_id is not None:
|
|
statement = statement.where(SoftwareJob.job_id == job_id)
|
|
rows = session.execute(statement).all()
|
|
prepared: list[tuple[SoftwareJob, Task, RepairPlan]] = []
|
|
for job, task in rows:
|
|
plan = build_plan(job, task)
|
|
if plan is None:
|
|
continue
|
|
print(
|
|
f"[INFO] job={job.job_id} move_dir={plan.move_legacy_dir} "
|
|
f"metadata_moves={len(plan.metadata_moves)} target={plan.output_dir}"
|
|
)
|
|
prepared.append((job, task, plan))
|
|
if not apply:
|
|
return len(prepared), 0
|
|
updated_artifacts = 0
|
|
for job, task, plan in prepared:
|
|
apply_files(plan)
|
|
updated_artifacts += update_rows(session, job, task, plan)
|
|
session.flush()
|
|
return len(prepared), updated_artifacts
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--apply", action="store_true", help="apply filesystem and database changes")
|
|
parser.add_argument("--job-id", type=UUID, help="limit repair to one software job")
|
|
args = parser.parse_args()
|
|
database_url = os.environ.get("ZCBOT_MIGRATION_DB_URL", "").strip()
|
|
if not database_url:
|
|
print("[ERR] ZCBOT_MIGRATION_DB_URL is required", file=sys.stderr)
|
|
return 2
|
|
engine = create_engine(database_url, pool_pre_ping=True, future=True)
|
|
try:
|
|
with Session(engine, future=True) as session:
|
|
try:
|
|
jobs, artifacts = repair(
|
|
session, apply=args.apply, job_id=args.job_id
|
|
)
|
|
if args.apply:
|
|
session.commit()
|
|
else:
|
|
session.rollback()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
engine.dispose()
|
|
mode = "applied" if args.apply else "validated"
|
|
print(f"[OK] {mode} jobs: {jobs}")
|
|
print(f"[OK] updated artifacts: {artifacts}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|