From b56573f15af0fa41520f82c5ccb6dae3f1f3d240 Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Fri, 9 Oct 2026 18:04:47 +0000 Subject: [PATCH 1/6] test: regression tests for #54 (types registered by ComplexHelper leak into the connection) --- tests/db/test_complex_types.py | 102 +++++++++++++++++++++++++++++++++ tests/db/test_crud.py | 32 +++++++++++ 2 files changed, 134 insertions(+) diff --git a/tests/db/test_complex_types.py b/tests/db/test_complex_types.py index 8671834..3083a99 100644 --- a/tests/db/test_complex_types.py +++ b/tests/db/test_complex_types.py @@ -287,3 +287,105 @@ async def test_complex_helper_auto_works_for_every_crud_helper(complex_types_tes # 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 + + +# --- Fixes #54: a ComplexHelper must not change how *other* code on the same connection reads enums/composites --- + + +async def _plain_read(con: psycopg.AsyncConnection): + """What unrelated code on `con` sees: an enum label, a composite and an enum array, read with a plain SELECT.""" + async with con.cursor() as cur: + await cur.execute("SELECT 'sad'::mood, ROW(1, 2)::dims, ARRAY['happy']::mood[]") + return await cur.fetchone() + + +@requires_podman +async def test_pg_helpers_do_not_change_how_the_connection_reads_enums_and_composites(complex_types_test_db): + from pgdevkit.db import pg_insert_many, pg_update, pg_update_dict, pg_upsert, pg_upsert_dict, pg_upsert_many_dict + + async with await _connect() as con: + await _setup_gadget(con) + table = ("public", "gadget") + before = await _plain_read(con) + assert before[0] == "sad" and isinstance(before[0], str) # the label, as a str + + async def check(what: str, call) -> None: + await call() + assert await _plain_read(con) == before, f"{what} changed how the connection reads enums/composites" + + for helper_arg in ("auto", ComplexHelper(con), ComplexHelper(con, cache=False)): + kw = {"complex_helper": helper_arg} + await check("pg_retrieve", lambda: pg_retrieve(con, Gadget, {"id": 1}, **kw)) + await check("pg_retrieve_many", lambda: pg_retrieve_many(con, Gadget, {}, limit=1, **kw)) + await check("pg_insert", lambda: pg_insert(con, table, {"id": 10, "mood": "sad", "dims": {"w": 1, "h": 2}}, **kw)) + await check("pg_insert (enum only)", lambda: pg_insert(con, table, {"id": 11, "mood": "happy"}, **kw)) + await check("pg_update_dict", lambda: pg_update_dict(con, table, {"id": 10, "dims": {"w": 5, "h": 5}}, ["id"], **kw)) + await check("pg_update", lambda: pg_update(con, Gadget(id=10, mood="happy", dims={"w": 6, "h": 6}), Gadget, **kw)) + await check("pg_upsert_dict", lambda: pg_upsert_dict(con, table, {"id": 10, "dims": {"w": 7, "h": 7}}, ["id"], **kw)) + await check("pg_upsert", lambda: pg_upsert(con, Gadget(id=12, mood="sad", dims={"w": 1, "h": 1}), Gadget, **kw)) + await check( + "pg_insert_many", + lambda: pg_insert_many(con, table, [{"id": 20, "dims": {"w": 1, "h": 1}}, {"id": 21, "dims": {"w": 2, "h": 2}}], **kw), + ) + await check( + "pg_upsert_many_dict", + lambda: pg_upsert_many_dict(con, table, [{"id": 20, "dims": {"w": 3, "h": 3}}, {"id": 22, "dims": None}], ["id"], **kw), + ) + await check( + "pg_upsert_many_dict(must_exist)", + lambda: pg_upsert_many_dict(con, table, [{"id": 20, "dims": {"w": 4, "h": 4}}], ["id"], must_exist=True, **kw), + ) + await con.execute("DELETE FROM gadget WHERE id >= 10") + + # the data written through the helper is right nevertheless + await pg_insert(con, table, {"id": 30, "mood": "sad", "dims": {"w": 8, "h": 9}}, complex_helper="auto") + got = await pg_retrieve(con, Gadget, {"id": 30}, complex_helper="auto") + assert got == Gadget(id=30, mood="sad", dims={"w": 8, "h": 9}, note=None) + assert await _plain_read(con) == before + + +@requires_podman +async def test_issue_54_snippet_enum_label_survives_pg_retrieve_many(complex_types_test_db): + async with await _connect() as con: + await _setup_gadget(con) + + async def plain() -> tuple: + return await (await con.execute("select 'happy'::mood")).fetchone() # type: ignore[return-value] + + assert await plain() == ("happy",) + await pg_retrieve_many(con, Gadget, {}, limit=1, complex_helper=ComplexHelper(con)) + assert await plain() == ("happy",) # not + row = await (await con.execute("select 'happy'::mood, ROW(1, 2)::dims")).fetchone() + assert row == ("happy", "(1,2)") # a composite is still its text form, not a namedtuple + + +@requires_podman +async def test_registration_is_scoped_to_the_cursor_that_writes(complex_types_test_db): + # the public way to use `recursive_convert` yourself: register on the cursor that runs the statement + async with await _connect() as con: + await con.execute("CREATE TYPE mood AS ENUM ('happy', 'sad')") + await con.execute("CREATE TYPE dims AS (w int, h mood)") + await con.execute("CREATE TABLE gadget (id serial PRIMARY KEY, moods mood[], d dims, d2 dims[])") + helper = ComplexHelper(con) + types = await helper.load_all_complex_types(("public", "gadget")) + async with con.cursor() as cur: + for info in types.values(): + await helper.register_on(cur, info) + await cur.execute( + "INSERT INTO gadget (moods, d, d2) VALUES (%s, %s, %s)", + ( + await helper.recursive_convert(["happy", "sad"], types["moods"], con), + await helper.recursive_convert({"w": 1, "h": "sad"}, types["d"], con), + await helper.recursive_convert([{"w": 2, "h": "happy"}], types["d2"], con), + ), + ) + await cur.execute("SELECT moods, d, d2 FROM gadget") + moods, d, d2 = await cur.fetchone() # type: ignore[misc] + assert [m.name for m in moods] == ["happy", "sad"] + assert (d.w, d.h.name) == (1, "sad") and (d2[0].w, d2[0].h.name) == (2, "happy") # nested types registered too + # a plain SELECT on the connection (or on another cursor) is untouched + assert await (await con.execute("SELECT moods, d FROM gadget")).fetchone() == ("{happy,sad}", "(1,sad)") + # a cursor that did not register the types can't dump what the helper converted + converted = await helper.recursive_convert({"w": 1, "h": "sad"}, types["d"], con) + with pytest.raises(psycopg.ProgrammingError): + await con.execute("INSERT INTO gadget (d) VALUES (%s)", (converted,)) diff --git a/tests/db/test_crud.py b/tests/db/test_crud.py index a5ca738..1ef3cb8 100644 --- a/tests/db/test_crud.py +++ b/tests/db/test_crud.py @@ -337,3 +337,35 @@ async def test_pg_retrieve_many_where_with_complex_helper(pool: PgPool): 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")] + + +@requires_podman +async def test_pooled_connection_is_unchanged_after_pg_helpers_with_complex_helper(pool: PgPool): + # Fixes #54: type registration used to stick to the physical connection and so to the next pool user. + single = PgPool(env_prefix=ENV_PREFIX, max_size=1, min_size=1) + await single.open() + try: + plain = "SELECT 'active'::gadget_status, ROW('a', 'b')::label_pair" + + async def plain_read() -> tuple[int, tuple]: + async with single.connection() as con: + row = await (await con.execute(plain)).fetchone() + return con.pgconn.backend_pid, row # type: ignore[return-value] + + pid, before = await plain_read() + assert before == ("active", "(a,b)") + + async with single.connection() as con: + assert con.pgconn.backend_pid == pid # one physical connection: the next user gets this one again + helper = ComplexHelper(con) + await pg_insert(con, ("public", "gadget"), {"status": "active"}, complex_helper=helper) + await pg_insert(con, ("public", "gizmo"), {"label": {"en": "Hi", "de": "Hallo"}}, complex_helper="auto") + await pg_upsert_many_dict(con, ("public", "gizmo"), [{"id": 1, "label": {"en": "a", "de": "b"}}], ["id"], complex_helper=helper) + assert [g.status for g in await pg_retrieve_many(con, Gadget, {}, complex_helper=helper)] == ["active"] + assert (await pg_retrieve(con, Gizmo, {"id": 1}, complex_helper="auto")) is not None + # inside the same checkout, too + assert await (await con.execute(plain)).fetchone() == before + + assert await plain_read() == (pid, before) # after release and re-acquire + finally: + await single.close() From 0121fedf80ef4ae257849dbfb4e5764b226e3b5f Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Fri, 9 Oct 2026 18:06:13 +0000 Subject: [PATCH 2/6] fix(db): ComplexHelper registers enum/composite adapters on the writing cursor only, never on the connection (Fixes #54) Bump to 0.15.1. Co-Authored-By: Claude Sonnet 5.5 --- README.md | 13 +++- pgdevkit/db/complex_types.py | 76 ++++++++++++------- pgdevkit/db/crud.py | 129 ++++++++++++++++++--------------- pgdevkit/testdb/schema.py | 3 + pyproject.toml | 2 +- skills/pgdevkit/SKILL.md | 11 ++- tests/db/test_complex_types.py | 8 +- tests/testdb/test_schema.py | 5 +- uv.lock | 2 +- 9 files changed, 152 insertions(+), 97 deletions(-) diff --git a/README.md b/README.md index 099a959..3fa0e53 100644 --- a/README.md +++ b/README.md @@ -446,8 +446,17 @@ Install with the `db` extra: `pip install pgdevkit[db]`. `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 + to any `pg_*` function that takes the argument. **Nothing is registered on the connection:** psycopg's type + registration is global to a connection (and so to the next user of a pooled one: a plain `SELECT` of an enum would + return psycopg's auto-numbered `Enum` member instead of the label, #54), so the helpers only register the + enum/composite adapters on the *cursor* that runs a write statement (insert/update/upsert, `*_many`), and only + when a dict/list value had to be converted. Reads (`pg_retrieve*`) select such columns as `to_jsonb(...)` and need + no registration. After any `pg_*` call the connection reads enums as label strings and composites as before. + (`RETURNING` rows of `pg_insert`/`pg_update_dict`/`pg_upsert_dict` come from that same cursor, so a column whose + value was converted comes back as the registered Python object, e.g. a composite as a namedtuple.) If you call + `recursive_convert` yourself, register the types on the cursor that sends the values: + `await helper.register_on(cur, info)` (passing the connection instead restores psycopg's global behaviour). + What it learns from the catalog (types and column types) 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 diff --git a/pgdevkit/db/complex_types.py b/pgdevkit/db/complex_types.py index d0e8a1f..dfdb77a 100644 --- a/pgdevkit/db/complex_types.py +++ b/pgdevkit/db/complex_types.py @@ -5,6 +5,9 @@ from typing import Any, Callable from psycopg import AsyncConnection +from psycopg import adapters as _psycopg_adapters +from psycopg.abc import AdaptContext +from psycopg.adapt import AdaptersMap from psycopg.rows import dict_row from psycopg.sql import Identifier from psycopg.types.composite import CompositeInfo, register_composite @@ -28,12 +31,11 @@ class JsonbArray: @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.""" + CompositeInfo/EnumInfo objects built from them (they carry that database's OIDs) and each table column's type. + Never "what is registered where": nothing is registered on the connection (see `ComplexHelper.register_on`).""" 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) @@ -52,6 +54,15 @@ def _cache_for(con: AsyncConnection) -> _TypeCache: return _TypeCache() +class _Scope: + """An `AdaptContext` of its own (a copy of psycopg's defaults, shared with nothing).""" + + connection = None + + def __init__(self) -> None: + self.adapters = AdaptersMap(_psycopg_adapters) + + class ComplexHelper: """psycopg adapter for PostgreSQL composite types, enums, and JSONB. @@ -65,7 +76,14 @@ class ComplexHelper: 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 + **It never registers anything on the connection.** psycopg's adapter registration is global for the connection + (and so for the next user of a pooled connection: a plain `SELECT` of an enum would start returning psycopg's + `Enum` members instead of the label, see issue #54). What `recursive_convert` returns is sent through a cursor + that has the types registered for itself only (`register_on(cursor, info)`; the `pg_*` helpers do this for the + cursor that runs their statement), and the connection behaves exactly as before. Reads need no registration at + all: the `pg_retrieve*` helpers select complex columns as `to_jsonb(...)`. + + What it learns from the catalog (the types and the columns' types) 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 @@ -104,19 +122,10 @@ def complex_types(self) -> dict[tuple[str, str], CompositeInfo | EnumInfo]: 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() @@ -154,9 +163,7 @@ async def _load_complex_type_from_colinfos(self, res: dict[str, Any] | None) -> return None udt_schema: str = res["udt_schema"] udt_name: str = res["udt_name"] - c = await self._get_complex_type(f"{udt_schema}.{udt_name}", res["is_enum"], self.con) - await self._recurse_register(c, self.con) - return c + return await self._get_complex_type(f"{udt_schema}.{udt_name}", res["is_enum"], self.con) async def load_all_complex_types( self, table_name: tuple[str, str], include_generated: bool = False @@ -236,24 +243,43 @@ async def _get_complex_type(self, name: str, is_enum: bool, con: AsyncConnection if (schema, type_name) not in self.complex_types: 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" + # `recursive_convert` needs the Python types (`ci.enum`, `ci.python_type`) that psycopg only creates when + # a type is registered: do that in a scratch scope that no connection or cursor ever sees. + scratch = _Scope() + if isinstance(ci, EnumInfo): + register_enum(ci, scratch) + else: + register_composite(ci, scratch) self.complex_types[(schema, type_name)] = ci return self.complex_types[(schema, type_name)] - async def _recurse_register(self, info: CompositeInfo | EnumInfo, con: AsyncConnection) -> None: + async def register_on(self, context: AdaptContext, info: ComplexTypeInfo) -> None: + """Register the psycopg adapters for `info` (and every composite/enum nested in it) in `context`, so that + what `recursive_convert` returns for it can be sent by statements run through that context, and values of + those types come back as those Python objects. + + Pass the **cursor** that runs the statement: psycopg gives each cursor its own copy of the adapters, so the + registration lasts as long as the cursor and nothing else on the connection (hence on a pooled one) is + affected. Passing the connection itself registers globally for that connection, for good (psycopg's + default); that is not what the `pg_*` helpers do. A no-op for JSONB (`info` being `Jsonb`/`JsonbArray`) and + for `None`.""" + if isinstance(info, (CompositeInfo, EnumInfo)): + await self._recurse_register(info, context, set()) + + async def _recurse_register(self, info: CompositeInfo | EnumInfo, context: AdaptContext, seen: set) -> None: assert self.system_complex_type_dict is not None, "System complex type dictionary not loaded" - if info not in self.registered: - if isinstance(info, EnumInfo): - register_enum(info, con) - else: - register_composite(info, con) - self.registered.add(info) + if info in seen: + return + seen.add(info) if isinstance(info, EnumInfo): + register_enum(info, context, enum=info.enum) return + register_composite(info, context, factory=info.python_type) for t in info.field_types: if t in self.system_complex_type_dict: name, typtype = self.system_complex_type_dict[t] - ci = await self._get_complex_type(name, typtype == "e", con) - await self._recurse_register(ci, con) + ci = await self._get_complex_type(name, typtype == "e", self.con) + await self._recurse_register(ci, context, seen) async def recursive_convert( self, diff --git a/pgdevkit/db/crud.py b/pgdevkit/db/crud.py index b957d7c..ae7fed6 100644 --- a/pgdevkit/db/crud.py +++ b/pgdevkit/db/crud.py @@ -3,6 +3,7 @@ from string.templatelib import Interpolation, Template from typing import Any, Callable, Literal, Mapping, Optional, Sequence, Type, TypeVar +from psycopg import AsyncCursor from psycopg.connection_async import AsyncConnection from psycopg.rows import dict_row from psycopg.sql import Literal as SqlLiteral @@ -44,12 +45,14 @@ async def _select_list( async def _convert_complex_values( con: AsyncConnection, + cur: AsyncCursor[Any], table_name: tuple[str, str], data: dict, 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) + psycopg types those columns need, and register those types on `cur` (the cursor that will run the + statement) only, never on the connection. A no-op (no extra query) unless `data` actually contains dict/list values.""" helper = _resolve_helper(con, complex_helper) if helper is None: @@ -61,12 +64,14 @@ async def _convert_complex_values( for k in candidate_keys: info = await helper.load_complex_type(table_name, k) if info is not None: + await helper.register_on(cur, info) converted[k] = await helper.recursive_convert(data[k], info, con) return converted async def _convert_complex_values_many( con: AsyncConnection, + cur: AsyncCursor[Any], table_name: tuple[str, str], rows: Sequence[dict], complex_helper: ComplexHelperArg, @@ -81,6 +86,8 @@ async def _convert_complex_values_many( complex_keys = {k for k, info in infos.items() if info is not None} if not complex_keys: return rows + for k in complex_keys: + await helper.register_on(cur, infos[k]) converted_rows = [] for row in rows: new_row = dict(row) @@ -225,13 +232,13 @@ async def pg_insert( 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) - query = SQL("INSERT INTO {tbl} ({cols}) VALUES ({vals}) RETURNING *").format( - tbl=Identifier(*table_name), - cols=SQL(", ").join(Identifier(k) for k in data), - vals=SQL(", ").join(Placeholder(k) for k in data), - ) async with con.cursor(row_factory=dict_row) as cur: + data = await _convert_complex_values(con, cur, table_name, data, complex_helper) + query = SQL("INSERT INTO {tbl} ({cols}) VALUES ({vals}) RETURNING *").format( + tbl=Identifier(*table_name), + cols=SQL(", ").join(Identifier(k) for k in data), + vals=SQL(", ").join(Placeholder(k) for k in data), + ) await cur.execute(query, data) row = await cur.fetchone() assert row is not None @@ -251,17 +258,17 @@ async def pg_update_dict( Pass `complex_helper` when the table has composite/enum columns among the values being set — omitted, this behaves exactly as before (plain values passed straight through to psycopg).""" - data = await _convert_complex_values(con, table_name, data, complex_helper) - set_parts = [ - SQL("{col} = {val}").format(col=Identifier(k), val=Placeholder(k)) for k in data if k not in primary_keys - ] - where_parts = [SQL("{col} = {val}").format(col=Identifier(pk), val=Placeholder(pk)) for pk in primary_keys] - query = SQL("UPDATE {tbl} SET {sets} WHERE {where} RETURNING *").format( - tbl=Identifier(*table_name), - sets=SQL(", ").join(set_parts), - where=SQL(" AND ").join(where_parts), - ) async with con.cursor() as cur: + data = await _convert_complex_values(con, cur, table_name, data, complex_helper) + set_parts = [ + SQL("{col} = {val}").format(col=Identifier(k), val=Placeholder(k)) for k in data if k not in primary_keys + ] + where_parts = [SQL("{col} = {val}").format(col=Identifier(pk), val=Placeholder(pk)) for pk in primary_keys] + query = SQL("UPDATE {tbl} SET {sets} WHERE {where} RETURNING *").format( + tbl=Identifier(*table_name), + sets=SQL(", ").join(set_parts), + where=SQL(" AND ").join(where_parts), + ) await cur.execute(query, data) return await cur.fetchone() @@ -284,19 +291,19 @@ async def pg_upsert_dict( 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) - fields = list(data) - updates = [SQL("{col} = EXCLUDED.{col}").format(col=Identifier(k)) for k in fields] - query = SQL( - "INSERT INTO {tbl} ({cols}) VALUES ({vals}) ON CONFLICT ({pks}) DO UPDATE SET {updates} RETURNING *" - ).format( - tbl=Identifier(*table_name), - cols=SQL(", ").join(Identifier(k) for k in fields), - vals=SQL(", ").join(Placeholder(k) for k in fields), - pks=SQL(", ").join(Identifier(pk) for pk in primary_keys), - updates=SQL(", ").join(updates), - ) async with con.cursor(row_factory=dict_row) as cur: + data = await _convert_complex_values(con, cur, table_name, data, complex_helper) + fields = list(data) + updates = [SQL("{col} = EXCLUDED.{col}").format(col=Identifier(k)) for k in fields] + query = SQL( + "INSERT INTO {tbl} ({cols}) VALUES ({vals}) ON CONFLICT ({pks}) DO UPDATE SET {updates} RETURNING *" + ).format( + tbl=Identifier(*table_name), + cols=SQL(", ").join(Identifier(k) for k in fields), + vals=SQL(", ").join(Placeholder(k) for k in fields), + pks=SQL(", ").join(Identifier(pk) for pk in primary_keys), + updates=SQL(", ").join(updates), + ) await cur.execute(query, data) row = await cur.fetchone() assert row is not None @@ -328,30 +335,34 @@ async def pg_upsert_many_dict( want a missing row to be a silent no-op rather than create one.""" if not data: return - data = await _convert_complex_values_many(con, table_name, data, complex_helper) - fields = list(data[0]) - if must_exist: - update_assignments = [ - SQL("{col} = {val}").format(col=Identifier(k), val=Placeholder(k)) for k in fields if k not in primary_keys - ] - target_eq = SQL(" AND ").join( - SQL("t.{col} = {val}").format(col=Identifier(pk), val=Placeholder(pk)) for pk in primary_keys - ) - query = SQL("UPDATE {tbl} t SET {updates} WHERE {target_eq}").format( - tbl=Identifier(*table_name), - updates=SQL(", ").join(update_assignments), - target_eq=target_eq, - ) - else: - updates = [SQL("{col} = EXCLUDED.{col}").format(col=Identifier(k)) for k in fields if k not in primary_keys] - query = SQL("INSERT INTO {tbl} ({cols}) VALUES ({vals}) ON CONFLICT ({pks}) DO UPDATE SET {updates}").format( - tbl=Identifier(*table_name), - cols=SQL(", ").join(Identifier(k) for k in fields), - vals=SQL(", ").join(Placeholder(k) for k in fields), - pks=SQL(", ").join(Identifier(pk) for pk in primary_keys), - updates=SQL(", ").join(updates), - ) async with con.cursor() as cur: + data = await _convert_complex_values_many(con, cur, table_name, data, complex_helper) + fields = list(data[0]) + if must_exist: + update_assignments = [ + SQL("{col} = {val}").format(col=Identifier(k), val=Placeholder(k)) + for k in fields + if k not in primary_keys + ] + target_eq = SQL(" AND ").join( + SQL("t.{col} = {val}").format(col=Identifier(pk), val=Placeholder(pk)) for pk in primary_keys + ) + query = SQL("UPDATE {tbl} t SET {updates} WHERE {target_eq}").format( + tbl=Identifier(*table_name), + updates=SQL(", ").join(update_assignments), + target_eq=target_eq, + ) + else: + updates = [SQL("{col} = EXCLUDED.{col}").format(col=Identifier(k)) for k in fields if k not in primary_keys] + query = SQL( + "INSERT INTO {tbl} ({cols}) VALUES ({vals}) ON CONFLICT ({pks}) DO UPDATE SET {updates}" + ).format( + tbl=Identifier(*table_name), + cols=SQL(", ").join(Identifier(k) for k in fields), + vals=SQL(", ").join(Placeholder(k) for k in fields), + pks=SQL(", ").join(Identifier(pk) for pk in primary_keys), + updates=SQL(", ").join(updates), + ) await cur.executemany(query, data) @@ -378,14 +389,14 @@ async def pg_insert_many( if not data: return dict_data = [d if isinstance(d, dict) else d.model_dump() for d in data] - dict_data = await _convert_complex_values_many(con, table_name, dict_data, complex_helper) - fields = list(dict_data[0]) - query = SQL("INSERT INTO {tbl} ({cols}) VALUES ({vals})").format( - tbl=Identifier(*table_name), - cols=SQL(", ").join(Identifier(k) for k in fields), - vals=SQL(", ").join(Placeholder(k) for k in fields), - ) async with con.cursor() as cur: + dict_data = await _convert_complex_values_many(con, cur, table_name, dict_data, complex_helper) + fields = list(dict_data[0]) + query = SQL("INSERT INTO {tbl} ({cols}) VALUES ({vals})").format( + tbl=Identifier(*table_name), + cols=SQL(", ").join(Identifier(k) for k in fields), + vals=SQL(", ").join(Placeholder(k) for k in fields), + ) await cur.executemany(query, dict_data) diff --git a/pgdevkit/testdb/schema.py b/pgdevkit/testdb/schema.py index 9cf70b6..0691726 100644 --- a/pgdevkit/testdb/schema.py +++ b/pgdevkit/testdb/schema.py @@ -260,6 +260,9 @@ async def _insert_test_data( await copy.write_row(tuple(row[c] for c in col_names)) return + for col in col_names: + # registered on this cursor only, never on the (possibly pooled) connection + await complex_helper.register_on(cur, complex_types.get(col)) for row in rows: for col in col_names: info = complex_types.get(col) diff --git a/pyproject.toml b/pyproject.toml index 6c31d27..f38552f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ packages = ["pgdevkit"] [project] name = "pgdevkit" -version = "0.15.0" +version = "0.15.1" 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 123efec..da030b0 100644 --- a/skills/pgdevkit/SKILL.md +++ b/skills/pgdevkit/SKILL.md @@ -268,7 +268,16 @@ when you pass `complex_helper="auto"` (or `complex_helper=ComplexHelper(con, nor 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 +**The helper never registers types on the connection** (psycopg's registration is connection-global and would stick +to a pooled connection: a plain `SELECT` of an enum would return psycopg's `Enum` member instead of the label, #54). +Writes register the enum/composite adapters on the one cursor that runs the statement, only when a dict/list value +needed converting; reads (`pg_retrieve*`) use `to_jsonb(...)` and register nothing. So after any `pg_*` call the +connection reads enums as label strings and composites as before. `RETURNING` rows of `pg_insert`/`pg_update_dict`/ +`pg_upsert_dict` come from that cursor: a converted column is the registered Python object (composite namedtuple). +Using `recursive_convert` yourself: `await helper.register_on(cur, info)` on your cursor first (the connection as +context would be a global registration again). + +What the helper learns from the catalog (types and column types) 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 diff --git a/tests/db/test_complex_types.py b/tests/db/test_complex_types.py index 3083a99..fa19f5c 100644 --- a/tests/db/test_complex_types.py +++ b/tests/db/test_complex_types.py @@ -53,8 +53,9 @@ async def test_array_of_enum_column_is_detected_and_converted(complex_types_test assert isinstance(info, EnumInfo) converted = await helper.recursive_convert(["happy", "sad"], info, con) - await con.execute("INSERT INTO gadget (id, moods) VALUES (1, %s)", (converted,)) async with con.cursor() as cur: + await helper.register_on(cur, info) # on this cursor only, never on the connection (#54) + await cur.execute("INSERT INTO gadget (id, moods) VALUES (1, %s)", (converted,)) await cur.execute("SELECT moods FROM gadget WHERE id = 1") (moods,) = await cur.fetchone() assert [m.name for m in moods] == ["happy", "sad"] @@ -127,7 +128,6 @@ async def test_jsonb_array_column_wraps_each_element_not_the_whole_list(complex_ 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 == {} @@ -385,7 +385,3 @@ async def test_registration_is_scoped_to_the_cursor_that_writes(complex_types_te assert (d.w, d.h.name) == (1, "sad") and (d2[0].w, d2[0].h.name) == (2, "happy") # nested types registered too # a plain SELECT on the connection (or on another cursor) is untouched assert await (await con.execute("SELECT moods, d FROM gadget")).fetchone() == ("{happy,sad}", "(1,sad)") - # a cursor that did not register the types can't dump what the helper converted - converted = await helper.recursive_convert({"w": 1, "h": "sad"}, types["d"], con) - with pytest.raises(psycopg.ProgrammingError): - await con.execute("INSERT INTO gadget (d) VALUES (%s)", (converted,)) diff --git a/tests/testdb/test_schema.py b/tests/testdb/test_schema.py index 5efbfd0..b7f5e83 100644 --- a/tests/testdb/test_schema.py +++ b/tests/testdb/test_schema.py @@ -137,8 +137,9 @@ async def test_apply_schema_seeds_composite_enum_and_jsonb_columns(schema_test_d async with con.cursor() as cur: await cur.execute("SELECT mood, size, tags FROM app.gadget WHERE id = 1") mood, size, tags = await cur.fetchone() - assert mood.name == "happy" - assert size == (10, 20) + # read back with a plain SELECT: seeding must not leave psycopg type registrations on the connection (#54) + assert mood == "happy" + assert size == "(10,20)" assert tags == ["small", "shiny"] diff --git a/uv.lock b/uv.lock index 3dcbd74..21dd842 100644 --- a/uv.lock +++ b/uv.lock @@ -451,7 +451,7 @@ wheels = [ [[package]] name = "pgdevkit" -version = "0.15.0" +version = "0.15.1" source = { editable = "." } dependencies = [ { name = "docker" }, From e177adaa76aed07887e02886c57c62c70364d581 Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Fri, 9 Oct 2026 18:08:33 +0000 Subject: [PATCH 3/6] fix(db): keep deprecated ComplexHelper.registered shim, reuse composite loader classes, docs wording, more tests Co-Authored-By: Claude Sonnet 5.5 --- README.md | 23 +++++++++-------- pgdevkit/db/complex_types.py | 47 ++++++++++++++++++++++++++++++---- pgdevkit/db/crud.py | 19 ++++++++++++-- pgdevkit/testdb/schema.py | 4 +-- skills/pgdevkit/SKILL.md | 19 ++++++++------ tests/db/test_complex_types.py | 27 ++++++++++++++++--- 6 files changed, 108 insertions(+), 31 deletions(-) diff --git a/README.md b/README.md index 3fa0e53..423b450 100644 --- a/README.md +++ b/README.md @@ -448,16 +448,19 @@ Install with the `db` extra: `pip install pgdevkit[db]`. `complex_helper=ComplexHelper(con)` (optionally with `normalizers`), or `complex_helper="auto"` for a default one, to any `pg_*` function that takes the argument. **Nothing is registered on the connection:** psycopg's type registration is global to a connection (and so to the next user of a pooled one: a plain `SELECT` of an enum would - return psycopg's auto-numbered `Enum` member instead of the label, #54), so the helpers only register the - enum/composite adapters on the *cursor* that runs a write statement (insert/update/upsert, `*_many`), and only - when a dict/list value had to be converted. Reads (`pg_retrieve*`) select such columns as `to_jsonb(...)` and need - no registration. After any `pg_*` call the connection reads enums as label strings and composites as before. - (`RETURNING` rows of `pg_insert`/`pg_update_dict`/`pg_upsert_dict` come from that same cursor, so a column whose - value was converted comes back as the registered Python object, e.g. a composite as a namedtuple.) If you call - `recursive_convert` yourself, register the types on the cursor that sends the values: - `await helper.register_on(cur, info)` (passing the connection instead restores psycopg's global behaviour). - What it learns from the catalog (types and column types) 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 + return psycopg's auto-numbered `Enum` member instead of the label, #54). The helpers therefore register the + enum/composite adapters only on the *cursor* that runs a write statement (`pg_insert`, `pg_update`/`pg_update_dict`, + `pg_upsert`/`pg_upsert_dict`, the `*_many` variants), and only for the columns where a dict/list value had to + be converted. Reads (`pg_retrieve*`) select such columns as `to_jsonb(...)` and need no registration. After any + `pg_*` call the connection reads enums as label strings and composites as text (e.g. `'(1,2)'`), as before the call. + `RETURNING` rows of `pg_insert`/`pg_update_dict`/`pg_upsert_dict` (and so `pg_update`/`pg_upsert`) come from that + same cursor: columns whose type, or a type nested in it, was registered on it come back as the registered Python + objects (a composite as a namedtuple, an enum as an enum member), all others as psycopg returns them by default. + If you call `recursive_convert` yourself, register the types on the cursor that sends the values, with the `info` + from `load_all_complex_types`/`load_complex_type`: `await helper.register_on(cur, info)`, then execute the converted + value on that same `cur` (passing the connection instead would restore psycopg's global behaviour). + What it learns from the catalog (types and column types) 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. diff --git a/pgdevkit/db/complex_types.py b/pgdevkit/db/complex_types.py index dfdb77a..ee40ec9 100644 --- a/pgdevkit/db/complex_types.py +++ b/pgdevkit/db/complex_types.py @@ -1,5 +1,6 @@ from __future__ import annotations +import warnings import weakref from dataclasses import dataclass, field from typing import Any, Callable @@ -54,6 +55,13 @@ def _cache_for(con: AsyncConnection) -> _TypeCache: return _TypeCache() +def _make_composite_object(values: Any, info: CompositeInfo) -> Any: + # One module-level function, so psycopg's cache of generated loader classes (keyed on it) hits on re-registration + # instead of growing by two classes per registration. + assert info.python_type is not None + return info.python_type(*values) + + class _Scope: """An `AdaptContext` of its own (a copy of psycopg's defaults, shared with nothing).""" @@ -67,8 +75,8 @@ class ComplexHelper: """psycopg adapter for PostgreSQL composite types, enums, and JSONB. Detects a table's non-scalar columns (composite types, enums, JSONB) and - converts plain dict/list Python values into the psycopg-registered types - those columns need, recursing into nested composite fields. + converts plain dict/list Python values into the Python objects (enum members, composite namedtuples) + psycopg sends for those columns, recursing into nested composite fields. `normalizers` lets a caller reshape a composite value before conversion, keyed by composite type name (e.g. a project with a `locale_labels` @@ -122,6 +130,27 @@ def complex_types(self) -> dict[tuple[str, str], CompositeInfo | EnumInfo]: def complex_types(self, value: dict[tuple[str, str], CompositeInfo | EnumInfo]) -> None: self._cache.complex_types = value + @property + def registered(self) -> set[CompositeInfo | EnumInfo]: + """Deprecated: nothing is registered on the connection any more, so this is always empty (a new set each + time). Use `register_on(cursor, info)` to register types for a statement.""" + warnings.warn( + "ComplexHelper.registered is deprecated: types are no longer registered on the connection; " + "use `await helper.register_on(cursor, info)` to register them on a cursor", + DeprecationWarning, + stacklevel=2, + ) + return set() + + @registered.setter + def registered(self, value: set[CompositeInfo | EnumInfo]) -> None: + warnings.warn( + "ComplexHelper.registered is deprecated and ignored (types are no longer registered on the connection; " + "see `register_on`)", + DeprecationWarning, + stacklevel=2, + ) + 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 @@ -249,7 +278,7 @@ async def _get_complex_type(self, name: str, is_enum: bool, con: AsyncConnection if isinstance(ci, EnumInfo): register_enum(ci, scratch) else: - register_composite(ci, scratch) + register_composite(ci, scratch, make_object=_make_composite_object) self.complex_types[(schema, type_name)] = ci return self.complex_types[(schema, type_name)] @@ -258,15 +287,23 @@ async def register_on(self, context: AdaptContext, info: ComplexTypeInfo) -> Non what `recursive_convert` returns for it can be sent by statements run through that context, and values of those types come back as those Python objects. + `info` is what `load_all_complex_types(table)` / `load_complex_type(table, column)` return for the column + you are writing, and the value `recursive_convert(value, info, con)` returns must be executed on the same + `context` (the cursor) that you registered it on. + Pass the **cursor** that runs the statement: psycopg gives each cursor its own copy of the adapters, so the registration lasts as long as the cursor and nothing else on the connection (hence on a pooled one) is affected. Passing the connection itself registers globally for that connection, for good (psycopg's default); that is not what the `pg_*` helpers do. A no-op for JSONB (`info` being `Jsonb`/`JsonbArray`) and for `None`.""" if isinstance(info, (CompositeInfo, EnumInfo)): + if self.system_complex_type_dict is None: + await self.load_complex_type_dict() await self._recurse_register(info, context, set()) - async def _recurse_register(self, info: CompositeInfo | EnumInfo, context: AdaptContext, seen: set) -> None: + async def _recurse_register( + self, info: CompositeInfo | EnumInfo, context: AdaptContext, seen: set[CompositeInfo | EnumInfo] + ) -> None: assert self.system_complex_type_dict is not None, "System complex type dictionary not loaded" if info in seen: return @@ -274,7 +311,7 @@ async def _recurse_register(self, info: CompositeInfo | EnumInfo, context: Adapt if isinstance(info, EnumInfo): register_enum(info, context, enum=info.enum) return - register_composite(info, context, factory=info.python_type) + register_composite(info, context, factory=info.python_type, make_object=_make_composite_object) for t in info.field_types: if t in self.system_complex_type_dict: name, typtype = self.system_complex_type_dict[t] diff --git a/pgdevkit/db/crud.py b/pgdevkit/db/crud.py index ae7fed6..79f584b 100644 --- a/pgdevkit/db/crud.py +++ b/pgdevkit/db/crud.py @@ -231,7 +231,12 @@ async def pg_insert( *, complex_helper: ComplexHelperArg = None, ) -> dict[str, Any]: - """Insert one row and return the full row (RETURNING *).""" + """Insert one row and return the full row (RETURNING *). + + A column of a composite/enum type (or an array of one) whose value was dict/list-converted by `complex_helper` + comes back as the Python object registered on this statement's cursor (a composite as a namedtuple, an enum as an + enum member); every other column comes back as plain psycopg returns it (an enum as its label, a composite as its + text, e.g. '(1,2)'). Nothing is registered on `con`.""" async with con.cursor(row_factory=dict_row) as cur: data = await _convert_complex_values(con, cur, table_name, data, complex_helper) query = SQL("INSERT INTO {tbl} ({cols}) VALUES ({vals}) RETURNING *").format( @@ -255,6 +260,11 @@ async def pg_update_dict( ) -> Any | None: """Update a row identified by primary_keys. Returns the raw row tuple. + Columns whose type, or a type nested in it, was registered on this statement's cursor to convert a value + come back as the registered Python objects (composites as namedtuples, enums as enum members), the others as + psycopg returns them by default (an enum as its label, a composite as text, e.g. '(1,2)'). Nothing is + registered on `con`. + Pass `complex_helper` when the table has composite/enum columns among the values being set — omitted, this behaves exactly as before (plain values passed straight through to psycopg).""" @@ -290,7 +300,12 @@ async def pg_upsert_dict( *, complex_helper: ComplexHelperArg = None, ) -> dict: - """INSERT ... ON CONFLICT ... DO UPDATE, returns the row as a dict.""" + """INSERT ... ON CONFLICT ... DO UPDATE, returns the row as a dict. + + Columns whose type, or a type nested in it, was registered on this statement's cursor to convert a value + come back as the registered Python objects (composites as namedtuples, enums as enum members), the others as + psycopg returns them by default (an enum as its label, a composite as text, e.g. '(1,2)'). Nothing is + registered on `con`.""" async with con.cursor(row_factory=dict_row) as cur: data = await _convert_complex_values(con, cur, table_name, data, complex_helper) fields = list(data) diff --git a/pgdevkit/testdb/schema.py b/pgdevkit/testdb/schema.py index 0691726..2e1692c 100644 --- a/pgdevkit/testdb/schema.py +++ b/pgdevkit/testdb/schema.py @@ -267,8 +267,8 @@ async def _insert_test_data( for col in col_names: info = complex_types.get(col) if info is not None and isinstance(row[col], (dict, list)): - # composite/enum/JSONB: needs psycopg-registered-type - # conversion. Anything else (plain scalars, and native + # composite/enum/JSONB: needs conversion to the + # Python objects registered on `cur` above. Anything else (plain scalars, and native # Postgres arrays like text[], which aren't "complex" — # psycopg already adapts a Python list to those natively) # is left untouched. diff --git a/skills/pgdevkit/SKILL.md b/skills/pgdevkit/SKILL.md index da030b0..2cbdfa9 100644 --- a/skills/pgdevkit/SKILL.md +++ b/skills/pgdevkit/SKILL.md @@ -262,20 +262,23 @@ See [`references/temporal-tables.md`](references/temporal-tables.md) for row-lev ### Custom Postgres types (composites, enums) `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. `pgdevkit.testdb.schema`'s test-data seeding uses it automatically; the `pg_*` CRUD helpers only do +columns and converts plain dict/list values into the Python objects psycopg sends +for those columns. `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 *`. **The helper never registers types on the connection** (psycopg's registration is connection-global and would stick to a pooled connection: a plain `SELECT` of an enum would return psycopg's `Enum` member instead of the label, #54). -Writes register the enum/composite adapters on the one cursor that runs the statement, only when a dict/list value -needed converting; reads (`pg_retrieve*`) use `to_jsonb(...)` and register nothing. So after any `pg_*` call the -connection reads enums as label strings and composites as before. `RETURNING` rows of `pg_insert`/`pg_update_dict`/ -`pg_upsert_dict` come from that cursor: a converted column is the registered Python object (composite namedtuple). -Using `recursive_convert` yourself: `await helper.register_on(cur, info)` on your cursor first (the connection as -context would be a global registration again). +Writes register the enum/composite adapters on the one cursor that runs the statement, only for the columns where a +dict/list value needed converting; reads (`pg_retrieve*`) use `to_jsonb(...)` and register nothing. So after any `pg_*` call the +connection reads enums as label strings and composites as text (`'(1,2)'`), as before the call. `RETURNING` rows of +`pg_insert`/`pg_update_dict`/`pg_upsert_dict` (hence `pg_update`/`pg_upsert`) come from that cursor: columns whose +type, or a type nested in it, was registered on it come back as the registered Python objects (a composite as a +namedtuple, an enum as an enum member), all others as psycopg returns them by default. +Using `recursive_convert` yourself: take `info` from `load_all_complex_types`/`load_complex_type`, call +`await helper.register_on(cur, info)` on your cursor, and execute the converted value on that same `cur` (the +connection as context would be a global registration again). What the helper learns from the catalog (types and column types) is cached **per connection**, so creating a `ComplexHelper(con)` per call (`pg_retrieve(con, M, pks, complex_helper=ComplexHelper(con))`) queries the diff --git a/tests/db/test_complex_types.py b/tests/db/test_complex_types.py index fa19f5c..402db8d 100644 --- a/tests/db/test_complex_types.py +++ b/tests/db/test_complex_types.py @@ -6,7 +6,18 @@ import pytest from psycopg.types.enum import EnumInfo -from pgdevkit.db import PostgresTableModel, pg_insert, pg_retrieve, pg_retrieve_many +from pgdevkit.db import ( + PostgresTableModel, + pg_insert, + pg_insert_many, + pg_retrieve, + pg_retrieve_many, + pg_update, + pg_update_dict, + pg_upsert, + pg_upsert_dict, + pg_upsert_many_dict, +) from pgdevkit.db.complex_types import ComplexHelper from pgdevkit.db.crud import _select_list from pgdevkit.testdb import constants @@ -128,6 +139,10 @@ async def test_jsonb_array_column_wraps_each_element_not_the_whole_list(complex_ def test_helper_attributes_stay_assignable(): helper = ComplexHelper(con=None) # type: ignore[arg-type] helper.complex_types = {} + with pytest.warns(DeprecationWarning, match="register_on"): + helper.registered = set() # deprecated and ignored: nothing is registered on the connection any more + with pytest.warns(DeprecationWarning, match="register_on"): + assert helper.registered == set() helper.system_complex_type_dict = {} assert helper.system_complex_type_dict == {} @@ -301,8 +316,6 @@ async def _plain_read(con: psycopg.AsyncConnection): @requires_podman async def test_pg_helpers_do_not_change_how_the_connection_reads_enums_and_composites(complex_types_test_db): - from pgdevkit.db import pg_insert_many, pg_update, pg_update_dict, pg_upsert, pg_upsert_dict, pg_upsert_many_dict - async with await _connect() as con: await _setup_gadget(con) table = ("public", "gadget") @@ -337,6 +350,12 @@ async def check(what: str, call) -> None: ) await con.execute("DELETE FROM gadget WHERE id >= 10") + # an enum *array* (a list converted to enum members) written through the helper + await con.execute("ALTER TABLE gadget ADD COLUMN moods mood[]") + helper_arg = ComplexHelper(con, cache=False) # sees the new column + await check("pg_insert (enum array)", lambda: pg_insert(con, table, {"id": 40, "moods": ["happy", "sad"]}, complex_helper=helper_arg)) + assert (await (await con.execute("SELECT moods FROM gadget WHERE id = 40")).fetchone()) == ("{happy,sad}",) + # the data written through the helper is right nevertheless await pg_insert(con, table, {"id": 30, "mood": "sad", "dims": {"w": 8, "h": 9}}, complex_helper="auto") got = await pg_retrieve(con, Gadget, {"id": 30}, complex_helper="auto") @@ -345,7 +364,7 @@ async def check(what: str, call) -> None: @requires_podman -async def test_issue_54_snippet_enum_label_survives_pg_retrieve_many(complex_types_test_db): +async def test_enum_label_survives_pg_retrieve_many_on_the_same_connection(complex_types_test_db): async with await _connect() as con: await _setup_gadget(con) From a7a40b232ec1b2a6e882c499a4251d478b260603 Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Fri, 9 Oct 2026 18:08:42 +0000 Subject: [PATCH 4/6] test: composite loader class cache does not grow per write Co-Authored-By: Claude Sonnet 5.5 --- tests/db/test_complex_types.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/db/test_complex_types.py b/tests/db/test_complex_types.py index 402db8d..cd5660d 100644 --- a/tests/db/test_complex_types.py +++ b/tests/db/test_complex_types.py @@ -404,3 +404,17 @@ async def test_registration_is_scoped_to_the_cursor_that_writes(complex_types_te assert (d.w, d.h.name) == (1, "sad") and (d2[0].w, d2[0].h.name) == (2, "happy") # nested types registered too # a plain SELECT on the connection (or on another cursor) is untouched assert await (await con.execute("SELECT moods, d FROM gadget")).fetchone() == ("{happy,sad}", "(1,sad)") + + +@requires_podman +async def test_repeated_writes_do_not_grow_psycopgs_loader_class_cache(complex_types_test_db): + from psycopg.types import composite as psycopg_composite + + async with await _connect() as con: + await _setup_gadget(con) + row = {"id": 1, "dims": {"w": 1, "h": 1}} + await pg_upsert_dict(con, ("public", "gadget"), row, ["id"], complex_helper="auto") + size = psycopg_composite._make_loader.cache_info().currsize + for i in range(5): + await pg_upsert_dict(con, ("public", "gadget"), {**row, "id": i + 2}, ["id"], complex_helper="auto") + assert psycopg_composite._make_loader.cache_info().currsize == size # one registration per write, no new classes From 489a2092a5ac7cecdf3a6f2a122877442d6655ac Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Fri, 9 Oct 2026 18:10:31 +0000 Subject: [PATCH 5/6] fix(db): RETURNING of pg_insert/pg_update_dict/pg_upsert_dict uses to_jsonb for complex columns (deterministic plain values); enum[] inside a composite is detected as an enum Co-Authored-By: Claude Sonnet 5.5 --- README.md | 6 ++--- pgdevkit/db/complex_types.py | 8 +++++-- pgdevkit/db/crud.py | 39 +++++++++++++++---------------- skills/pgdevkit/SKILL.md | 6 ++--- tests/db/test_complex_types.py | 42 ++++++++++++++++++++++++++++++++-- tests/db/test_crud.py | 9 +++----- 6 files changed, 74 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index 423b450..3e1a5d7 100644 --- a/README.md +++ b/README.md @@ -453,9 +453,9 @@ Install with the `db` extra: `pip install pgdevkit[db]`. `pg_upsert`/`pg_upsert_dict`, the `*_many` variants), and only for the columns where a dict/list value had to be converted. Reads (`pg_retrieve*`) select such columns as `to_jsonb(...)` and need no registration. After any `pg_*` call the connection reads enums as label strings and composites as text (e.g. `'(1,2)'`), as before the call. - `RETURNING` rows of `pg_insert`/`pg_update_dict`/`pg_upsert_dict` (and so `pg_update`/`pg_upsert`) come from that - same cursor: columns whose type, or a type nested in it, was registered on it come back as the registered Python - objects (a composite as a namedtuple, an enum as an enum member), all others as psycopg returns them by default. + `RETURNING` rows of `pg_insert`/`pg_update_dict`/`pg_upsert_dict` (and so `pg_update`/`pg_upsert`) select such + columns as `to_jsonb(...)` too, so with a helper they are plain Python values (a composite as a dict, an enum as its + label) whichever columns were written; without one it is a plain `RETURNING *`. If you call `recursive_convert` yourself, register the types on the cursor that sends the values, with the `info` from `load_all_complex_types`/`load_complex_type`: `await helper.register_on(cur, info)`, then execute the converted value on that same `cur` (passing the connection instead would restore psycopg's global behaviour). diff --git a/pgdevkit/db/complex_types.py b/pgdevkit/db/complex_types.py index ee40ec9..fca83d8 100644 --- a/pgdevkit/db/complex_types.py +++ b/pgdevkit/db/complex_types.py @@ -163,7 +163,10 @@ async def load_complex_type_dict(self) -> None: await cur.execute(""" SELECT t.oid, pg_catalog.format_type ( t.oid, NULL ) AS obj_name, - t.typtype + -- an array type reports its element type's typtype, so an enum[] is known to be an enum + CASE WHEN t.typcategory = 'A' + THEN ( SELECT et.typtype FROM pg_catalog.pg_type et WHERE et.oid = t.typelem ) + ELSE t.typtype END AS typtype FROM pg_catalog.pg_type t JOIN pg_catalog.pg_namespace n ON n.oid = t.typnamespace @@ -215,7 +218,8 @@ async def load_all_complex_types( from information_schema.columns c left join enum_types e on e.enum_schema=c.udt_schema and (e.enum_name=c.udt_name or (c.data_type='ARRAY' and c.udt_name='_'||e.enum_name)) - where table_schema=%(schema)s and table_name = %(tbl)s and (is_generated <> 'ALWAYS' or %(include_generated)s)""" + where table_schema=%(schema)s and table_name = %(tbl)s and (is_generated <> 'ALWAYS' or %(include_generated)s) + order by c.ordinal_position""" async with self.con.cursor(row_factory=dict_row) as cur: await cur.execute( colquery, diff --git a/pgdevkit/db/crud.py b/pgdevkit/db/crud.py index 79f584b..42f5238 100644 --- a/pgdevkit/db/crud.py +++ b/pgdevkit/db/crud.py @@ -27,7 +27,7 @@ def _resolve_helper(con: AsyncConnection, complex_helper: ComplexHelperArg) -> C async def _select_list( con: AsyncConnection, table_name: tuple[str, str], complex_helper: ComplexHelperArg ) -> Composable: - """Column list for a SELECT, wrapping composite/enum/JSONB columns in + """Column list for a SELECT or RETURNING, 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.""" helper = _resolve_helper(con, complex_helper) @@ -231,18 +231,18 @@ async def pg_insert( *, complex_helper: ComplexHelperArg = None, ) -> dict[str, Any]: - """Insert one row and return the full row (RETURNING *). + """Insert one row and return the full row (`RETURNING *`). - A column of a composite/enum type (or an array of one) whose value was dict/list-converted by `complex_helper` - comes back as the Python object registered on this statement's cursor (a composite as a namedtuple, an enum as an - enum member); every other column comes back as plain psycopg returns it (an enum as its label, a composite as its - text, e.g. '(1,2)'). Nothing is registered on `con`.""" + With a `complex_helper`, composite/enum/JSONB columns are returned as `to_jsonb(...)`, exactly like + `pg_retrieve` does: plain Python values (a composite as a dict, an enum as its label), whichever columns were + written. Without one it is a plain `RETURNING *` as psycopg reads it.""" async with con.cursor(row_factory=dict_row) as cur: data = await _convert_complex_values(con, cur, table_name, data, complex_helper) - query = SQL("INSERT INTO {tbl} ({cols}) VALUES ({vals}) RETURNING *").format( + query = SQL("INSERT INTO {tbl} ({cols}) VALUES ({vals}) RETURNING {returning}").format( tbl=Identifier(*table_name), cols=SQL(", ").join(Identifier(k) for k in data), vals=SQL(", ").join(Placeholder(k) for k in data), + returning=await _select_list(con, table_name, complex_helper), ) await cur.execute(query, data) row = await cur.fetchone() @@ -258,26 +258,25 @@ async def pg_update_dict( *, complex_helper: ComplexHelperArg = None, ) -> Any | None: - """Update a row identified by primary_keys. Returns the raw row tuple. + """Update a row identified by primary_keys. Returns the raw row tuple (in column order). - Columns whose type, or a type nested in it, was registered on this statement's cursor to convert a value - come back as the registered Python objects (composites as namedtuples, enums as enum members), the others as - psycopg returns them by default (an enum as its label, a composite as text, e.g. '(1,2)'). Nothing is - registered on `con`. + With a `complex_helper`, composite/enum/JSONB columns in it are plain Python values, like `pg_retrieve` returns + them (a composite as a dict, an enum as its label), whichever columns were written; without one, plain psycopg + behaviour (`RETURNING *`). Pass `complex_helper` when the table has composite/enum columns among the - values being set — omitted, this behaves exactly as before (plain values - passed straight through to psycopg).""" + values being set — omitted, values are passed straight through to psycopg.""" async with con.cursor() as cur: data = await _convert_complex_values(con, cur, table_name, data, complex_helper) set_parts = [ SQL("{col} = {val}").format(col=Identifier(k), val=Placeholder(k)) for k in data if k not in primary_keys ] where_parts = [SQL("{col} = {val}").format(col=Identifier(pk), val=Placeholder(pk)) for pk in primary_keys] - query = SQL("UPDATE {tbl} SET {sets} WHERE {where} RETURNING *").format( + query = SQL("UPDATE {tbl} SET {sets} WHERE {where} RETURNING {returning}").format( tbl=Identifier(*table_name), sets=SQL(", ").join(set_parts), where=SQL(" AND ").join(where_parts), + returning=await _select_list(con, table_name, complex_helper), ) await cur.execute(query, data) return await cur.fetchone() @@ -302,22 +301,22 @@ async def pg_upsert_dict( ) -> dict: """INSERT ... ON CONFLICT ... DO UPDATE, returns the row as a dict. - Columns whose type, or a type nested in it, was registered on this statement's cursor to convert a value - come back as the registered Python objects (composites as namedtuples, enums as enum members), the others as - psycopg returns them by default (an enum as its label, a composite as text, e.g. '(1,2)'). Nothing is - registered on `con`.""" + With a `complex_helper`, composite/enum/JSONB columns in it are plain Python values, like `pg_retrieve` returns + them (a composite as a dict, an enum as its label), whichever columns were written; without one, plain psycopg + behaviour (`RETURNING *`).""" async with con.cursor(row_factory=dict_row) as cur: data = await _convert_complex_values(con, cur, table_name, data, complex_helper) fields = list(data) updates = [SQL("{col} = EXCLUDED.{col}").format(col=Identifier(k)) for k in fields] query = SQL( - "INSERT INTO {tbl} ({cols}) VALUES ({vals}) ON CONFLICT ({pks}) DO UPDATE SET {updates} RETURNING *" + "INSERT INTO {tbl} ({cols}) VALUES ({vals}) ON CONFLICT ({pks}) DO UPDATE SET {updates} RETURNING {returning}" ).format( tbl=Identifier(*table_name), cols=SQL(", ").join(Identifier(k) for k in fields), vals=SQL(", ").join(Placeholder(k) for k in fields), pks=SQL(", ").join(Identifier(pk) for pk in primary_keys), updates=SQL(", ").join(updates), + returning=await _select_list(con, table_name, complex_helper), ) await cur.execute(query, data) row = await cur.fetchone() diff --git a/skills/pgdevkit/SKILL.md b/skills/pgdevkit/SKILL.md index 2cbdfa9..e5ae837 100644 --- a/skills/pgdevkit/SKILL.md +++ b/skills/pgdevkit/SKILL.md @@ -273,9 +273,9 @@ to a pooled connection: a plain `SELECT` of an enum would return psycopg's `Enum Writes register the enum/composite adapters on the one cursor that runs the statement, only for the columns where a dict/list value needed converting; reads (`pg_retrieve*`) use `to_jsonb(...)` and register nothing. So after any `pg_*` call the connection reads enums as label strings and composites as text (`'(1,2)'`), as before the call. `RETURNING` rows of -`pg_insert`/`pg_update_dict`/`pg_upsert_dict` (hence `pg_update`/`pg_upsert`) come from that cursor: columns whose -type, or a type nested in it, was registered on it come back as the registered Python objects (a composite as a -namedtuple, an enum as an enum member), all others as psycopg returns them by default. +`pg_insert`/`pg_update_dict`/`pg_upsert_dict` (hence `pg_update`/`pg_upsert`) also use `to_jsonb(...)` for such +columns: with a helper they are plain Python values (a composite as a dict, an enum as its label) whichever columns +were written; without one it is a plain `RETURNING *`. Using `recursive_convert` yourself: take `info` from `load_all_complex_types`/`load_complex_type`, call `await helper.register_on(cur, info)` on your cursor, and execute the converted value on that same `cur` (the connection as context would be a global registration again). diff --git a/tests/db/test_complex_types.py b/tests/db/test_complex_types.py index cd5660d..8e3c586 100644 --- a/tests/db/test_complex_types.py +++ b/tests/db/test_complex_types.py @@ -237,7 +237,7 @@ async def retrieve(): 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] + assert row["dims"] == {"w": 1, "h": 2} # type: ignore[index] # RETURNING gives plain values, like pg_retrieve @requires_podman @@ -292,7 +292,7 @@ async def test_complex_helper_auto_works_for_every_crud_helper(complex_types_tes 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] + assert row["mood"] == "sad" and row["dims"] == {"w": 1, "h": 2} # type: ignore[index] 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") @@ -418,3 +418,41 @@ async def test_repeated_writes_do_not_grow_psycopgs_loader_class_cache(complex_t for i in range(5): await pg_upsert_dict(con, ("public", "gadget"), {**row, "id": i + 2}, ["id"], complex_helper="auto") assert psycopg_composite._make_loader.cache_info().currsize == size # one registration per write, no new classes + + +@requires_podman +async def test_returning_rows_are_plain_values_whichever_columns_were_written(complex_types_test_db): + async with await _connect() as con: + await _setup_gadget(con) + table = ("public", "gadget") + full = {"id": 1, "mood": "happy", "dims": {"w": 3, "h": 4}, "note": "a"} + for helper in ("auto", ComplexHelper(con)): + # only the composite written: the enum (not converted) and the composite come back the same way ... + r1 = await pg_upsert_dict(con, table, {"id": 1, "dims": {"w": 3, "h": 4}}, ["id"], complex_helper=helper) + # ... as when only the enum is written, or both, or neither + r2 = await pg_upsert_dict(con, table, {"id": 1, "mood": "happy"}, ["id"], complex_helper=helper) + r3 = await pg_upsert_dict(con, table, {"id": 1, "mood": "happy", "dims": {"w": 3, "h": 4}}, ["id"], complex_helper=helper) + r4 = await pg_upsert_dict(con, table, {"id": 1, "note": "a"}, ["id"], complex_helper=helper) + r5 = await pg_insert(con, table, {"id": 9, "note": "z"}, complex_helper=helper) + await con.execute("DELETE FROM gadget WHERE id = 9") + assert r1 == r2 == r3 == r4 == full + assert r5 == {"id": 9, "mood": None, "dims": None, "note": "z"} + u1 = await pg_update_dict(con, table, {"id": 1, "dims": {"w": 3, "h": 4}}, ["id"], complex_helper=helper) + u2 = await pg_update_dict(con, table, {"id": 1, "note": "a"}, ["id"], complex_helper=helper) + assert u1 == u2 == (1, "happy", {"w": 3, "h": 4}, "a") # the raw tuple, in column order + # without a helper nothing changes: RETURNING * as psycopg reads it + assert await pg_upsert_dict(con, table, {"id": 1, "note": "a"}, ["id"]) == { + "id": 1, "mood": "happy", "dims": "(3,4)", "note": "a" + } + + +@requires_podman +async def test_composite_whose_field_is_an_enum_array_is_converted(complex_types_test_db): + async with await _connect() as con: + await con.execute("CREATE TYPE mood AS ENUM ('happy', 'sad')") + await con.execute("CREATE TYPE onlyarr AS (ms mood[])") + await con.execute("CREATE TABLE holder (id int PRIMARY KEY, a onlyarr)") + row = await pg_insert(con, ("public", "holder"), {"id": 1, "a": {"ms": ["sad"]}}, complex_helper="auto") + assert row == {"id": 1, "a": {"ms": ["sad"]}} + got = await ComplexHelper(con).load_all_complex_types(("public", "holder")) + assert isinstance(got["a"], psycopg.types.composite.CompositeInfo) diff --git a/tests/db/test_crud.py b/tests/db/test_crud.py index 1ef3cb8..fd6f930 100644 --- a/tests/db/test_crud.py +++ b/tests/db/test_crud.py @@ -184,10 +184,8 @@ async def test_insert_many_accepts_model_instances(pool: PgPool): async def test_insert_retrieve_composite_column_with_normalizer(pool: PgPool): # label_pair mirrors a real project's locale-labels composite type: a # normalizer backfills a missing locale before the value is converted - # into the psycopg-registered composite type for INSERT. `RETURNING *` - # gives back the psycopg-generated composite instance directly (since - # register_composite() also wires up decoding); pg_retrieve's to_jsonb() - # wrapping instead unwraps it back into a plain dict. + # into the psycopg composite type for INSERT. Both RETURNING and + # pg_retrieve select complex columns as to_jsonb(), i.e. a plain dict. def backfill_de(value: dict) -> dict: if not value.get("de"): value = {**value, "de": f"[{value['en']}]"} @@ -196,8 +194,7 @@ def backfill_de(value: dict) -> dict: async with pool.connection() as con: helper = ComplexHelper(con, normalizers={"label_pair": backfill_de}) inserted = await pg_insert(con, ("public", "gizmo"), {"label": {"en": "Hello"}}, complex_helper=helper) - assert inserted["label"].en == "Hello" - assert inserted["label"].de == "[Hello]" + assert inserted["label"] == {"en": "Hello", "de": "[Hello]"} # RETURNING: plain values, like pg_retrieve gizmo_id = inserted["id"] fetched = await pg_retrieve(con, Gizmo, {"id": gizmo_id}, complex_helper=helper) From e336fad0b76dd1a3973a36001eb669b889929b25 Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Fri, 9 Oct 2026 19:27:40 +0000 Subject: [PATCH 6/6] docs: re-wrap --- README.md | 4 ++-- skills/pgdevkit/SKILL.md | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 3e1a5d7..d803cbf 100644 --- a/README.md +++ b/README.md @@ -460,8 +460,8 @@ Install with the `db` extra: `pip install pgdevkit[db]`. from `load_all_complex_types`/`load_complex_type`: `await helper.register_on(cur, info)`, then execute the converted value on that same `cur` (passing the connection instead would restore psycopg's global behaviour). What it learns from the catalog (types and column types) 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 + 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 diff --git a/skills/pgdevkit/SKILL.md b/skills/pgdevkit/SKILL.md index e5ae837..e1d7674 100644 --- a/skills/pgdevkit/SKILL.md +++ b/skills/pgdevkit/SKILL.md @@ -271,8 +271,8 @@ converted and reads are a plain `SELECT *`. **The helper never registers types on the connection** (psycopg's registration is connection-global and would stick to a pooled connection: a plain `SELECT` of an enum would return psycopg's `Enum` member instead of the label, #54). Writes register the enum/composite adapters on the one cursor that runs the statement, only for the columns where a -dict/list value needed converting; reads (`pg_retrieve*`) use `to_jsonb(...)` and register nothing. So after any `pg_*` call the -connection reads enums as label strings and composites as text (`'(1,2)'`), as before the call. `RETURNING` rows of +dict/list value needed converting; reads (`pg_retrieve*`) use `to_jsonb(...)` and register nothing. So after any +`pg_*` call the connection reads enums as label strings and composites as text (`'(1,2)'`), as before the call. `RETURNING` rows of `pg_insert`/`pg_update_dict`/`pg_upsert_dict` (hence `pg_update`/`pg_upsert`) also use `to_jsonb(...)` for such columns: with a helper they are plain Python values (a composite as a dict, an enum as its label) whichever columns were written; without one it is a plain `RETURNING *`.