pablan/backend/app/ingestion/queue.py

"""Postgres-backed background queue.

One asyncio loop in the app lifespan claims jobs via
SELECT … FOR UPDATE SKIP LOCKED. The claim transaction stays open while the
handler runs: a crash rolls everything back and the job remains pending and
claimable after restart — handler writes are atomic with job completion.
Failure bookkeeping (attempts, backoff, last_error) happens in a follow-up
transaction. Single worker per process; the loop moves into a worker
container unchanged when scale demands it (Variant B).
"""

import asyncio
import logging
from collections.abc import Awaitable, Callable
from datetime import UTC, datetime, timedelta
from typing import Any

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from app.config import get_settings
from app.db import async_session_factory
from app.log import safe_error
from app.metrics import metrics
from app.models import Job, JobStatus

logger = logging.getLogger("pablan.queue")

JobHandler = Callable[[AsyncSession, Job], Awaitable[None]]

_HANDLERS: dict[str, JobHandler] = {}

MAX_ATTEMPTS = 5
BACKOFF_BASE_SECONDS = 30.0  # 30s, 1m, 2m, 4m between retries


def job_handler(job_type: str) -> Callable[[JobHandler], JobHandler]:
    def register(fn: JobHandler) -> JobHandler:
        _HANDLERS[job_type] = fn
        return fn

    return register


def backoff_delay(attempts: int) -> timedelta:
    return timedelta(seconds=BACKOFF_BASE_SECONDS * 2 ** (attempts - 1))


async def enqueue(
    db: AsyncSession,
    job_type: str,
    payload: dict[str, Any] | None = None,
    run_after: datetime | None = None,
) -> Job:
    job = Job(type=job_type, payload=payload or {})
    if run_after is not None:
        job.run_after = run_after
    db.add(job)
    await db.flush()
    return job


async def process_one(
    session_factory: async_sessionmaker[AsyncSession] = async_session_factory,
) -> bool:
    """Claim and process a single due job. Returns True if one was processed."""
    started = asyncio.get_running_loop().time()
    async with session_factory() as db:
        job = (
            await db.execute(
                select(Job)
                .where(Job.status == JobStatus.pending, Job.run_after <= func.now())
                .order_by(Job.run_after)
                .limit(1)
                .with_for_update(skip_locked=True)
            )
        ).scalar_one_or_none()
        if job is None:
            await db.rollback()
            return False

        job_id, job_type, attempts_before = job.id, job.type, job.attempts
        try:
            handler = _HANDLERS.get(job_type)
            if handler is None:
                raise LookupError(f"no handler registered for job type {job_type!r}")
            await handler(db, job)
            job.status = JobStatus.done
            job.attempts = attempts_before + 1
            await db.commit()
        except Exception as exc:
            await db.rollback()
            await _record_failure(session_factory, job_id, exc)
            duration = asyncio.get_running_loop().time() - started
            metrics.inc("jobs_processed_total", {"type": job_type, "status": "failed"})
            metrics.observe("job_seconds", duration, {"type": job_type})
            logger.warning(
                "job failed",
                extra={
                    "event": "job_failed",
                    "job_id": str(job_id),
                    "job_type": job_type,
                    "attempt": attempts_before + 1,
                    "error": safe_error(exc),
                },
            )
            return True

    duration = asyncio.get_running_loop().time() - started
    metrics.inc("jobs_processed_total", {"type": job_type, "status": "done"})
    metrics.observe("job_seconds", duration, {"type": job_type})
    logger.info(
        "job done",
        extra={
            "event": "job_done",
            "job_id": str(job_id),
            "job_type": job_type,
            "attempt": attempts_before + 1,
            "duration_ms": round(duration * 1000),
        },
    )
    return True


async def _record_failure(
    session_factory: async_sessionmaker[AsyncSession],
    job_id: Any,
    exc: Exception,
) -> None:
    async with session_factory() as db:
        job = await db.get(Job, job_id, with_for_update=True)
        if job is None:  # pragma: no cover — job deleted underneath us
            return
        job.attempts += 1
        job.last_error = safe_error(exc, limit=500)
        if job.attempts >= MAX_ATTEMPTS:
            job.status = JobStatus.failed
            metrics.inc("jobs_exhausted_total", {"type": job.type})
        else:
            job.status = JobStatus.pending
            job.run_after = datetime.now(UTC) + backoff_delay(job.attempts)
            metrics.inc("jobs_retried_total", {"type": job.type})
        await db.commit()


async def _update_depth_gauge(
    session_factory: async_sessionmaker[AsyncSession],
) -> None:
    async with session_factory() as db:
        depth = (
            await db.execute(
                select(func.count(Job.id)).where(Job.status == JobStatus.pending)
            )
        ).scalar_one()
    metrics.set_gauge("jobs_queue_depth", float(depth))


async def run_queue(
    stop_event: asyncio.Event,
    session_factory: async_sessionmaker[AsyncSession] = async_session_factory,
) -> None:
    poll_seconds = get_settings().job_poll_seconds
    logger.info("job queue started", extra={"event": "queue_started"})
    while not stop_event.is_set():
        worked = False
        try:
            worked = await process_one(session_factory)
            await _update_depth_gauge(session_factory)
        except Exception as exc:
            logger.error(
                "queue iteration failed",
                extra={"event": "queue_error", "error": safe_error(exc)},
            )
        if not worked:
            try:
                await asyncio.wait_for(stop_event.wait(), timeout=poll_seconds)
            except TimeoutError:
                pass
    logger.info("job queue stopped", extra={"event": "queue_stopped"})