diff --git a/pgdevkit/db/crud.py b/pgdevkit/db/crud.py index f8f05fa..56728f4 100644 --- a/pgdevkit/db/crud.py +++ b/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( diff --git a/pyproject.toml b/pyproject.toml index 58c3222..1b86fb4 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/db/test_crud.py b/tests/db/test_crud.py index a0992e0..5556070 100644 --- a/tests/db/test_crud.py +++ b/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"} diff --git a/uv.lock b/uv.lock index c610b87..acd2bab 100644 --- a/uv.lock +++ b/uv.lock @@ -313,7 +313,7 @@ wheels = [ [[package]] name = "pgdevkit" -version = "0.3.0" +version = "0.3.1" source = { editable = "." } dependencies = [ { name = "docker" },