fix: address PR #836 review — forward task_producer to MCP contexts + cleanups
🔴 nc_get_vector_sync_status reported pending=0 for INGEST_QUEUE=postgres: the AppContext/OAuthAppContext per-session yields snapshotted the stream fields but never forwarded task_producer, so lifespan_ctx.task_producer was always None. Convert task_producer to a @property that reads _vector_sync_state live (like eviction_task_group), removing the snapshot field so the yields can't drop it. Add a regression test pinning the contract on both contexts. 🟡 Remove the unused _RECLAIM_TASK_NAME constant. 🟡 get_procrastinate_conninfo: warn + document that DATABASE_URL query params (application_name, connect_timeout, …) are dropped. 🟡 worker: open the procrastinate App once — apply_ingest_queue_schema gains manage_connection=False so the worker reuses its own open connector instead of a redundant open/close before run_worker_async. 🟢 Clarify the apply-schema broad-except comment (non-race errors re-raise) and document the deliberate Any typing in ingest_status.get_ingest_pending. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
63e073c224
commit
cfdef3c2c5
@@ -355,10 +355,18 @@ class AppContext:
|
||||
storage: "RefreshTokenStorage | None" = None
|
||||
document_send_stream: MemoryObjectSendStream | None = None
|
||||
document_receive_stream: MemoryObjectReceiveStream | None = None
|
||||
task_producer: "TaskProducer | None" = None
|
||||
shutdown_event: anyio.Event | None = None
|
||||
scanner_wake_event: anyio.Event | None = None
|
||||
|
||||
@property
|
||||
def task_producer(self) -> "TaskProducer | None":
|
||||
# Read dynamically from the module-level singleton (like
|
||||
# eviction_task_group) rather than snapshotting at yield time — that way
|
||||
# a session can't observe a stale ``None`` and the per-session yields
|
||||
# can't forget to forward it (the bug this property replaces). The
|
||||
# vector-sync status tool reads this for postgres-backend job counts.
|
||||
return _vector_sync_state.task_producer
|
||||
|
||||
@property
|
||||
def eviction_task_group(self) -> TaskGroup | None:
|
||||
# Read dynamically from the module-level singleton instead of
|
||||
@@ -381,10 +389,14 @@ class OAuthAppContext:
|
||||
server_client_id: str | None = None # MCP server's OAuth client ID (static or DCR)
|
||||
document_send_stream: MemoryObjectSendStream | None = None
|
||||
document_receive_stream: MemoryObjectReceiveStream | None = None
|
||||
task_producer: "TaskProducer | None" = None
|
||||
shutdown_event: anyio.Event | None = None
|
||||
scanner_wake_event: anyio.Event | None = None
|
||||
|
||||
@property
|
||||
def task_producer(self) -> "TaskProducer | None":
|
||||
# See AppContext.task_producer for rationale.
|
||||
return _vector_sync_state.task_producer
|
||||
|
||||
@property
|
||||
def eviction_task_group(self) -> TaskGroup | None:
|
||||
# See AppContext.eviction_task_group for rationale.
|
||||
@@ -606,8 +618,8 @@ async def app_lifespan_basic(server: FastMCP) -> AsyncIterator[AppContext]:
|
||||
document_receive_stream=_vector_sync_state.document_receive_stream,
|
||||
shutdown_event=_vector_sync_state.shutdown_event,
|
||||
scanner_wake_event=_vector_sync_state.scanner_wake_event,
|
||||
# eviction_task_group is exposed via @property (reads
|
||||
# _vector_sync_state at access time, not snapshot).
|
||||
# task_producer and eviction_task_group are exposed via @property
|
||||
# (read _vector_sync_state at access time, not snapshot).
|
||||
)
|
||||
finally:
|
||||
logger.info("Shutting down BasicAuth session")
|
||||
@@ -1239,8 +1251,8 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
|
||||
document_receive_stream=_vector_sync_state.document_receive_stream,
|
||||
shutdown_event=_vector_sync_state.shutdown_event,
|
||||
scanner_wake_event=_vector_sync_state.scanner_wake_event,
|
||||
# eviction_task_group is exposed via @property (reads
|
||||
# _vector_sync_state at access time, not snapshot).
|
||||
# task_producer and eviction_task_group are exposed via
|
||||
# @property (read _vector_sync_state at access time).
|
||||
)
|
||||
finally:
|
||||
logger.info("Shutting down MCP server")
|
||||
|
||||
@@ -326,9 +326,12 @@ def worker(concurrency: int | None):
|
||||
app = get_procrastinate_app()
|
||||
|
||||
async def _run() -> None:
|
||||
# Defensive apply (the always-on API pod is the authoritative applier).
|
||||
await apply_ingest_queue_schema(app)
|
||||
# Open the connector pool once and reuse it for both the defensive
|
||||
# schema apply (the always-on API pod is the authoritative applier) and
|
||||
# the worker loop — manage_connection=False avoids a redundant
|
||||
# open/close cycle on startup.
|
||||
async with app.open_async():
|
||||
await apply_ingest_queue_schema(app, manage_connection=False)
|
||||
click.echo(
|
||||
f"Ingest worker started: queue={INGEST_QUEUE_NAME} concurrency={workers}"
|
||||
)
|
||||
|
||||
@@ -1437,6 +1437,13 @@ def get_procrastinate_conninfo(database_url: str | None = None) -> str:
|
||||
(set ``DATABASE_URL`` to a ``+psycopg`` URL) rather than rewriting the driver
|
||||
in code — see charts repo, not this repo.
|
||||
|
||||
Only the host/port/dbname/user/password components are forwarded; any
|
||||
``?key=value`` query parameters on the URL (e.g. ``application_name``,
|
||||
``connect_timeout``) are **dropped** — TLS is set separately via
|
||||
:func:`_pg_ssl_params`, and SQLAlchemy-specific query options don't map
|
||||
cleanly to libpq keywords. A warning is logged when params are dropped so
|
||||
operators aren't surprised.
|
||||
|
||||
Raises ``ValueError`` for a non-Postgres URL — procrastinate is Postgres-only.
|
||||
"""
|
||||
from psycopg.conninfo import make_conninfo # noqa: PLC0415
|
||||
@@ -1449,6 +1456,13 @@ def get_procrastinate_conninfo(database_url: str | None = None) -> str:
|
||||
f"got driver {url.drivername!r}"
|
||||
)
|
||||
|
||||
if url.query:
|
||||
logging.getLogger(__name__).warning(
|
||||
"Dropping DATABASE_URL query parameters not forwarded to the "
|
||||
"procrastinate connector: %s",
|
||||
", ".join(sorted(url.query)),
|
||||
)
|
||||
|
||||
params: dict[str, str] = {}
|
||||
if url.host:
|
||||
params["host"] = url.host
|
||||
|
||||
@@ -36,6 +36,12 @@ async def get_ingest_pending(
|
||||
) -> IngestPending:
|
||||
"""Compute outstanding ingest work for the configured queue backend.
|
||||
|
||||
``task_producer`` and ``document_receive_stream`` are intentionally typed
|
||||
``Any``: they're duck-typed across backends. Only ``ProcrastinateTaskProducer``
|
||||
exposes ``job_counts`` (the ``TaskProducer`` protocol doesn't), and the memory
|
||||
backend reads the anyio stream's ``statistics()`` — so no single concrete type
|
||||
or Protocol fits both branches, and we probe with ``hasattr`` instead.
|
||||
|
||||
Never raises — a status surface must stay available even if the queue is
|
||||
unreachable; failures degrade to ``pending=0``.
|
||||
"""
|
||||
|
||||
@@ -54,7 +54,6 @@ 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
|
||||
@@ -217,8 +216,8 @@ async def _ingest_schema_present(app: App) -> bool:
|
||||
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).
|
||||
async def _apply_ingest_queue_schema_open(app: App) -> None:
|
||||
"""Apply the ingest-queue schema on an already-open connector (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
|
||||
@@ -230,12 +229,7 @@ async def apply_ingest_queue_schema(app: App | None = None) -> None:
|
||||
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
|
||||
@@ -243,14 +237,35 @@ async def apply_ingest_queue_schema(app: App | None = None) -> None:
|
||||
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.
|
||||
# The apply runs in a single transaction, so any failure rolls back
|
||||
# atomically (no partial schema). The only benign case is losing the
|
||||
# create race to another pod — confirmed by re-checking presence. Any
|
||||
# other failure (network, auth, …) leaves the schema absent, so this
|
||||
# branch re-raises it rather than masking it.
|
||||
if await _ingest_schema_present(app):
|
||||
logger.info("Ingest queue schema applied concurrently by another pod")
|
||||
return
|
||||
raise
|
||||
|
||||
|
||||
async def apply_ingest_queue_schema(
|
||||
app: App | None = None, *, manage_connection: bool = True
|
||||
) -> None:
|
||||
"""Create procrastinate's tables on a fresh database (apply-if-absent).
|
||||
|
||||
By default opens a short-lived connection, so it is safe to call standalone
|
||||
from the CLI ``db upgrade`` path. Pass ``manage_connection=False`` when the
|
||||
caller already holds an open connector (the ``worker`` command opens the App
|
||||
once and reuses it) to avoid a redundant open/close cycle.
|
||||
"""
|
||||
app = app or get_procrastinate_app()
|
||||
if not manage_connection:
|
||||
await _apply_ingest_queue_schema_open(app)
|
||||
return
|
||||
async with app.open_async():
|
||||
await _apply_ingest_queue_schema_open(app)
|
||||
|
||||
|
||||
# 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")
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
"""Regression test for the lifespan context `task_producer` exposure (Deck #183).
|
||||
|
||||
`nc_get_vector_sync_status` reads `lifespan_ctx.task_producer` for postgres-backend
|
||||
job counts. It was previously a snapshot dataclass field the per-session yields
|
||||
forgot to populate, so the tool always reported `pending=0` on the postgres
|
||||
backend. It is now a `@property` that reads the module singleton live (like
|
||||
`eviction_task_group`); these tests pin that contract.
|
||||
"""
|
||||
|
||||
from typing import cast
|
||||
|
||||
import pytest
|
||||
|
||||
import nextcloud_mcp_server.app as app_module
|
||||
from nextcloud_mcp_server.app import AppContext, OAuthAppContext
|
||||
from nextcloud_mcp_server.client import NextcloudClient
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
|
||||
def test_app_context_task_producer_reads_vector_sync_state(monkeypatch):
|
||||
sentinel = object()
|
||||
monkeypatch.setattr(app_module._vector_sync_state, "task_producer", sentinel)
|
||||
ctx = AppContext(client=cast(NextcloudClient, None))
|
||||
assert ctx.task_producer is sentinel
|
||||
|
||||
|
||||
def test_oauth_app_context_task_producer_reads_vector_sync_state(monkeypatch):
|
||||
sentinel = object()
|
||||
monkeypatch.setattr(app_module._vector_sync_state, "task_producer", sentinel)
|
||||
ctx = OAuthAppContext(
|
||||
nextcloud_host="https://example.test", token_verifier=object()
|
||||
)
|
||||
assert ctx.task_producer is sentinel
|
||||
Reference in New Issue
Block a user