feat(fastapi): CancelOnDisconnectRoute - cancel the handler (and its query) when the client disconnects

README.md

@@ -417,7 +417,7 @@ Install with the `db` extra: `pip install pgdevkit[db]`.
   with `set_default_pool()` (a `ConnectionSource` such as `PgPool`, or a plain callable returning an
   async connection context manager). Set the `cancel` `asyncio.Event` to abort the query on the server
   (`QueryCanceled`; to stream a result as JSON from FastAPI see `PostgresJsonResponse`, which watches for the
-  disconnect itself); cancelling the awaiting task cancels it there too, and
+  disconnect itself, and for buffered endpoints see `CancelOnDisconnectRoute`); cancelling the awaiting task cancels it there too, and
   `CancelledError` propagates as usual. `statement_timeout=<seconds>` lets Postgres abort the query after that
   long (`QueryCanceled`, "canceling statement due to statement timeout"); it is applied per call with `SET LOCAL`
   semantics, so it never outlives the call (a borrowed `con` gets its previous setting back); see
@@ -511,14 +511,16 @@ async with pool.connection() as con:
     await pg_upsert(con, Widget(id=1, name="thing"), Widget)
 ```
 
-## `pgdevkit.fastapi` — streaming JSON responses
+## `pgdevkit.fastapi` — streaming JSON responses and cancel-on-disconnect
 
 Requires the `fastapi` extra (`pip install pgdevkit[fastapi]`, FastAPI >= 0.118).
 
 `PostgresJsonResponse` runs a query and streams its rows to the client as a JSON array. Postgres builds the JSON
 (`row_to_json`), so there is no row-by-row Python serialization and no pydantic round trip, and memory use stays flat
 for big results. **If the client goes away, the query is cancelled in Postgres** instead of running to completion.
-(`fetch_all(..., cancel=event)` cancels a buffered query too, but you wire the event to the disconnect yourself.)
+(`fetch_all(..., cancel=event)` cancels a buffered query too; for `fetch_all`-based endpoints put the route on
+`CancelOnDisconnectRoute` ([below](#cancelling-other-endpoints-when-the-client-leaves)) instead of wiring the event
+yourself.)
 
 ```python
 from pgdevkit.db import PgPool, set_default_pool
@@ -571,6 +573,49 @@ big results, can use up the pool and make every other request wait for a connect
 reverse proxy and size the pool accordingly. Disconnect behaviour can only be tested over a real socket
 (`httpx.ASGITransport` never sends `http.disconnect`); see `tests/test_pg_json_granian.py`.
 
+### Cancelling other endpoints when the client leaves
+
+`PostgresJsonResponse` is the only thing that notices a disconnect by itself. ASGI servers such as Granian do **not**
+cancel a request handler when the client goes away, so a handler built on `fetch_all` (or a raw cursor, Meilisearch,
+...) runs to the end and its result is thrown away, which defeats a frontend `AbortController`. Put such read
+endpoints on a router with `CancelOnDisconnectRoute`: it runs the handler as a task and cancels it on
+`http.disconnect` (in the request's own task, with an `asyncio.timeout` as the cancel scope, as
+`PostgresJsonResponse` does), and psycopg/pgdevkit then cancel the running statement on the server (see `fetch_all`'s
+`cancel`).
+
+```python
+from fastapi import APIRouter
+from pgdevkit.db import fetch_all
+from pgdevkit.fastapi import CancelOnDisconnectRoute
+
+cancellable = APIRouter(route_class=CancelOnDisconnectRoute)  # or router.add_api_route(..., route_class_override=...)
+
+@cancellable.post("/customers/overview")  # a POST that only reads
+async def overview(payload: Filter) -> list[Row]:
+    return await fetch_all("select ...", {"q": payload.q})  # aborted in Postgres when the client leaves
+
+app.include_router(cancellable)
+```
+
+- **Opt-in, reads only.** A cancelled write is rolled back at best and half-applied at worst. Don't select routes by
+  HTTP method: search and grid endpoints are often POSTs.
+- **Everything inside the route is cancelled**, dependencies included (authentication, audit logging, counters), and
+  the client gets no response, so `BackgroundTasks` registered by the handler never run. Don't rely on side effects
+  that must always happen in an opted-in route: shield them or do them in a middleware. A `Depends` with `yield` sees
+  an `HTTPException` (status 499) after the cancellation, so a transaction rolls back instead of committing.
+- The client that left gets a 499 nobody reads; the middleware stack sees a normal response (no
+  `No response returned`).
+- The request body is **not** buffered: the handler reads it from a `Request` whose `receive()` is fed by a reader
+  task that also spots the disconnect. That reader runs at most a few chunks ahead of the handler, so a route that
+  never reads a large body can't pin it in memory, but a disconnect is then only noticed once the body is consumed.
+- **Shared work must survive one caller's cancellation.** Await shared tasks and futures (single-flight caches,
+  "refresh once, many wait" loaders) through `asyncio.shield`; otherwise one aborted request cancels the load every
+  other request is waiting on.
+- A returned `StreamingResponse` is not covered (Starlette watches for the disconnect there only on servers with ASGI
+  spec < 2.4; `PostgresJsonResponse` always does). Threads (`asyncio.to_thread`) and
+  external services (Databricks statements, HTTP calls) keep running; only the awaiting stops.
+- Like the above, disconnects only show up over a real socket: see `tests/test_pg_json_granian.py` for a Granian test.
+
 ### Statement timeout
 
 `statement_timeout=<seconds>` (optional, also on `fetch_all`) lets **Postgres** abort the query after that long. It

