Add complex_helper support to pg_update_dict/pg_update

pgdevkit/db/crud.py

@@ -157,8 +157,15 @@ async def pg_update_dict(
     table_name: tuple[str, str],
     data: dict,
     primary_keys: Sequence[str],
+    *,
+    complex_helper: ComplexHelper | None = 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.
+
+    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
     ]
@@ -173,9 +180,13 @@ async def pg_update_dict(
         return await cur.fetchone()
 
 
-async def pg_update(con: AsyncConnection, data: T, data_type: type[T]) -> Any | None:
+async def pg_update(
+    con: AsyncConnection, data: T, data_type: type[T], *, complex_helper: ComplexHelper | None = None
+) -> Any | None:
     """Update a typed model instance."""
-    return await pg_update_dict(con, data_type.get_table_name(), data.model_dump(), data_type.get_primary_key())
+    return await pg_update_dict(
+        con, data_type.get_table_name(), data.model_dump(), data_type.get_primary_key(), complex_helper=complex_helper
+    )
 
 
 async def pg_upsert_dict(

pyproject.toml

@@ -11,7 +11,7 @@ packages = ["pgdevkit"]
 
 [project]
 name = "pgdevkit"
-version = "0.3.0"
+version = "0.3.1"
 description = "A helper for developing with Postgres"
 readme = "README.md"
 requires-python = ">=3.14"

tests/db/test_crud.py

@@ -16,6 +16,7 @@
     pg_retrieve,
     pg_retrieve_many,
     pg_update,
+    pg_update_dict,
     pg_upsert,
     pg_upsert_many,
     pg_upsert_many_dict,
@@ -60,6 +61,20 @@ def get_primary_key() -> Sequence[str]:
         return ["id"]
 
 
+class Gadget(PostgresTableModel):
+    model_config = ConfigDict(from_attributes=True)
+    id: int
+    status: str
+
+    @staticmethod
+    def get_table_name() -> tuple[str, str]:
+        return ("public", "gadget")
+
+    @staticmethod
+    def get_primary_key() -> Sequence[str]:
+        return ["id"]
+
+
 @pytest.fixture
 async def pool(monkeypatch: pytest.MonkeyPatch):
     ensure_container()
@@ -70,6 +85,8 @@ async def pool(monkeypatch: pytest.MonkeyPatch):
         con.execute("CREATE TABLE widget (id serial PRIMARY KEY, name text NOT NULL)")
         con.execute("CREATE TYPE label_pair AS (en text, de text)")
         con.execute("CREATE TABLE gizmo (id serial PRIMARY KEY, label label_pair NOT NULL)")
+        con.execute("CREATE TYPE gadget_status AS ENUM ('active', 'inactive')")
+        con.execute("CREATE TABLE gadget (id serial PRIMARY KEY, status gadget_status NOT NULL)")
 
     monkeypatch.setenv(f"{ENV_PREFIX}HOST", constants.HOST)
     monkeypatch.setenv(f"{ENV_PREFIX}PORT", str(constants.PORT))
@@ -184,3 +201,70 @@ def backfill_de(value: dict) -> dict:
         fetched = await pg_retrieve(con, Gizmo, {"id": gizmo_id}, complex_helper=helper)
         assert fetched is not None
         assert fetched.label == {"en": "Hello", "de": "[Hello]"}
+
+
+@requires_podman
+async def test_update_dict_without_complex_helper_still_works_on_plain_columns(pool: PgPool):
+    # complex_helper is a new optional kwarg -- omitting it (the pre-existing call
+    # shape) must behave exactly as before for tables with no composite/enum columns.
+    async with pool.connection() as con:
+        inserted = await pg_insert(con, ("public", "widget"), {"name": "sprocket"})
+        widget_id = inserted["id"]
+
+        await pg_update_dict(con, ("public", "widget"), {"id": widget_id, "name": "renamed"}, ["id"])
+
+        fetched = await pg_retrieve(con, Widget, {"id": widget_id})
+        assert fetched is not None
+        assert fetched.name == "renamed"
+
+
+@requires_podman
+async def test_update_dict_converts_enum_column_with_complex_helper(pool: PgPool):
+    # psycopg already binds a plain str to an enum-typed UPDATE target column
+    # fine on its own (Postgres resolves it via the usual assignment cast), so
+    # this isn't a failing-without/passing-with case the way composite columns
+    # are below -- it just confirms complex_helper doesn't break the enum path
+    # and round-trips correctly end to end.
+    async with pool.connection() as con:
+        helper = ComplexHelper(con)
+        inserted = await pg_insert(con, ("public", "gadget"), {"status": "active"}, complex_helper=helper)
+        gadget_id = inserted["id"]
+
+        await pg_update_dict(
+            con, ("public", "gadget"), {"id": gadget_id, "status": "inactive"}, ["id"], complex_helper=helper
+        )
+        fetched = await pg_retrieve(con, Gadget, {"id": gadget_id}, complex_helper=helper)
+        assert fetched is not None
+        assert fetched.status == "inactive"
+
+        # pg_update (typed-model variant) threads complex_helper through the same way.
+        fetched.status = "active"
+        await pg_update(con, fetched, Gadget, complex_helper=helper)
+        refetched = await pg_retrieve(con, Gadget, {"id": gadget_id}, complex_helper=helper)
+        assert refetched is not None
+        assert refetched.status == "active"
+
+
+@requires_podman
+async def test_update_dict_composite_column_requires_complex_helper(pool: PgPool):
+    # Unlike enums, psycopg has no built-in way to adapt a plain Python dict
+    # to a composite column type -- this is the real case pg_update_dict's
+    # complex_helper support fixes.
+    async with pool.connection() as con:
+        helper = ComplexHelper(con)
+        inserted = await pg_insert(con, ("public", "gizmo"), {"label": {"en": "Hello", "de": "Hallo"}}, complex_helper=helper)
+        gizmo_id = inserted["id"]
+
+        with pytest.raises(psycopg.ProgrammingError):
+            await pg_update_dict(con, ("public", "gizmo"), {"id": gizmo_id, "label": {"en": "Bye", "de": "Tschuess"}}, ["id"])
+
+        await pg_update_dict(
+            con,
+            ("public", "gizmo"),
+            {"id": gizmo_id, "label": {"en": "Bye", "de": "Tschuess"}},
+            ["id"],
+            complex_helper=helper,
+        )
+        fetched = await pg_retrieve(con, Gizmo, {"id": gizmo_id}, complex_helper=helper)
+        assert fetched is not None
+        assert fetched.label == {"en": "Bye", "de": "Tschuess"}