diff --git a/README.md b/README.md index 5f4e9a1..099a959 100644 --- a/README.md +++ b/README.md @@ -405,22 +405,57 @@ Install with the `db` extra: `pip install pgdevkit[db]`. password — same host-based auto-detection as `pgdb compare`'s `--entra-user`. For Lakebase hosts, also set the `{env_prefix}DATABRICKS_WORKSPACE_HOST` and `{env_prefix}DATABRICKS_INSTANCE` - env vars. + env vars. Values (password, `dsn_params`) are quoted properly, so spaces, backslashes and quotes are fine. + `min_size`, `timeout` (how long to wait for a free connection before `PoolTimeout`), `configure` (async hook run on + every new connection, e.g. to set session timeouts or `application_name`) and `check` are passed through to the + psycopg pool; `max_size` defaults to 40 and `min_size` to 4 (or `max_size` if smaller). - **`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 + 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 `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). `query` must be a literal string, a sqlglot + semantics, so it never outlives the call (a borrowed `con` gets its previous setting back); see + [Statement timeout](#statement-timeout) for what it does to a timeout the connection already has. `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`, + `exp.Placeholder` + `params`). A literal `%` in a sqlglot expression (`LIKE 'a%'`, the modulo operator) works with + or without `params`: it is escaped for psycopg when `params` is given (plain strings and `psycopg.sql` composables + are passed on as is, so there you write `%%` yourself when you pass `params`). +- **`fetch_one`, `fetch_scalar`, `execute`** — siblings of `fetch_all` with the same `con`/`pool`/`cancel`/ + `statement_timeout` keywords and the same query types (including t-strings, which take no `params`), so there is no + need for hand-written cursor code. `fetch_one(query, params, *, model=None, row_mapper=None)` returns the first row + (dict, `model` or `row_mapper` result) or `None`, `fetch_scalar(query, params)` the first column of the first row or + `None` (`await fetch_scalar("SELECT name FROM widget WHERE id = %(id)s", {"id": 1})`), and `execute(query, params)` + the row count of a write without `RETURNING` (`-1` for statements without one, like DDL). Without `con`, `execute` + runs on a pooled connection that commits when it returns (true for `PgPool`/`psycopg_pool`; a custom connection + source decides for itself); on your `con` nothing is committed. +- **`readonly_transaction(con)`** — async context manager that runs its block in a read-only transaction: + `async with pool.connection() as con, readonly_transaction(con): await fetch_all(sql, con=con)`. + Postgres rejects writes with `ReadOnlySqlTransaction`; the transaction is committed when the block ends normally + and rolled back on error, and `read_only` is **always** put back afterwards (normally to `None`), so a connection + that goes back to a shared pool is never left read-only. `con` must not be inside a transaction already (a fresh + pooled connection is fine). **It is a safety net, not a sandbox for untrusted SQL:** text without bound parameters + uses the simple query protocol, where `COMMIT; INSERT ...` ends the read-only transaction, TEMP tables stay + writable and `SET` statements stick to the connection. For LLM-generated SQL also connect as a role that only has + `SELECT` privileges, and consider `statement_timeout=at_most(...)`. +- **`ComplexHelper`** — handles composite, enum and JSONB columns for the CRUD functions: pass + `complex_helper=ComplexHelper(con)` (optionally with `normalizers`), or `complex_helper="auto"` for a default one, + to any `pg_*` function that takes the argument. What it learns from the catalog (types, column types, adapter + registrations) is cached per connection, so a helper created per call queries the catalog only the first time a + connection is used (the SELECT/INSERT itself is then the only query). The cache cannot see later DDL, nor DDL that + was rolled back after it had been cached: call `helper.clear_cache()` after `CREATE TYPE`/`ALTER TYPE`/column + changes or rollbacks of those (typical for tests that create types inside a rolled-back transaction on a reused + connection). `ComplexHelper(con, cache=False)` gives the helper a private cache and the pre-cache behaviour. +- **CRUD functions** — `pg_retrieve`, `pg_retrieve_many` (equality `filters`, plus optional `where=` as a t-string or + a `psycopg.sql` composable with `params=`, `order_by=` as a column name, a *list* of names or + `(name, "asc"|"desc")` pairs, or a composable, and `limit=`; prefer a t-string, since in a composable a literal `%` + must be written `%%` whenever filters or params are present), `pg_insert`, `pg_insert_many`, `pg_update`, `pg_update_dict`, `pg_upsert`, `pg_upsert_dict`, `pg_upsert_many`, `pg_upsert_many_dict`, `pg_delete`, `pg_delete_dict` — typed (`TableModel`-based) or dict-based CRUD against a @@ -519,7 +554,9 @@ 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 +help, see below). With `PgPool` (40 connections by default; `max_size=`) that means a few slow readers, or many parallel streams of +big results, can use up the pool and make every other request wait for a connection. 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 @@ -527,9 +564,15 @@ help, see below). Enforce a write timeout in the reverse proxy and size the pool `statement_timeout=` (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 +the client gets a `504 {"error": "Query timed out"}` (the key is `error_key`, see Errors); 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. +A plain value *replaces* whatever timeout the connection already has, in both directions: with a database-level +safety net of 600 s, a per-call `statement_timeout=120` shortens it for that call, but `statement_timeout=900` loosens +it. To make a per-call value only ever tighten it, wrap it: `statement_timeout=at_most(120)` +(`from pgdevkit.db import at_most`) uses `min(current, 120)`, where a current timeout of 0 (none) counts as unlimited. +`at_most` works anywhere `statement_timeout` does, `PostgresJsonResponse` included. + 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. @@ -545,12 +588,27 @@ class AppJsonResponse(PostgresJsonResponse): expose_errors = IS_DEV ``` +FastAPI's own error handlers answer with `{"detail": ...}` instead. If a frontend already handles that shape, set +`error_key = "detail"` in the subclass and these responses use it too. + 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 +Timestamps: Postgres' `row_to_json` writes a `timestamptz` as `2026-01-01T10:00:00+00:00` (with the session time +zone's offset, e.g. `+01:00`), where Pydantic/FastAPI serialize UTC values as `...Z`. Both are valid ISO-8601 and +parse to the same instant, but code that compares or displays the raw strings sees a difference; normalize in the +frontend, or in the query (`to_char(...)`). + +`NaN`/`Infinity`: JSON has no such numbers, so `row_to_json` writes a `float`/`numeric` NaN or Infinity as the +*string* `"NaN"`/`"Infinity"`. A frontend expecting numbers must handle those strings, or the query should map them +(`nullif(x, 'NaN')`). + +Behind Starlette's `BaseHTTPMiddleware` a failure after the first byte has been reported to end as a clean `200` with a +truncated body: the missing closing `]` is then the only signal (clients fail to parse it), the status cannot carry it. + 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. diff --git a/pgdevkit/db/__init__.py b/pgdevkit/db/__init__.py index dcddef4..4c0a1e8 100644 --- a/pgdevkit/db/__init__.py +++ b/pgdevkit/db/__init__.py @@ -16,19 +16,35 @@ pg_upsert_many, pg_upsert_many_dict, ) -from .fetch import ConnectionSource, SqlQuery, fetch_all, set_default_pool +from .fetch import ( + ConnectionSource, + PoolLike, + SqlQuery, + at_most, + execute, + fetch_all, + fetch_one, + fetch_scalar, + set_default_pool, +) from .loader import SqlLoader from .model import PostgresTableModel, TableModel +from .transaction import readonly_transaction __all__ = [ "ComplexHelper", "ConnectionSource", "PgPool", + "PoolLike", "PostgresTableModel", "SqlLoader", "SqlQuery", "TableModel", + "at_most", + "execute", "fetch_all", + "fetch_one", + "fetch_scalar", "pg_delete", "pg_delete_dict", "pg_insert", @@ -41,5 +57,6 @@ "pg_upsert_dict", "pg_upsert_many", "pg_upsert_many_dict", + "readonly_transaction", "set_default_pool", ] diff --git a/pgdevkit/db/complex_types.py b/pgdevkit/db/complex_types.py index e9bef41..d0e8a1f 100644 --- a/pgdevkit/db/complex_types.py +++ b/pgdevkit/db/complex_types.py @@ -1,5 +1,7 @@ from __future__ import annotations +import weakref +from dataclasses import dataclass, field from typing import Any, Callable from psycopg import AsyncConnection @@ -23,6 +25,33 @@ class JsonbArray: ComplexTypeInfo = CompositeInfo | EnumInfo | type[Jsonb] | type[JsonbArray] | None +@dataclass +class _TypeCache: + """What a `ComplexHelper` has learned about one connection's database: the catalog's complex types, the + CompositeInfo/EnumInfo objects built from them (they carry that database's OIDs), which of those are registered + on the connection (an adapter registration lives and dies with the connection), and each table column's type.""" + + system_complex_type_dict: dict[Any, tuple[str, str]] | None = None + complex_types: dict[tuple[str, str], CompositeInfo | EnumInfo] = field(default_factory=dict) + registered: set[CompositeInfo | EnumInfo] = field(default_factory=set) + columns: dict[tuple[tuple[str, str], bool], dict[str, ComplexTypeInfo]] = field(default_factory=dict) + column: dict[tuple[tuple[str, str], str], ComplexTypeInfo] = field(default_factory=dict) + shared: bool = False # whether other helpers on the same connection use it too (then column lookups are cached) + + +# One cache per live connection, so helpers created per call (`pg_retrieve(..., complex_helper=ComplexHelper(con))`) +# share it, and it goes away with the connection: pooled connections keep theirs for as long as the pool keeps them. +# Per connection rather than per process on purpose: the cached OIDs belong to one database. +_caches: weakref.WeakKeyDictionary[AsyncConnection, _TypeCache] = weakref.WeakKeyDictionary() + + +def _cache_for(con: AsyncConnection) -> _TypeCache: + try: + return _caches.setdefault(con, _TypeCache(shared=True)) + except TypeError: # not something a weak reference can point to (e.g. a stand-in for a connection): no sharing + return _TypeCache() + + class ComplexHelper: """psycopg adapter for PostgreSQL composite types, enums, and JSONB. @@ -35,24 +64,62 @@ class ComplexHelper: composite type that needs locale-key backfilling before it's built) — this is intentionally the only project-specific extension point; nothing else about a project's types is hardcoded here. + + What it learns from the catalog (the types, the columns' types, which ones are registered on the connection) is + cached per connection, for as long as that connection lives: a helper created for every call, on a connection + that is reused (e.g. from a pool), only queries the catalog the first time. The cache can't see schema changes + made afterwards (`CREATE TYPE`, `ALTER TYPE ... ADD VALUE`, altered columns), nor DDL that was rolled back after + it had been cached (a type's OID is then gone): call `clear_cache()` after such changes. `cache=False` gives the + helper a private cache instead, which behaves as before the shared cache existed (types are remembered per helper + instance, column lookups always ask the database). """ def __init__( self, con: AsyncConnection, normalizers: dict[str, Callable[[dict], dict]] | None = None, + *, + cache: bool = True, ) -> None: self.con = con - self.system_complex_type_dict: dict[Any, tuple[str, str]] | None = None - # Instance-scoped: a CompositeInfo/EnumInfo carries OIDs from a - # specific connection/database, so caching it on the class (shared - # across every connection) would leak stale OIDs across databases - # that happen to reuse the same type name — exactly the case for - # pgdevkit's per-worktree isolated test databases. - self.complex_types: dict[tuple[str, str], CompositeInfo | EnumInfo] = {} - self.registered: set[CompositeInfo | EnumInfo] = set() + # Never class-scoped: a CompositeInfo/EnumInfo carries OIDs from a specific connection/database, so + # sharing it across every connection would leak stale OIDs across databases that happen to reuse the + # same type name — exactly the case for pgdevkit's per-worktree isolated test databases. + self._cache = _cache_for(con) if cache else _TypeCache() self._normalizers = normalizers or {} + @property + def system_complex_type_dict(self) -> dict[Any, tuple[str, str]] | None: + return self._cache.system_complex_type_dict + + @system_complex_type_dict.setter + def system_complex_type_dict(self, value: dict[Any, tuple[str, str]] | None) -> None: + self._cache.system_complex_type_dict = value + + @property + def complex_types(self) -> dict[tuple[str, str], CompositeInfo | EnumInfo]: + return self._cache.complex_types + + @complex_types.setter + def complex_types(self, value: dict[tuple[str, str], CompositeInfo | EnumInfo]) -> None: + self._cache.complex_types = value + + @property + def registered(self) -> set[CompositeInfo | EnumInfo]: + return self._cache.registered + + @registered.setter + def registered(self, value: set[CompositeInfo | EnumInfo]) -> None: + self._cache.registered = value + + def clear_cache(self) -> None: + """Forget what was learned about this connection's database (see the class docstring).""" + self._cache.system_complex_type_dict = None + self._cache.complex_types.clear() + self._cache.registered.clear() + self._cache.columns.clear() + self._cache.column.clear() + async def load_complex_type_dict(self) -> None: async with self.con.cursor(row_factory=dict_row) as cur: await cur.execute(""" @@ -94,6 +161,9 @@ async def _load_complex_type_from_colinfos(self, res: dict[str, Any] | None) -> async def load_all_complex_types( self, table_name: tuple[str, str], include_generated: bool = False ) -> dict[str, ComplexTypeInfo]: + cache_key = (table_name, include_generated) + if self._cache.shared and cache_key in self._cache.columns: + return dict(self._cache.columns[cache_key]) if self.system_complex_type_dict is None: await self.load_complex_type_dict() colquery = """ @@ -120,9 +190,15 @@ async def load_all_complex_types( }, ) res = await cur.fetchall() - return {r["column_name"]: await self._load_complex_type_from_colinfos(r) for r in res} + types = {r["column_name"]: await self._load_complex_type_from_colinfos(r) for r in res} + if types and self._cache.shared: # (an empty result may just be a table that doesn't exist yet) + self._cache.columns[cache_key] = types + return dict(types) async def load_complex_type(self, table_name: tuple[str, str], col_name: str) -> ComplexTypeInfo: + cache_key = (table_name, col_name) + if self._cache.shared and cache_key in self._cache.column: + return self._cache.column[cache_key] if self.system_complex_type_dict is None: await self.load_complex_type_dict() colquery = """ @@ -146,7 +222,10 @@ async def load_complex_type(self, table_name: tuple[str, str], col_name: str) -> {"schema": table_name[0], "tbl": table_name[1], "col": col_name}, ) res = await cur.fetchone() - return await self._load_complex_type_from_colinfos(res) + info = await self._load_complex_type_from_colinfos(res) + if res and self._cache.shared: # (a missing column may just not exist yet) + self._cache.column[cache_key] = info + return info async def _get_complex_type(self, name: str, is_enum: bool, con: AsyncConnection) -> CompositeInfo | EnumInfo: if name.endswith("[]"): @@ -154,13 +233,9 @@ async def _get_complex_type(self, name: str, is_enum: bool, con: AsyncConnection schema, type_name = name.split(".") if "." in name else ("public", name) if type_name.startswith("_"): # the array type in PostgreSQL starts with an underscore type_name = type_name[1:] - if is_enum: - ci = await EnumInfo.fetch(con, Identifier(schema, type_name)) - assert ci is not None, f"Enum type {name} not found in database" - self.complex_types[(schema, type_name)] = ci if (schema, type_name) not in self.complex_types: - ci = await CompositeInfo.fetch(con, Identifier(schema, type_name)) - assert ci is not None, f"Complex type {name} not found in database" + ci = await (EnumInfo if is_enum else CompositeInfo).fetch(con, Identifier(schema, type_name)) + assert ci is not None, f"{'Enum' if is_enum else 'Complex'} type {name} not found in database" self.complex_types[(schema, type_name)] = ci return self.complex_types[(schema, type_name)] diff --git a/pgdevkit/db/connection.py b/pgdevkit/db/connection.py index 533c88a..2891722 100644 --- a/pgdevkit/db/connection.py +++ b/pgdevkit/db/connection.py @@ -2,8 +2,11 @@ import asyncio import os -from typing import Literal +from collections.abc import Awaitable, Callable +from typing import Any, Literal +from psycopg import AsyncConnection +from psycopg.conninfo import make_conninfo from psycopg_pool import AsyncConnectionPool, AsyncNullConnectionPool from ..connection import detect_provider, get_azure_postgres_password, is_azure_postgres_host @@ -26,6 +29,12 @@ class PgPool: auto-includes `prepare_threshold=None` too, since PgBouncer transaction mode doesn't support server-side prepared statements. Pass an explicit `use_null_pool`/`connection_kwargs` to override either. + + `min_size`, `timeout` (seconds to wait for a free connection before the pool raises `PoolTimeout`), + `configure` (an async callable run on every new connection, e.g. to set session timeouts) and `check` (the + health check run when a connection is handed out; the default is the pool class' `check_connection`) are passed + through to the underlying psycopg pool when given (`min_size` defaults to 4, or `max_size` if that is smaller; a + null pool takes no `min_size`). """ def __init__( @@ -40,6 +49,10 @@ def __init__( use_null_pool: bool | Literal["auto"] = "auto", connection_kwargs: dict | None = None, dsn_params: dict[str, str] | None = None, + min_size: int | None = None, + timeout: float | None = None, + configure: Callable[[AsyncConnection], Awaitable[None]] | None = None, + check: Callable[[AsyncConnection], Awaitable[None]] | None = None, ) -> None: self._env_prefix = env_prefix self._max_size = max_size @@ -50,6 +63,12 @@ def __init__( self._use_null_pool = use_null_pool self._connection_kwargs = connection_kwargs self._dsn_params = dsn_params or {} + # only what was given: the psycopg pool's own defaults apply to the rest + self._pool_options: dict[str, Any] = { + k: v for k, v in {"timeout": timeout, "configure": configure}.items() if v is not None + } + self._min_size = min_size + self._check = check self._pool: AsyncConnectionPool | AsyncNullConnectionPool | None = None def _is_azure_postgres(self, host: str) -> bool: @@ -77,10 +96,8 @@ async def _dsn(self) -> str: exclude_interactive_browser_credential=self._exclude_interactive_browser_credential, ) - dsn = f"host={host} port={port} dbname={dbname} user={user} password={password}" - for key, value in self._dsn_params.items(): - dsn += f" {key}={value}" - return dsn + # make_conninfo quotes values (a password with a space, backslash or quote would break a hand-built string) + return make_conninfo(host=host, port=port, dbname=dbname, user=user, password=password, **self._dsn_params) async def open(self) -> None: if self._pool is None: @@ -93,13 +110,17 @@ async def open(self) -> None: connection_kwargs = {"prepare_threshold": None} if is_azure_postgres else {} max_lifetime = 3600.0 if self._max_lifetime is None else self._max_lifetime if use_null_pool: + null_options: dict[str, Any] = dict(self._pool_options) + if self._min_size is not None: + null_options["min_size"] = self._min_size self._pool = AsyncNullConnectionPool( conninfo=self._dsn, open=False, max_size=self._max_size, max_lifetime=max_lifetime, - check=AsyncNullConnectionPool.check_connection, + check=self._check or AsyncNullConnectionPool.check_connection, kwargs=connection_kwargs, + **null_options, ) else: self._pool = AsyncConnectionPool( @@ -107,8 +128,11 @@ async def open(self) -> None: open=False, max_size=self._max_size, max_lifetime=max_lifetime, - check=AsyncConnectionPool.check_connection, + check=self._check or AsyncConnectionPool.check_connection, kwargs=connection_kwargs, + # (psycopg's default min_size of 4 would reject a max_size below 4) + min_size=min(4, self._max_size) if self._min_size is None else self._min_size, + **self._pool_options, ) if not self._pool._opened: await self._pool.open() diff --git a/pgdevkit/db/crud.py b/pgdevkit/db/crud.py index 56728f4..b957d7c 100644 --- a/pgdevkit/db/crud.py +++ b/pgdevkit/db/crud.py @@ -1,10 +1,12 @@ from __future__ import annotations -from typing import Any, Callable, Mapping, Optional, Sequence, Type, TypeVar +from string.templatelib import Interpolation, Template +from typing import Any, Callable, Literal, Mapping, Optional, Sequence, Type, TypeVar from psycopg.connection_async import AsyncConnection from psycopg.rows import dict_row -from psycopg.sql import SQL, Composable, Identifier, Placeholder +from psycopg.sql import Literal as SqlLiteral +from psycopg.sql import SQL, Composable, Composed, Identifier, Placeholder from pydantic import BaseModel from .complex_types import ComplexHelper @@ -12,16 +14,25 @@ T = TypeVar("T", bound=PostgresTableModel) +# `complex_helper=` of every function below: a `ComplexHelper` you configured (e.g. with `normalizers`), `"auto"` for +# a default one on the connection (what it learns is cached per connection), or None for plain psycopg behaviour. +type ComplexHelperArg = ComplexHelper | Literal["auto"] | None + + +def _resolve_helper(con: AsyncConnection, complex_helper: ComplexHelperArg) -> ComplexHelper | None: + return ComplexHelper(con) if complex_helper == "auto" else complex_helper + async def _select_list( - con: AsyncConnection, table_name: tuple[str, str], complex_helper: ComplexHelper | None + con: AsyncConnection, table_name: tuple[str, str], complex_helper: ComplexHelperArg ) -> Composable: """Column list for a SELECT, wrapping composite/enum/JSONB columns in `to_jsonb(...)` so psycopg gets back plain Python values. Falls back to `SELECT *` (no extra query) when no ComplexHelper is given.""" - if complex_helper is None: + helper = _resolve_helper(con, complex_helper) + if helper is None: return SQL("*") - complex_types = await complex_helper.load_all_complex_types(table_name, include_generated=True) + complex_types = await helper.load_all_complex_types(table_name, include_generated=True) if not complex_types: return SQL("*") parts = [ @@ -35,21 +46,22 @@ async def _convert_complex_values( con: AsyncConnection, table_name: tuple[str, str], data: dict, - complex_helper: ComplexHelper | None, + complex_helper: ComplexHelperArg, ) -> dict: """Convert dict/list values destined for composite/enum columns into the psycopg-registered types those columns need. A no-op (no extra query) unless `data` actually contains dict/list values.""" - if complex_helper is None: + helper = _resolve_helper(con, complex_helper) + if helper is None: return data candidate_keys = [k for k, v in data.items() if isinstance(v, (dict, list))] if not candidate_keys: return data converted = dict(data) for k in candidate_keys: - info = await complex_helper.load_complex_type(table_name, k) + info = await helper.load_complex_type(table_name, k) if info is not None: - converted[k] = await complex_helper.recursive_convert(data[k], info, con) + converted[k] = await helper.recursive_convert(data[k], info, con) return converted @@ -57,14 +69,15 @@ async def _convert_complex_values_many( con: AsyncConnection, table_name: tuple[str, str], rows: Sequence[dict], - complex_helper: ComplexHelper | None, + complex_helper: ComplexHelperArg, ) -> Sequence[dict]: - if complex_helper is None or not rows: + helper = _resolve_helper(con, complex_helper) + if helper is None or not rows: return rows candidate_keys = {k for row in rows for k, v in row.items() if isinstance(v, (dict, list))} if not candidate_keys: return rows - infos = {k: await complex_helper.load_complex_type(table_name, k) for k in candidate_keys} + infos = {k: await helper.load_complex_type(table_name, k) for k in candidate_keys} complex_keys = {k for k, info in infos.items() if info is not None} if not complex_keys: return rows @@ -73,7 +86,7 @@ async def _convert_complex_values_many( new_row = dict(row) for k in complex_keys: if isinstance(new_row.get(k), (dict, list)): - new_row[k] = await complex_helper.recursive_convert(new_row[k], infos[k], con) + new_row[k] = await helper.recursive_convert(new_row[k], infos[k], con) converted_rows.append(new_row) return converted_rows @@ -83,13 +96,14 @@ async def pg_retrieve( data_type: Type[T], pks: dict, *, - complex_helper: ComplexHelper | None = None, + complex_helper: ComplexHelperArg = None, ) -> T | None: """Fetch a single row by primary key(s). Pass `complex_helper` (a `ComplexHelper`, optionally configured with - `normalizers`) when the table has composite/enum columns; omitted, this - behaves exactly like a plain `SELECT *`.""" + `normalizers`, or `"auto"` for a default one) when the table has + composite/enum columns; omitted, this behaves exactly like a plain + `SELECT *`.""" async with con.cursor(row_factory=dict_row) as cur: table_name = data_type.get_table_name() select_cols = await _select_list(con, table_name, complex_helper) @@ -103,29 +117,101 @@ async def pg_retrieve( return data_type(**row) if row else None +type OrderBy = str | Sequence[str | tuple[str, Literal["asc", "desc"]]] | Composable + + +def _order_by_sql(order_by: OrderBy | None) -> Composable | None: + """` ORDER BY ...` for `pg_retrieve_many(order_by=...)`: column names (quoted as identifiers, optionally + as `(name, "desc")`), or a `psycopg.sql` composable for anything else (an expression, `NULLS LAST`, ...).""" + if order_by is None: + return None + if isinstance(order_by, Composable): + return SQL(" ORDER BY {}").format(order_by) + if isinstance(order_by, tuple) and len(order_by) == 2 and str(order_by[1]).lower() in ("asc", "desc"): + raise ValueError(f"order_by={order_by!r} would order by two columns; use a list: [{order_by!r}]") + items = [order_by] if isinstance(order_by, str) else list(order_by) + parts: list[Composable] = [] + for item in items: + column, direction = (item, "asc") if isinstance(item, str) else item + if direction.lower() not in ("asc", "desc"): + raise ValueError(f"order_by direction must be 'asc' or 'desc', got {direction!r}") + parts.append(SQL("{} DESC" if direction.lower() == "desc" else "{} ASC").format(Identifier(column))) + return SQL(" ORDER BY {}").format(SQL(", ").join(parts)) if parts else None + + async def pg_retrieve_many( con: AsyncConnection, data_type: Type[T], filters: dict, *, from_dict: Optional[Callable[[Mapping], T]] = None, - complex_helper: ComplexHelper | None = None, + complex_helper: ComplexHelperArg = None, + where: Composable | Template | None = None, + params: Mapping[str, Any] | None = None, + order_by: OrderBy | None = None, + limit: int | None = None, ) -> Sequence[T]: - """Fetch multiple rows matching all filter key=value pairs.""" + """Fetch multiple rows matching all filter key=value pairs. + + Beyond those equality `filters` (`{}` for none), the query can be narrowed and shaped with: + + - `where=`: more conditions, ANDed with the filters -- a t-string (`t"price > {minimum} AND name LIKE {pattern}"`, + values are bound for you; prefer it), or a `psycopg.sql` composable, whose own `sql.Placeholder("name")`s are + bound from `params=`. In a composable a literal `%` must be written `%%` whenever `filters` or `params` are + non-empty (psycopg's rule); bind patterns as values instead (`name LIKE {pattern}`). + - `order_by=`: a column name, a *list* of them (each a name or `(name, "asc" | "desc")`, quoted as + identifiers; a bare `("name", "desc")` is rejected as ambiguous), or a composable for anything fancier. + - `limit=`: at most that many rows. + + Without those it behaves exactly as before. For joins or other shapes use `fetch_all` instead.""" + # validate before the select list may cost a catalog round trip + tail: list[Composable] = [] + if (order := _order_by_sql(order_by)) is not None: + tail.append(order) + if limit is not None: + if isinstance(limit, bool) or not isinstance(limit, int) or limit < 0: + raise ValueError(f"limit must be a non-negative integer, got {limit!r}") + tail.append(SQL(" LIMIT {}").format(SqlLiteral(limit))) + if isinstance(where, Template) and params: + raise TypeError("A t-string `where` carries its own values; don't pass `params` with it.") + table_name = data_type.get_table_name() + select_cols = await _select_list(con, table_name, complex_helper) + query: Template | Composed + query_params: Mapping[str, Any] | None + if isinstance(where, Template): + parts: list[str | Interpolation] = [ + "SELECT ", + Interpolation(select_cols, "cols", None, "q"), + " FROM ", + Interpolation(Identifier(*table_name), "tbl", None, "i"), + " WHERE ", + ] + for k, v in filters.items(): + parts += [Interpolation(Identifier(k), k, None, "i"), " = ", Interpolation(v, k, None, ""), " AND "] + parts += ["(", Interpolation(where, "where", None, "q"), ")"] + parts += [Interpolation(t, "tail", None, "q") for t in tail] + query, query_params = Template(*parts), None + else: + conditions: list[Composable] = [ + SQL("{col} = {val}").format(col=Identifier(k), val=Placeholder(k)) for k in filters + ] + if where is not None: + conditions.append(SQL("({})").format(where)) + query_params = {**filters} + for name, value in (params or {}).items(): + if name in query_params: + raise ValueError(f"`params` key {name!r} collides with a filter of the same name") + query_params[name] = value + query = SQL("SELECT {cols} FROM {tbl}").format(cols=select_cols, tbl=Identifier(*table_name)) + if conditions: + query = SQL("{} WHERE {}").format(query, SQL(" AND ").join(conditions)) + query = Composed([query, *tail]) + query_params = query_params or None # an empty dict would make psycopg parse `%` for nothing async with con.cursor(row_factory=dict_row) as cur: - table_name = data_type.get_table_name() - select_cols = await _select_list(con, table_name, complex_helper) - if filters: - query = SQL("SELECT {cols} FROM {tbl} WHERE {where}").format( - cols=select_cols, - tbl=Identifier(*table_name), - where=SQL(" AND ").join( - SQL("{col} = {val}").format(col=Identifier(k), val=Placeholder(k)) for k in filters - ), - ) + if isinstance(query, Template): + await cur.execute(query) else: - query = SQL("SELECT {cols} FROM {tbl}").format(cols=select_cols, tbl=Identifier(*table_name)) - await cur.execute(query, filters) + await cur.execute(query, query_params) rows = await cur.fetchall() fn = from_dict or (lambda d: data_type(**d)) return [fn(r) for r in rows] @@ -136,7 +222,7 @@ async def pg_insert( table_name: tuple[str, str], data: dict, *, - complex_helper: ComplexHelper | None = None, + complex_helper: ComplexHelperArg = None, ) -> dict[str, Any]: """Insert one row and return the full row (RETURNING *).""" data = await _convert_complex_values(con, table_name, data, complex_helper) @@ -158,7 +244,7 @@ async def pg_update_dict( data: dict, primary_keys: Sequence[str], *, - complex_helper: ComplexHelper | None = None, + complex_helper: ComplexHelperArg = None, ) -> Any | None: """Update a row identified by primary_keys. Returns the raw row tuple. @@ -181,7 +267,7 @@ async def pg_update_dict( async def pg_update( - con: AsyncConnection, data: T, data_type: type[T], *, complex_helper: ComplexHelper | None = None + con: AsyncConnection, data: T, data_type: type[T], *, complex_helper: ComplexHelperArg = None ) -> Any | None: """Update a typed model instance.""" return await pg_update_dict( @@ -195,7 +281,7 @@ async def pg_upsert_dict( data: dict, primary_keys: Sequence[str], *, - complex_helper: ComplexHelper | None = None, + complex_helper: ComplexHelperArg = None, ) -> dict: """INSERT ... ON CONFLICT ... DO UPDATE, returns the row as a dict.""" data = await _convert_complex_values(con, table_name, data, complex_helper) @@ -218,7 +304,7 @@ async def pg_upsert_dict( async def pg_upsert( - con: AsyncConnection, data: T, data_type: type[T], *, complex_helper: ComplexHelper | None = None + con: AsyncConnection, data: T, data_type: type[T], *, complex_helper: ComplexHelperArg = None ) -> dict: """Upsert a typed model instance.""" return await pg_upsert_dict( @@ -233,7 +319,7 @@ async def pg_upsert_many_dict( primary_keys: Sequence[str], *, must_exist: bool = False, - complex_helper: ComplexHelper | None = None, + complex_helper: ComplexHelperArg = None, ) -> None: """Batch upsert — one round-trip via executemany. @@ -270,7 +356,7 @@ async def pg_upsert_many_dict( async def pg_upsert_many( - con: AsyncConnection, data: Sequence[T], data_type: type[T], *, complex_helper: ComplexHelper | None = None + con: AsyncConnection, data: Sequence[T], data_type: type[T], *, complex_helper: ComplexHelperArg = None ) -> None: await pg_upsert_many_dict( con, @@ -286,7 +372,7 @@ async def pg_insert_many( table_name: tuple[str, str], data: Sequence[dict | BaseModel], *, - complex_helper: ComplexHelper | None = None, + complex_helper: ComplexHelperArg = None, ) -> None: """Batch insert — no RETURNING, one round-trip via executemany.""" if not data: diff --git a/pgdevkit/db/fetch.py b/pgdevkit/db/fetch.py index d554fc6..b3b94e8 100644 --- a/pgdevkit/db/fetch.py +++ b/pgdevkit/db/fetch.py @@ -4,14 +4,15 @@ import logging import math import re +import uuid 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 typing import TYPE_CHECKING, Any, Literal, LiteralString, Protocol, cast, overload from psycopg.connection_async import AsyncConnection from psycopg.errors import QueryCanceled -from psycopg.rows import dict_row +from psycopg.rows import dict_row, tuple_row from psycopg.sql import Composable from pydantic import BaseModel @@ -37,17 +38,30 @@ class ConnectionSource(Protocol): def connection(self) -> AbstractAsyncContextManager[AsyncConnection]: ... -_default_source: ConnectionSource | None = None +# What `set_default_pool()` takes: a `ConnectionSource` (has `.connection()`, like `PgPool`), or a plain +# zero-argument callable returning an async connection context manager (e.g. a module-level pool's bound method). +type PoolLike = ConnectionSource | Callable[[], AbstractAsyncContextManager[AsyncConnection]] +_default_source: PoolLike | 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 + +def set_default_pool(pool: PoolLike | None) -> None: + """Register where the `fetch_*()`/`execute()` helpers get their connection when called without `con`/`pool` + (typically the app's `PgPool`, once at startup). Pass `None` to clear it. Besides a `ConnectionSource` (anything + with a `.connection()` method), a plain callable returning an async connection context manager is accepted, for + a pool that is resolved lazily or has another method name: `set_default_pool(lambda: get_pool().acquire())` + (only here: the `pool=` keyword of the helpers takes a `ConnectionSource`). This is process-wide state meant for single-database apps: with several databases/tenants pass `pool=` explicitly.""" global _default_source _default_source = pool +def _connection_cm(source: PoolLike) -> AbstractAsyncContextManager[AsyncConnection]: + if hasattr(source, "connection"): + return cast(ConnectionSource, source).connection() + return cast("Callable[[], AbstractAsyncContextManager[AsyncConnection]]", source)() + + @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 @@ -57,21 +71,25 @@ async def _acquire(con: AsyncConnection | None, pool: ConnectionSource | None) - raise TypeError("Pass either `con` or `pool`, not both.") yield con return - source = pool if pool is not None else _default_source + source: PoolLike | None = 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: + async with _connection_cm(source) as borrowed: yield borrowed _PLACEHOLDER_NAME = re.compile(r"\w+") -def _render(query: SqlQuery) -> Any: +def _render(query: SqlQuery, params: QueryParams = None, *, caller: str = "fetch_all") -> Any: + """Turn `query` into something psycopg executes. A sqlglot expression becomes postgres-dialect SQL; when + `params` is not None, psycopg also parses `%` in the text, so every literal `%` of the expression (in a + string literal such as `LIKE '50%'`, a quoted identifier, the modulo operator) is doubled -- but not the + `%(name)s`/`%s` placeholders sqlglot renders for `exp.Placeholder` nodes.""" if isinstance(query, (str, Composable, Template)): return query if not hasattr(query, "sql"): - raise TypeError(f"fetch_all: unsupported query type {type(query).__name__}") + raise TypeError(f"{caller}: 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 @@ -79,10 +97,30 @@ def _render(query: SqlQuery) -> Any: 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") + raise ValueError(f"{caller}: 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") + raise ValueError(f"{caller}: invalid placeholder name {node.name!r}") + if params is None: # psycopg doesn't look at `%` without params + return query.sql(dialect="postgres") + return _escape_percent(query, exp) + + +def _escape_percent(query: Any, exp: Any) -> str: + """Render a sqlglot expression with every literal `%` doubled, but not the `%(name)s` placeholders. The + placeholders are rendered as unique tokens first and put back afterwards: telling `%(id)s` apart from a + literal `'%(id)s'` in the rendered text would be guesswork.""" + nonce = uuid.uuid4().hex + placeholders: list[str] = [] + + def to_token(node: Any) -> Any: + if not isinstance(node, exp.Placeholder): + return node + placeholders.append(node.sql(dialect="postgres")) + return exp.Var(this=f"pgdevkit_ph_{nonce}_{len(placeholders) - 1}_") # (a transform works on a copy) + + rendered: str = query.transform(to_token).sql(dialect="postgres").replace("%", "%%") + # one pass, with a closing delimiter on the token: `_1_` must not match the start of `_10_` + return re.sub(rf"pgdevkit_ph_{nonce}_(\d+)_", lambda m: placeholders[int(m.group(1))], rendered) def _check_statement_timeout(statement_timeout: float | None) -> None: @@ -90,34 +128,65 @@ def _check_statement_timeout(statement_timeout: float | None) -> None: raise ValueError(f"statement_timeout must be a positive number of seconds, got {statement_timeout}") +class _AtMost(float): + """A `statement_timeout` that may only tighten the connection's current one (see `at_most`).""" + + +def at_most(seconds: float) -> float: + """Use as `statement_timeout=at_most(120)`: like a plain number of seconds, except that it never *loosens* the + timeout the connection already has. A plain `statement_timeout` replaces the current one in both directions + (a per-call 900 s overrides a database-level safety net of 600 s); with `at_most` the effective timeout is + `min(current, seconds)`, where a current setting of 0 (none) counts as unlimited, i.e. `seconds` applies.""" + return _AtMost(seconds) + + +_SET_TIMEOUT = "SELECT current_setting('statement_timeout'), set_config('statement_timeout', %(ms)s, true)" +_SET_TIMEOUT_AT_MOST = """ +SELECT current_setting('statement_timeout'), set_config('statement_timeout', ( + SELECT CASE WHEN setting::bigint > 0 THEN least(setting::bigint, %(ms)s::bigint) ELSE %(ms)s::bigint END + FROM pg_settings WHERE name = 'statement_timeout' +)::text, true)""" + + @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.""" + `None` leaves the connection's own setting alone; an `at_most()` value never raises it.""" if seconds is None: yield return + # `pg_settings.setting` of statement_timeout is in milliseconds, 0 = disabled + query = _SET_TIMEOUT_AT_MOST if isinstance(seconds, _AtMost) else _SET_TIMEOUT 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)))}, - ) + cur = await c.execute(query, {"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]}) +type _Mode = Literal["all", "one", "scalar", "count"] + + 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: + c: AsyncConnection, query: Any, params: QueryParams, statement_timeout: float | None, mode: _Mode +) -> Any: + row_factory = dict_row if mode in ("all", "one") else tuple_row + async with _statement_timeout(c, statement_timeout), c.cursor(row_factory=row_factory) as cur: await cur.execute(query, params) - return await cur.fetchall() + if mode == "all": + return await cur.fetchall() + if mode == "count": + return cur.rowcount + row = await cur.fetchone() + if mode == "scalar": + first = cast("tuple[Any, ...] | None", row) + return first[0] if first is not None else None + return row async def _abort(c: AsyncConnection, work: asyncio.Future[Any]) -> None: @@ -129,7 +198,7 @@ async def _abort(c: AsyncConnection, work: asyncio.Future[Any]) -> None: try: await c.cancel_safe() except Exception: - logger.warning("fetch_all: cancelling the query failed; closing the connection", exc_info=True) + logger.warning("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() @@ -153,25 +222,29 @@ async def _abort_uninterruptibly(c: AsyncConnection, work: asyncio.Future[Any]) raise asyncio.CancelledError -async def _fetch_dicts( +async def _run( 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]]: + mode: _Mode, + caller: str, +) -> Any: + """The one place that acquires a connection, runs the statement and wires up `cancel`/the task's own + cancellation; `mode` only says what is read from the cursor (see `_run_query`).""" if cancel is not None and cancel.is_set(): - raise QueryCanceled("fetch_all: cancelled before the query started") + raise QueryCanceled(f"{caller}: cancelled before the query started") async with _acquire(con, pool) as c: - work = asyncio.ensure_future(_run_query(c, query, params, statement_timeout)) + work = asyncio.ensure_future(_run_query(c, query, params, statement_timeout, mode)) 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`") + raise QueryCanceled(f"{caller}: 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 @@ -184,6 +257,33 @@ async def _fetch_dicts( stop.cancel() +def _prepare( + caller: str, + query: SqlQuery, + params: QueryParams, + model: type[BaseModel] | None, + row_mapper: Callable[[dict[str, Any]], Any] | None, + statement_timeout: float | None, +) -> Any: + """Validate the arguments shared by every helper and render the query.""" + 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) + return _render(query, params, caller=caller) + + +def _convert( + row: dict[str, Any], model: type[BaseModel] | None, row_mapper: Callable[[dict[str, Any]], Any] | None +) -> Any: + if model is not None: + return model.model_validate(row) + if row_mapper is not None: + return row_mapper(row) + return row + + @overload async def fetch_all[M: BaseModel]( query: SqlQueryNoTemplate, @@ -272,7 +372,8 @@ async def fetch_all( cancel: asyncio.Event | None = None, statement_timeout: float | None = None, ) -> list[Any]: - """Run one query and return every row. + """Run one query and return every row. (`fetch_one`, `fetch_scalar` and `execute` are the siblings for one row, + one value, and a write whose result you don't read; they take the same connection/cancel/timeout keywords.) 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. @@ -281,7 +382,8 @@ async def fetch_all( 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. + Without it, the connection's own `statement_timeout` (default: none) applies. A plain value replaces that + setting in both directions; `statement_timeout=at_most(seconds)` never loosens it. 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 @@ -299,17 +401,187 @@ async def fetch_all( 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 literal `%` in one (`LIKE '50%'`, the modulo operator) is fine with or without `params`; it is escaped for + psycopg when needed. (A plain string or `psycopg.sql` composable is passed on as is: with `params`, write + `%%` there.) 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 + rendered = _prepare("fetch_all", query, params, model, row_mapper, statement_timeout) + rows = await _run(rendered, params, con, pool, cancel, statement_timeout, "all", "fetch_all") + return [_convert(row, model, row_mapper) for row in rows] + + +@overload +async def fetch_one[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, +) -> M | None: ... + + +@overload +async def fetch_one[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, +) -> T | None: ... + + +@overload +async def fetch_one( + 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, +) -> dict[str, Any] | None: ... + + +@overload +async def fetch_one[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, +) -> M | None: ... + + +@overload +async def fetch_one[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, +) -> T | None: ... + + +@overload +async def fetch_one( + 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, +) -> dict[str, Any] | None: ... + + +async def fetch_one( + 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, +) -> Any | None: + """Run one query and return its first row, or `None` if it returned none (further rows are ignored: add + `LIMIT 1`/`ORDER BY` yourself). Same arguments and semantics as `fetch_all`: the row is a dict, validated into + `model`, or converted by `row_mapper`.""" + rendered = _prepare("fetch_one", query, params, model, row_mapper, statement_timeout) + row = await _run(rendered, params, con, pool, cancel, statement_timeout, "one", "fetch_one") + return None if row is None else _convert(row, model, row_mapper) + + +@overload +async def fetch_scalar( + query: SqlQueryNoTemplate, + params: QueryParams = None, + *, + con: AsyncConnection | None = None, + pool: ConnectionSource | None = None, + cancel: asyncio.Event | None = None, + statement_timeout: float | None = None, +) -> Any | None: ... + + +@overload +async def fetch_scalar( + query: Template, + *, + con: AsyncConnection | None = None, + pool: ConnectionSource | None = None, + cancel: asyncio.Event | None = None, + statement_timeout: float | None = None, +) -> Any | None: ... + + +@overload +async def execute( + query: SqlQueryNoTemplate, + params: QueryParams = None, + *, + con: AsyncConnection | None = None, + pool: ConnectionSource | None = None, + cancel: asyncio.Event | None = None, + statement_timeout: float | None = None, +) -> int: ... + + +@overload +async def execute( + query: Template, + *, + con: AsyncConnection | None = None, + pool: ConnectionSource | None = None, + cancel: asyncio.Event | None = None, + statement_timeout: float | None = None, +) -> int: ... + + +async def fetch_scalar( + query: SqlQuery, + params: QueryParams = None, + *, + con: AsyncConnection | None = None, + pool: ConnectionSource | None = None, + cancel: asyncio.Event | None = None, + statement_timeout: float | None = None, +) -> Any | None: + """Run one query and return the first column of its first row, or `None` if it returned no row (a NULL + value is `None` too: use `fetch_one` to tell them apart). Same arguments and semantics as `fetch_all`, + minus `model`/`row_mapper`. The statement must return a result set (e.g. `SELECT count(*) ...`, + or an `INSERT ... RETURNING id`).""" + rendered = _prepare("fetch_scalar", query, params, None, None, statement_timeout) + return await _run(rendered, params, con, pool, cancel, statement_timeout, "scalar", "fetch_scalar") + + +async def execute( + query: SqlQuery, + params: QueryParams = None, + *, + con: AsyncConnection | None = None, + pool: ConnectionSource | None = None, + cancel: asyncio.Event | None = None, + statement_timeout: float | None = None, +) -> int: + """Run one statement whose result you don't read (`INSERT`/`UPDATE`/`DELETE`/DDL without `RETURNING`) and + return its row count (`-1` for statements that have none, like DDL). Same arguments and semantics as + `fetch_all`. Without `con`, the statement runs on a pooled connection that commits when the call returns; on + a `con` you pass, nothing is committed: that stays your transaction's business.""" + rendered = _prepare("execute", query, params, None, None, statement_timeout) + return await _run(rendered, params, con, pool, cancel, statement_timeout, "count", "execute") diff --git a/pgdevkit/db/transaction.py b/pgdevkit/db/transaction.py new file mode 100644 index 0000000..4bffbee --- /dev/null +++ b/pgdevkit/db/transaction.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +from psycopg.connection_async import AsyncConnection + + +@asynccontextmanager +async def readonly_transaction(con: AsyncConnection) -> AsyncIterator[AsyncConnection]: + """Run a block in a read-only transaction on `con`: Postgres refuses writes with `ReadOnlySqlTransaction`. + + ```python + async with pool.connection() as con, readonly_transaction(con): + rows = await fetch_all(generated_sql, con=con) + ``` + + The transaction is committed when the block ends normally and rolled back if it raises. Afterwards `con`'s + `read_only` setting is *always* put back to what it was before (normally `None`, the server default), because + the connection usually goes back to a shared pool, where a leaked `read_only=True` would turn an unrelated + write into `ReadOnlySqlTransaction`. + + This is a safety net, **not a sandbox for untrusted SQL**. Text without bound parameters goes to Postgres + through the simple query protocol, where several statements are allowed: `COMMIT; INSERT ...` ends the + read-only transaction and writes. TEMP tables stay writable, and `SET`/`SET search_path` persist on the + connection. To run untrusted SQL, also connect as a role that only has SELECT privileges. + + `con` must not be inside a transaction already (psycopg raises `ProgrammingError` when `read_only` is changed + then): a connection fresh from `pool.connection()` is fine, one that ran a statement without + `commit()`/`rollback()` (non-autocommit) is not. + """ + previous = con.read_only + await con.set_read_only(True) + try: + async with con.transaction(): + yield con + finally: + if not (con.closed or con.broken): # a dead connection can't leak into a pool, and can't be asked + await con.set_read_only(previous) diff --git a/pgdevkit/fastapi/json_response.py b/pgdevkit/fastapi/json_response.py index 492c094..c0d86dc 100644 --- a/pgdevkit/fastapi/json_response.py +++ b/pgdevkit/fastapi/json_response.py @@ -58,18 +58,22 @@ class PostgresJsonResponse(StreamingResponse): ``statement_timeout`` (seconds) lets Postgres abort the query after that long, as in ``fetch_all``: before the first byte that is a 504 ``{"error": "Query timed out"}``, later the response is cut short like any other error. It limits query execution only: a client that stops reading is *not* freed by it (Postgres defers the - cancel while blocked writing), so keep a write timeout in the reverse proxy. Needs a connection that is not in - autocommit mode. + cancel while blocked writing), so keep a write timeout in the reverse proxy: a pooled connection stays checked + out until the client has read the whole body, so slow readers can exhaust the pool. Needs a connection that is not + in autocommit mode. ``statement_timeout=at_most(seconds)`` never loosens a timeout the connection already has. Errors before the first byte become a JSON ``{"error": ...}`` response (status 500, or the status of a raised - ``HTTPException``). After that the status line is gone: the error is logged and re-raised, and the response + ``HTTPException``); FastAPI's own handlers use ``{"detail": ...}``, so a subclass can set ``error_key = "detail"`` + to keep an existing frontend's error handling working. After that the status line is gone: the error is logged and re-raised, and the response ends without its closing ``]``, so clients can't mistake it for a complete array. Subclass and set ``expose_errors = True`` (e.g. in dev/test) to include ``str(error)`` in the 500 response. - Postgres sends rows in 8 kB buffers, so a slow query with small rows delivers its first bytes late. + Postgres sends rows in 8 kB buffers, so a slow query with small rows delivers its first bytes late. A + ``timestamptz`` is written as ``...+00:00`` (``row_to_json``), where FastAPI/Pydantic emit ``...Z``. """ expose_errors: bool = False + error_key: str = "error" # the key of the message in the JSON body of an error response def __init__( self, @@ -91,7 +95,7 @@ def __init__( _check_statement_timeout(statement_timeout) 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.") - rendered = _render(query) + rendered = _render(query, params) self.query = rendered if query_produces_json else _as_json_rows(rendered) self.params = params self.con = con @@ -169,7 +173,7 @@ async def _send_error(self, send: Send, err: Exception) -> None: status, message = err.status_code, err.detail else: status, message = 500, str(err) if self.expose_errors else "Internal Server Error" - body = json.dumps({"error": message}).encode() + body = json.dumps({self.error_key: message}).encode() headers = [(b"content-type", b"application/json"), (b"content-length", str(len(body)).encode())] if isinstance(err, HTTPException) and err.headers: headers += [(k.lower().encode(), v.encode()) for k, v in err.headers.items()] diff --git a/pyproject.toml b/pyproject.toml index fdf4eed..6c31d27 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ packages = ["pgdevkit"] [project] name = "pgdevkit" -version = "0.14.2" +version = "0.15.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 ebb64cf..123efec 100644 --- a/skills/pgdevkit/SKILL.md +++ b/skills/pgdevkit/SKILL.md @@ -4,7 +4,7 @@ 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 and - query helpers (`pgdevkit.db`: `pg_*`, `fetch_all`, `PgPool`, `SqlLoader`), streaming a query as JSON from a + 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, and database-in-source skills — pgdevkit is a real dependency now, not @@ -120,9 +120,15 @@ pool = PgPool(env_prefix="APP_POSTGRES_") # matches ensure_testdb()'s {env_pref async def startup(): await pool.open() - set_default_pool(pool) # lets `fetch_all()` borrow from it when no `con=`/`pool=` is given + set_default_pool(pool) # lets `fetch_all()` & co. borrow from it when no `con=`/`pool=` is given ``` +`set_default_pool()` takes anything with a `.connection()` method (`PgPool`), or a plain callable returning an async +connection context manager (typed `PoolLike`; e.g. `set_default_pool(lambda: get_pool().connection())`; only the +default pool accepts a callable, the `pool=` keyword wants a `ConnectionSource`). `await pool.open()` is idempotent but +still has to run once at startup; there is no lazy auto-open. `PgPool` also takes `min_size=`, `timeout=` (acquire +timeout), `configure=` (async hook per new connection) and `check=`; its password and `dsn_params` are quoted for you. + ### Models ```python @@ -164,7 +170,7 @@ from pgdevkit.db import pg_retrieve, pg_insert, pg_upsert, pg_delete | Helper | Purpose | |--------|---------| | `pg_retrieve` | Fetch single row by PK | -| `pg_retrieve_many` | Fetch rows matching filter dict | +| `pg_retrieve_many` | Fetch rows matching a filter dict; optional `where=` (t-string / `psycopg.sql` + `params=`), `order_by=`, `limit=` | | `pg_insert` | Insert one row, `RETURNING *` | | `pg_update` / `pg_update_dict` | Update by PK | | `pg_upsert` / `pg_upsert_dict` | `INSERT ... ON CONFLICT ... DO UPDATE` | @@ -172,7 +178,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 run it with `fetch_all` (below) in a repository method. +Use these for simple CRUD. `pg_retrieve_many(con, Model, {"status": "open"}, where=t"price > {minimum}", order_by=[("created", "desc")], limit=50)` covers a bit more than equality; for joins, aggregations or anything bigger write a dedicated `.sql` file and run it with `fetch_all` (below) in a repository method. ### Loading `.sql` files @@ -206,10 +212,10 @@ 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 +### `fetch_all`, `fetch_one`, `fetch_scalar`, `execute` — custom queries without cursor boilerplate ```python -from pgdevkit.db import fetch_all, set_default_pool +from pgdevkit.db import execute, fetch_all, fetch_one, fetch_scalar, set_default_pool set_default_pool(pool) # once, at startup -- the connection source when none is passed @@ -217,13 +223,18 @@ users = await fetch_all(sql.load_sql("users", "list_active_users"), {"limit": 10 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) + +user = await fetch_one(sql.load_sql("users", "get_user_by_id"), {"id": 1}, model=UserSummary) # UserSummary | None +email = await fetch_scalar("SELECT email FROM users WHERE id = %(id)s", {"id": 1}) # str | None +deleted = await execute("DELETE FROM sessions WHERE user_id = %(id)s", {"id": 1}) # row count ``` +- Siblings with the same keywords (`con=`, `pool=`, `cancel=`, `statement_timeout=`, t-string queries): `fetch_one(...)` returns the first row or `None` (`model=`/`row_mapper=` work as in `fetch_all`; add `LIMIT 1` yourself), `fetch_scalar(...)` the first column of the first row or `None` (e.g. `await fetch_scalar("SELECT count(*) FROM users")`), `execute(...)` the row count of a write without `RETURNING` (`await execute("DELETE FROM sessions WHERE expires < now()")`; on a pooled connection it commits, on your `con=` it does not; with `PgPool`, a custom connection source decides for itself). - `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). +- `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. -- `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`. +- `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. 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 @@ -252,10 +263,32 @@ See [`references/temporal-tables.md`](references/temporal-tables.md) for row-lev `pgdevkit.db.complex_types.ComplexHelper` detects a table's composite/enum/JSONB columns and converts plain dict/list values into the psycopg-registered types -those columns need — used automatically by both `pgdevkit.testdb.schema`'s -test-data seeding and `pgdevkit.db.crud`'s CRUD helpers. Pass a `normalizers` -dict (keyed by composite type name) if a project needs to reshape a value -before conversion (e.g. backfilling missing locale keys). +those columns need. `pgdevkit.testdb.schema`'s test-data seeding uses it automatically; the `pg_*` CRUD helpers only do +when you pass `complex_helper="auto"` (or `complex_helper=ComplexHelper(con, normalizers={...})` to reshape a value +before conversion, keyed by composite type name, e.g. backfilling missing locale keys); without it values are not +converted and reads are a plain `SELECT *`. + +What the helper learns from the catalog (types, column types, adapter registrations) is cached **per connection**, so +creating a `ComplexHelper(con)` per call (`pg_retrieve(con, M, pks, complex_helper=ComplexHelper(con))`) queries the +catalog only on that connection's first use, not every time. The cache cannot see later schema changes +(`CREATE TYPE`, `ALTER TYPE ... ADD VALUE`, new columns), nor DDL that was rolled back after it had been cached (test +fixtures that create types inside a rolled-back transaction on a reused connection): call `helper.clear_cache()` after +such DDL. `ComplexHelper(con, cache=False)` gives the helper a private cache (the pre-cache behaviour). + +### Read-only transactions — `readonly_transaction` + +```python +from pgdevkit.db import fetch_all, readonly_transaction + +async with pool.connection() as con, readonly_transaction(con): + rows = await fetch_all(sql_text, con=con) # a write raises psycopg.errors.ReadOnlySqlTransaction +``` + +It sets `read_only` before the transaction (committed on success, rolled back on error) and **always** puts it back in +a `finally`, so a pooled connection is never handed on read-only. `con` must not be inside a transaction already +(psycopg refuses to change `read_only` then). **A safety net, not a sandbox:** SQL without bound parameters may contain +several statements, and `COMMIT; INSERT ...` leaves the read-only transaction; TEMP tables stay writable. For untrusted +(e.g. LLM-written) SQL also use a role with only `SELECT` privileges. ### Streaming JSON from FastAPI — `pgdevkit.fastapi` @@ -275,7 +308,9 @@ async def articles(lng: str) -> PostgresJsonResponse: - Query: from a `.sql` file via `SqlLoader` like any non-trivial query (or `psycopg.sql`, sqlglot or a t-string for dynamic SQL), values via `params` (`%(name)s`; none with a t-string, it carries its own). That it is a literal is enforced by the type checker only: never build it from user input. Data-modifying CTE or custom JSON: `query_produces_json=True` (one text column per row). - `responses={200: {"model": ...}}` keeps the OpenAPI schema typed for the generated frontend client; `response_model` is ignored and rows are **not validated**. - Connections as in `fetch_all`: default pool, `pool=`, or `con=`. Prefer the pool: a `con` you opened must outlive the response (`Depends` with `yield`, never `async with ... as conn: return PostgresJsonResponse(q, con=conn)`), and a disconnect aborts its transaction. -- Errors show as `{"error": "Internal Server Error"}` (500); `expose_errors = IS_DEV` in a subclass shows the text (leaks SQL/values: never in prod). An error after the first byte leaves the array unterminated, so clients fail to parse it. +- Errors show as `{"error": "Internal Server Error"}` (500), not FastAPI's `{"detail": ...}`: set `error_key = "detail"` in a subclass if the frontend expects that. `expose_errors = IS_DEV` in a subclass shows the text (leaks SQL/values: never in prod). An error after the first byte leaves the array unterminated, so clients fail to parse it. +- A pooled connection is held until the client has read the whole body, so slow readers (and many parallel streams) can exhaust the pool (`PgPool` defaults to 40): size it for that (`PgPool(max_size=...)`) and enforce a proxy write timeout. +- `timestamptz` columns come out as `2026-01-01T10:00:00+00:00` (Postgres' `row_to_json`, with the session time zone's offset), where FastAPI/Pydantic emit `...Z` for UTC: equivalent for any ISO-8601 parser, but a frontend comparing strings sees a difference. `NaN`/`Infinity` come out as the strings `"NaN"`/`"Infinity"`. - `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. @@ -354,7 +389,9 @@ 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 -- [ ] 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` +- [ ] Custom read queries use `fetch_all(...)`/`fetch_one(...)`/`fetch_scalar(...)` and writes without `RETURNING` use `execute(...)` (`model=` where the shape is stable; pass `con=` inside a transaction); results mapped to a Pydantic model; table-mapped models extend `PostgresTableModel` +- [ ] 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) - [ ] 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` diff --git a/skills/pgdevkit/references/dynamic-sql.md b/skills/pgdevkit/references/dynamic-sql.md index bd13c52..a0f926f 100644 --- a/skills/pgdevkit/references/dynamic-sql.md +++ b/skills/pgdevkit/references/dynamic-sql.md @@ -25,7 +25,7 @@ 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): +`fetch_all`/`fetch_one`/`fetch_scalar`/`execute` accept t-strings and `psycopg.sql` composables too (so does `pg_retrieve_many(..., where=t"...")`) (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`):** diff --git a/tests/db/test_complex_types.py b/tests/db/test_complex_types.py index 7488882..8671834 100644 --- a/tests/db/test_complex_types.py +++ b/tests/db/test_complex_types.py @@ -1,9 +1,12 @@ from __future__ import annotations +from typing import Sequence + import psycopg import pytest from psycopg.types.enum import EnumInfo +from pgdevkit.db import PostgresTableModel, pg_insert, pg_retrieve, pg_retrieve_many from pgdevkit.db.complex_types import ComplexHelper from pgdevkit.db.crud import _select_list from pgdevkit.testdb import constants @@ -14,11 +17,11 @@ def _admin_dsn() -> str: - return f"postgresql://{constants.USER}:{constants.PASSWORD}@{constants.HOST}:{constants.PORT}/postgres" + return constants.conninfo("postgres") def _db_dsn() -> str: - return f"postgresql://{constants.USER}:{constants.PASSWORD}@{constants.HOST}:{constants.PORT}/{TEST_DB}" + return constants.conninfo(TEST_DB) @pytest.fixture @@ -121,6 +124,14 @@ async def test_jsonb_array_column_wraps_each_element_not_the_whole_list(complex_ assert metadata == metadata_value +def test_helper_attributes_stay_assignable(): + helper = ComplexHelper(con=None) # type: ignore[arg-type] + helper.complex_types = {} + helper.registered = set() + helper.system_complex_type_dict = {} + assert helper.system_complex_type_dict == {} + + def test_complex_types_cache_is_per_instance_not_shared(): # Regression: complex_types used to be a mutable class attribute, so a # CompositeInfo/EnumInfo (which carries OIDs from one specific @@ -136,3 +147,143 @@ def test_complex_types_cache_is_per_instance_not_shared(): assert helper_b.complex_types == {} assert helper_a.complex_types is not helper_b.complex_types + + +class _CountingCursor(psycopg.AsyncCursor): + async def execute(self, *args, **kwargs): + type(self.connection).queries += 1 + return await super().execute(*args, **kwargs) + + +async def _connect() -> psycopg.AsyncConnection: + """A connection that counts the statements run through its cursors (every psycopg helper uses them).""" + + class CountingConnection(psycopg.AsyncConnection): + queries = 0 + + return await CountingConnection.connect(_db_dsn(), autocommit=True, cursor_factory=_CountingCursor) + + +class Gadget(PostgresTableModel): + id: int + mood: str | None = None + dims: dict | None = None + note: str | None = None + + @staticmethod + def get_table_name() -> tuple[str, str]: + return ("public", "gadget") + + @staticmethod + def get_primary_key() -> Sequence[str]: + return ["id"] + + +async def _setup_gadget(con: psycopg.AsyncConnection) -> None: + await con.execute("CREATE TYPE mood AS ENUM ('happy', 'sad')") + await con.execute("CREATE TYPE dims AS (w int, h int)") + await con.execute("CREATE TABLE gadget (id int PRIMARY KEY, mood mood, dims dims, note text)") + await con.execute("INSERT INTO gadget VALUES (1, 'happy', ROW(3, 4), 'a'), (2, 'sad', NULL, 'b')") + + +async def _queries_of(con, call) -> tuple[int, object]: + before = type(con).queries + result = await call() + return type(con).queries - before, result + + +@requires_podman +async def test_pg_retrieve_does_not_requery_the_catalog_for_every_call(complex_types_test_db): + async with await _connect() as con: + await _setup_gadget(con) + + async def retrieve(): + return await pg_retrieve(con, Gadget, {"id": 1}, complex_helper=ComplexHelper(con)) # a new helper per call + + first_count, first = await _queries_of(con, retrieve) + second_count, second = await _queries_of(con, retrieve) + third_count, third = await _queries_of(con, retrieve) + assert first == second == third == Gadget(id=1, mood="happy", dims={"w": 3, "h": 4}, note="a") + assert first_count > 1 # type catalog + columns + the enum/composite lookups + the SELECT itself + assert (second_count, third_count) == (1, 1) # just the SELECT + + # the other read path, and writes with composite values, benefit too + count, rows = await _queries_of( + con, lambda: pg_retrieve_many(con, Gadget, {"mood": "sad"}, complex_helper=ComplexHelper(con)) + ) + assert count == 1 and [r.id for r in rows] == [2] # type: ignore[union-attr] + count, _ = await _queries_of( + con, + lambda: pg_insert(con, ("public", "gadget"), {"id": 3, "mood": "sad", "dims": {"w": 1, "h": 2}}, complex_helper=ComplexHelper(con)), + ) + assert count > 1 # first conversion of values on this connection: looks the two columns up + count, row = await _queries_of( + con, + lambda: pg_insert(con, ("public", "gadget"), {"id": 4, "mood": "sad", "dims": {"w": 1, "h": 2}}, complex_helper=ComplexHelper(con)), + ) + assert count == 1 + assert row["dims"].w == 1 # type: ignore[index] + + +@requires_podman +async def test_cache_can_be_disabled_or_cleared(complex_types_test_db): + async with await _connect() as con: + await _setup_gadget(con) + first_count, _ = await _queries_of( + con, lambda: pg_retrieve(con, Gadget, {"id": 1}, complex_helper=ComplexHelper(con)) + ) + uncached_count, _ = await _queries_of( + con, lambda: pg_retrieve(con, Gadget, {"id": 1}, complex_helper=ComplexHelper(con, cache=False)) + ) + assert uncached_count == first_count + helper = ComplexHelper(con) + assert (await _queries_of(con, lambda: pg_retrieve(con, Gadget, {"id": 1}, complex_helper=helper)))[0] == 1 + helper.clear_cache() + cleared_count, _ = await _queries_of(con, lambda: pg_retrieve(con, Gadget, {"id": 1}, complex_helper=helper)) + assert cleared_count == first_count + # ... which is how a schema change gets picked up + await con.execute("ALTER TABLE gadget ADD COLUMN extra mood") + stale = await ComplexHelper(con).load_all_complex_types(("public", "gadget"), include_generated=True) + assert "extra" not in stale + helper.clear_cache() + fresh = await ComplexHelper(con).load_all_complex_types(("public", "gadget"), include_generated=True) + assert "extra" in fresh + # cache=False never caches column lookups, not even within one helper + private = ComplexHelper(con, cache=False) + assert "extra" in await private.load_all_complex_types(("public", "gadget"), include_generated=True) + await con.execute("ALTER TABLE gadget ADD COLUMN extra2 mood") + assert "extra2" in await private.load_all_complex_types(("public", "gadget"), include_generated=True) + + +@requires_podman +async def test_the_cache_is_per_connection_and_skips_missing_tables(complex_types_test_db): + async with await _connect() as con_a, await _connect() as con_b: + await _setup_gadget(con_a) + assert await ComplexHelper(con_a).load_all_complex_types(("public", "nope")) == {} # not cached + await con_a.execute("CREATE TABLE nope (id int, m mood)") + assert set(await ComplexHelper(con_a).load_all_complex_types(("public", "nope"))) == {"id", "m"} + helper_a, helper_b = ComplexHelper(con_a), ComplexHelper(con_b) + assert helper_a.complex_types is not helper_b.complex_types + await helper_a.load_all_complex_types(("public", "gadget")) + assert helper_b.system_complex_type_dict is None # b has learned nothing from a + assert ComplexHelper(con_a).complex_types is helper_a.complex_types # a's helpers share + + +@requires_podman +async def test_complex_helper_auto_works_for_every_crud_helper(complex_types_test_db): + from pgdevkit.db import pg_insert_many, pg_update_dict, pg_upsert_dict + + async with await _connect() as con: + await _setup_gadget(con) + table = ("public", "gadget") + row = await pg_insert(con, table, {"id": 5, "mood": "sad", "dims": {"w": 1, "h": 2}}, complex_helper="auto") + assert row["mood"] == "sad" and row["dims"].w == 1 # type: ignore[attr-defined] + await pg_update_dict(con, table, {"id": 5, "dims": {"w": 9, "h": 9}}, ["id"], complex_helper="auto") + await pg_upsert_dict(con, table, {"id": 5, "mood": "happy"}, ["id"], complex_helper="auto") + await pg_insert_many(con, table, [{"id": 6, "mood": "sad", "dims": {"w": 1, "h": 1}}, {"id": 7, "mood": "happy", "dims": {"w": 2, "h": 2}}], complex_helper="auto") + got = await pg_retrieve(con, Gadget, {"id": 5}, complex_helper="auto") + assert got == Gadget(id=5, mood="happy", dims={"w": 9, "h": 9}, note=None) + assert {g.id for g in await pg_retrieve_many(con, Gadget, {"mood": "happy"}, complex_helper="auto")} == {1, 5, 7} + # one helper per call, yet the catalog was asked once: the next call is just the statement itself + count, _ = await _queries_of(con, lambda: pg_retrieve(con, Gadget, {"id": 1}, complex_helper="auto")) + assert count == 1 diff --git a/tests/db/test_connection.py b/tests/db/test_connection.py index 6223d12..194a0ec 100644 --- a/tests/db/test_connection.py +++ b/tests/db/test_connection.py @@ -1,7 +1,10 @@ from __future__ import annotations import pytest +from psycopg import AsyncConnection +from psycopg.conninfo import conninfo_to_dict +from pgdevkit.db import fetch_scalar from pgdevkit.db.connection import PgPool @@ -137,3 +140,110 @@ def fake_get_lakebase_password(workspace_host, instance_name): "user=alice@example.com password=LAKEBASE_TOKEN" ) assert calls == [("https://adb-123.azuredatabricks.net", "myinstance")] + + +@pytest.mark.parametrize("password", ["p ss", "back\\slash", "it's", "a=b", "two spaces\\ and 'quote'", "", "ünï"]) +async def test_dsn_quotes_the_password(monkeypatch, password): + monkeypatch.setenv(f"{ENV_PREFIX}HOST", "localhost") + monkeypatch.setenv(f"{ENV_PREFIX}PORT", "5432") + monkeypatch.setenv(f"{ENV_PREFIX}DB", "mydb") + monkeypatch.setenv(f"{ENV_PREFIX}USER", "myuser") + monkeypatch.setenv(f"{ENV_PREFIX}PASSWORD", password) + + info = conninfo_to_dict(await PgPool(env_prefix=ENV_PREFIX)._dsn()) + assert info == {"host": "localhost", "port": "5432", "dbname": "mydb", "user": "myuser", "password": password} + + +async def test_dsn_quotes_extra_params_and_lets_them_override(monkeypatch): + monkeypatch.setenv(f"{ENV_PREFIX}HOST", "localhost") + monkeypatch.setenv(f"{ENV_PREFIX}PORT", "5432") + monkeypatch.setenv(f"{ENV_PREFIX}DB", "mydb") + monkeypatch.setenv(f"{ENV_PREFIX}USER", "myuser") + monkeypatch.setenv(f"{ENV_PREFIX}PASSWORD", "pw") + + pool = PgPool(env_prefix=ENV_PREFIX, dsn_params={"application_name": "my app\\'s", "sslmode": "require"}) + info = conninfo_to_dict(await pool._dsn()) + assert info["application_name"] == "my app\\'s" + assert info["sslmode"] == "require" + + +async def test_azure_token_is_quoted_too(monkeypatch): + monkeypatch.setenv(f"{ENV_PREFIX}HOST", "myserver.postgres.database.azure.com") + monkeypatch.setenv(f"{ENV_PREFIX}PORT", "5432") + monkeypatch.setenv(f"{ENV_PREFIX}DB", "mydb") + monkeypatch.setattr("pgdevkit.db.connection.get_azure_postgres_password", lambda **kwargs: "tok en\\") + info = conninfo_to_dict(await PgPool(env_prefix=ENV_PREFIX, entra_user="alice@example.com")._dsn()) + assert info["password"] == "tok en\\" + + +def _set_local_env(monkeypatch, dsn: str) -> None: + info = conninfo_to_dict(dsn) + for key, field in (("HOST", "host"), ("PORT", "port"), ("DB", "dbname"), ("USER", "user"), ("PASSWORD", "password")): + monkeypatch.setenv(f"PGDEVKIT_CONNLIVE_{key}", info.get(field, "")) + + +async def test_pool_options_are_passed_through(monkeypatch): + monkeypatch.setenv(f"{ENV_PREFIX}HOST", "localhost") + monkeypatch.setenv(f"{ENV_PREFIX}PORT", "5432") + monkeypatch.setenv(f"{ENV_PREFIX}DB", "mydb") + monkeypatch.setenv(f"{ENV_PREFIX}USER", "myuser") + monkeypatch.setenv(f"{ENV_PREFIX}PASSWORD", "pw") + + async def configure(con: AsyncConnection) -> None: ... + + async def check(con: AsyncConnection) -> None: ... + + pool = PgPool(env_prefix=ENV_PREFIX, max_size=7, min_size=2, timeout=3.5, configure=configure, check=check) + await pool.open() + try: + raw = pool.raw_pool + assert (raw.min_size, raw.max_size, raw.timeout) == (2, 7, 3.5) + assert raw._configure is configure + assert raw._check is check + finally: + await pool.close() + # untouched defaults: the library's own check, no hook + plain = PgPool(env_prefix=ENV_PREFIX) + await plain.open() + try: + assert plain.raw_pool._configure is None + assert plain.raw_pool._check is not None + finally: + await plain.close() + + +async def test_configure_hook_and_timeout_work_on_a_live_pool(monkeypatch, postgres_dsn): + _set_local_env(monkeypatch, postgres_dsn) + + async def configure(con: AsyncConnection) -> None: + await con.execute("SET application_name = 'pgdevkit_configured'") + await con.commit() + + pool = PgPool(env_prefix="PGDEVKIT_CONNLIVE_", max_size=1, timeout=0.3, configure=configure) + await pool.open() + try: + assert await fetch_scalar("SHOW application_name", pool=pool) == "pgdevkit_configured" + from psycopg_pool import PoolTimeout + + async with pool.connection(): # the only connection is taken: the next borrower times out + with pytest.raises(PoolTimeout): + await fetch_scalar("SELECT 1", pool=pool) + finally: + await pool.close() + + +async def test_pool_options_are_passed_to_a_null_pool_too(monkeypatch): + monkeypatch.setenv(f"{ENV_PREFIX}HOST", "myserver.postgres.database.azure.com") + monkeypatch.setenv(f"{ENV_PREFIX}PORT", "5432") + monkeypatch.setenv(f"{ENV_PREFIX}DB", "mydb") + monkeypatch.setenv(f"{ENV_PREFIX}USER", "myuser") + monkeypatch.setenv(f"{ENV_PREFIX}PASSWORD", "pw") + + async def configure(con: AsyncConnection) -> None: ... + + pool = PgPool(env_prefix=ENV_PREFIX, timeout=2.5, configure=configure) + await pool.open() + try: + assert (pool.raw_pool.timeout, pool.raw_pool._configure) == (2.5, configure) + finally: + await pool.close() diff --git a/tests/db/test_crud.py b/tests/db/test_crud.py index 5556070..a5ca738 100644 --- a/tests/db/test_crud.py +++ b/tests/db/test_crud.py @@ -4,6 +4,8 @@ import psycopg import pytest +from psycopg.sql import SQL, Placeholder +from psycopg.sql import Literal as SqlLiteral from pydantic import ConfigDict from pgdevkit.db import ( @@ -268,3 +270,70 @@ async def test_update_dict_composite_column_requires_complex_helper(pool: PgPool fetched = await pg_retrieve(con, Gizmo, {"id": gizmo_id}, complex_helper=helper) assert fetched is not None assert fetched.label == {"en": "Bye", "de": "Tschuess"} + + +@requires_podman +async def test_pg_retrieve_many_where_order_by_and_limit(pool: PgPool): + async with pool.connection() as con: + await pg_insert_many(con, ("public", "widget"), [{"name": n} for n in ("cog", "sprocket", "spring", "50% off", "gear")]) + + def names(rows) -> list[str]: + return [r.name for r in rows] + + # unchanged: equality filters only + assert names(await pg_retrieve_many(con, Widget, {"name": "cog"})) == ["cog"] + + # order_by: a column, several, a direction, or a composable + assert names(await pg_retrieve_many(con, Widget, {}, order_by="name")) == ["50% off", "cog", "gear", "spring", "sprocket"] + assert names(await pg_retrieve_many(con, Widget, {}, order_by=[("name", "desc")], limit=2)) == ["sprocket", "spring"] + assert names(await pg_retrieve_many(con, Widget, {}, order_by=["name"], limit=0)) == [] + by_length = SQL("length(name) DESC, name") + assert names(await pg_retrieve_many(con, Widget, {}, order_by=by_length, limit=2)) == ["sprocket", "50% off"] + + # where as a t-string: values are bound, a literal % is no problem, filters are ANDed in + pattern = "sp%" + assert names(await pg_retrieve_many(con, Widget, {}, where=t"name LIKE {pattern}", order_by="name")) == ["spring", "sprocket"] + assert names(await pg_retrieve_many(con, Widget, {"name": "spring"}, where=t"name LIKE {pattern}")) == ["spring"] + assert names(await pg_retrieve_many(con, Widget, {}, where=t"name LIKE '50%' OR name = {'cog'}", order_by="name")) == ["50% off", "cog"] + minimum = 0 + assert len(await pg_retrieve_many(con, Widget, {}, where=t"id > {minimum}", limit=3)) == 3 + assert await pg_retrieve_many(con, Widget, {}, where=t"false") == [] + + # where as a composable (+ params), also together with filters + like = SQL("name LIKE {p}").format(p=Placeholder("pattern")) + assert names(await pg_retrieve_many(con, Widget, {}, where=like, params={"pattern": "sp%"}, order_by="name")) == ["spring", "sprocket"] + assert names(await pg_retrieve_many(con, Widget, {"name": "cog"}, where=like, params={"pattern": "c%"})) == ["cog"] + literal = SQL("name LIKE {p}").format(p=SqlLiteral("50%")) # no params at all: % is not parsed + assert names(await pg_retrieve_many(con, Widget, {}, where=literal)) == ["50% off"] + + with pytest.raises(ValueError, match="collides"): + await pg_retrieve_many(con, Widget, {"name": "cog"}, where=like, params={"name": "x", "pattern": "x"}) + with pytest.raises(TypeError, match="t-string"): + await pg_retrieve_many(con, Widget, {}, where=t"true", params={"x": 1}) + with pytest.raises(ValueError, match="use a list"): + await pg_retrieve_many(con, Widget, {}, order_by=("name", "desc")) # type: ignore[arg-type] + assert names(await pg_retrieve_many(con, Widget, {}, order_by=[("name", "DESC")], limit=1)) == ["sprocket"] + # `params` without `where` is simply unused; a missing `where` placeholder is psycopg's error + assert names(await pg_retrieve_many(con, Widget, {"name": "cog"}, params={"unused": 1})) == ["cog"] + with pytest.raises(ValueError, match="direction"): + await pg_retrieve_many(con, Widget, {}, order_by=[("name", "sideways")]) # type: ignore[list-item] + for bad in (-1, True, "3", 1.5): + with pytest.raises(ValueError, match="limit"): + await pg_retrieve_many(con, Widget, {}, limit=bad) # type: ignore[arg-type] + # a column name is an identifier, never SQL + with pytest.raises(psycopg.errors.UndefinedColumn): + await pg_retrieve_many(con, Widget, {}, order_by="name; DROP TABLE widget") + await con.rollback() + + +@requires_podman +async def test_pg_retrieve_many_where_with_complex_helper(pool: PgPool): + async with pool.connection() as con: + await pg_insert_many(con, ("public", "gadget"), [{"status": "active"}, {"status": "inactive"}, {"status": "active"}]) + rows = await pg_retrieve_many( + con, Gadget, {"status": "active"}, complex_helper=ComplexHelper(con), order_by=[("id", "desc")], limit=1 + ) + assert [(r.id, r.status) for r in rows] == [(3, "active")] + status = "inactive" + rows = await pg_retrieve_many(con, Gadget, {}, complex_helper=ComplexHelper(con), where=t"status = {status}") + assert [(r.id, r.status) for r in rows] == [(2, "inactive")] diff --git a/tests/db/test_fetch.py b/tests/db/test_fetch.py index aa9ffa1..63e5428 100644 --- a/tests/db/test_fetch.py +++ b/tests/db/test_fetch.py @@ -11,7 +11,8 @@ from pydantic import BaseModel from sqlglot import exp, select -from pgdevkit.db import PgPool, fetch_all, set_default_pool +from pgdevkit.db import PgPool, execute, fetch_all, fetch_one, fetch_scalar, set_default_pool +from pgdevkit.db.fetch import _render from pgdevkit.testdb import constants from pgdevkit.testdb.container import ensure_container from tests.testdb.conftest import RUN_SUFFIX, requires_podman @@ -251,3 +252,150 @@ async def broken_cancel(self, **_kwargs) -> None: ) await observer.close() await con.close() + + +# --- fetch_one / fetch_scalar / execute --------------------------------------------------------- + + +async def test_siblings_validate_their_arguments_before_touching_a_connection(): + with pytest.raises(TypeError, match="model"): + await fetch_one("SELECT 1", model=Widget, row_mapper=dict) # type: ignore[call-overload] + value = 1 + for fn in (fetch_one, fetch_scalar, execute): + with pytest.raises(TypeError, match="t-string"): + await fn(t"SELECT {value}", {"x": 1}) # type: ignore[call-overload] + with pytest.raises(ValueError, match="statement_timeout"): + await fn("SELECT 1", statement_timeout=0) + with pytest.raises(RuntimeError, match="set_default_pool"): + await fn("SELECT 1") + with pytest.raises(TypeError, match=fn.__name__): + await fn(123, pool=_NoPool()) # type: ignore[call-overload] + cancel = asyncio.Event() + cancel.set() + with pytest.raises(QueryCanceled, match=fn.__name__): + await fn("SELECT 1", cancel=cancel, pool=_NoPool()) + + +@requires_podman +async def test_fetch_one_returns_first_row_or_none(pool: PgPool): + assert await fetch_one("SELECT id, name FROM widget ORDER BY id", pool=pool) == {"id": 1, "name": "sprocket"} + assert await fetch_one("SELECT id FROM widget WHERE id = %(id)s", {"id": 99}, pool=pool) is None + assert await fetch_one("SELECT id, name FROM widget WHERE id = 2", model=Widget, pool=pool) == Widget(id=2, name="cog") + assert await fetch_one("SELECT name FROM widget WHERE id = 2", row_mapper=lambda r: r["name"].upper(), pool=pool) == "COG" + assert await fetch_one("SELECT name FROM widget WHERE id = 99", model=Widget, pool=pool) is None + widget_id = 1 + assert await fetch_one(t"SELECT name FROM widget WHERE id = {widget_id}", pool=pool) == {"name": "sprocket"} + query = select("name").from_("widget").where(exp.column("id").eq(exp.Placeholder(this="id"))) + assert await fetch_one(query, {"id": 2}, pool=pool) == {"name": "cog"} + + +@requires_podman +async def test_fetch_scalar_returns_first_column_of_first_row_or_none(pool: PgPool): + set_default_pool(pool) + assert await fetch_scalar("SELECT count(*) FROM widget") == 2 + assert await fetch_scalar("SELECT name, id FROM widget ORDER BY id DESC") == "cog" + assert await fetch_scalar("SELECT name FROM widget WHERE id = %(id)s", {"id": 99}) is None + assert await fetch_scalar("SELECT NULL::int") is None + assert await fetch_scalar("SELECT 1 AS a, 2 AS a") == 1 # first column, even if names repeat + assert await fetch_scalar(t"SELECT {5}::int + 1") == 6 + + +@requires_podman +async def test_execute_returns_rowcount_and_commits_on_a_pooled_connection(pool: PgPool): + assert await execute("UPDATE widget SET name = upper(name)", pool=pool) == 2 + assert await execute("DELETE FROM widget WHERE id = %(id)s", {"id": 1}, pool=pool) == 1 + assert await execute("DELETE FROM widget WHERE id = %(id)s", {"id": 1}, pool=pool) == 0 + new_id = 7 + assert await execute(t"INSERT INTO widget VALUES ({new_id}, 'gear')", pool=pool) == 1 + assert await fetch_all("SELECT id, name FROM widget ORDER BY id", pool=pool) == [ + {"id": 2, "name": "COG"}, + {"id": 7, "name": "gear"}, + ] + assert await execute("CREATE TABLE extra (id int)", pool=pool) == -1 + + +@requires_podman +async def test_execute_on_a_given_connection_leaves_the_transaction_to_the_caller(pool: PgPool): + async with pool.connection() as con: + assert await execute("INSERT INTO widget VALUES (3, 'gear')", con=con) == 1 + assert con.info.transaction_status == psycopg.pq.TransactionStatus.INTRANS + await con.rollback() + assert await fetch_scalar("SELECT count(*) FROM widget", pool=pool) == 2 + + +@requires_podman +async def test_siblings_support_cancel_and_statement_timeout(pool: PgPool): + cancel = asyncio.Event() + task = asyncio.ensure_future(fetch_scalar("SELECT pg_sleep(30)", pool=pool, cancel=cancel)) + await _wait_for_active_sleeps(pool, 1) + cancel.set() + with pytest.raises(QueryCanceled): + await task + await _wait_for_active_sleeps(pool, 0) + with pytest.raises(QueryCanceled, match="statement timeout"): + await fetch_one("SELECT pg_sleep(30)", pool=pool, statement_timeout=0.3) + with pytest.raises(QueryCanceled, match="statement timeout"): + await execute("SELECT pg_sleep(30)", pool=pool, statement_timeout=0.3) + assert await fetch_scalar("SELECT 1", pool=pool, statement_timeout=5) == 1 + + +@requires_podman +async def test_default_pool_may_be_a_plain_callable(pool: PgPool): + set_default_pool(lambda: pool.connection()) + assert await fetch_scalar("SELECT count(*) FROM widget") == 2 + assert await fetch_one("SELECT name FROM widget WHERE id = 1") == {"name": "sprocket"} + + +# --- literal `%` in sqlglot expressions ---------------------------------------------------------- + + +def test_render_leaves_percent_alone_without_params(): + query = select("a").from_("t").where(exp.column("a").like("50%")) + assert _render(query) == "SELECT a FROM t WHERE a LIKE '50%'" + assert _render(query, None) == "SELECT a FROM t WHERE a LIKE '50%'" + + +@pytest.mark.parametrize("params", [{}, {"id": 1}]) +def test_render_doubles_literal_percent_but_not_placeholders_when_params_are_given(params): + query = ( + select(exp.column("50%"), exp.column("a") % 2) + .from_("t") + .where(exp.column("a").like("50%")) + .where(exp.column("b").eq("%(id)s")) # a literal that merely looks like a placeholder + .where(exp.column("c").eq(exp.Placeholder(this="id"))) + ) + assert _render(query, params) == ( + "SELECT \"50%%\", a %% 2 FROM t WHERE (a LIKE '50%%' AND b = '%%(id)s') AND c = %(id)s" + ) + # the query object itself is left untouched + assert "pgdevkit_ph_" not in query.sql(dialect="postgres") + assert query.sql(dialect="postgres").count("%(id)s") == 2 # the placeholder and the literal + + +def test_render_keeps_ten_or_more_placeholders_apart(): + # (placeholder token _1 must not be mistaken for the start of _10) + names = [f"p{i}" for i in range(25)] + query = select(*[exp.Placeholder(this=n) for n in names]).where(exp.column("a").like("x%")) + assert _render(query, {}) == "SELECT " + ", ".join(f"%({n})s" for n in names) + " WHERE a LIKE 'x%%'" + + +def test_render_percent_escaping_leaves_other_query_types_alone(): + assert _render("SELECT '50%%'", {}) == "SELECT '50%%'" + composed = sql.SQL("SELECT {}").format(sql.Literal("50%")) + assert _render(composed, {}) is composed + + +@requires_podman +async def test_percent_in_sqlglot_string_literals_works_with_and_without_params(pool: PgPool): + like = select("name").from_("widget").where(exp.column("name").like("sp%")) + assert await fetch_all(like, pool=pool) == [{"name": "sprocket"}] + assert await fetch_all(like, {}, pool=pool) == [{"name": "sprocket"}] + with_placeholder = like.where(exp.column("id").eq(exp.Placeholder(this="id"))) + assert await fetch_all(with_placeholder, {"id": 1}, pool=pool) == [{"name": "sprocket"}] + assert await fetch_all(with_placeholder, {"id": 2}, pool=pool) == [] + assert await fetch_scalar(select(exp.Literal.string("50%").as_("pct")), {}, pool=pool) == "50%" + assert await fetch_scalar(select(exp.Literal.string("50%").as_("pct")), pool=pool) == "50%" + # the modulo operator, and a literal that looks like a placeholder + modulo = select((exp.column("id") % 2).as_("odd")).from_("widget").where(exp.column("id").eq(exp.Placeholder(this="id"))) + assert await fetch_scalar(modulo, {"id": 1}, pool=pool) == 1 + assert await fetch_scalar(select(exp.Literal.string("%(id)s")), {}, pool=pool) == "%(id)s" diff --git a/tests/db/test_fetch_statement_timeout.py b/tests/db/test_fetch_statement_timeout.py index 19d51a8..976db35 100644 --- a/tests/db/test_fetch_statement_timeout.py +++ b/tests/db/test_fetch_statement_timeout.py @@ -9,7 +9,7 @@ from psycopg.pq import TransactionStatus from psycopg_pool import AsyncConnectionPool -from pgdevkit.db import fetch_all +from pgdevkit.db import at_most, fetch_all, fetch_scalar SLOW = "SELECT pg_sleep(30)" @@ -59,3 +59,48 @@ async def test_autocommit_connection_works_and_stays_usable(postgres_dsn: str) - async def test_statement_timeout_must_be_positive(seconds: float) -> None: with pytest.raises(ValueError, match="statement_timeout"): await fetch_all("SELECT 1", statement_timeout=seconds) + + +async def _effective(pool: AsyncConnectionPool, timeout: float, current: str) -> int: + """The timeout (ms) in force inside a call with `statement_timeout=timeout`, given a connection whose own + setting is `current`.""" + async with pool.connection() as con: + await con.execute("SELECT set_config('statement_timeout', %(current)s, false)", {"current": current}) + await con.commit() + try: + return await fetch_scalar( + "SELECT setting::bigint FROM pg_settings WHERE name = 'statement_timeout'", + con=con, + statement_timeout=timeout, + ) + finally: + await con.execute("SELECT set_config('statement_timeout', '0', false)") + await con.commit() + + +async def test_at_most_never_loosens_the_connections_timeout(pool: AsyncConnectionPool) -> None: + # a plain value replaces the current one in both directions ... + assert await _effective(pool, 900, "600s") == 900_000 + assert await _effective(pool, 120, "600s") == 120_000 + # ... at_most() only tightens it + assert await _effective(pool, at_most(900), "600s") == 600_000 + assert await _effective(pool, at_most(120), "600s") == 120_000 + assert await _effective(pool, at_most(120), "0") == 120_000 # no timeout set: the requested one applies + assert await _effective(pool, at_most(0.0004), "600s") == 1 # sub-millisecond requests round up, as usual + + +async def test_at_most_restores_the_previous_setting_and_still_aborts(pool: AsyncConnectionPool) -> None: + async with pool.connection() as con: + await con.execute("SELECT set_config('statement_timeout', '7s', false)") + assert await fetch_all("SELECT 1 AS one", con=con, statement_timeout=at_most(2)) == [{"one": 1}] + assert await fetch_all("SHOW statement_timeout", con=con) == [{"statement_timeout": "7s"}] + await con.execute("SELECT set_config('statement_timeout', '0', false)") + with pytest.raises(QueryCanceled, match="statement timeout"): + await fetch_all(SLOW, pool=pool, statement_timeout=at_most(0.3)) + assert await fetch_all("SHOW statement_timeout", pool=pool) == [{"statement_timeout": "0"}] + + +async def test_at_most_is_validated_like_a_plain_timeout() -> None: + assert at_most(1.5) == 1.5 + with pytest.raises(ValueError, match="statement_timeout"): + await fetch_all("SELECT 1", statement_timeout=at_most(0)) diff --git a/tests/db/test_transaction.py b/tests/db/test_transaction.py new file mode 100644 index 0000000..39a3c78 --- /dev/null +++ b/tests/db/test_transaction.py @@ -0,0 +1,106 @@ +"""readonly_transaction(): writes are refused inside, and the setting never leaks into the pool.""" + +from collections.abc import AsyncIterator + +import pytest +from psycopg import AsyncConnection, ProgrammingError +from psycopg.errors import ReadOnlySqlTransaction +from psycopg.pq import TransactionStatus +from psycopg_pool import AsyncConnectionPool + +from pgdevkit.db import fetch_all, fetch_scalar, readonly_transaction + +WRITE = "CREATE TABLE pgdevkit_readonly_probe (id int)" +CLEANUP = "DROP TABLE IF EXISTS pgdevkit_readonly_probe" + + +@pytest.fixture +async def pool(postgres_dsn: str) -> AsyncIterator[AsyncConnectionPool]: + # one connection only: every borrower gets the very same (possibly tainted) connection back + async with AsyncConnectionPool(postgres_dsn, min_size=1, max_size=1, open=False) as pool: + yield pool + async with pool.connection() as con: + await con.execute(CLEANUP) + + +async def test_reads_work_and_writes_are_refused(pool: AsyncConnectionPool) -> None: + async with pool.connection() as con: + async with readonly_transaction(con) as ro: + assert ro is con + assert await fetch_scalar("SHOW transaction_read_only", con=con) == "on" + assert await fetch_all("SELECT 1 AS one", con=con) == [{"one": 1}] + with pytest.raises(ReadOnlySqlTransaction): + await con.execute(WRITE) + + +async def test_a_pooled_connection_is_not_left_read_only(pool: AsyncConnectionPool) -> None: + async with pool.connection() as con: + async with readonly_transaction(con): + await con.execute("SELECT 1") + assert con.read_only is None + # the same connection again: an unrelated write must work + async with pool.connection() as con: + assert await fetch_scalar("SHOW transaction_read_only", con=con) == "off" + await con.execute(WRITE) + await con.execute(CLEANUP) + + +async def test_read_only_is_reset_when_the_block_raises(pool: AsyncConnectionPool) -> None: + async with pool.connection() as con: + with pytest.raises(RuntimeError, match="boom"): + async with readonly_transaction(con): + await con.execute("SELECT 1") + raise RuntimeError("boom") + assert con.read_only is None + assert con.info.transaction_status == TransactionStatus.IDLE + async with pool.connection() as con: + await con.execute(WRITE) + await con.execute(CLEANUP) + + +async def test_read_only_is_reset_after_a_refused_write(pool: AsyncConnectionPool) -> None: + async with pool.connection() as con: + with pytest.raises(ReadOnlySqlTransaction): + async with readonly_transaction(con): + await con.execute(WRITE) + assert con.read_only is None + assert con.info.transaction_status == TransactionStatus.IDLE + await con.execute(WRITE) # allowed again, same connection + await con.execute(CLEANUP) + + +async def test_previous_read_only_setting_is_restored(pool: AsyncConnectionPool) -> None: + async with pool.connection() as con: + await con.set_read_only(True) # e.g. a connection to a replica + async with readonly_transaction(con): + await con.execute("SELECT 1") + assert con.read_only is True + await con.set_read_only(None) + + +async def test_refuses_a_connection_that_is_in_a_transaction(pool: AsyncConnectionPool) -> None: + async with pool.connection() as con: + await con.execute("SELECT 1") # starts a transaction (no autocommit) + with pytest.raises(ProgrammingError, match="read_only"): + async with readonly_transaction(con): + pass + assert con.read_only is None + await con.rollback() + + +async def test_works_on_an_autocommit_connection(postgres_dsn: str) -> None: + async with await AsyncConnection.connect(postgres_dsn, autocommit=True) as con: + async with readonly_transaction(con): + assert await fetch_scalar("SHOW transaction_read_only", con=con) == "on" + with pytest.raises(ReadOnlySqlTransaction): + await con.execute(WRITE) + assert con.read_only is None + assert await fetch_scalar("SHOW transaction_read_only", con=con) == "off" + + +async def test_closed_connection_does_not_mask_the_original_error(postgres_dsn: str) -> None: + con = await AsyncConnection.connect(postgres_dsn) + with pytest.raises(RuntimeError, match="boom"): + async with readonly_transaction(con): + await con.close() + raise RuntimeError("boom") diff --git a/tests/test_pg_json_response.py b/tests/test_pg_json_response.py index 8ffdb5d..b8f1df2 100644 --- a/tests/test_pg_json_response.py +++ b/tests/test_pg_json_response.py @@ -13,6 +13,7 @@ from psycopg.rows import dict_row from psycopg.sql import SQL, Literal from psycopg_pool import AsyncConnectionPool +from sqlglot import exp, select from starlette.exceptions import HTTPException from pgdevkit.db import fetch, set_default_pool @@ -24,6 +25,10 @@ class VerboseResponse(PostgresJsonResponse): expose_errors = True +class DetailResponse(PostgresJsonResponse): + error_key = "detail" + + class ForbiddenPool: def connection(self) -> AbstractAsyncContextManager[AsyncConnection]: raise HTTPException(status_code=403, detail="no access", headers={"WWW-Authenticate": "Bearer"}) @@ -122,6 +127,24 @@ async def no_pool_route(): async def stacked(): return VerboseResponse("select 1 as a) s; select 2 as b from (select 1", pool=pool) + @app.get("/syntax-error-detail") + async def syntax_error_detail(): + return DetailResponse("selec 1", pool=pool) + + @app.get("/forbidden-detail") + async def forbidden_detail(): + return DetailResponse("select 1", pool=ForbiddenPool()) + + @app.get("/sqlglot-like") + async def sqlglot_like(prefix: str): + query = ( + select("name") + .from_("(select 'a%' as name union all select 'ab' union all select 'b') s") + .where(exp.column("name").like("a%")) + .where(exp.column("name").neq(exp.Placeholder(this="prefix"))) + ) + return PostgresJsonResponse(query, {"prefix": prefix}, pool=pool) + @app.get("/syntax-error") async def syntax_error(): return PostgresJsonResponse("selec 1", pool=pool) @@ -275,3 +298,18 @@ async def test_t_string_with_query_produces_json(client: httpx.AsyncClient) -> N def test_t_string_takes_no_params(pool: AsyncConnectionPool) -> None: with pytest.raises(TypeError, match="t-string"): PostgresJsonResponse(t"select {1}", {"x": 1}, pool=pool) + + +async def test_literal_percent_in_a_sqlglot_query_with_params(client: httpx.AsyncClient) -> None: + response = await client.get("/sqlglot-like", params={"prefix": "ab"}) + assert response.json() == [{"name": "a%"}] + + +async def test_error_key_can_be_switched_to_detail(client: httpx.AsyncClient) -> None: + response = await client.get("/syntax-error-detail") + assert response.status_code == 500 + assert response.json() == {"detail": "Internal Server Error"} + response = await client.get("/forbidden-detail") + assert response.status_code == 403 + assert response.json() == {"detail": "no access"} + assert response.headers["www-authenticate"] == "Bearer" diff --git a/tests/test_pg_json_statement_timeout.py b/tests/test_pg_json_statement_timeout.py index 04c95c5..100a0e3 100644 --- a/tests/test_pg_json_statement_timeout.py +++ b/tests/test_pg_json_statement_timeout.py @@ -9,6 +9,7 @@ from psycopg.errors import QueryCanceled from psycopg_pool import AsyncConnectionPool +from pgdevkit.db import at_most from pgdevkit.fastapi import PostgresJsonResponse from tests._asgi import FakeClient @@ -88,3 +89,16 @@ async def test_autocommit_connection_is_refused_with_a_clear_error(postgres_dsn: def test_statement_timeout_must_be_positive(pool: AsyncConnectionPool, seconds: float) -> None: with pytest.raises(ValueError, match="statement_timeout"): PostgresJsonResponse("select 1", pool=pool, statement_timeout=seconds) + + +async def test_at_most_never_loosens_a_caller_owned_connections_timeout(postgres_dsn: str) -> None: + async with await AsyncConnection.connect(postgres_dsn) as conn: + await conn.execute("select set_config('statement_timeout', '600s', false)") + await conn.commit() + query = "select setting::bigint as ms from pg_settings where name = 'statement_timeout'" + for timeout, expected in ((at_most(900), 600_000), (at_most(120), 120_000), (900, 900_000)): + client = FakeClient() + await client.call(PostgresJsonResponse(query, con=conn, statement_timeout=timeout)) + assert json.loads(client.body) == [{"ms": expected}] + await conn.execute("select set_config('statement_timeout', '0', false)") + await conn.commit() diff --git a/uv.lock b/uv.lock index 3ef308b..3dcbd74 100644 --- a/uv.lock +++ b/uv.lock @@ -451,7 +451,7 @@ wheels = [ [[package]] name = "pgdevkit" -version = "0.14.2" +version = "0.15.0" source = { editable = "." } dependencies = [ { name = "docker" },