pgdevkit/fastapi/__init__.py

@@ -1,3 +1,4 @@
+from .cancel_on_disconnect import CancelOnDisconnectRoute
 from .json_response import PostgresJsonResponse
 
-__all__ = ["PostgresJsonResponse"]
+__all__ = ["CancelOnDisconnectRoute", "PostgresJsonResponse"]

pgdevkit/fastapi/cancel_on_disconnect.py

@@ -0,0 +1,132 @@
+"""Cancel a FastAPI handler when its client goes away, so the database work it started stops too.
+
+Requires the ``fastapi`` extra: ``pip install pgdevkit[fastapi]``.
+
+ASGI servers such as Granian do not cancel a request handler when the client disconnects: the handler runs to the
+end and its result is thrown away. ``PostgresJsonResponse`` watches for the disconnect itself, but ``fetch_all`` and
+friends (or any other awaitable work, raw psycopg cursors included) only stop when the awaiting task is cancelled.
+``CancelOnDisconnectRoute`` closes that gap: it runs the route handler as a task and cancels it on
+``http.disconnect``; psycopg and pgdevkit then cancel the running statement on the server.
+"""
+
+import asyncio
+import logging
+from collections.abc import Callable, Coroutine
+from typing import Any
+
+from fastapi import HTTPException, Request, Response
+from fastapi.routing import APIRoute
+from starlette.requests import ClientDisconnect
+from starlette.types import Message, Receive
+
+from .json_response import CLIENT_CLOSED_REQUEST
+
+logger = logging.getLogger(__name__)
+
+# Body chunks read ahead of the handler: enough to never stall it, too few to buffer an upload in memory.
+_MAX_PENDING_MESSAGES = 8
+
+
+class CancelOnDisconnectRoute(APIRoute):
+    """An ``APIRoute`` that cancels its handler as soon as the client disconnects.
+
+    Opt in per router, never for the whole app, and only for routes that are safe to abort halfway: reads, searches,
+    exports. A cancelled write is rolled back at best and half-applied at worst::
+
+        cancellable = APIRouter(route_class=CancelOnDisconnectRoute)
+
+        @cancellable.post("/customers/overview")  # a POST that only reads
+        async def overview(payload: Filter) -> Page:
+            return await fetch_all(...)  # aborted in Postgres when the client leaves
+
+        app.include_router(cancellable)
+
+    A single route can opt in with ``router.add_api_route(..., route_class_override=CancelOnDisconnectRoute)``
+    (``APIRouter`` only; the ``@router.get`` decorators don't take it).
+
+    Things to know:
+
+    * **Everything inside the route is cancelled**, dependencies included (authentication, audit logging, counters)
+      and so is ``BackgroundTasks`` registration, since the client gets no response. Don't rely on side effects that
+      must always happen; shield them or do them in a middleware. A ``Depends`` with ``yield`` sees the 499
+      ``HTTPException`` after the cancellation, so its cleanup rolls back instead of committing.
+    * The request body is not buffered: the handler reads it as usual, from a ``Request`` whose ``receive()`` is fed
+      by a reader task that also spots the disconnect. While the handler isn't reading a large body, the reader
+      stops after a few chunks, so a disconnect is only noticed once the handler has consumed the body.
+    * Work shared with other requests must survive the cancellation of one of them: await shared tasks and futures
+      through ``asyncio.shield`` (cache single-flight loaders, background refreshes), otherwise one aborted request
+      cancels the load every other request is waiting on.
+    * A returned ``StreamingResponse`` is not covered (the handler has finished by the time it is sent): Starlette
+      watches for the disconnect there only on servers with ASGI spec < 2.4, ``PostgresJsonResponse`` always does.
+    * Threads (``asyncio.to_thread``) and external services (Databricks statements, HTTP calls) keep running after
+      the cancellation; only awaiting them stops.
+    """
+
+    def get_route_handler(self) -> Callable[[Request], Coroutine[Any, Any, Response]]:
+        handler = super().get_route_handler()
+
+        async def cancel_on_disconnect(request: Request) -> Response:
+            # The handler runs in the request's own task, as without this route (context variables, dependencies with
+            # `yield`, an outer `asyncio.timeout` all behave as usual); an asyncio timeout is the cancel scope, as in
+            # PostgresJsonResponse. The disconnect pulls its deadline to "now".
+            loop = asyncio.get_running_loop()
+            receiver = _Receiver(request.receive)
+            request._receive = receiver.receive  # the same Request object, so a cached body is seen by everyone
+            reader = asyncio.create_task(receiver.read())
+
+            async def cancel_when_gone(cancel_scope: asyncio.Timeout) -> None:
+                await receiver.disconnected.wait()
+                cancel_scope.reschedule(loop.time())
+
+            try:
+                async with asyncio.timeout(None) as cancel_scope:
+                    watcher = asyncio.create_task(cancel_when_gone(cancel_scope))
+                    try:
+                        return await handler(request)
+                    finally:
+                        watcher.cancel()
+            except TimeoutError:
+                if not cancel_scope.expired():  # the handler's own timeout
+                    raise
+            except ClientDisconnect:  # it was reading the body when the client left
+                pass
+            finally:
+                reader.cancel()
+            logger.debug("client disconnected, cancelled %s %s", request.method, request.url.path)
+            # Raised, not returned: the dependencies' exit code (e.g. a transaction) must see a failure, not a success.
+            raise HTTPException(status_code=CLIENT_CLOSED_REQUEST, detail="Client closed the request")
+
+        return cancel_on_disconnect
+
+
+class _Receiver:
+    """Gives the handler's ``Request`` its own ``receive()`` while one reader task watches for the disconnect."""
+
+    def __init__(self, receive: Receive) -> None:
+        self._receive = receive
+        self._messages: asyncio.Queue[Message | Exception] = asyncio.Queue(maxsize=_MAX_PENDING_MESSAGES)
+        self._gone: Message | None = None
+        self.disconnected = asyncio.Event()
+
+    async def read(self) -> None:
+        """Pump the server's messages into the queue until the client is gone (the queue's limit is the back-pressure)."""
+        try:
+            while True:
+                message = await self._receive()
+                if message["type"] == "http.disconnect":
+                    self.disconnected.set()
+                await self._messages.put(message)
+                if message["type"] == "http.disconnect":
+                    return
+        except Exception as exc:  # noqa: BLE001 - a broken server: the handler sees it, the client is not treated as gone
+            await self._messages.put(exc)
+
+    async def receive(self) -> Message:
+        if self._gone is not None:  # the disconnect is final, whoever asks again
+            return self._gone
+        message = await self._messages.get()
+        if isinstance(message, Exception):
+            raise message
+        if message["type"] == "http.disconnect":
+            self._gone = message
+        return message

