pgdevkit 0.12.0__tar.gz → 0.13.0__tar.gz

{pgdevkit-0.12.0 → pgdevkit-0.13.0}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.5
 Name: pgdevkit
-Version: 0.12.0
+Version: 0.13.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: psycopg[binary]>=3.3.0
 Requires-Dist: sqlglot>=30.11.0
 Provides-Extra: azure
 Requires-Dist: azure-identity>=1.19.0; extra == 'azure'
@@ -427,6 +427,17 @@ 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)`**
+  — 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`); cancelling the awaiting task cancels it there too, and
+  `CancelledError` propagates as usual. `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`,

{pgdevkit-0.12.0 → pgdevkit-0.13.0}/README.md

@@ -406,6 +406,17 @@ 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)`**
+  — 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`); cancelling the awaiting task cancels it there too, and
+  `CancelledError` propagates as usual. `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`,

{pgdevkit-0.12.0 → pgdevkit-0.13.0}/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.13.0/pgdevkit/db/fetch.py

@@ -0,0 +1,269 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import re
+from collections.abc import AsyncIterator, Callable, Mapping, Sequence
+from contextlib import AbstractAsyncContextManager, 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")
+
+
+async def _run_query(c: AsyncConnection, query: Any, params: QueryParams) -> list[dict[str, Any]]:
+    async with 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,
+) -> 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))
+        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,
+) -> 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,
+) -> 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,
+) -> 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,
+) -> 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,
+) -> 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,
+) -> 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,
+) -> 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 `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 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.")
+    rows = await _fetch_dicts(_render(query), params, con, pool, cancel)
+    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.12.0 → pgdevkit-0.13.0}/pyproject.toml

@@ -11,13 +11,13 @@ packages = ["pgdevkit"]
 
 [project]
 name = "pgdevkit"
-version = "0.12.0"
+version = "0.13.0"
 description = "A helper for developing with Postgres"
 readme = "README.md"
 requires-python = ">=3.14"
 dependencies = [
     "docker>=7.1.0",
-    "psycopg[binary]>=3.2.0",
+    "psycopg[binary]>=3.3.0",
     "sqlglot>=30.11.0",
 ]
 
@@ -45,7 +45,7 @@ dev = [
     "ty>=0.0.59",
 ]
 test = [
-    "psycopg[binary]>=3.2.0",
+    "psycopg[binary]>=3.3.0",
     "pytest>=9.1.0",
     "pytest-asyncio>=0.24.0",
     "pytest-env>=1.1.0",

{pgdevkit-0.12.0 → pgdevkit-0.13.0}/skills/pgdevkit/SKILL.md

@@ -3,8 +3,8 @@ name: pgdevkit
 plugin: coding
 description: >
   Use pgdevkit for any PostgreSQL work in a Python project: a local
-  Docker/Podman test database (`pgdb testdb`), importable ORM-free CRUD
-  helpers (`pgdevkit.db`), and the `database/`-folder schema-as-code
+  Docker/Podman test database (`pgdb testdb`), importable ORM-free CRUD and
+  query helpers (`pgdevkit.db`: `pg_*`, `fetch_all`, `PgPool`, `SqlLoader`), 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
@@ -33,7 +33,7 @@ Core rules for every piece of database code in a project using pgdevkit:
 - **Inline SQL** — trivial queries of **4 lines or fewer** may be written inline in Python. Anything with JOINs, subqueries, CTEs, aggregations, or multiple conditions lives in its own `.sql` file.
 - **Named parameters** — always `%(name)s` style, never positional `%s`.
 - **The `database/` folder is the source of truth for the schema** — see [docs/database-layout.md](../../docs/database-layout.md) for the layer/object-type/file-naming conventions.
-- **Result mapping** — every query result maps to a Pydantic model; table-mapped models extend `pgdevkit.db.PostgresTableModel`.
+- **Result mapping** — every query result maps to a Pydantic model (or is deliberately returned as dicts / through `row_mapper`); table-mapped models extend `pgdevkit.db.PostgresTableModel`.
 
 ---
 
@@ -113,12 +113,13 @@ SQL files live under `db/queries/<topic>/`. Every custom query gets its own file
 
 ```python
 # db/connection.py
-from pgdevkit.db import PgPool
+from pgdevkit.db import PgPool, set_default_pool
 
 pool = PgPool(env_prefix="APP_POSTGRES_")  # matches ensure_testdb()'s {env_prefix}POSTGRES_* vars
 
 async def startup():
     await pool.open()
