o2o-plagiarism-ai/app/jobs/store.py

58 lines
1.6 KiB
Python

"""배치 잡 in-memory 스토어. 운영 시 Redis/PostgreSQL 등으로 교체."""
from __future__ import annotations
import threading
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Literal
from app.api.schemas import DetectResponse
JobStatus = Literal["queued", "running", "completed", "failed"]
@dataclass
class Job:
job_id: str
status: JobStatus
total: int
processed: int = 0
created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
finished_at: datetime | None = None
results: list[DetectResponse] = field(default_factory=list)
error: str | None = None
class JobStore:
def __init__(self):
self._jobs: dict[str, Job] = {}
self._lock = threading.Lock()
def create(self, total: int) -> Job:
job = Job(job_id=str(uuid.uuid4()), status="queued", total=total)
with self._lock:
self._jobs[job.job_id] = job
return job
def get(self, job_id: str) -> Job | None:
with self._lock:
return self._jobs.get(job_id)
def update(self, job_id: str, **fields) -> None:
with self._lock:
job = self._jobs.get(job_id)
if not job:
return
for key, value in fields.items():
setattr(job, key, value)
def append_result(self, job_id: str, result: DetectResponse) -> None:
with self._lock:
job = self._jobs.get(job_id)
if not job:
return
job.results.append(result)
job.processed = len(job.results)