fix(postgres): use .transaction method instead of managing our own

Author: cpcloud

.github/workflows/ibis-backends.yml

@@ -141,6 +141,7 @@ jobs:
               - default-libmysqlclient-dev
           - name: postgres
             title: PostgreSQL
+            serial: true
             extras:
               - --extra postgres
               - --extra geospatial
@@ -150,6 +151,7 @@ jobs:
               - libgeos-dev
           - name: postgres
             title: PostgreSQL + Torch
+            serial: true
             extras:
               - --extra postgres
               - --extra geospatial
@@ -319,6 +321,7 @@ jobs:
             backend:
               name: postgres
               title: PostgreSQL
+              serial: true
               extras:
                 - --extra postgres
                 - --extra geospatial
@@ -339,6 +342,7 @@ jobs:
             backend:
               name: postgres
               title: PostgreSQL + Torch
+              serial: true
               extras:
                 - --extra postgres
                 - --extra geospatial

ibis/backends/postgres/__init__.py

@@ -34,6 +34,7 @@
     import pandas as pd
     import polars as pl
     import pyarrow as pa
+    from typing_extensions import Self
 
 
 class NatDumper(psycopg.adapt.Dumper):
@@ -131,23 +132,17 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
             name, schema=schema, columns=True, placeholder="%s"
         )
 
-        with self.begin() as cur:
-            cur.execute(create_stmt_sql)
-            cur.executemany(sql, data)
+        with self.begin() as cursor:
+            cursor.execute(create_stmt_sql)
+            cursor.executemany(sql, data)
 
     @contextlib.contextmanager
-    def begin(self):
-        con = self.con
-        cursor = con.cursor()
-        try:
+    def begin(self, *, name: str = "", withhold: bool = False):
+        with (
+            (con := self.con).transaction(),
+            con.cursor(name=name, withhold=withhold) as cursor,
+        ):
             yield cursor