pgdevkit/fastapi/json_response.py

@@ -4,10 +4,11 @@
 """
 
 import asyncio
+import contextlib
 import json
 import logging
-from contextlib import aclosing
 from collections.abc import AsyncGenerator, Mapping
+from contextlib import aclosing
 from string.templatelib import Template
 from typing import Any, cast
 
@@ -34,6 +35,8 @@
 
 logger = logging.getLogger(__name__)
 
+CLIENT_CLOSED_REQUEST = 499  # nginx's code for a client that left; nobody reads it
+
 _CANCEL_TIMEOUT_SECONDS = 5.0
 # libpq < 17 (a system libpq; psycopg's binary wheels bundle a newer one) has no chunked rows mode.
 _CHUNKED_ROWS = pq.version() >= 170000
@@ -115,17 +118,29 @@ async def watch_for_disconnect(cancel_scope: asyncio.Timeout) -> None:
                 pass
             cancel_scope.reschedule(loop.time())  # i.e. cancel the stream now
 
+        started = False
+
+        async def tracking_send(message: Message) -> None:
+            nonlocal started
+            started = started or message["type"] == "http.response.start"
+            await send(message)
+
         try:
             async with asyncio.timeout(None) as cancel_scope:
                 watcher = asyncio.create_task(watch_for_disconnect(cancel_scope))
                 try:
-                    await self.stream_response(send)
+                    await self.stream_response(tracking_send)
                 finally:
                     watcher.cancel()
         except TimeoutError:
             if not cancel_scope.expired():
                 raise
-            # the client is gone and the query has been cancelled: nobody is left to answer
+            # The client is gone and the query has been cancelled: nobody is left to answer. A middleware in front
+            # (BaseHTTPMiddleware) still expects a response, though, or it fails with "No response returned.".
+            if not started:
+                with contextlib.suppress(OSError, ClientDisconnect):
+                    await send(self._start_message(CLIENT_CLOSED_REQUEST, []))
+                    await send({"type": "http.response.body", "body": b"", "more_body": False})
         if self.background is not None:
             await self.background()
 

pyproject.toml

@@ -11,7 +11,7 @@ packages = ["pgdevkit"]
 
 [project]
 name = "pgdevkit"
-version = "0.15.2"
+version = "0.16.0"
 description = "A helper for developing with Postgres"
 readme = "README.md"
 requires-python = ">=3.14"