"""Procrastinate-backed ingest queue — the Postgres ``TaskProducer`` + worker (Deck #183). This replaces NATS JetStream and the old Postgres-queue stub. The MCP server now owns *both* sides of ingest: - **Producer** (API role / scanner) — :class:`ProcrastinateTaskProducer.send` *defers* one ``ingest:process_document`` job per changed document into the per-tenant Postgres (the same app DB; procrastinate manages its own tables). - **Consumer** (worker role) — ``nextcloud-mcp-server worker`` runs :func:`procrastinate.App.run_worker`, which drains the ``ingest`` queue and invokes the existing :func:`process_document` pipeline. Design notes: - **No execution ``lock``, only ``queueing_lock``.** procrastinate does NOT auto-reclaim ``doing`` jobs, so a per-doc execution lock would permanently deadlock a document if a worker crashed mid-job. The Qdrant upsert is idempotent (deterministic ``uuid5`` point IDs), so a concurrent/re-run is harmless; ``queueing_lock`` (partial-unique on ``status='todo'``) is enough to dedupe enqueues, and :func:`reclaim_stalled_ingest_jobs` retries jobs orphaned in ``doing`` by a crash. - procrastinate is Postgres-only and uses asyncio; ``anyio`` runs natively on the asyncio backend, so the worker can call the anyio-based pipeline directly. - Tasks are defined on a :class:`procrastinate.Blueprint` so the connector is decoupled from the task registry: production binds a real :class:`PsycopgConnector`; unit tests bind ``testing.InMemoryConnector``. """ from __future__ import annotations import logging from dataclasses import asdict from datetime import datetime, timezone from types import TracebackType from typing import TYPE_CHECKING from procrastinate import App, Blueprint, JobContext, PsycopgConnector, RetryStrategy from procrastinate.connector import BaseConnector from procrastinate.exceptions import AlreadyEnqueued from ...config import get_procrastinate_conninfo, get_settings from ..scanner import DocumentTask if TYPE_CHECKING: from ...client import NextcloudClient logger = logging.getLogger(__name__) # Single queue for document ingest. KEDA scales the worker Deployment on the # depth of this queue (``SELECT count(*) FROM procrastinate_jobs WHERE # queue_name='ingest' AND status='todo'``). INGEST_QUEUE_NAME = "ingest" # Blueprint namespace → registered task names are prefixed ``ingest:``. _NAMESPACE = "ingest" INGEST_TASK_NAME = f"{_NAMESPACE}:process_document" _RECLAIM_TASK_NAME = f"{_NAMESPACE}:reclaim_stalled_jobs" # A crashed worker leaves its job in ``doing``; reclaim it once its (per-worker) # heartbeat is this many seconds stale. Sized well above the longest expected # ``process_document`` (PDF render + embedding) so a slow-but-live worker — whose # heartbeat stays current during a long job — is never reclaimed out from under # itself. _STALLED_AFTER_SECONDS = 300 # Tasks are defined as plain functions and registered onto a *fresh* Blueprint # per app (see _build_ingest_blueprint). procrastinate's add_tasks_from mutates # the blueprint's task names in place (namespace prefixing), so a single shared # Blueprint cannot be added to more than one App — which the tests (in-memory + # real Postgres) and any re-init path require. async def process_document_task( *, user_id: str, doc_id: str, doc_type: str, operation: str, modified_at: int, file_path: str | None = None, metadata: dict[str, int | str] | None = None, etag: str | None = None, owner_id: str | None = None, ) -> None: """Worker entry: rebuild the DocumentTask, resolve creds, run the pipeline.""" # Local imports avoid a heavy import chain at blueprint-definition time # (this module is also imported by the API pod just to defer jobs). from ..oauth_sync import NotProvisionedError # noqa: PLC0415 from ..processor import process_document # noqa: PLC0415 task = DocumentTask( user_id=user_id, doc_id=doc_id, doc_type=doc_type, operation=operation, modified_at=modified_at, file_path=file_path, metadata=metadata, etag=etag, owner_id=owner_id, ) try: nc_client = await _resolve_client(user_id) except NotProvisionedError: # A deprovisioned user must not pin a worker slot retrying forever. # Finish the job as a no-op; the next scan re-enqueues once the user # re-provisions an app password. Other errors (transient DB/network) # propagate so procrastinate's retry strategy handles them. logger.warning( "ingest.skip_no_credentials user=%s doc=%s:%s", user_id, doc_type, doc_id ) return try: # Durable retry is procrastinate's job; disable the in-process loop. await process_document(task, nc_client, max_retries=1) finally: await nc_client.close() async def reclaim_stalled_ingest_jobs(context: JobContext, timestamp: int) -> None: """Re-queue ingest jobs orphaned in ``doing`` by a crashed worker. procrastinate prunes dead *workers* but does not reset their in-flight jobs; without this they'd sit in ``doing`` forever. ``timestamp`` is procrastinate's periodic-run marker (unused). """ manager = context.app.job_manager retry_at = datetime.now(tz=timezone.utc) reclaimed = 0 for job in await manager.get_stalled_jobs( queue=INGEST_QUEUE_NAME, seconds_since_heartbeat=_STALLED_AFTER_SECONDS ): if job.id is None: continue await manager.retry_job_by_id_async(job_id=job.id, retry_at=retry_at) reclaimed += 1 if reclaimed: logger.warning("ingest.reclaimed_stalled_jobs count=%d", reclaimed) async def _resolve_client(user_id: str) -> NextcloudClient: """Build an authenticated NextcloudClient for ``user_id`` in the worker. Single-user BasicAuth uses the shared env credentials; every multi-user mode resolves the user's locally-stored app password (BasicAuth). """ from ...client import NextcloudClient # noqa: PLC0415 from ...config_validators import AuthMode, detect_auth_mode # noqa: PLC0415 settings = get_settings() if detect_auth_mode(settings) == AuthMode.SINGLE_USER_BASIC: return NextcloudClient.from_env() from ..oauth_sync import get_user_client_basic_auth # noqa: PLC0415 host = settings.nextcloud_host if not host: raise ValueError("NEXTCLOUD_HOST is required for multi-user ingest") return await get_user_client_basic_auth(user_id, host) def _build_ingest_blueprint() -> Blueprint: """Create a fresh Blueprint with the ingest tasks registered. Fresh per call because ``add_tasks_from`` mutates the blueprint's task names (namespace prefixing), so the same Blueprint cannot be reused across Apps. """ bp = Blueprint() # Durable retry owned by the queue (survives worker crashes); the in-process # retry loop in process_document is disabled on this path via max_retries=1. bp.task( name="process_document", queue=INGEST_QUEUE_NAME, retry=RetryStrategy(max_attempts=5, exponential_wait=4), )(process_document_task) reclaim = bp.task( name="reclaim_stalled_jobs", queue=INGEST_QUEUE_NAME, pass_context=True )(reclaim_stalled_ingest_jobs) bp.periodic(cron="*/5 * * * *", periodic_id="reclaim_stalled_ingest")(reclaim) return bp def build_app(connector: BaseConnector) -> App: """Build an App for the given connector with the ingest tasks registered. Shared by production (:func:`get_procrastinate_app`) and tests (which pass a ``testing.InMemoryConnector``). """ app = App(connector=connector) app.add_tasks_from(_build_ingest_blueprint(), namespace=_NAMESPACE) return app def build_app_for_url(database_url: str) -> App: """Build an App bound to an explicit Postgres URL (for the CLI, which may target a ``--database-url`` that differs from the ``DATABASE_URL`` env).""" return build_app( PsycopgConnector(conninfo=get_procrastinate_conninfo(database_url)) ) _app: App | None = None def get_procrastinate_app() -> App: """Process-wide procrastinate App bound to the Postgres app database.""" global _app if _app is None: _app = build_app(PsycopgConnector(conninfo=get_procrastinate_conninfo())) return _app async def _ingest_schema_present(app: App) -> bool: row = await app.connector.execute_query_one_async( "SELECT to_regclass('procrastinate_jobs') IS NOT NULL AS present" ) return bool(row["present"]) async def apply_ingest_queue_schema(app: App | None = None) -> None: """Create procrastinate's tables on a fresh database (apply-if-absent). procrastinate's ``schema.sql`` uses bare ``CREATE TYPE``/``CREATE TABLE`` (not ``IF NOT EXISTS``), so it errors if re-applied — it is meant to run once on a fresh DB. We skip when ``procrastinate_jobs`` already exists; *version* upgrades use procrastinate's own migration files (operator-run, a lineage independent of the app's Alembic schema). Safe to call concurrently across rolling-update pods without an advisory lock: Postgres DDL is transactional and procrastinate applies the whole schema in one transaction, so a pod that loses the race rolls back cleanly and we treat the resulting error as benign once the schema is present. Opens a short-lived connection, so it is safe to call from the CLI (``db upgrade`` / worker startup). """ app = app or get_procrastinate_app() async with app.open_async(): if await _ingest_schema_present(app): logger.debug("ingest queue schema already present; skipping apply") return try: await app.schema_manager.apply_schema_async() logger.info("Applied procrastinate ingest queue schema") except Exception: # A racing pod likely committed the schema while our transaction # rolled back atomically. Benign iff the schema is now present. if await _ingest_schema_present(app): logger.info("Ingest queue schema applied concurrently by another pod") return raise # Job-status keys procrastinate flattens into each list_queues row (alongside # ``name`` and ``jobs_count``). ``aborting`` is legacy/unused since v3.0.0. _JOB_STATUSES = ("todo", "doing", "succeeded", "failed", "cancelled", "aborted") async def get_ingest_job_counts(app: App | None = None) -> dict[str, int]: """Return ingest job counts by status (``todo``/``doing``/``failed``/…). Reads procrastinate's per-queue stats via the manager API (not hand-written SQL) so a future schema bump doesn't silently break the status surface. The manager flattens its per-status ``stats`` into top-level row keys, so we read the known status keys directly. Assumes the app's connector is already open. """ app = app or get_procrastinate_app() counts: dict[str, int] = {} for row in await app.job_manager.list_queues_async(queue=INGEST_QUEUE_NAME): for status in _JOB_STATUSES: if status in row: counts[status] = counts.get(status, 0) + int(row[status]) return counts def _doc_queueing_lock(task: DocumentTask) -> str: """Per-document enqueue-dedup key (partial-unique on ``status='todo'``).""" return f"{task.user_id}:{task.doc_type}:{task.doc_id}" class ProcrastinateTaskProducer: """``TaskProducer`` that defers ingest jobs into Postgres via procrastinate. The App's connector pool is owned by the server lifespan (opened once, closed on shutdown), so ``clone``/``aenter``/``aexit``/``aclose`` are no-ops — there is no per-handle resource like the memory stream's clones. """ def __init__(self, app: App): self._app = app @classmethod async def connect(cls) -> ProcrastinateTaskProducer: app = get_procrastinate_app() await app.open_async() return cls(app) async def send(self, task: DocumentTask, /) -> None: key = _doc_queueing_lock(task) deferrer = self._app.configure_task(INGEST_TASK_NAME, queueing_lock=key) try: await deferrer.defer_async(**asdict(task)) except AlreadyEnqueued: # A todo job already exists for this doc; the next periodic scan # re-evaluates freshness (placeholder/Qdrant modified_at only # advances after a successful index), so this is not a lost update. logger.debug("ingest.already_enqueued key=%s", key) async def job_counts(self) -> dict[str, int]: """Ingest job counts by status (for the vector-sync status surface).""" return await get_ingest_job_counts(self._app) def clone(self) -> ProcrastinateTaskProducer: return self async def __aenter__(self) -> ProcrastinateTaskProducer: return self async def __aexit__( self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None, ) -> None: return None # The bare suppression marker silences S7503 (async method without await): # ``async def`` is required by the TaskProducer protocol; per-handle close is # a no-op (the pool is owned by the lifespan, drained once on shutdown). async def aclose(self) -> None: # NOSONAR return None async def drain(self) -> None: """Close the shared connector pool (lifespan shutdown only).""" await self._app.close_async()