-        except Exception:
-            con.rollback()
-            raise
-        else:
-            con.commit()
-        finally:
-            cursor.close()
 
     def _fetch_from_cursor(
         self, cursor: psycopg.Cursor, schema: sch.Schema
@@ -190,6 +185,7 @@ def do_connect(
         port: int = 5432,
         database: str | None = None,
         schema: str | None = None,
+        autocommit: bool = True,
         **kwargs: Any,
     ) -> None:
         """Create an Ibis client connected to PostgreSQL database.
@@ -208,6 +204,8 @@ def do_connect(
             Database to connect to
         schema
             PostgreSQL schema to use. If `None`, use the default `search_path`.
+        autocommit
+            Whether or not to autocommit
         kwargs
             Additional keyword arguments to pass to the backend client connection.
 
@@ -252,6 +250,7 @@ def do_connect(
             password=password,
             dbname=database,
             options=(f"-csearch_path={schema}" * (schema is not None)) or None,
+            autocommit=autocommit,
             **kwargs,
         )
 
@@ -276,8 +275,23 @@ def from_connection(cls, con: psycopg.Connection, /) -> Backend:
         return new_backend
 
     def _post_connect(self) -> None:
-        with self.begin() as cur:
-            cur.execute("SET TIMEZONE = UTC")
+        import psycopg.types
+        import psycopg.types.hstore
+
+        try:
+            # try to load hstore
+            with self.begin() as cursor:
+                cursor.execute("CREATE EXTENSION IF NOT EXISTS hstore")
+            psycopg.types.hstore.register_hstore(
+                psycopg.types.TypeInfo.fetch(self.con, "hstore"), self.con
+            )
+        except psycopg.Error as e:
+            warnings.warn(f"Failed to load hstore extension: {e}")
+        except TypeError:
+            pass
+
+        with self.begin() as cursor:
+            cursor.execute("SET TIMEZONE = UTC")
 
     @property
     def _session_temp_db(self) -> str | None:
@@ -706,23 +720,12 @@ def drop_table(
             pass
 
     @contextlib.contextmanager
-    def _safe_raw_sql(self, *args, **kwargs):
-        with contextlib.closing(self.raw_sql(*args, **kwargs)) as result:
-            yield result
+    def _safe_raw_sql(self, query: str | sg.Expression, **kwargs: Any):
+        with contextlib.suppress(AttributeError):
+            query = query.sql(dialect=self.dialect)
 
-    def _register_hstore(self, cursor: psycopg.Cursor) -> None:
-        import psycopg.types
-        import psycopg.types.hstore
-
-        try:
-            # try to load hstore
-            psycopg.types.hstore.register_hstore(
-                psycopg.types.TypeInfo.fetch(self.con, "hstore"), cursor
-            )
-        except psycopg.Error as e:
-            warnings.warn(f"Failed to load hstore extension: {e}")
-        except TypeError:
-            pass
+        with self.begin() as cursor:
+            yield cursor.execute(query, **kwargs)
 
     def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
         with contextlib.suppress(AttributeError):
@@ -731,16 +734,12 @@ def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
         con = self.con
         cursor = con.cursor()
 
-        self._register_hstore(cursor)
-
         try:
             cursor.execute(query, **kwargs)
         except Exception:
-            con.rollback()
             cursor.close()
             raise
         else:
-            con.commit()
             return cursor
 
     @util.experimental
@@ -757,33 +756,22 @@ def to_pyarrow_batches(
         import pandas as pd
         import pyarrow as pa
 
-        def _batches(*, schema: pa.Schema, con: psycopg.Connection, query: str):
+        def _batches(self: Self, *, schema: pa.Schema, query: str):
             columns = schema.names
             # server-side cursors need to be uniquely named
-            with con.cursor(name=util.gen_name("postgres_cursor")) as cursor:
-                self._register_hstore(cursor)
-
-                try:
-                    cur = cursor.execute(query)
-                except Exception:
-                    con.rollback()
-                    raise
-                else:
-                    try:
-                        while batch := cur.fetchmany(chunk_size):
-                            yield pa.RecordBatch.from_pandas(
-                                pd.DataFrame(batch, columns=columns), schema=schema
-                            )
-                    except Exception:
-                        con.rollback()
-                        raise
-                    else:
-                        con.commit()
+            with self.begin(
+                name=util.gen_name("postgres_cursor"), withhold=True
+            ) as cursor:
+                cursor.execute(query)
+                while batch := cursor.fetchmany(chunk_size):
+                    yield pa.RecordBatch.from_pandas(
+                        pd.DataFrame(batch, columns=columns), schema=schema
+                    )
 
         self._run_pre_execute_hooks(expr)
 
         schema = expr.as_table().schema().to_pyarrow()
         query = self.compile(expr, limit=limit, params=params)
         return pa.RecordBatchReader.from_batches(
-            schema, _batches(schema=schema, con=self.con, query=query)
+            schema, _batches(self, schema=schema, query=query)
         )

ibis/backends/postgres/tests/conftest.py

@@ -64,6 +64,7 @@ def connect(*, tmpdir, worker_id, **kw):
             user=PG_USER,
             password=PG_PASS,
             database=IBIS_TEST_POSTGRES_DB,
+            application_name="Ibis test suite",
             **kw,
         )
 

ibis/backends/postgres/tests/test_client.py

@@ -255,7 +255,10 @@ def test_kwargs_passthrough_in_connect():
     con = ibis.connect(
         "postgresql://postgres:postgres@localhost:5432/ibis_testing?sslmode=allow"
     )
-    assert con.current_catalog == "ibis_testing"
+    try:
+        assert con.current_catalog == "ibis_testing"
+    finally:
+        con.disconnect()
 
 
 def test_port():

ibis/backends/tests/test_client.py

@@ -819,8 +819,10 @@ def test_unsigned_integer_type(con, temp_table):
 )
 def test_connect_url(url):
     con = ibis.connect(url)
-    one = ibis.literal(1)
-    assert con.execute(one) == 1
+    try:
+        assert con.execute(ibis.literal(1)) == 1
+    finally:
+        con.disconnect()
 
 
 @pytest.mark.parametrize(
@@ -1295,7 +1297,11 @@ def test_set_backend_url(url, monkeypatch):
     monkeypatch.setattr(ibis.options, "default_backend", None)
     name = url.split("://")[0]
     ibis.set_backend(url)
-    assert ibis.get_backend().name == name
+    con = ibis.get_backend()
+    try:
+        assert con.name == name
+    finally:
+        con.disconnect()
 
 
 @pytest.mark.notyet(