diff --git a/README.md b/README.md index e14c89e..86e4d2c 100644 --- a/README.md +++ b/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=` 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=` (optional, also on `fetch_all`) lets **Postgres** abort the query after that long. It diff --git a/pgdevkit/fastapi/__init__.py b/pgdevkit/fastapi/__init__.py index a86932f..dc19d16 100644 --- a/pgdevkit/fastapi/__init__.py +++ b/pgdevkit/fastapi/__init__.py @@ -1,3 +1,4 @@ +from .cancel_on_disconnect import CancelOnDisconnectRoute from .json_response import PostgresJsonResponse -__all__ = ["PostgresJsonResponse"] +__all__ = ["CancelOnDisconnectRoute", "PostgresJsonResponse"] diff --git a/pgdevkit/fastapi/cancel_on_disconnect.py b/pgdevkit/fastapi/cancel_on_disconnect.py new file mode 100644 index 0000000..0b32815 --- /dev/null +++ b/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 diff --git a/pgdevkit/fastapi/json_response.py b/pgdevkit/fastapi/json_response.py index c0d86dc..318491f 100644 --- a/pgdevkit/fastapi/json_response.py +++ b/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() diff --git a/pyproject.toml b/pyproject.toml index a505512..01369e1 100644 --- a/pyproject.toml +++ b/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" diff --git a/skills/pgdevkit/SKILL.md b/skills/pgdevkit/SKILL.md index 61c37e0..d460281 100644 --- a/skills/pgdevkit/SKILL.md +++ b/skills/pgdevkit/SKILL.md @@ -6,7 +6,8 @@ description: > Docker/Podman test database (`pgdb testdb`), importable ORM-free CRUD and query helpers (`pgdevkit.db`: `pg_*`, `fetch_all`/`fetch_one`/`fetch_scalar`/`execute`, `readonly_transaction`, `PgPool`, `SqlLoader`), streaming a query as JSON from a FastAPI endpoint with query cancellation on client disconnect (`pgdevkit.fastapi.PostgresJsonResponse`: large grid - endpoints), and the `database/`-folder schema-as-code convention. Supersedes the old postgres-test-setup, postgres-best-practices, + endpoints; `pgdevkit.fastapi.CancelOnDisconnectRoute`: cancel plain `fetch_all`/cursor read handlers when the client + aborts, e.g. a frontend AbortController), and the `database/`-folder schema-as-code convention. Supersedes the old postgres-test-setup, postgres-best-practices, and database-in-source skills — pgdevkit is a real dependency now, not copy-pasted reference files. Use whenever the user wants to set up a local test Postgres, write or review psycopg code, add a table/view/function to @@ -233,7 +234,7 @@ deleted = await execute("DELETE FROM sessions WHERE user_id = %(id)s", {"id": 1} - `model=` validates each row into a Pydantic model; `row_mapper=` maps each dict row to anything else (exclusive with `model`). - Passing `con=` runs on *your* connection: `fetch_all` never commits, closes or releases it. Without it, a connection is borrowed from `pool=` (or the `set_default_pool()` pool) and released afterwards. - `statement_timeout=` (optional) lets Postgres abort the query after that long, raising `QueryCanceled`; it is applied with `SET LOCAL` semantics for this call only (a borrowed `con` gets its previous value back). A plain value *replaces* the connection's own setting in both directions, so a per-call 900 s overrides a database-level safety net of 600 s; `statement_timeout=at_most(120)` (`from pgdevkit.db import at_most`) never loosens it: the effective timeout is `min(current, 120)`. -- `cancel=` takes an `asyncio.Event`: set it (e.g. when the HTTP client disconnects) and the running query is cancelled on the server, raising `psycopg.errors.QueryCanceled`. If the awaiting task is cancelled instead, the server-side query is cancelled too and `CancelledError` propagates as usual (not `QueryCanceled`), so no query keeps running unattended. Don't use `cancel=` on a `con` that concurrent tasks share (it aborts whatever statement is running on it); if the cancel request itself can't be delivered, the connection is closed. After `QueryCanceled` a borrowed `con` is usable again once you `await con.rollback()` (autocommit connections need nothing). To stream a whole result as JSON from FastAPI (it wires the disconnect itself) see `PostgresJsonResponse` below. +- `cancel=` takes an `asyncio.Event`: set it (e.g. when the HTTP client disconnects) and the running query is cancelled on the server, raising `psycopg.errors.QueryCanceled`. If the awaiting task is cancelled instead, the server-side query is cancelled too and `CancelledError` propagates as usual (not `QueryCanceled`), so no query keeps running unattended. Don't use `cancel=` on a `con` that concurrent tasks share (it aborts whatever statement is running on it); if the cancel request itself can't be delivered, the connection is closed. After `QueryCanceled` a borrowed `con` is usable again once you `await con.rollback()` (autocommit connections need nothing). To stream a whole result as JSON from FastAPI (it wires the disconnect itself) see `PostgresJsonResponse` below; to cancel a plain `fetch_all` endpoint when the client leaves, see `CancelOnDisconnectRoute` below. - `query` is typed `SqlQuery`: a literal string (`SqlLoader.load_sql()`), a sqlglot expression (any `exp.Expr`, so Select, Union, Subquery and other `exp.Query` values are fine), a `psycopg.sql` composable or a t-string. A plain `str` fails type checking on purpose, so don't build SQL with f-strings/concatenation. A sqlglot expression is only as safe as the strings it was built from (its builders parse plain strings as SQL): pass user values as `exp.Placeholder` + `params`, never as literals/raw text. A literal `%` in a sqlglot expression (`LIKE 'a%'`, modulo) just works, with or without `params`; in a plain string or `psycopg.sql` composable it is psycopg's rule: write `%%` when you pass `params`. A t-string carries its own values, so don't also pass `params`. ### Dynamic SQL @@ -326,6 +327,29 @@ async def articles(lng: str) -> PostgresJsonResponse: - `statement_timeout=` (optional, as in `fetch_all`): Postgres aborts a slow query; before the first byte the client gets a 504. It does not free a client that stopped reading (use a proxy write timeout), and needs a non-autocommit connection. - It does no authorization; scope the query yourself. Full reference: the `pgdevkit.fastapi` section of the README. +### Cancelling `fetch_all` endpoints when the client leaves + +Granian does not cancel a request handler on disconnect, so a frontend abort (debounced search, grid filter change) only stops `PostgresJsonResponse` endpoints. Put other read endpoints — including POSTs that only read, e.g. grid/search — on a router with `CancelOnDisconnectRoute`: the handler task is cancelled on `http.disconnect` and psycopg cancels the statement in Postgres. + +```python +from fastapi import APIRouter +from pgdevkit.fastapi import CancelOnDisconnectRoute + +cancellable = APIRouter(route_class=CancelOnDisconnectRoute) # single route: router.add_api_route(..., route_class_override=CancelOnDisconnectRoute) + +@cancellable.post("/customers/overview") # reads only +async def overview(payload: Filter) -> list[Row]: + return await fetch_all(sql.load_sql("customers", "overview"), payload.model_dump()) + +app.include_router(cancellable) +``` + +- **Reads only, opt-in.** Never for writes; don't select by HTTP method (searches are POSTs). +- **Everything in the route is cancelled**, dependencies included (auth, audit logging), and `BackgroundTasks` never run; a `yield` dependency sees an `HTTPException` (499) and rolls back. Don't put must-happen side effects in such a route. +- The client that left gets a 499 nobody reads. The body is not buffered (the handler gets a `Request` fed by a reader task). +- Shared work (single-flight caches, "refresh once, many wait" loaders) must be awaited through `asyncio.shield`, or one aborted request cancels it for all waiters. +- `to_thread` work and external services (Databricks statements, HTTP calls) keep running; only awaiting them stops. Returned `StreamingResponse`s aren't covered (use `PostgresJsonResponse`). + ### SQL formatting ```bash @@ -405,6 +429,7 @@ Reports drift between the `database/` `.sql` files and the actual schema — tab - [ ] Untrusted SQL runs under a SELECT-only role (plus `readonly_transaction(con)` and an `at_most(...)` timeout as safety nets, not as a sandbox) - [ ] `set_default_pool()` at startup; `helper.clear_cache()` after DDL on a long-lived connection when using `complex_helper=` - [ ] Large/grid reads from FastAPI may use `pgdevkit.fastapi.PostgresJsonResponse` (streams, cancels on disconnect, no model validation) +- [ ] Other read endpoints built on `fetch_all` that should stop when the client leaves use `APIRouter(route_class=CancelOnDisconnectRoute)` (never writes) - [ ] No `LATERAL JOIN` — use a CTE that groups/aggregates first, then joins it - [ ] `..sql` files (e.g. `.prod.sql`) are skipped by `pgdb testdb` unless it's run with a matching `--env` - [ ] Every table (and non-obvious column) has a `COMMENT ON`, placed in the object's own `.sql` file diff --git a/tests/_granian_app.py b/tests/_granian_app.py index 56c125a..3dc0b9b 100644 --- a/tests/_granian_app.py +++ b/tests/_granian_app.py @@ -4,10 +4,11 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager -from fastapi import FastAPI +from fastapi import APIRouter, FastAPI, Request from psycopg_pool import AsyncConnectionPool -from pgdevkit.fastapi import PostgresJsonResponse +from pgdevkit.db import fetch_all +from pgdevkit.fastapi import CancelOnDisconnectRoute, PostgresJsonResponse pool = AsyncConnectionPool(os.environ["PGJSON_TEST_DSN"], min_size=1, max_size=2, open=False) @@ -21,6 +22,32 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: app = FastAPI(lifespan=lifespan) +@app.middleware("http") # a BaseHTTPMiddleware, as in our apps: it must cope with a handler that was cancelled +async def passthrough(request: Request, call_next): + return await call_next(request) + + +cancellable = APIRouter(route_class=CancelOnDisconnectRoute) + + +@cancellable.post("/cancellable-sleep") +async def cancellable_sleep() -> list[dict]: + return await fetch_all("select pg_sleep(30) as slept", pool=pool) + + +@cancellable.post("/cancellable-ok") +async def cancellable_ok() -> list[dict]: + return await fetch_all("select 1 as one", pool=pool) + + +app.include_router(cancellable) + + +@app.post("/plain-sleep") # not opted in: runs to completion even if the client is gone +async def plain_sleep() -> list[dict]: + return await fetch_all("select pg_sleep(3) as slept", pool=pool) + + @app.get("/ok") async def ok(): return PostgresJsonResponse("select i from generate_series(1, 3) i", pool=pool) diff --git a/tests/test_cancel_on_disconnect.py b/tests/test_cancel_on_disconnect.py new file mode 100644 index 0000000..94afeaa --- /dev/null +++ b/tests/test_cancel_on_disconnect.py @@ -0,0 +1,332 @@ +"""CancelOnDisconnectRoute cancels the handler (and with it the Postgres query) when the client goes away.""" + +import asyncio +import contextvars +import json +import uuid +from collections.abc import AsyncIterator + +import pytest +from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request +from fastapi.exceptions import RequestValidationError +from fastapi.responses import JSONResponse +from psycopg_pool import AsyncConnectionPool +from pydantic import BaseModel + +from pgdevkit.db import fetch_all +from pgdevkit.fastapi import CancelOnDisconnectRoute +from tests._asgi import active_queries, wait_until + + +class Filter(BaseModel): + n: int + + +current_user: contextvars.ContextVar[str] = contextvars.ContextVar("current_user") + + +class HttpClient: + """Drives an ASGI app by hand: sends one request, then stays connected until `disconnect` is set.""" + + def __init__(self, method: str = "GET", path: str = "/", body: bytes | list[bytes] = b"") -> None: + self.method, self.path = method, path + self.chunks = [body] if isinstance(body, bytes) else list(body) + self.chunks_read = 0 + self.disconnect = asyncio.Event() + self.sent: list[dict] = [] + self.gone_at_start = False # the client disconnected before the body arrived + self.gone_after_body = False # ... right after it + + async def receive(self) -> dict: + if self.gone_at_start: + return {"type": "http.disconnect"} + if self.chunks_read < max(len(self.chunks), 1): + chunk = self.chunks[self.chunks_read] if self.chunks else b"" + self.chunks_read += 1 + return {"type": "http.request", "body": chunk, "more_body": self.chunks_read < len(self.chunks)} + if not self.gone_after_body: + await self.disconnect.wait() + return {"type": "http.disconnect"} + + async def send(self, message: dict) -> None: + self.sent.append(message) + + async def call(self, app: FastAPI) -> None: + scope = { + "type": "http", + "asgi": {"spec_version": "2.3"}, + "http_version": "1.1", + "method": self.method, + "path": self.path, + "raw_path": self.path.encode(), + "query_string": b"", + "headers": [(b"content-type", b"application/json")] if self.chunks and self.method == "POST" else [], + "client": ("testclient", 1), + "server": ("testserver", 80), + "scheme": "http", + "root_path": "", + } + await app(scope, self.receive, self.send) + + @property + def status(self) -> int | None: + return next((m["status"] for m in self.sent if m["type"] == "http.response.start"), None) + + @property + def body(self) -> bytes: + return b"".join(m.get("body", b"") for m in self.sent if m["type"] == "http.response.body") + + +@pytest.fixture +def app_name() -> str: + return f"cod-{uuid.uuid4().hex[:12]}" # isolates this test's backends in pg_stat_activity + + +@pytest.fixture +async def pool(postgres_dsn: str, app_name: str) -> AsyncIterator[AsyncConnectionPool]: + kwargs = {"application_name": app_name} + async with AsyncConnectionPool(postgres_dsn, min_size=1, max_size=2, kwargs=kwargs, open=False) as pool: + yield pool + + +@pytest.fixture +def events() -> list[str]: + return [] + + +@pytest.fixture +def app(pool: AsyncConnectionPool, events: list[str]) -> FastAPI: + cancellable = APIRouter(route_class=CancelOnDisconnectRoute) + plain = APIRouter() + + @cancellable.get("/sleep") + async def sleep() -> list[dict]: + try: + return await fetch_all("select pg_sleep(30) as slept", pool=pool) + except asyncio.CancelledError: + events.append("cancelled") + raise + + @plain.get("/plain-sleep") + async def plain_sleep() -> list[dict]: + return await fetch_all("select pg_sleep(1.5) as slept", pool=pool) + + @cancellable.get("/ok") + async def ok() -> list[dict]: + return await fetch_all("select 1 as one", pool=pool) + + @cancellable.post("/echo") + async def echo(payload: dict) -> dict: # the body was buffered for the watcher: the handler must still get it + return payload + + @cancellable.post("/size") # reads a (large) body in the handler + async def size(request: Request) -> dict: + return {"size": len(await request.body())} + + async def deny() -> None: + raise HTTPException(status_code=401, detail="no") + + @cancellable.post("/denied", dependencies=[Depends(deny)]) # no body parameters, rejected before any body read + async def denied() -> None: ... + + async def transaction() -> AsyncIterator[None]: + try: + yield + except BaseException as exc: + events.append(f"dependency saw {type(exc).__name__}") + raise + events.append("dependency saw success") + + @cancellable.get("/in-transaction", dependencies=[Depends(transaction)]) + async def in_transaction() -> list[dict]: + return await fetch_all("select pg_sleep(30) as slept", pool=pool) + + @cancellable.get("/slow-cleanup") + async def slow_cleanup() -> None: + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + events.append("cleaning up") + await asyncio.sleep(30) + + @cancellable.get("/instant") + async def instant() -> dict: + return {"done": True} + + async def with_context() -> AsyncIterator[None]: # set before, reset after the handler: needs one context + token = current_user.set("me") + try: + yield + finally: + current_user.reset(token) + + @cancellable.get("/context", dependencies=[Depends(with_context)]) + async def context() -> dict: + return {"user": current_user.get()} + + @cancellable.post("/filter") + async def filter_(payload: Filter) -> Filter: + return payload + + @cancellable.get("/forbidden") + async def forbidden() -> None: + raise HTTPException(status_code=403, detail="no") + + @cancellable.get("/boom") + async def boom() -> None: + raise RuntimeError("boom") + + app = FastAPI() + + @app.exception_handler(RequestValidationError) + async def log_the_body(request: Request, exc: RequestValidationError) -> JSONResponse: + return JSONResponse({"detail": "invalid", "body": (await request.body()).decode()}, status_code=422) + + app.include_router(cancellable) + app.include_router(plain) + return app + + +async def query_running(dsn: str, app_name: str) -> bool: + return await active_queries(dsn, app_name) == 1 + + +async def query_stopped(dsn: str, app_name: str) -> bool: + return await active_queries(dsn, app_name) == 0 + + +async def test_disconnect_cancels_the_handler_and_the_query( + app: FastAPI, events: list[str], postgres_dsn: str, app_name: str +) -> None: + client = HttpClient(path="/sleep") + task = asyncio.create_task(client.call(app)) + assert await wait_until(lambda: query_running(postgres_dsn, app_name)) + + client.disconnect.set() + await asyncio.wait_for(task, timeout=5) + + assert events == ["cancelled"] + assert await wait_until(lambda: query_stopped(postgres_dsn, app_name), timeout=3) + assert client.status == 499 + + +async def test_a_plain_route_is_not_cancelled(app: FastAPI, postgres_dsn: str, app_name: str) -> None: + client = HttpClient(path="/plain-sleep") + task = asyncio.create_task(client.call(app)) + assert await wait_until(lambda: query_running(postgres_dsn, app_name)) + + client.disconnect.set() + await asyncio.wait_for(task, timeout=5) + + assert client.status == 200 # opt-in: it ran to completion + + +async def test_normal_request_and_body_still_reach_the_handler(app: FastAPI) -> None: + ok = HttpClient(path="/ok") + await ok.call(app) + assert (ok.status, ok.body) == (200, b'[{"one":1}]') + + echo = HttpClient("POST", "/echo", b'{"a": 1}') + await echo.call(app) + assert (echo.status, echo.body) == (200, b'{"a":1}') + + +async def test_errors_of_the_handler_are_handled_as_usual(app: FastAPI) -> None: + forbidden = HttpClient(path="/forbidden") + await forbidden.call(app) + assert forbidden.status == 403 + + boom = HttpClient(path="/boom") + with pytest.raises(RuntimeError, match="boom"): + await boom.call(app) + + +async def test_cancelling_the_request_cancels_the_handler_and_the_query( + app: FastAPI, events: list[str], postgres_dsn: str, app_name: str +) -> None: + client = HttpClient(path="/sleep") + task = asyncio.create_task(client.call(app)) + assert await wait_until(lambda: query_running(postgres_dsn, app_name)) + + task.cancel() + with pytest.raises(asyncio.CancelledError): # an outer cancellation is never swallowed + await task + + assert events == ["cancelled"] + assert await wait_until(lambda: query_stopped(postgres_dsn, app_name), timeout=3) + + +async def test_no_task_is_left_behind(app: FastAPI) -> None: + before = len(asyncio.all_tasks()) + client = HttpClient(path="/ok") + await client.call(app) + await asyncio.sleep(0) + assert len(asyncio.all_tasks()) <= before + + +async def test_client_gone_before_the_handler_starts_is_not_an_error(app: FastAPI, events: list[str]) -> None: + client = HttpClient("POST", "/size", b"x") + client.gone_at_start = True + await client.call(app) # no ClientDisconnect / 500 / traceback + assert client.status == 499 + assert events == [] + + +async def test_a_large_body_is_streamed_to_the_handler_not_cut_off(app: FastAPI) -> None: + client = HttpClient("POST", "/size", [b"x" * 65536] * 50) + await client.call(app) + assert (client.status, client.body) == (200, b'{"size":3276800}') + + +async def test_an_unread_body_is_not_buffered(app: FastAPI) -> None: + """A route that rejects the request before reading the body must not pull the whole upload into memory.""" + client = HttpClient("POST", "/denied", [b"x" * 65536] * 100) + await client.call(app) + assert client.status == 401 + assert client.chunks_read <= 12, f"{client.chunks_read} of 100 chunks were read ahead" + + +async def test_yield_dependencies_see_a_failure_not_a_success_after_a_disconnect( + app: FastAPI, events: list[str], postgres_dsn: str, app_name: str +) -> None: + client = HttpClient(path="/in-transaction") + task = asyncio.create_task(client.call(app)) + assert await wait_until(lambda: query_running(postgres_dsn, app_name)) + + client.disconnect.set() + await asyncio.wait_for(task, timeout=5) + + assert events == ["dependency saw HTTPException"] + assert client.status == 499 + + +async def test_a_cancellation_during_the_handlers_cleanup_is_not_swallowed(app: FastAPI, events: list[str]) -> None: + client = HttpClient(path="/slow-cleanup") + task = asyncio.create_task(client.call(app)) + await asyncio.sleep(0.2) + client.disconnect.set() + assert await wait_until(lambda: events == ["cleaning up"], timeout=3) + + task.cancel() # e.g. server shutdown or a timeout, while the handler is still cleaning up + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=3) + + +async def test_a_handler_that_finished_wins_over_a_disconnect_that_arrives_with_it(app: FastAPI) -> None: + client = HttpClient(path="/instant") + client.gone_after_body = True + await client.call(app) + assert (client.status, client.body) == (200, b'{"done":true}') + + +async def test_context_variables_and_yield_dependencies_work_as_without_the_route(app: FastAPI) -> None: + client = HttpClient(path="/context") + await client.call(app) # reset() in another context than set() would raise ValueError: 500 + assert (client.status, client.body) == (200, b'{"user":"me"}') + + +async def test_exception_handlers_can_still_read_the_body(app: FastAPI) -> None: + client = HttpClient("POST", "/filter", b'{"n": "x"}') + await asyncio.wait_for(client.call(app), timeout=3) # used to wait for a body the reader had already taken + assert client.status == 422 + assert json.loads(client.body)["body"] == '{"n": "x"}' diff --git a/tests/test_pg_json_cancel.py b/tests/test_pg_json_cancel.py index aedf86e..130bf17 100644 --- a/tests/test_pg_json_cancel.py +++ b/tests/test_pg_json_cancel.py @@ -15,8 +15,7 @@ from psycopg_pool import AsyncConnectionPool from starlette.requests import ClientDisconnect -from pgdevkit.fastapi import PostgresJsonResponse -from pgdevkit.fastapi import json_response +from pgdevkit.fastapi import PostgresJsonResponse, json_response from pgdevkit.fastapi.json_response import _cancel_if_running from tests._asgi import FakeClient, active_queries, wait_until @@ -56,7 +55,11 @@ async def test_disconnect_before_the_first_byte_cancels_the_query( client.disconnect.set() await asyncio.wait_for(task, timeout=5) # finishes quietly: nobody is left to answer - assert client.sent == [] + # nobody is left to read it, but a BaseHTTPMiddleware in front would fail with "No response returned." otherwise + assert [(m["type"], m.get("status")) for m in client.sent] == [ + ("http.response.start", 499), + ("http.response.body", None), + ] assert await wait_until(lambda: query_stopped(postgres_dsn, app_name), timeout=3) assert time.monotonic() - started < 5, "the 30s query must not run to completion" async with pool.connection() as conn: # the pool replaced the connection we closed diff --git a/tests/test_pg_json_granian.py b/tests/test_pg_json_granian.py index ca0545e..def2e99 100644 --- a/tests/test_pg_json_granian.py +++ b/tests/test_pg_json_granian.py @@ -25,11 +25,16 @@ def app_name() -> str: @pytest.fixture(scope="module") -def server(postgres_dsn: str, app_name: str, tmp_path_factory: pytest.TempPathFactory) -> Iterator[str]: +def server_log(tmp_path_factory: pytest.TempPathFactory) -> Path: + return tmp_path_factory.mktemp("granian") / "granian.log" + + +@pytest.fixture(scope="module") +def server(postgres_dsn: str, app_name: str, server_log: Path) -> Iterator[str]: with socket.socket() as s: s.bind(("127.0.0.1", 0)) port = s.getsockname()[1] - log = tmp_path_factory.mktemp("granian") / "granian.log" + log = server_log env = {**os.environ, "PGJSON_TEST_DSN": f"{postgres_dsn} application_name={app_name}"} with log.open("w") as log_file: proc = subprocess.Popen( @@ -65,7 +70,10 @@ async def test_normal_response(server: str) -> None: assert (await client.get(f"{server}/ok")).json() == [{"i": 1}, {"i": 2}, {"i": 3}] -async def test_aborted_client_cancels_the_query(server: str, postgres_dsn: str, app_name: str) -> None: +async def test_aborted_client_cancels_the_query( + server: str, server_log: Path, postgres_dsn: str, app_name: str +) -> None: + log_start = len(server_log.read_text()) async def query_running() -> bool: return await active_queries(postgres_dsn, app_name) == 1 @@ -82,6 +90,7 @@ async def client_gives_up() -> None: await request assert await wait_until(query_stopped, timeout=5), "the query kept running after the client left" + assert "Traceback" not in server_log.read_text()[log_start:] # e.g. BaseHTTPMiddleware: "No response returned." async def test_error_after_the_first_byte_ends_the_response_without_a_valid_json_array(server: str) -> None: @@ -96,3 +105,45 @@ async def test_error_after_the_first_byte_ends_the_response_without_a_valid_json with pytest.raises(json.JSONDecodeError): response.json() assert (await client.get(f"{server}/ok")).status_code == 200 # the server and its pool are fine + + +async def abort_after(server: str, path: str, seconds: float) -> None: + async with httpx.AsyncClient(timeout=seconds) as client: + with pytest.raises(httpx.ReadTimeout): # drops the connection + await client.post(f"{server}{path}") + + +async def test_cancel_on_disconnect_route_cancels_the_query_when_the_client_leaves( + server: str, server_log: Path, postgres_dsn: str, app_name: str +) -> None: + async def running() -> bool: + return await active_queries(postgres_dsn, app_name) == 1 + + async def stopped() -> bool: + return await active_queries(postgres_dsn, app_name) == 0 + + log_start = len(server_log.read_text()) # the log is shared with tests that fail on purpose + request = asyncio.create_task(abort_after(server, "/cancellable-sleep", 2)) + assert await wait_until(running, timeout=5), "the query never started" + await request + + assert await wait_until(stopped, timeout=5), "the query kept running after the client left" + async with httpx.AsyncClient() as client: # server, pool and middleware are fine afterwards + assert (await client.post(f"{server}/cancellable-ok")).json() == [{"one": 1}] + assert "Traceback" not in server_log.read_text()[log_start:] + + +async def test_a_route_that_did_not_opt_in_runs_to_completion(server: str, postgres_dsn: str, app_name: str) -> None: + async def running() -> bool: + return await active_queries(postgres_dsn, app_name) == 1 + + request = asyncio.create_task(abort_after(server, "/plain-sleep", 1)) + assert await wait_until(running, timeout=5) + await request + await asyncio.sleep(0.5) + assert await active_queries(postgres_dsn, app_name) == 1, "a plain route must not be cancelled" + assert await wait_until(lambda: _idle(postgres_dsn, app_name), timeout=6) + + +async def _idle(dsn: str, app_name: str) -> bool: + return await active_queries(dsn, app_name) == 0 diff --git a/uv.lock b/uv.lock index 9889250..0c8e19d 100644 --- a/uv.lock +++ b/uv.lock @@ -451,7 +451,7 @@ wheels = [ [[package]] name = "pgdevkit" -version = "0.15.2" +version = "0.16.0" source = { editable = "." } dependencies = [ { name = "docker" },