pgdevkit 0.12.0__tar.gz → 0.14.2__tar.gz

pgdevkit-0.12.0/README.md → pgdevkit-0.14.2/PKG-INFO

@@ -1,3 +1,28 @@
+Metadata-Version: 2.5
+Name: pgdevkit
+Version: 0.14.2
+Summary: A helper for developing with Postgres
+Requires-Python: >=3.14
+Requires-Dist: docker>=7.1.0
+Requires-Dist: psycopg[binary]>=3.3.0
+Requires-Dist: sqlglot>=30.11.0
+Provides-Extra: azure
+Requires-Dist: azure-identity>=1.19.0; extra == 'azure'
+Provides-Extra: cli
+Requires-Dist: rich>=13.0.0; extra == 'cli'
+Requires-Dist: tqdm>=4.66.0; extra == 'cli'
+Requires-Dist: typer>=0.26.7; extra == 'cli'
+Provides-Extra: db
+Requires-Dist: psycopg-pool>=3.3.0; extra == 'db'
+Requires-Dist: pydantic>=2.0; extra == 'db'
+Provides-Extra: fastapi
+Requires-Dist: fastapi>=0.118.0; extra == 'fastapi'
+Requires-Dist: psycopg-pool>=3.3.0; extra == 'fastapi'
+Requires-Dist: pydantic>=2.0; extra == 'fastapi'
+Provides-Extra: mssql
+Requires-Dist: mssql-python>=1.0.0; extra == 'mssql'
+Description-Content-Type: text/markdown
+
 # pgdevkit
 
 A helper for developing with Postgres.
@@ -406,6 +431,20 @@ Install with the `db` extra: `pip install pgdevkit[db]`.
   `--entra-user`. For Lakebase hosts, also set the
   `{env_prefix}DATABRICKS_WORKSPACE_HOST` and `{env_prefix}DATABRICKS_INSTANCE`
   env vars.
+- **`fetch_all(query, params=None, *, model=None, row_mapper=None, con=None, pool=None, cancel=None, statement_timeout=None)`**
+  — run one custom query and get every row back as dicts, as validated
+  Pydantic `model` instances, or through a `row_mapper`. Pass `con` to run on a
+  connection you already hold (it is borrowed: never committed, closed or
+  released); otherwise a connection is taken from `pool` or the pool registered
+  with `set_default_pool()`. 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
+  `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). `query` must be a literal string, a sqlglot
+  expression, a `psycopg.sql` composable or a t-string — never an f-string (a sqlglot
+  expression is only as safe as the strings it was built from: pass user values as
+  `exp.Placeholder` + `params`).
 - **CRUD functions** — `pg_retrieve`, `pg_retrieve_many`, `pg_insert`,
   `pg_insert_many`, `pg_update`, `pg_update_dict`, `pg_upsert`,
   `pg_upsert_dict`, `pg_upsert_many`, `pg_upsert_many_dict`, `pg_delete`,
@@ -450,6 +489,97 @@ async with pool.connection() as con:
     await pg_upsert(con, Widget(id=1, name="thing"), Widget)
 ```
 