+    set_default_pool(pool)  # lets `fetch_all()` borrow from it when no `con=`/`pool=` is given
 ```
 
 ### Models
@@ -170,7 +171,7 @@ from pgdevkit.db import pg_retrieve, pg_insert, pg_upsert, pg_delete
 | `pg_insert_many` | Batch insert via `executemany` |
 | `pg_delete` / `pg_delete_dict` | Delete by PK, returns deleted row |
 
-Use these for simple CRUD. For custom `WHERE` clauses, joins, aggregations, or ordering, write a dedicated `.sql` file and a repository method.
+Use these for simple CRUD. For custom `WHERE` clauses, joins, aggregations, or ordering, write a dedicated `.sql` file and run it with `fetch_all` (below) in a repository method.
 
 ### Loading `.sql` files
 
@@ -184,8 +185,7 @@ sql = SqlLoader(Path(__file__).parent / "queries")
 
 ```python
 # db/repositories/user_repository.py
-from psycopg.rows import dict_row
-from pgdevkit.db import pg_retrieve, pg_delete
+from pgdevkit.db import fetch_all, pg_retrieve, pg_delete
 from db.connection import pool
 from db.loader import sql
 from models.user_models import UserRow, UserSummary
@@ -196,11 +196,7 @@ class UserRepository:
             return await pg_retrieve(conn, UserRow, {"id": user_id})
 
     async def list_active(self, limit: int = 100) -> list[UserSummary]:
-        async with pool.connection() as conn:
-            async with conn.cursor(row_factory=dict_row) as cur:
-                await cur.execute(sql.load_sql("users", "list_active_users"), {"limit": limit})
-                rows = await cur.fetchall()
-        return [UserSummary.model_validate(r) for r in rows]
+        return await fetch_all(sql.load_sql("users", "list_active_users"), {"limit": limit}, model=UserSummary)
 
     async def delete(self, user: UserRow) -> UserRow | None:
         async with pool.connection() as conn:
@@ -209,6 +205,24 @@ class UserRepository:
 
 Named parameters, `%(name)s` style, dict argument — never positional `%s`, never f-strings or `str.format()` for SQL text.
 
+### `fetch_all` — custom queries without cursor boilerplate
+
+```python
+from pgdevkit.db import fetch_all, set_default_pool
+
+set_default_pool(pool)  # once, at startup -- the connection source when none is passed
+
+users = await fetch_all(sql.load_sql("users", "list_active_users"), {"limit": 10}, model=UserSummary)  # list[UserSummary]
+rows = await fetch_all(sql.load_sql("users", "list_active_users"), {"limit": 10})                      # list[dict]
+async with pool.connection() as conn:  # inside a transaction you already hold
+    users = await fetch_all(query, params, model=UserSummary, con=conn)
+```
+
+- `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.
+- `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).
+- `query` is typed `SqlQuery`: a literal string (`SqlLoader.load_sql()`), a sqlglot expression, 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, and write a literal `%` as `%%` when you also pass `params`. A t-string carries its own values, so don't also pass `params`.
+
 ### Dynamic SQL
 
 Avoid it whenever possible — a static `.sql` file is always clearer. See [`references/dynamic-sql.md`](references/dynamic-sql.md) for t-string templates (3.14+) and `psycopg.sql` (< 3.14) when column/table names genuinely vary at runtime.
@@ -316,7 +330,7 @@ Reports drift between the `database/` `.sql` files and the actual schema — tab
 - [ ] Simple CRUD uses `pgdevkit.db`'s `pg_*` helpers; custom queries use `.sql` files loaded via `SqlLoader`
 - [ ] Inline SQL only for trivial queries ≤ 4 lines; anything with JOINs/CTEs/aggregations/subqueries uses a `.sql` file
 - [ ] All parameters use `%(name)s` style with a dict argument
-- [ ] Results mapped to a Pydantic model; table-mapped models extend `PostgresTableModel`
+- [ ] Custom read queries use `fetch_all(...)` (`model=` where the shape is stable; pass `con=` inside a transaction); results mapped to a Pydantic model; table-mapped models extend `PostgresTableModel`
 - [ ] No `LATERAL JOIN` — use a CTE that groups/aggregates first, then joins it
 - [ ] `.<env>.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

{pgdevkit-0.12.0 → pgdevkit-0.13.0}/skills/pgdevkit/references/dynamic-sql.md

@@ -25,6 +25,9 @@ T-strings look like f-strings but are evaluated by psycopg — values are always
 await cur.execute(t"SELECT * FROM users WHERE id = {user_id}")
 ```
 
+`fetch_all` accepts t-strings and `psycopg.sql` composables too (don't also pass `params` with a t-string):
+`await fetch_all(t"SELECT {column:i} FROM users WHERE active = {active}", model=UserRow)`.
+
 **Dynamic identifier (`:i`):**
 
 ```python