191 lines
7.1 KiB
Python
191 lines
7.1 KiB
Python
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
|
|
|
|
@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)
|