+## `pgdevkit.fastapi` — streaming JSON responses
+
+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.)
+
+```python
+from pgdevkit.db import PgPool, set_default_pool
+from pgdevkit.fastapi import PostgresJsonResponse
+
+pool = PgPool(env_prefix="APP_POSTGRES_")
+set_default_pool(pool)  # once at startup, shared with fetch_all
+
+@router.get("/articles", responses={200: {"model": list[ArticleOut]}})
+async def articles(lng: str) -> PostgresJsonResponse:
+    return PostgresJsonResponse(
+        "select id, name from articles where lng = %(lng)s order by name",
+        {"lng": lng},
+    )
+```
+
+The `responses=` entry keeps the OpenAPI schema (and a generated frontend client) typed; `response_model` is ignored
+for a `Response`, and the rows are *not* validated against the model.
+
+- **Connections** work as in `fetch_all`: the pool registered with `set_default_pool()`, an explicit `pool=` (anything
+  with `.connection()`, e.g. `PgPool`), or a `con=` you opened yourself. A pooled connection is acquired when streaming
+  starts and released when it ends. A `con` must stay open until the response has been sent, i.e. come from a `Depends`
+  with `yield` (FastAPI >= 0.118); `async with pool.connection() as conn: return PostgresJsonResponse(q, con=conn)`
+  closes it too early. Prefer the pool.
+- **Query** — a literal string, a `psycopg.sql` composable, a sqlglot expression or a t-string, with values bound
+  through `params` (`%(name)s`; a t-string carries its own values and takes no `params`). Passing only literals
+  is enforced by the type checker, **not at runtime**: the query is wrapped as a subquery, never parsed, so never
+  build it from user input. For a data-modifying CTE, or custom JSON, pass `query_produces_json=True` and return one
+  *text* column per row (a `json` column or NULL fails). A trailing `;` is fine in a string, a `;` followed by a
+  comment is not.
+- **Options** — `batch_size` (rows per chunk, >= 1, default 1000; don't take it from user input),
+  `statement_timeout` (seconds, see below), `status_code`, `headers` (e.g. `{"Cache-Control": "max-age=3600"}`).
+
+### Cancellation
+
+The response watches for `http.disconnect` itself (Starlette only does that for ASGI spec < 2.4), also while the
+query is still running and nothing has been sent yet. It then cancels the query with `cancel_safe()`, shielded from
+the cancellation that is tearing the request down, and closes the connection if it was left mid-query: a pool opens a
+fresh one instead of reusing it. Only a query that is still running is cancelled, so a late cancel can never
+hit the next statement on a reused connection. A frontend `AbortController` / closed tab therefore frees the database.
+
+With `con=`, a disconnect aborts the query *and your transaction*, unlike `fetch_all`: the connection is afterwards
+either left in an aborted transaction (idle, with autocommit) or closed, depending on whether the stream was waiting for
+Postgres or for the client.
+Don't use it afterwards, and don't count on earlier writes in that transaction.
+
+Connections are held while the client reads, so a client that reads very slowly pins one (`statement_timeout` does not
+help, see below). Enforce a write timeout in the 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`.
+
+### Statement timeout
+
+`statement_timeout=<seconds>` (optional, also on `fetch_all`) lets **Postgres** abort the query after that long. It
+is set with `SET LOCAL` semantics for this one query: transaction-scoped, so it can't leak into the pool and works
+behind PgBouncer's transaction pooling, and a `con` you passed gets its previous setting back. Before the first byte
+the client gets a `504 {"error": "Query timed out"}`; after that the response is cut short like any other error. It
+needs a connection that is not in autocommit mode. Without it, the connection's own `statement_timeout` applies.
+
+It limits *query execution*, nothing else: Postgres defers a cancel while it is blocked writing to a client that
+has stopped reading, so a stalled reader keeps its connection regardless (a backend was observed staying `active`
+in `ClientWrite` well past a 1 s timeout). For that, use a write timeout in the reverse proxy.
+
+### Errors
+
+An error before the first byte becomes a JSON `{"error": ...}` response (500, or the status and headers of an
+`HTTPException` raised by a custom pool); the text is hidden unless you subclass and set `expose_errors = True`
+(it contains SQL fragments and parameter values: never in production):
+
+```python
+class AppJsonResponse(PostgresJsonResponse):
+    expose_errors = IS_DEV
+```
+
+After the first byte the status line is gone: the error is logged and re-raised, and the response ends **without
+the closing `]`**, so clients fail to parse it rather than accepting a silently shortened array. (Granian, which our
+apps run on, ends such a response cleanly with the truncated body.)
+
+### Notes
+
+It does no authorization: scope the query yourself. The first byte is sent once `batch_size` rows (or the whole
+result) are available, so lower `batch_size` for earlier output; Postgres itself also sends rows in 8 kB buffers,
+so a slow query with tiny rows delivers late either way.
+
 ## Releasing
 
 Bump `version` in `pyproject.toml` as part of your PR, same as any other

pgdevkit-0.12.0/PKG-INFO → pgdevkit-0.14.2/README.md

@@ -1,24 +1,3 @@
-Metadata-Version: 2.5
-Name: pgdevkit
-Version: 0.12.0
-Summary: A helper for developing with Postgres
-Requires-Python: >=3.14
-Requires-Dist: docker>=7.1.0
-Requires-Dist: psycopg[binary]>=3.2.0
-Requires-Dist: sqlglot>=30.11.0
-Provides-Extra: azure
-Requires-Dist: azure-identity>=1.19.0; extra == 'azure'
-Provides-Extra: cli
-Requires-Dist: rich>=13.0.0; extra == 'cli'
-Requires-Dist: tqdm>=4.66.0; extra == 'cli'
-Requires-Dist: typer>=0.26.7; extra == 'cli'
-Provides-Extra: db
-Requires-Dist: psycopg-pool>=3.3.0; extra == 'db'
-Requires-Dist: pydantic>=2.0; extra == 'db'
-Provides-Extra: mssql
-Requires-Dist: mssql-python>=1.0.0; extra == 'mssql'
-Description-Content-Type: text/markdown
-
 # pgdevkit
 
 A helper for developing with Postgres.
@@ -427,6 +406,20 @@ Install with the `db` extra: `pip install pgdevkit[db]`.
   `--entra-user`. For Lakebase hosts, also set the
   `{env_prefix}DATABRICKS_WORKSPACE_HOST` and `{env_prefix}DATABRICKS_INSTANCE`
   env vars.
+- **`fetch_all(query, params=None, *, model=None, row_mapper=None, con=None, pool=None, cancel=None, statement_timeout=None)`**
+  — run one custom query and get every row back as dicts, as validated
+  Pydantic `model` instances, or through a `row_mapper`. Pass `con` to run on a
+  connection you already hold (it is borrowed: never committed, closed or
+  released); otherwise a connection is taken from `pool` or the pool registered
+  with `set_default_pool()`. 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
+  `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). `query` must be a literal string, a sqlglot
+  expression, a `psycopg.sql` composable or a t-string — never an f-string (a sqlglot
+  expression is only as safe as the strings it was built from: pass user values as
+  `exp.Placeholder` + `params`).
 - **CRUD functions** — `pg_retrieve`, `pg_retrieve_many`, `pg_insert`,
   `pg_insert_many`, `pg_update`, `pg_update_dict`, `pg_upsert`,
   `pg_upsert_dict`, `pg_upsert_many`, `pg_upsert_many_dict`, `pg_delete`,
