pgdevkit 0.14.2__tar.gz → 0.15.0__tar.gz

{pgdevkit-0.14.2 → pgdevkit-0.15.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: pgdevkit
-Version: 0.14.2
+Version: 0.15.0
 Summary: A helper for developing with Postgres
 Requires-Python: >=3.14
 Requires-Dist: docker>=7.1.0
@@ -430,22 +430,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=<seconds>` lets Postgres abort the query after that
   long (`QueryCanceled`, "canceling statement due to statement timeout"); it is applied per call with `SET LOCAL`
-  semantics, so it never outlives the call (a borrowed `con` gets its previous setting back). `query` must be a literal string, a sqlglot
+  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
@@ -544,7 +579,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
@@ -552,9 +589,15 @@ help, see below). Enforce a write timeout in the reverse proxy and size the pool
 `statement_timeout=<seconds>` (optional, also on `fetch_all`) lets **Postgres** abort the query after that long. It
 is set with `SET LOCAL` semantics for this one query: transaction-scoped, so it can't leak into the pool and works
 behind PgBouncer's transaction pooling, and a `con` you passed gets its previous setting back. Before the first byte
-the client gets a `504 {"error": "Query timed out"}`; after that the response is cut short like any other error. It
+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.
@@ -570,12 +613,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.

{pgdevkit-0.14.2 → pgdevkit-0.15.0}/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=<seconds>` lets Postgres abort the query after that
   long (`QueryCanceled`, "canceling statement due to statement timeout"); it is applied per call with `SET LOCAL`
-  semantics, so it never outlives the call (a borrowed `con` gets its previous setting back). `query` must be a literal string, a sqlglot
+  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=<seconds>` (optional, also on `fetch_all`) lets **Postgres** abort the query after that long. It
 is set with `SET LOCAL` semantics for this one query: transaction-scoped, so it can't leak into the pool and works
 behind PgBouncer's transaction pooling, and a `con` you passed gets its previous setting back. Before the first byte
-the client gets a `504 {"error": "Query timed out"}`; after that the response is cut short like any other error. It
+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.

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

@@ -16,19 +16,35 @@ from .crud import (
     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 @@ __all__ = [
     "pg_upsert_dict",
     "pg_upsert_many",
     "pg_upsert_many_dict",
+    "readonly_transaction",
     "set_default_pool",
 ]

{pgdevkit-0.14.2 → pgdevkit-0.15.0}/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 @@ class ComplexHelper:
     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 @@ class ComplexHelper:
                 },
             )
             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 @@ class ComplexHelper:
                 {"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 @@ class ComplexHelper:
         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)]
 

{pgdevkit-0.14.2 → pgdevkit-0.15.0}/pgdevkit/db/connection.py

@@ -2,8 +2,11 @@ from __future__ import annotations
 
 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 @@ class PgPool:
         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 @@ class PgPool:
         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 @@ class PgPool:
                     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 @@ class PgPool:
                 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 @@ class PgPool:
                     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()