"""Job state store for long-running catalogue jobs (sync, enrich, pair, build). These jobs used to hold an HTTP connection open for minutes. This module lets a job run in the background while a client polls `get(job_id)` for progress. State lives in Postgres rather than Redis. Redis was the first choice -- job state is ephemeral and a TTL is free there -- but the deployment could not reach it reliably, and a control plane that fails whenever its own store is unreachable is worse than one that shares a database already proven to work. Postgres also gives two things Redis did not: a job survives a restart, and the "one job per tenant per kind" rule becomes a unique index the database enforces rather than a lock the application has to remember to release. """ import json import logging import time import uuid from datetime import datetime, timedelta, timezone import anyio from psycopg2 import errors from psycopg2.extras import Json, RealDictCursor from app.services.infra.database import get_master_db_connection logger = logging.getLogger(__name__) STALE_AFTER_SECONDS = 300 JOB_TTL_SECONDS = 86400 ACTIVE_STATUSES = ("queued", "running") class JobAlreadyRunning(Exception): """Raised when a tenant already has a job of this kind in flight. Two concurrent runs of the same kind (e.g. two `pair` jobs) would both rewrite the same downstream table, so `create` refuses the second one. """ def __init__(self, job_id: str): self.job_id = job_id super().__init__(f"Job already running: {job_id}") def ensure_table() -> None: conn = get_master_db_connection() try: with conn.cursor() as cur: cur.execute(""" CREATE TABLE IF NOT EXISTS catalog_jobs ( job_id TEXT PRIMARY KEY, tenant_id TEXT NOT NULL, kind TEXT NOT NULL, status TEXT NOT NULL, percent INT NOT NULL DEFAULT 0, step TEXT NOT NULL DEFAULT '', result JSONB, error TEXT, created_at TIMESTAMPTZ NOT NULL DEFAULT now(), updated_at TIMESTAMPTZ NOT NULL DEFAULT now() ) """) # The lock, enforced by the database rather than by remembering to # release one. A second active job for the same tenant and kind # cannot be inserted at all, so there is no window between checking # and taking it. cur.execute(""" CREATE UNIQUE INDEX IF NOT EXISTS catalog_jobs_active_idx ON catalog_jobs (tenant_id, kind) WHERE status IN ('queued', 'running') """) cur.execute( "CREATE INDEX IF NOT EXISTS catalog_jobs_tenant_idx " "ON catalog_jobs (tenant_id, created_at DESC)") conn.commit() except Exception: conn.rollback() logger.error("Could not ensure catalog_jobs table", exc_info=True) raise finally: conn.close() def _row_to_state(row: dict) -> dict: status = row["status"] updated_at = row["updated_at"] age = (datetime.now(timezone.utc) - updated_at).total_seconds() # A job dies with the process that runs it, leaving the row stuck on # "running" forever. Age it out here rather than trusting what is stored, # or a client would poll "running" indefinitely. if status == "running" and age > STALE_AFTER_SECONDS: status = "lost" return { "job_id": row["job_id"], "tenant_id": row["tenant_id"], "kind": row["kind"], "status": status, "percent": int(row["percent"]), "step": row["step"] or "", "result": row["result"], "error": row["error"] or None, "created_at": row["created_at"].timestamp(), "updated_at": updated_at.timestamp(), } def _create(tenant_id: str, kind: str) -> str: ensure_table() job_id = f"job_{uuid.uuid4().hex}" stale_before = datetime.now(timezone.utc) - timedelta(seconds=STALE_AFTER_SECONDS) conn = get_master_db_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: # Redis released a stuck lock by letting the key expire. Postgres # has no TTL, so a job whose process died would hold the unique # index forever -- retire it first, or one crash locks a tenant out # of this job kind permanently. cur.execute(""" UPDATE catalog_jobs SET status = 'lost', updated_at = now() WHERE tenant_id = %s AND kind = %s AND status IN ('queued', 'running') AND updated_at < %s """, (tenant_id, kind, stale_before)) cur.execute( "DELETE FROM catalog_jobs WHERE created_at < %s", (datetime.now(timezone.utc) - timedelta(seconds=JOB_TTL_SECONDS),)) try: cur.execute(""" INSERT INTO catalog_jobs (job_id, tenant_id, kind, status) VALUES (%s, %s, %s, 'queued') """, (job_id, tenant_id, kind)) except errors.UniqueViolation: conn.rollback() with conn.cursor(cursor_factory=RealDictCursor) as retry: retry.execute(""" SELECT job_id FROM catalog_jobs WHERE tenant_id = %s AND kind = %s AND status IN ('queued', 'running') LIMIT 1 """, (tenant_id, kind)) existing = retry.fetchone() raise JobAlreadyRunning(existing["job_id"] if existing else job_id) conn.commit() except JobAlreadyRunning: raise except Exception: conn.rollback() raise finally: conn.close() return job_id async def create(tenant_id: str, kind: str) -> str: return await anyio.to_thread.run_sync(lambda: _create(tenant_id, kind)) def tick(job_id: str, percent: int, step: str) -> None: """Best-effort progress report. Never raises. Deliberately synchronous: the job runs in a worker thread, so it can write its own progress rather than handing a coroutine back to the event loop. This runs alongside work that is already spending real money on model calls -- a database hiccup here must not abort an enrichment run. """ try: conn = get_master_db_connection() try: with conn.cursor() as cur: cur.execute(""" UPDATE catalog_jobs SET status = 'running', percent = %s, step = %s, updated_at = now() WHERE job_id = %s """, (int(percent), step, job_id)) conn.commit() finally: conn.close() except Exception: logger.warning("Progress update failed for job %s", job_id, exc_info=True) def _finish(job_id: str, result: dict = None, error: str = None) -> None: conn = get_master_db_connection() try: with conn.cursor() as cur: # A failed job keeps the percentage it reached. Claiming 100 for a # job that died at 30 tells a merchant the work finished. cur.execute(""" UPDATE catalog_jobs SET status = %s, percent = CASE WHEN %s IS NULL THEN 100 ELSE percent END, result = COALESCE(%s, result), error = %s, updated_at = now() WHERE job_id = %s """, ("failed" if error else "done", error, Json(result) if result is not None else None, error, job_id)) conn.commit() except Exception: conn.rollback() logger.error("Could not finish job %s", job_id, exc_info=True) finally: conn.close() async def finish(job_id: str, result: dict = None, error: str = None) -> None: await anyio.to_thread.run_sync(lambda: _finish(job_id, result, error)) def _get(job_id: str): conn = get_master_db_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute("SELECT * FROM catalog_jobs WHERE job_id = %s", (job_id,)) row = cur.fetchone() finally: conn.close() return _row_to_state(dict(row)) if row else None async def get(job_id: str): return await anyio.to_thread.run_sync(lambda: _get(job_id)) def _recent(tenant_id: str, limit: int) -> list: conn = get_master_db_connection() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute(""" SELECT * FROM catalog_jobs WHERE tenant_id = %s ORDER BY created_at DESC LIMIT %s """, (tenant_id, limit)) rows = cur.fetchall() finally: conn.close() return [_row_to_state(dict(row)) for row in rows] async def recent(tenant_id: str, limit: int = 20) -> list: return await anyio.to_thread.run_sync(lambda: _recent(tenant_id, limit)) def _clear_tenant(tenant_id: str) -> None: conn = get_master_db_connection() try: with conn.cursor() as cur: cur.execute("DELETE FROM catalog_jobs WHERE tenant_id = %s", (tenant_id,)) conn.commit() finally: conn.close() async def clear_tenant(tenant_id: str) -> None: """Test-only teardown.""" await anyio.to_thread.run_sync(lambda: _clear_tenant(tenant_id))