@@ -471,6 +464,97 @@ async with pool.connection() as con:
     await pg_upsert(con, Widget(id=1, name="thing"), Widget)
 ```
 
+## `pgdevkit.fastapi` — streaming JSON responses
+
+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.)
+
+```python
+from pgdevkit.db import PgPool, set_default_pool
+from pgdevkit.fastapi import PostgresJsonResponse
+
+pool = PgPool(env_prefix="APP_POSTGRES_")
+set_default_pool(pool)  # once at startup, shared with fetch_all
+
+@router.get("/articles", responses={200: {"model": list[ArticleOut]}})
+async def articles(lng: str) -> PostgresJsonResponse:
+    return PostgresJsonResponse(
+        "select id, name from articles where lng = %(lng)s order by name",
+        {"lng": lng},
+    )
+```
+
+The `responses=` entry keeps the OpenAPI schema (and a generated frontend client) typed; `response_model` is ignored
+for a `Response`, and the rows are *not* validated against the model.
+
+- **Connections** work as in `fetch_all`: the pool registered with `set_default_pool()`, an explicit `pool=` (anything
+  with `.connection()`, e.g. `PgPool`), or a `con=` you opened yourself. A pooled connection is acquired when streaming
+  starts and released when it ends. A `con` must stay open until the response has been sent, i.e. come from a `Depends`
+  with `yield` (FastAPI >= 0.118); `async with pool.connection() as conn: return PostgresJsonResponse(q, con=conn)`
+  closes it too early. Prefer the pool.
+- **Query** — a literal string, a `psycopg.sql` composable, a sqlglot expression or a t-string, with values bound
+  through `params` (`%(name)s`; a t-string carries its own values and takes no `params`). Passing only literals
+  is enforced by the type checker, **not at runtime**: the query is wrapped as a subquery, never parsed, so never
+  build it from user input. For a data-modifying CTE, or custom JSON, pass `query_produces_json=True` and return one
+  *text* column per row (a `json` column or NULL fails). A trailing `;` is fine in a string, a `;` followed by a
+  comment is not.
+- **Options** — `batch_size` (rows per chunk, >= 1, default 1000; don't take it from user input),
+  `statement_timeout` (seconds, see below), `status_code`, `headers` (e.g. `{"Cache-Control": "max-age=3600"}`).
+
+### Cancellation
+
+The response watches for `http.disconnect` itself (Starlette only does that for ASGI spec < 2.4), also while the
+query is still running and nothing has been sent yet. It then cancels the query with `cancel_safe()`, shielded from
+the cancellation that is tearing the request down, and closes the connection if it was left mid-query: a pool opens a
+fresh one instead of reusing it. Only a query that is still running is cancelled, so a late cancel can never
+hit the next statement on a reused connection. A frontend `AbortController` / closed tab therefore frees the database.
+
+With `con=`, a disconnect aborts the query *and your transaction*, unlike `fetch_all`: the connection is afterwards
+either left in an aborted transaction (idle, with autocommit) or closed, depending on whether the stream was waiting for
+Postgres or for the client.
+Don't use it afterwards, and don't count on earlier writes in that transaction.
+
+Connections are held while the client reads, so a client that reads very slowly pins one (`statement_timeout` does not
+help, see below). Enforce a write timeout in the 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`.
+
+### Statement timeout
+
+`statement_timeout=<seconds>` (optional, also on `fetch_all`) lets **Postgres** abort the query after that long. It
+is set with `SET LOCAL` semantics for this one query: transaction-scoped, so it can't leak into the pool and works
+behind PgBouncer's transaction pooling, and a `con` you passed gets its previous setting back. Before the first byte
+the client gets a `504 {"error": "Query timed out"}`; after that the response is cut short like any other error. It
+needs a connection that is not in autocommit mode. Without it, the connection's own `statement_timeout` applies.
+
+It limits *query execution*, nothing else: Postgres defers a cancel while it is blocked writing to a client that
+has stopped reading, so a stalled reader keeps its connection regardless (a backend was observed staying `active`
+in `ClientWrite` well past a 1 s timeout). For that, use a write timeout in the reverse proxy.
+
+### Errors
+
+An error before the first byte becomes a JSON `{"error": ...}` response (500, or the status and headers of an
+`HTTPException` raised by a custom pool); the text is hidden unless you subclass and set `expose_errors = True`
+(it contains SQL fragments and parameter values: never in production):
+
+```python
+class AppJsonResponse(PostgresJsonResponse):
+    expose_errors = IS_DEV
+```
+
+After the first byte the status line is gone: the error is logged and re-raised, and the response ends **without
+the closing `]`**, so clients fail to parse it rather than accepting a silently shortened array. (Granian, which our
+apps run on, ends such a response cleanly with the truncated body.)
+
+### Notes
+
+It does no authorization: scope the query yourself. The first byte is sent once `batch_size` rows (or the whole
+result) are available, so lower `batch_size` for earlier output; Postgres itself also sends rows in 8 kB buffers,
+so a slow query with tiny rows delivers late either way.
+
 ## Releasing
 
 Bump `version` in `pyproject.toml` as part of your PR, same as any other

