from __future__ import annotations import uuid from datetime import datetime, timezone from collections.abc import Callable from typing import Any from app.core.errors import AppError from app.models import Job from app.schemas.job import JobCreate, JobRead class JobService: VALID_STATUSES = {"queued", "running", "success", "failed"} @staticmethod def run_sync_job( db, project_id: uuid.UUID, job_type: str, parameters: dict[str, Any] | None, operation: Callable[[], Any], input_dataset_id: uuid.UUID | None = None, ) -> dict[str, Any]: created = JobService.create_job( db, JobCreate( job_type=job_type, project_id=project_id, input_dataset_id=input_dataset_id, parameters_json=JobService._coerce_payload(parameters), ), ) try: JobService.mark_running(db, created.id) result = operation() output_dataset_id = None if isinstance(result, uuid.UUID): output_dataset_id = result result = {"output_dataset_id": str(result)} if isinstance(result, dict): candidate_output_dataset_id = result.get("output_dataset_id") if isinstance(candidate_output_dataset_id, str): try: output_dataset_id = uuid.UUID(candidate_output_dataset_id) except ValueError: output_dataset_id = None elif isinstance(candidate_output_dataset_id, uuid.UUID): output_dataset_id = candidate_output_dataset_id if isinstance(result, dict): job = JobService.mark_success(db, created.id, result=result, output_dataset_id=output_dataset_id) else: job = JobService.mark_success(db, created.id, result={"result": result}, output_dataset_id=output_dataset_id) job_payload = job.model_dump() if isinstance(job_payload.get("output_dataset_id"), uuid.UUID): job_payload["output_dataset_id"] = str(job_payload["output_dataset_id"]) result_json = job_payload.get("result_json") if isinstance(result_json, dict): if isinstance(result_json.get("output_dataset_id"), uuid.UUID): result_json["output_dataset_id"] = str(result_json["output_dataset_id"]) job_payload["result_json"] = result_json return job_payload except AppError as exc: failed = JobService.mark_failed( db, created.id, error_message=exc.message, details={"code": exc.code, "details": exc.details}, ) payload = failed.model_dump() if isinstance(payload.get("output_dataset_id"), uuid.UUID): payload["output_dataset_id"] = str(payload["output_dataset_id"]) result_json = payload.get("result_json") if isinstance(result_json, dict): if isinstance(result_json.get("output_dataset_id"), uuid.UUID): result_json["output_dataset_id"] = str(result_json["output_dataset_id"]) payload["result_json"] = result_json raise except Exception: # An unexpected error must never leave the job stuck in "running". try: db.rollback() except Exception: pass try: JobService.mark_failed( db, created.id, error_message="Unexpected internal error during synchronous job execution", details={"code": "JOB_INTERNAL_ERROR"}, ) except Exception: pass raise @staticmethod def _coerce_payload(payload: dict[str, Any] | None) -> dict[str, Any]: return dict(payload or {}) @staticmethod def create_job(db, payload: JobCreate) -> JobRead: job = Job( id=uuid.uuid4(), job_type=payload.job_type, status="queued", project_id=payload.project_id, dataset_id=payload.dataset_id, input_dataset_id=payload.input_dataset_id, output_dataset_id=payload.output_dataset_id, parameters_json=JobService._coerce_payload(payload.parameters_json), result_json=None, error_message=None, ) db.add(job) db.commit() db.refresh(job) return JobRead.model_validate(job) @staticmethod def mark_running(db, job_id: uuid.UUID) -> JobRead: job = JobService._get_job(db, job_id) job.status = "running" job.started_at = datetime.now(timezone.utc) job.error_message = None db.add(job) db.commit() db.refresh(job) return JobRead.model_validate(job) @staticmethod def mark_success( db, job_id: uuid.UUID, result: dict[str, Any] | None = None, output_dataset_id: uuid.UUID | None = None, ) -> JobRead: job = JobService._get_job(db, job_id) job.status = "success" job.finished_at = datetime.now(timezone.utc) if output_dataset_id is not None: job.output_dataset_id = output_dataset_id job.result_json = result job.error_message = None db.add(job) db.commit() db.refresh(job) return JobRead.model_validate(job) @staticmethod def mark_failed(db, job_id: uuid.UUID, error_message: str, details: dict[str, Any] | None = None) -> JobRead: job = JobService._get_job(db, job_id) job.status = "failed" job.finished_at = datetime.now(timezone.utc) if details: job.result_json = details job.error_message = error_message db.add(job) db.commit() db.refresh(job) return JobRead.model_validate(job) @staticmethod def get_job(db, job_id: uuid.UUID) -> JobRead: return JobRead.model_validate(JobService._get_job(db, job_id)) @staticmethod def get_job_status(db, job_id: uuid.UUID) -> dict: job = JobService._get_job(db, job_id) return { "id": job.id, "project_id": str(job.project_id), "status": job.status, "error_message": job.error_message, "started_at": job.started_at, "finished_at": job.finished_at, "result_json": job.result_json, } @staticmethod def list_jobs( db, project_id: uuid.UUID | None = None, dataset_id: uuid.UUID | None = None, limit: int = 50, offset: int = 0, ) -> tuple[list[JobRead], int]: query = db.query(Job) if project_id is not None: query = query.filter(Job.project_id == project_id) if dataset_id is not None: query = query.filter((Job.dataset_id == dataset_id) | (Job.input_dataset_id == dataset_id) | (Job.output_dataset_id == dataset_id)) total = query.count() rows = query.order_by(Job.created_at.desc()).offset(offset).limit(limit).all() return [JobRead.model_validate(row) for row in rows], total @staticmethod def _get_job(db, job_id: uuid.UUID) -> Job: job = db.get(Job, job_id) if not job: raise AppError(code="JOB_NOT_FOUND", message="Job not found", status_code=404) return job @staticmethod def validate_status(status: str) -> None: if status not in JobService.VALID_STATUSES: raise AppError(code="INVALID_JOB_STATUS", message="Invalid job status", status_code=400)