97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
"""DB-backed tracking for long-running admin operations.
|
|
|
|
Jobs are persisted to the ``jobs`` table rather than kept in a module-level
|
|
dict, so that a status-poll request landing on a *different* worker
|
|
process/replica than the one running the job can still see its progress —
|
|
required once the app runs with more than one uvicorn worker or replica
|
|
(a module-level dict is process-local and would report "not found" for a job
|
|
started elsewhere, even while that job is still running and writing to the
|
|
shared database).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import uuid
|
|
from collections.abc import Awaitable, Callable
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
from sqlalchemy import delete
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.database import AsyncSessionLocal
|
|
from app.models.job import Job
|
|
|
|
_MAX_FINISHED_AGE = timedelta(hours=1)
|
|
|
|
|
|
async def get(db: AsyncSession, job_id: str) -> Job | None:
|
|
return await db.get(Job, job_id)
|
|
|
|
|
|
async def start(kind: str, run: Callable[[str], Awaitable[None]]) -> Job:
|
|
"""Create a job row, launch ``run(job_id)`` in the background, and return
|
|
the job immediately.
|
|
|
|
``run`` is responsible for calling ``mark_done``/``mark_error`` to record
|
|
the outcome, including on 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.
|
|
"""
|
|
async with AsyncSessionLocal() as db:
|
|
await _evict_stale(db)
|
|
job = Job(id=uuid.uuid4().hex, kind=kind, status="running", message="Working…")
|
|
db.add(job)
|
|
await db.commit()
|
|
|
|
async def _runner() -> None:
|
|
try:
|
|
await run(job.id)
|
|
except Exception as exc:
|
|
await mark_error(job.id, f"{kind} failed: {exc}")
|
|
|
|
asyncio.create_task(_runner())
|
|
return job
|
|
|
|
|
|
async def mark_done(job_id: str, message: str, result: dict | None = None) -> None:
|
|
await _update(job_id, status="done", message=message, result=result)
|
|
|
|
|
|
async def mark_error(job_id: str, message: str) -> None:
|
|
await _update(job_id, status="error", message=message)
|
|
|
|
|
|
async def set_message(job_id: str, message: str) -> None:
|
|
await _update(job_id, message=message)
|
|
|
|
|
|
async def _update(
|
|
job_id: str,
|
|
*,
|
|
status: str | None = None,
|
|
message: str | None = None,
|
|
result: dict | None = None,
|
|
) -> None:
|
|
async with AsyncSessionLocal() as db:
|
|
job = await db.get(Job, job_id)
|
|
if job is None:
|
|
return
|
|
if status is not None:
|
|
job.status = status
|
|
if status in ("done", "error"):
|
|
job.finished_at = datetime.now(UTC)
|
|
if message is not None:
|
|
job.message = message
|
|
if result is not None:
|
|
job.result = result
|
|
await db.commit()
|
|
|
|
|
|
async def _evict_stale(db: AsyncSession) -> None:
|
|
cutoff = datetime.now(UTC) - _MAX_FINISHED_AGE
|
|
await db.execute(
|
|
delete(Job).where(Job.finished_at.is_not(None), Job.finished_at < cutoff)
|
|
)
|
|
await db.commit()
|