Rejigging async jobs

app/services/jobs.py

@@ -1,11 +1,12 @@
-"""Lightweight in-process tracking for long-running admin operations.
+"""DB-backed 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.
+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
@@ -13,62 +14,83 @@ from __future__ import annotations
 import asyncio
 import uuid
 from collections.abc import Awaitable, Callable
-from dataclasses import dataclass, field
-from datetime import UTC, datetime
+from datetime import UTC, datetime, timedelta
 
-_MAX_FINISHED_AGE_SECONDS = 3600
+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)
 
 
-@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
+async def get(db: AsyncSession, job_id: str) -> Job | None:
+    return await db.get(Job, job_id)
 
 
-_jobs: dict[str, Job] = {}
+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.
 
-
-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.
+    ``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.
     """
-    _evict_stale()
-    job = Job(id=uuid.uuid4().hex, kind=kind)
-    _jobs[job.id] = job
+    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)
+            await run(job.id)
         except Exception as exc:
-            job.status = "error"
-            job.message = f"{kind} failed: {exc}"
-        finally:
-            job.finished_at = datetime.now(UTC)
+            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()