{pgdevkit-0.12.0 → pgdevkit-0.14.2}/pgdevkit/db/__init__.py

@@ -16,15 +16,19 @@ from .crud import (
     pg_upsert_many,
     pg_upsert_many_dict,
 )
+from .fetch import ConnectionSource, SqlQuery, fetch_all, set_default_pool
 from .loader import SqlLoader
 from .model import PostgresTableModel, TableModel
 
 __all__ = [
     "ComplexHelper",
+    "ConnectionSource",
     "PgPool",
     "PostgresTableModel",
     "SqlLoader",
+    "SqlQuery",
     "TableModel",
+    "fetch_all",
     "pg_delete",
     "pg_delete_dict",
     "pg_insert",
@@ -37,4 +41,5 @@ __all__ = [
     "pg_upsert_dict",
     "pg_upsert_many",
     "pg_upsert_many_dict",
+    "set_default_pool",
 ]

pgdevkit-0.14.2/pgdevkit/db/fetch.py

@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import math
+import re
+from collections.abc import AsyncIterator, Callable, Mapping, Sequence
+from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager, suppress
+from string.templatelib import Template
+from typing import TYPE_CHECKING, Any, LiteralString, Protocol, overload
+
+from psycopg.connection_async import AsyncConnection
+from psycopg.errors import QueryCanceled
+from psycopg.rows import dict_row
+from psycopg.sql import Composable
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sqlglot import exp  # imported lazily: only needed to type sqlglot queries (rendering is duck-typed)
+
+# Same params shape psycopg's Cursor.execute() accepts: a %(name)s-style mapping,
+# or a positional %s-style sequence.
+QueryParams = Mapping[str, Any] | Sequence[Any] | None
+
+# The only things a query may be: a literal string (e.g. from `SqlLoader.load_sql()`), a sqlglot
+# expression, or psycopg's own safe-composition types (`psycopg.sql` / a t-string). A plain `str`
+# is rejected by the type checker on purpose, so user-controlled text can't be concatenated in.
+type SqlQueryNoTemplate = LiteralString | exp.Expression | Composable
+type SqlQuery = SqlQueryNoTemplate | Template
+
+
+class ConnectionSource(Protocol):
+    """Anything that hands out an async connection context manager -- `PgPool` fits."""
+
+    def connection(self) -> AbstractAsyncContextManager[AsyncConnection]: ...
+
+
+_default_source: ConnectionSource | None = None
+
+
+def set_default_pool(pool: ConnectionSource | None) -> None:
+    """Register where `fetch_all()` gets its connection when called without `con`/`pool`
+    (typically the app's `PgPool`, once at startup). Pass `None` to clear it. This is process-wide
+    state meant for single-database apps: with several databases/tenants pass `pool=` explicitly."""
+    global _default_source
+    _default_source = pool
+
+
+@asynccontextmanager
+async def _acquire(con: AsyncConnection | None, pool: ConnectionSource | None) -> AsyncIterator[AsyncConnection]:
+    """Yield `con` untouched (the caller owns it: no commit/close/return-to-pool here), or a
+    connection borrowed from `pool` / the default pool for the duration of the block."""
+    if con is not None:
+        if pool is not None:
+            raise TypeError("Pass either `con` or `pool`, not both.")
+        yield con
+        return
+    source = pool if pool is not None else _default_source
+    if source is None:
+        raise RuntimeError("No connection: pass `con=`/`pool=`, or register a pool with `set_default_pool()`.")
+    async with source.connection() as borrowed:
+        yield borrowed
+
+
+_PLACEHOLDER_NAME = re.compile(r"\w+")
+
+
+def _render(query: SqlQuery) -> Any:
+    if isinstance(query, (str, Composable, Template)):
+        return query
+    if not hasattr(query, "sql"):
+        raise TypeError(f"fetch_all: unsupported query type {type(query).__name__}")
+    # Anything else is a sqlglot expression (no sqlglot import needed to know that; it is loaded
+    # already). Its postgres dialect renders `exp.Placeholder("id")` as `%(id)s`.
+    from sqlglot import exp
+
+    for node in query.walk():
+        # sqlglot renders these verbatim, so they could smuggle arbitrary SQL into the statement.
+        if isinstance(node, exp.Command):
+            raise ValueError("fetch_all: raw sqlglot Command nodes are not allowed")
+        if isinstance(node, exp.Placeholder) and node.name and not _PLACEHOLDER_NAME.fullmatch(node.name):
+            raise ValueError(f"fetch_all: invalid placeholder name {node.name!r}")
+    return query.sql(dialect="postgres")
+
+
+def _check_statement_timeout(statement_timeout: float | None) -> None:
+    if statement_timeout is not None and statement_timeout <= 0:
+        raise ValueError(f"statement_timeout must be a positive number of seconds, got {statement_timeout}")
+
+
+@asynccontextmanager
+async def _statement_timeout(c: AsyncConnection, seconds: float | None) -> AsyncIterator[None]:
+    """Run the block under Postgres' `statement_timeout`, set with `set_config(..., is_local => true)`: it is
+    scoped to the current transaction (so it can't leak into a pool, and it is safe behind PgBouncer's
+    transaction pooling) and put back afterwards for a borrowed connection that stays in its transaction.
+    `None` leaves the connection's own setting alone."""
+    if seconds is None:
+        yield
+        return
+    async with AsyncExitStack() as stack:
+        if c.autocommit:
+            await stack.enter_async_context(c.transaction())  # a local setting needs a transaction to live in
+        cur = await c.execute(
+            "SELECT current_setting('statement_timeout'), set_config('statement_timeout', %(ms)s, true)",
+            {"ms": str(max(1, math.ceil(seconds * 1000)))},
+        )
+        row = await cur.fetchone()
+        assert row is not None
+        yield
+        await c.execute("SELECT set_config('statement_timeout', %(previous)s, true)", {"previous": row[0]})
+
+
+async def _run_query(
+    c: AsyncConnection, query: Any, params: QueryParams, statement_timeout: float | None
+) -> list[dict[str, Any]]:
+    async with _statement_timeout(c, statement_timeout), c.cursor(row_factory=dict_row) as cur:
+        await cur.execute(query, params)
+        return await cur.fetchall()
+
+
+async def _abort(c: AsyncConnection, work: asyncio.Future[Any]) -> None:
+    """Cancel `work`'s query on the server and wait until it has unwound, leaving `c` idle (an
+    aborted transaction, to be rolled back by its owner). If the cancel request itself can't be
+    delivered, the connection's state is unknown, so it is closed (the server abandons the query as
+    soon as it next talks to the client) and a pool discards it instead of reusing it."""
+    if not work.done():
+        try:
+            await c.cancel_safe()
+        except Exception:
+            logger.warning("fetch_all: cancelling the query failed; closing the connection", exc_info=True)
+            with suppress(Exception):
+                await c.close()  # first: a cancelled-but-still-busy connection would make `work` wait out the query
+            work.cancel()
+            await asyncio.gather(work, return_exceptions=True)
+            return
+    await asyncio.gather(work, return_exceptions=True)  # always consume `work`'s outcome
+
+
+async def _abort_uninterruptibly(c: AsyncConnection, work: asyncio.Future[Any]) -> None:
+    """`_abort`, but a further cancellation of the awaiting task can't cut it short (which would
+    leave `work` running); it is re-raised once the abort is complete."""
+    task = asyncio.ensure_future(_abort(c, work))
+    interrupted = False
+    while not task.done():
+        try:
+            await asyncio.shield(task)
+        except asyncio.CancelledError:
+            interrupted = True
+    task.result()
+    if interrupted:
+        raise asyncio.CancelledError
+
+
+async def _fetch_dicts(
+    query: Any,  # already rendered (see `_render`)
+    params: QueryParams,
+    con: AsyncConnection | None,
+    pool: ConnectionSource | None,
+    cancel: asyncio.Event | None,
+    statement_timeout: float | None,
+) -> list[dict[str, Any]]:
+    if cancel is not None and cancel.is_set():
+        raise QueryCanceled("fetch_all: cancelled before the query started")
+    async with _acquire(con, pool) as c:
+        work = asyncio.ensure_future(_run_query(c, query, params, statement_timeout))
+        stop = asyncio.ensure_future(cancel.wait()) if cancel is not None else None
+        try:
+            await asyncio.wait([work, *([stop] if stop else [])], return_when=asyncio.FIRST_COMPLETED)
+            if not work.done():  # `cancel` was set while the query runs: abort it on the server
+                await _abort_uninterruptibly(c, work)
+            if work.cancelled():
+                raise QueryCanceled("fetch_all: cancelled via `cancel`")
+            return work.result()  # raises psycopg.errors.QueryCanceled if it was cancelled
+        except asyncio.CancelledError:
+            # The awaiting task itself was cancelled: abandoning the query client-side would leave it
+            # running on the server (and the connection mid-query), so cancel it there and let it
+            # unwind first -- the connection is then idle for its owner/the pool.
+            await _abort_uninterruptibly(c, work)
+            raise
+        finally:
+            if stop is not None:
+                stop.cancel()
+
+
+@overload
+async def fetch_all[M: BaseModel](
+    query: SqlQueryNoTemplate,
+    params: QueryParams = None,
+    *,
+    model: type[M],
+    con: AsyncConnection | None = None,
+    pool: ConnectionSource | None = None,
+    cancel: asyncio.Event | None = None,
+    statement_timeout: float | None = None,
+) -> list[M]: ...
+
+
+@overload
+async def fetch_all[T](
+    query: SqlQueryNoTemplate,
+    params: QueryParams = None,
+    *,
+    row_mapper: Callable[[dict[str, Any]], T],
+    con: AsyncConnection | None = None,
+    pool: ConnectionSource | None = None,
+    cancel: asyncio.Event | None = None,
+    statement_timeout: float | None = None,
+) -> list[T]: ...
+
+
+@overload
+async def fetch_all(
+    query: SqlQueryNoTemplate,
+    params: QueryParams = None,
+    *,
+    model: None = None,
+    row_mapper: None = None,
+    con: AsyncConnection | None = None,
+    pool: ConnectionSource | None = None,
+    cancel: asyncio.Event | None = None,
+    statement_timeout: float | None = None,
+) -> list[dict[str, Any]]: ...
+
+
+@overload
+async def fetch_all[M: BaseModel](
+    query: Template,
+    *,
+    model: type[M],
+    con: AsyncConnection | None = None,
+    pool: ConnectionSource | None = None,
+    cancel: asyncio.Event | None = None,
+    statement_timeout: float | None = None,
+) -> list[M]: ...
+
+
+@overload
+async def fetch_all[T](
+    query: Template,
+    *,
+    row_mapper: Callable[[dict[str, Any]], T],
+    con: AsyncConnection | None = None,
+    pool: ConnectionSource | None = None,
+    cancel: asyncio.Event | None = None,
+    statement_timeout: float | None = None,
+) -> list[T]: ...
+
+
+@overload
+async def fetch_all(
+    query: Template,
+    *,
+    model: None = None,
+    row_mapper: None = None,
+    con: AsyncConnection | None = None,
+    pool: ConnectionSource | None = None,
+    cancel: asyncio.Event | None = None,
+    statement_timeout: float | None = None,
+) -> list[dict[str, Any]]: ...
+
+
+async def fetch_all(
+    query: SqlQuery,
+    params: QueryParams = None,
+    *,
+    model: type[BaseModel] | None = None,
+    row_mapper: Callable[[dict[str, Any]], Any] | None = None,
+    con: AsyncConnection | None = None,
+    pool: ConnectionSource | None = None,
+    cancel: asyncio.Event | None = None,
+    statement_timeout: float | None = None,
+) -> list[Any]:
+    """Run one query and return every row.
+
+    Rows come back as plain dicts, or, with `model=` (a Pydantic model class), validated into
+    that model; `row_mapper=` converts each dict row into any other shape. The two are exclusive.
+
+    Pass `statement_timeout` (seconds) to let Postgres abort the query after that long: `fetch_all` then
+    raises `psycopg.errors.QueryCanceled` ("canceling statement due to statement timeout"). It is applied
+    with `SET LOCAL` semantics, i.e. for this query's transaction only; a borrowed `con` in autocommit mode
+    gets a short transaction around the query, and any other borrowed `con` gets its previous setting back.
+    Without it, the connection's own `statement_timeout` (default: none) applies.
+
+    Pass `con` to run on a connection you already hold (e.g. inside a transaction): it is
+    borrowed, never committed, closed or returned to a pool. Without it, a connection is taken
+    from `pool` or the pool registered via `set_default_pool()` and released afterwards.
+
+    Set `cancel` (an `asyncio.Event`, e.g. when the HTTP client disconnects; to stream a large result
+    from FastAPI see `pgdevkit.fastapi.PostgresJsonResponse`, which does that wiring itself) to abort a running
+    query on the server: `fetch_all` then raises `psycopg.errors.QueryCanceled`. If the awaiting
+    task is cancelled instead, the query is cancelled server-side too and `CancelledError`
+    propagates, so the query never keeps running unattended. `cancel` aborts whatever statement is
+    running on the connection, so don't use it on a `con` that concurrent tasks share. If the
+    cancel request itself can't be delivered, the connection is closed.
+
+    `query` must be a literal string, a sqlglot expression, a `psycopg.sql` composable or a
+    t-string -- never text built with f-strings/concatenation (a plain `str` fails type checking).
+    A sqlglot expression is only as safe as the strings it was built from (the builders parse plain
+    strings as SQL): pass user values as `exp.Placeholder` + `params`, never as literals/raw text.
+    In one, write a literal `%` as `%%` when you also pass `params`.
+    A t-string already carries its values, so it must not be combined with `params`.
+    """
+    if model is not None and row_mapper is not None:
+        raise TypeError("Pass either `model` or `row_mapper`, not both.")
+    if isinstance(query, Template) and params is not None:
+        raise TypeError("A t-string query carries its own values; don't pass `params` with it.")
+    _check_statement_timeout(statement_timeout)
+    rows = await _fetch_dicts(_render(query), params, con, pool, cancel, statement_timeout)
+    if model is not None:
+        return [model.model_validate(row) for row in rows]
+    if row_mapper is not None:
+        return [row_mapper(row) for row in rows]
+    return rows

pgdevkit-0.14.2/pgdevkit/fastapi/__init__.py

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