feat: replace NATS ingest with procrastinate Postgres queue (#183)
Re-architect document ingest from the shared NATS-glued document-processor to a per-tenant, in-process model owned by nextcloud-mcp-server (Deck #183). The MCP server now owns both sides of ingest: - Producer (api role): the scanner defers one job per changed document into the app's Postgres via procrastinate (queueing_lock dedup; no execution lock, so a crashed worker can't deadlock a doc — Qdrant upserts are idempotent). - Consumer (worker role): `nextcloud-mcp-server worker` drains the queue and runs the existing process_document pipeline; a periodic task reclaims jobs orphaned in `doing` by a crash. INGEST_QUEUE selects the transport (auto: postgres when DATABASE_URL is Postgres, else the in-process anyio queue for SQLite/dev). procrastinate manages its own tables (applied on a fresh DB at startup and by `db upgrade`). The vector-sync status surface reads job counts from Postgres in postgres mode. procrastinate + psycopg3 ship in the [postgres] extra; the app's own engine still uses asyncpg (driver unification is a follow-up handled in the rendered Helm chart). NATS JetStream, the Postgres-queue stub, the bus status subscriber, and nats-py are removed. BREAKING CHANGE: the external-NATS-ingest env vars are removed (INGEST_MODE, STATUS_BACKEND, INGEST_BUS_URL, INGEST_BUS_NUM_REPLICAS, FACT_EVENT_EMITTER). Use INGEST_QUEUE (memory|postgres) and the `worker` command instead. TENANT_ID is retained (no longer NATS-subject-charset-validated). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Chris Coutinho
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
b91af923d2
commit
21b7922bac
@@ -1,7 +1,7 @@
|
||||
"""Ingest-path ports & adapters (design §10, hexagonal)."""
|
||||
"""Ingest-path ports & adapters (design §10, hexagonal; Deck #183)."""
|
||||
|
||||
from .factory import build_external_producer
|
||||
from .factory import build_producer
|
||||
from .memory import MemoryTaskProducer
|
||||
from .ports import TaskProducer
|
||||
|
||||
__all__ = ["MemoryTaskProducer", "TaskProducer", "build_external_producer"]
|
||||
__all__ = ["MemoryTaskProducer", "TaskProducer", "build_producer"]
|
||||
|
||||
@@ -1,16 +1,17 @@
|
||||
"""Composition root for the ingest producer (design §10).
|
||||
"""Composition root for the ingest producer (Deck #183).
|
||||
|
||||
``build_external_producer`` is called by the lifespan only when
|
||||
``INGEST_MODE=external``; local mode uses the in-memory stream directly (it
|
||||
already satisfies :class:`TaskProducer`). The transport under ``external`` is
|
||||
selected from the ``INGEST_BUS_URL`` scheme (``nats://`` now, ``postgres://``
|
||||
later) so moving the external processor to Postgres needs no new INGEST_MODE.
|
||||
The transport is selected from ``INGEST_QUEUE``:
|
||||
|
||||
- ``postgres`` → :class:`ProcrastinateTaskProducer`, which defers jobs into the
|
||||
per-tenant Postgres for the out-of-process ``worker`` role to drain.
|
||||
- ``memory`` (SQLite/dev default) → the in-process anyio stream, built inline by
|
||||
the server lifespan (it owns both the send and receive ends), so it is not
|
||||
produced here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from ...config import Settings
|
||||
from .ports import TaskProducer
|
||||
@@ -18,43 +19,19 @@ from .ports import TaskProducer
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _transport_for(url: str) -> str:
|
||||
scheme = urlsplit(url).scheme.lower()
|
||||
if scheme.startswith("postgres"):
|
||||
return "postgres"
|
||||
if not scheme.startswith("nats"):
|
||||
logger.warning(
|
||||
"INGEST_BUS_URL scheme %r is neither nats:// nor postgres://; "
|
||||
"defaulting to the NATS transport",
|
||||
scheme,
|
||||
)
|
||||
return "nats"
|
||||
async def build_producer(settings: Settings) -> TaskProducer:
|
||||
"""Build the Postgres (procrastinate) ingest producer.
|
||||
|
||||
|
||||
async def build_external_producer(settings: Settings) -> TaskProducer:
|
||||
"""Build the external-ingest producer for the configured transport.
|
||||
|
||||
Precondition: ``settings.ingest_mode == "external"`` (so __post_init__ has
|
||||
guaranteed ``ingest_bus_url`` and ``tenant_id`` are set).
|
||||
Precondition: ``settings.ingest_queue == "postgres"`` (the memory transport
|
||||
is constructed inline by the lifespan because it needs the paired receive
|
||||
stream for the in-process processor pool).
|
||||
"""
|
||||
# Defence-in-depth (robust under ``python -O``, which strips asserts):
|
||||
# __post_init__ already guarantees these when ingest_mode == external.
|
||||
if settings.ingest_bus_url is None or settings.tenant_id is None:
|
||||
if settings.ingest_queue != "postgres":
|
||||
raise ValueError(
|
||||
"build_external_producer requires INGEST_BUS_URL and TENANT_ID "
|
||||
"(guaranteed by Settings validation when INGEST_MODE=external)"
|
||||
"build_producer is only for INGEST_QUEUE=postgres; the memory "
|
||||
f"transport is built inline by the lifespan (got {settings.ingest_queue!r})"
|
||||
)
|
||||
|
||||
transport = _transport_for(settings.ingest_bus_url)
|
||||
if transport == "postgres":
|
||||
from .postgres import PostgresTaskProducer # noqa: PLC0415
|
||||
from .procrastinate import ProcrastinateTaskProducer # noqa: PLC0415
|
||||
|
||||
return await PostgresTaskProducer.connect(settings)
|
||||
|
||||
from .nats import NatsTaskProducer # noqa: PLC0415
|
||||
|
||||
return await NatsTaskProducer.connect(
|
||||
url=settings.ingest_bus_url,
|
||||
tenant_id=settings.tenant_id,
|
||||
num_replicas=settings.ingest_bus_num_replicas,
|
||||
)
|
||||
return await ProcrastinateTaskProducer.connect()
|
||||
|
||||
@@ -1,173 +0,0 @@
|
||||
"""NATS JetStream ``TaskProducer`` — external ingest transport (design §3.4).
|
||||
|
||||
Publishes ``mcp.ingest.requested.{tenant_id}`` for the external
|
||||
document-processor to consume. Translates the in-process ``DocumentTask`` into
|
||||
the wire ``IngestMessage`` schema (mirrored in astrolabe-cloud-website's
|
||||
``bus/messages.py``), with the JetStream ``Nats-Msg-Id`` dedup header per §3.4.
|
||||
|
||||
This server is only the *producer* on this transport; the document-processor
|
||||
owns the consumer. ``nats-py`` is imported lazily so deployments that never
|
||||
enable external ingest don't pay the import.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
from types import TracebackType
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ...canonical import canonical_json
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..scanner import DocumentTask
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
STREAM_NAME = "mcp"
|
||||
INGEST_SUBJECT_PREFIX = "mcp.ingest.requested"
|
||||
|
||||
|
||||
def warn_if_insecure_nats_url(url: str) -> None:
|
||||
"""Log a warning when the bus URL is not TLS-encrypted.
|
||||
|
||||
``nats://`` (and ``ws://``) carry tenant document metadata in cleartext;
|
||||
production deployments should use ``tls://`` (or ``wss://``). We connect
|
||||
regardless — this is an operator alert, not a hard failure.
|
||||
"""
|
||||
scheme = url.split("://", 1)[0].lower()
|
||||
if scheme not in ("tls", "wss"):
|
||||
logger.warning(
|
||||
"NATS bus URL uses unencrypted transport (scheme=%s://); "
|
||||
"use tls:// in production to protect document metadata in transit",
|
||||
scheme,
|
||||
)
|
||||
|
||||
|
||||
def _modified_at_rfc3339(modified_at: int) -> str:
|
||||
"""DocumentTask.modified_at is an epoch int (0 for deletes)."""
|
||||
return datetime.fromtimestamp(int(modified_at), tz=timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _content_hash(task: DocumentTask) -> str:
|
||||
"""etag is the change-detection token; fall back to modified_at when it is
|
||||
absent (e.g. deletes, or sources whose etag we don't thread through).
|
||||
|
||||
TODO(follow-up, PR #814 review): thread etags for file / deck_card /
|
||||
news_item scans too (only note scans pass etag today). Until then their
|
||||
JetStream Nats-Msg-Id dedup keys off modified_at, which misses content
|
||||
changes that leave modified_at unchanged (e.g. a file move/rename).
|
||||
"""
|
||||
return task.etag or str(task.modified_at)
|
||||
|
||||
|
||||
def msg_id(tenant_id: str, doc_id: str, modified_at_rfc3339: str) -> str:
|
||||
"""JetStream dedup header per §3.4. SHA-256 over canonical JSON (NOT the
|
||||
BLAKE2b helper) — it is an opaque external header, not a stored field."""
|
||||
return hashlib.sha256(
|
||||
canonical_json(
|
||||
{
|
||||
"tenant_id": tenant_id,
|
||||
"doc_id": doc_id,
|
||||
"modified_at": modified_at_rfc3339,
|
||||
}
|
||||
)
|
||||
).hexdigest()
|
||||
|
||||
|
||||
class NatsTaskProducer:
|
||||
"""Publishes ingest requests to NATS JetStream."""
|
||||
|
||||
def __init__(self, nc: Any, js: Any, tenant_id: str):
|
||||
self._nc = nc
|
||||
self._js = js
|
||||
self.tenant_id = tenant_id
|
||||
|
||||
@classmethod
|
||||
async def connect(
|
||||
cls, *, url: str, tenant_id: str, num_replicas: int = 1
|
||||
) -> NatsTaskProducer:
|
||||
import nats # noqa: PLC0415 (lazy: optional dependency for external mode)
|
||||
|
||||
warn_if_insecure_nats_url(url)
|
||||
nc = await nats.connect(url)
|
||||
js = nc.jetstream()
|
||||
await cls._ensure_stream(js, num_replicas)
|
||||
logger.info("Connected NATS ingest producer: url=%s, tenant=%s", url, tenant_id)
|
||||
return cls(nc, js, tenant_id)
|
||||
|
||||
@staticmethod
|
||||
async def _ensure_stream(js: Any, num_replicas: int) -> None:
|
||||
# noqa: PLC0415 — nats.js types are only importable once nats-py is present.
|
||||
from nats.js.api import RetentionPolicy, StreamConfig # noqa: PLC0415
|
||||
|
||||
config = StreamConfig(
|
||||
name=STREAM_NAME,
|
||||
subjects=["mcp.>"],
|
||||
retention=RetentionPolicy.LIMITS,
|
||||
num_replicas=num_replicas,
|
||||
)
|
||||
try:
|
||||
await js.add_stream(config=config)
|
||||
logger.info("nats.stream_created stream=%s", STREAM_NAME)
|
||||
except Exception as exc:
|
||||
# add_stream is idempotent in spirit but errors when the stream
|
||||
# already exists; treat as benign (mirrors the processor's
|
||||
# ensure_stream). A genuinely broken broker surfaces on publish.
|
||||
logger.info("nats.stream_exists_or_unavailable detail=%s", exc)
|
||||
|
||||
def ingest_message(self, task: DocumentTask) -> dict[str, Any]:
|
||||
"""DocumentTask → wire IngestMessage dict (mirrors the sibling schema)."""
|
||||
return {
|
||||
"tenant_id": self.tenant_id,
|
||||
"doc_id": task.doc_id,
|
||||
"content_hash": _content_hash(task),
|
||||
"modified_at": _modified_at_rfc3339(task.modified_at),
|
||||
"doc_type": task.doc_type,
|
||||
"operation": task.operation,
|
||||
"user_id": task.user_id,
|
||||
"file_path": task.file_path,
|
||||
}
|
||||
|
||||
async def send(self, task: DocumentTask) -> None:
|
||||
message = self.ingest_message(task)
|
||||
subject = f"{INGEST_SUBJECT_PREFIX}.{self.tenant_id}"
|
||||
headers = {
|
||||
"Nats-Msg-Id": msg_id(self.tenant_id, task.doc_id, message["modified_at"])
|
||||
}
|
||||
await self._js.publish(subject, canonical_json(message), headers=headers)
|
||||
|
||||
# The scanner/oauth_sync use the producer as a clone-able async context
|
||||
# manager (memory-stream semantics). The bus connection is owned by the
|
||||
# lifespan, so cloning shares it and __aexit__ is a no-op (close happens via
|
||||
# aclose() on shutdown).
|
||||
def clone(self) -> NatsTaskProducer:
|
||||
return self
|
||||
|
||||
async def __aenter__(self) -> NatsTaskProducer:
|
||||
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, but this handle
|
||||
# close is a genuine no-op.
|
||||
async def aclose(self) -> None: # NOSONAR
|
||||
# Per-handle close (e.g. a per-user scanner clone exiting). The bus
|
||||
# connection is shared and owned by the lifespan, so this is a no-op;
|
||||
# the connection is torn down once via ``drain()`` on shutdown.
|
||||
return None
|
||||
|
||||
async def drain(self) -> None:
|
||||
"""Drain + close the shared NATS connection (lifespan shutdown only)."""
|
||||
try:
|
||||
await self._nc.drain()
|
||||
except Exception:
|
||||
logger.warning("NATS drain on shutdown failed", exc_info=True)
|
||||
@@ -1,18 +1,18 @@
|
||||
"""Ingest-path ports (design §10, hexagonal).
|
||||
"""Ingest-path ports (design §10, hexagonal; Deck #183).
|
||||
|
||||
A ``TaskProducer`` is where the scanner + webhook receiver send a
|
||||
``DocumentTask``. The transport behind it is swappable:
|
||||
|
||||
- the in-process anyio ``MemoryObjectSendStream`` (local ingest — the default),
|
||||
- ``NatsTaskProducer`` (external ingest → the document-processor), and
|
||||
- a future Postgres-queue producer (seam only; the *external* processor owns the
|
||||
consume side — see ``postgres.py``).
|
||||
- ``MemoryTaskProducer`` over the in-process anyio ``MemoryObjectSendStream``
|
||||
(``INGEST_QUEUE=memory`` — the SQLite/dev default), and
|
||||
- ``ProcrastinateTaskProducer`` (``INGEST_QUEUE=postgres``), which defers jobs
|
||||
into the per-tenant Postgres for the out-of-process ``worker`` role to drain.
|
||||
|
||||
The protocol is exactly the surface the scanner/oauth_sync already use on the
|
||||
memory stream (``send`` + ``clone`` + ``async with``), so both adapters drop in
|
||||
with only a type-annotation change at the call sites. There is intentionally NO
|
||||
consumer port: the MCP server's only in-process consumer is the memory stream;
|
||||
when ingest is external the document-processor is the consumer, not this server.
|
||||
with only a type-annotation change at the call sites. There is no consumer port:
|
||||
in memory mode the in-process processor pool is the consumer; in postgres mode
|
||||
the procrastinate worker is.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
"""Postgres-queue ``TaskProducer`` — documented seam, not implemented.
|
||||
|
||||
The external document-processor may later drain a Postgres-backed queue instead
|
||||
of NATS to limit NATS operational overhead. Processing stays *external*; only
|
||||
the transport changes — so on this server it would be a drop-in producer swap.
|
||||
The consume side + the queue-table migration belong to that processor-side
|
||||
refactor (cross-repo), NOT here. This stub exists so the transport value and the
|
||||
``TaskProducer`` Protocol conformance are testable today.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import TracebackType
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..scanner import DocumentTask
|
||||
|
||||
_NOT_IMPLEMENTED = (
|
||||
"Postgres ingest transport is a documented seam. The external "
|
||||
"document-processor owns the Postgres-drain refactor (transport swap only; "
|
||||
"processing stays external). Use INGEST_BUS_URL=nats://… for now."
|
||||
)
|
||||
|
||||
|
||||
class PostgresTaskProducer:
|
||||
@classmethod
|
||||
async def connect(cls, settings: Any) -> PostgresTaskProducer:
|
||||
raise NotImplementedError(_NOT_IMPLEMENTED)
|
||||
|
||||
async def send(self, task: DocumentTask) -> None: # pragma: no cover
|
||||
raise NotImplementedError(_NOT_IMPLEMENTED)
|
||||
|
||||
def clone(self) -> PostgresTaskProducer: # pragma: no cover
|
||||
return self
|
||||
|
||||
async def __aenter__(self) -> PostgresTaskProducer: # pragma: no cover
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc: BaseException | None,
|
||||
tb: TracebackType | None,
|
||||
) -> None: # pragma: no cover
|
||||
return None
|
||||
|
||||
# The bare suppression marker silences S7503 (async method without await):
|
||||
# ``async def`` is required by the TaskProducer protocol; this stub is a
|
||||
# no-op until the Postgres transport lands.
|
||||
async def aclose(self) -> None: # NOSONAR # pragma: no cover
|
||||
return None
|
||||
@@ -0,0 +1,335 @@
|
||||
"""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()
|
||||
@@ -1,195 +0,0 @@
|
||||
"""Status surface for ingest jobs (design §10.1, ``STATUS_BACKEND``).
|
||||
|
||||
- ``local``: in-process job state — the memory-stream buffer (today's behavior,
|
||||
read directly by the status endpoint).
|
||||
- ``bus``: a background subscriber consumes
|
||||
``mcp.document.{ready,failed,reparsed}.{tenant_id}`` into a bounded in-process
|
||||
:class:`StatusStore` that the status endpoint / ``nc_get_vector_sync_status``
|
||||
read.
|
||||
|
||||
**Honest constraint (design §10.2 / decision):** MCP progress notifications
|
||||
(``ctx.report_progress``) can only be emitted inside an *active tool-call
|
||||
request*; a background subscriber has no ``ctx`` and the MCP SDK exposes no
|
||||
out-of-band push. So "surface events as MCP progress notifications" is delivered
|
||||
via this store (polled by the status endpoint / a tool), not an unsolicited
|
||||
server push. True server-initiated progress / SSE is a follow-up — the
|
||||
``on_event`` callback seam is left in place for it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections import OrderedDict
|
||||
from typing import TYPE_CHECKING, Any, Callable
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import anyio
|
||||
from anyio.abc import TaskStatus
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Terminal/intermediate document states carried on mcp.document.* subjects.
|
||||
_VALID_STATES = {"ready", "failed", "reparsed"}
|
||||
|
||||
|
||||
class StatusStore:
|
||||
"""Bounded LRU of recent document states keyed by ``doc_id``."""
|
||||
|
||||
def __init__(self, max_size: int = 10_000):
|
||||
self._entries: OrderedDict[str, dict[str, Any]] = OrderedDict()
|
||||
self._max = max_size
|
||||
|
||||
def record(
|
||||
self,
|
||||
doc_id: str,
|
||||
state: str,
|
||||
*,
|
||||
content_hash: str | None = None,
|
||||
transitioned_at: str | None = None,
|
||||
) -> None:
|
||||
self._entries[doc_id] = {
|
||||
"state": state,
|
||||
"content_hash": content_hash,
|
||||
"transitioned_at": transitioned_at,
|
||||
}
|
||||
self._entries.move_to_end(doc_id)
|
||||
while len(self._entries) > self._max:
|
||||
self._entries.popitem(last=False)
|
||||
|
||||
def get(self, doc_id: str) -> dict[str, Any] | None:
|
||||
return self._entries.get(doc_id)
|
||||
|
||||
def counts(self) -> dict[str, int]:
|
||||
out: dict[str, int] = {}
|
||||
for entry in self._entries.values():
|
||||
out[entry["state"]] = out.get(entry["state"], 0) + 1
|
||||
return out
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._entries)
|
||||
|
||||
|
||||
def state_from_subject(subject: str) -> str | None:
|
||||
"""``mcp.document.<state>.<tenant_id>`` → ``<state>`` (or None if unknown)."""
|
||||
parts = subject.split(".")
|
||||
if len(parts) >= 4 and parts[0] == "mcp" and parts[1] == "document":
|
||||
state = parts[2]
|
||||
if state in _VALID_STATES:
|
||||
return state
|
||||
return None
|
||||
|
||||
|
||||
class NatsStatusSubscriber:
|
||||
"""Consumes ``mcp.document.*.{tenant_id}`` into a :class:`StatusStore`."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
nc: Any,
|
||||
js: Any,
|
||||
tenant_id: str,
|
||||
store: StatusStore,
|
||||
on_event: Callable[[str, str], None] | None = None,
|
||||
):
|
||||
self._nc = nc
|
||||
self._js = js
|
||||
self.tenant_id = tenant_id
|
||||
self.store = store
|
||||
# on_event(doc_id, state) — seam for a future SSE / progress bridge.
|
||||
self._on_event = on_event
|
||||
|
||||
def handle_message(self, subject: str, data: bytes) -> None:
|
||||
"""Parse one status message into the store. Unit-testable without NATS."""
|
||||
import json # noqa: PLC0415
|
||||
|
||||
state = state_from_subject(subject)
|
||||
if state is None:
|
||||
logger.warning("status.unknown_subject subject=%s", subject)
|
||||
return
|
||||
try:
|
||||
payload = json.loads(data)
|
||||
doc_id = payload["doc_id"]
|
||||
except Exception:
|
||||
logger.warning("status.bad_message subject=%s", subject, exc_info=True)
|
||||
return
|
||||
self.store.record(
|
||||
doc_id,
|
||||
state,
|
||||
content_hash=payload.get("content_hash"),
|
||||
transitioned_at=payload.get("transitioned_at"),
|
||||
)
|
||||
if self._on_event is not None:
|
||||
self._on_event(doc_id, state)
|
||||
|
||||
@classmethod
|
||||
async def connect(
|
||||
cls, *, url: str, tenant_id: str, store: StatusStore
|
||||
) -> NatsStatusSubscriber:
|
||||
import nats # noqa: PLC0415
|
||||
|
||||
from .nats import warn_if_insecure_nats_url # noqa: PLC0415
|
||||
|
||||
warn_if_insecure_nats_url(url)
|
||||
nc = await nats.connect(url)
|
||||
js = nc.jetstream()
|
||||
return cls(nc, js, tenant_id, store)
|
||||
|
||||
async def run(
|
||||
self,
|
||||
shutdown_event: anyio.Event,
|
||||
*,
|
||||
task_status: TaskStatus | None = None,
|
||||
) -> None:
|
||||
"""Durable pull-consumer loop. Requires a live broker (integration)."""
|
||||
import anyio # noqa: PLC0415
|
||||
import nats.errors # noqa: PLC0415
|
||||
|
||||
subject = f"mcp.document.*.{self.tenant_id}"
|
||||
# Signal "task running" *before* the first (fallible) subscribe: bus
|
||||
# status is a non-critical observability path, so a broker that isn't
|
||||
# ready at startup should retry below rather than crash the lifespan.
|
||||
# ``started()`` therefore means "the subscriber loop is running", not
|
||||
# "the subscription succeeded".
|
||||
if task_status is not None:
|
||||
task_status.started()
|
||||
|
||||
sub = None
|
||||
while not shutdown_event.is_set():
|
||||
if sub is None:
|
||||
try:
|
||||
sub = await self._js.pull_subscribe(
|
||||
subject, durable=f"mcp-status-{self.tenant_id}"
|
||||
)
|
||||
except Exception:
|
||||
# Broker not ready / transient connect error: back off and
|
||||
# retry the subscribe instead of giving up.
|
||||
logger.warning(
|
||||
"NATS status subscribe failed; retrying", exc_info=True
|
||||
)
|
||||
await anyio.sleep(5)
|
||||
continue
|
||||
try:
|
||||
msgs = await sub.fetch(batch=16, timeout=5)
|
||||
except nats.errors.TimeoutError:
|
||||
# Expected when idle: no messages within the fetch window. Loop
|
||||
# straight back to re-check shutdown — no log, no extra sleep.
|
||||
continue
|
||||
except Exception:
|
||||
# Real broker error (disconnect, auth failure, stream deleted):
|
||||
# drop the (possibly dead) subscription, back off, and
|
||||
# re-subscribe on the next iteration rather than hot-spinning.
|
||||
logger.warning(
|
||||
"NATS status subscriber fetch failed; re-subscribing",
|
||||
exc_info=True,
|
||||
)
|
||||
sub = None
|
||||
await anyio.sleep(5)
|
||||
continue
|
||||
for msg in msgs:
|
||||
self.handle_message(msg.subject, msg.data)
|
||||
await msg.ack()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
try:
|
||||
await self._nc.drain()
|
||||
except Exception:
|
||||
logger.warning("NATS status subscriber drain failed", exc_info=True)
|
||||
Reference in New Issue
Block a user