75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
"""Lightweight in-process tracking for long-running admin operations.
|
|
|
|
Jobs live in a module-level dict, so this only works correctly for a single
|
|
OS process (the app's current deployment: one uvicorn worker, no multi-process
|
|
or multi-replica setup — see ``entrypoint.sh`` / ``docker-compose.yml``). If
|
|
the app is ever scaled to multiple workers/replicas, this needs to move to
|
|
shared storage (a DB table, Redis) so a poll request landing on a different
|
|
process can still see the job.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import uuid
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass, field
|
|
from datetime import UTC, datetime
|
|
|
|
_MAX_FINISHED_AGE_SECONDS = 3600
|
|
|
|
|
|
@dataclass
|
|
class Job:
|
|
id: str
|
|
kind: str
|
|
status: str = "running" # running | done | error
|
|
message: str = "Working…"
|
|
result: object | None = None
|
|
created_at: datetime = field(default_factory=lambda: datetime.now(UTC))
|
|
finished_at: datetime | None = None
|
|
|
|
|
|
_jobs: dict[str, Job] = {}
|
|
|
|
|
|
def get(job_id: str) -> Job | None:
|
|
return _jobs.get(job_id)
|
|
|
|
|
|
def _evict_stale() -> None:
|
|
now = datetime.now(UTC)
|
|
stale = [
|
|
job_id
|
|
for job_id, job in _jobs.items()
|
|
if job.finished_at
|
|
and (now - job.finished_at).total_seconds() > _MAX_FINISHED_AGE_SECONDS
|
|
]
|
|
for job_id in stale:
|
|
del _jobs[job_id]
|
|
|
|
|
|
def start(kind: str, run: Callable[[Job], Awaitable[None]]) -> Job:
|
|
"""Create a job, launch ``run(job)`` in the background, and return it immediately.
|
|
|
|
``run`` is responsible for updating ``job.status``/``job.message``/``job.result``
|
|
to reflect the outcome, including catching its own errors. Any exception that
|
|
escapes anyway is caught here as a last resort so the job doesn't hang forever
|
|
in "running" state.
|
|
"""
|
|
_evict_stale()
|
|
job = Job(id=uuid.uuid4().hex, kind=kind)
|
|
_jobs[job.id] = job
|
|
|
|
async def _runner() -> None:
|
|
try:
|
|
await run(job)
|
|
except Exception as exc:
|
|
job.status = "error"
|
|
job.message = f"{kind} failed: {exc}"
|
|
finally:
|
|
job.finished_at = datetime.now(UTC)
|
|
|
|
asyncio.create_task(_runner())
|
|
return job
|