feat(postgres): use psycopg rather than psycopg2

.github/renovate.json

@@ -69,7 +69,7 @@
     },
     {
       "addLabels": ["postgres"],
-      "matchPackageNames": ["/psycopg2/", "/postgres/"]
+      "matchPackageNames": ["/psycopg/", "/postgres/"]
     },
     {
       "addLabels": ["druid"],
@@ -89,7 +89,7 @@
     },
     {
       "addLabels": ["risingwave"],
-      "matchPackageNames": ["/risingwave/"]
+      "matchPackageNames": ["/psycopg/", "/risingwave/"]
     },
     {
       "addLabels": ["snowflake"],

conda/environment-arm64-flink.yml

@@ -26,7 +26,7 @@ dependencies:
   - pins >=0.8.2
   - uv>=0.4.29
   - polars >=1,<2
-  - psycopg2 >=2.8.4
+  - psycopg >= 3.2.0
   - pyarrow =11.0.0
   - pyarrow-tests
   - pyarrow-hotfix >=0.4

conda/environment-arm64.yml

@@ -26,7 +26,7 @@ dependencies:
   - pins >=0.8.2
   - uv>=0.4.29
   - polars >=1,<2
-  - psycopg2 >=2.8.4
+  - psycopg >= 3.2.0
   - pyarrow >=10.0.1
   - pyarrow-tests
   - pyarrow-hotfix >=0.4

conda/environment.yml

@@ -27,7 +27,7 @@ dependencies:
   - pip
   - uv>=0.4.29
   - polars >=1,<2
-  - psycopg2 >=2.8.4
+  - psycopg >= 3.2.0
   - pyarrow >=10.0.1
   - pyarrow-hotfix >=0.4
   - pydata-google-auth

ibis/backends/postgres/__init__.py

@@ -31,7 +31,7 @@
 
     import pandas as pd
     import polars as pl
-    import psycopg2
+    import psycopg
     import pyarrow as pa
 
 
@@ -90,8 +90,6 @@
         return self.connect(**kwargs)
 
     def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
-        from psycopg2.extras import execute_batch
-
         schema = op.schema
         if null_columns := [col for col, dtype in schema.items() if dtype.is_null()]:
             raise exc.IbisTypeError(
@@ -129,7 +127,7 @@
 
         with self.begin() as cur:
             cur.execute(create_stmt_sql)
-            execute_batch(cur, sql, data, 128)
+            cur.executemany(sql, data)
 
     @contextlib.contextmanager
     def begin(self):
@@ -145,14 +143,16 @@
         finally:
             cursor.close()
 
-    def _fetch_from_cursor(self, cursor, schema: sch.Schema) -> pd.DataFrame:
+    def _fetch_from_cursor(
+        self, cursor: psycopg.Cursor, schema: sch.Schema
+    ) -> pd.DataFrame:
         import pandas as pd
 
         from ibis.backends.postgres.converter import PostgresPandasData
 
         try:
             df = pd.DataFrame.from_records(
-                cursor, columns=schema.names, coerce_float=True
+                cursor.fetchall(), columns=schema.names, coerce_float=True
             )
         except Exception:
             # clean up the cursor if we fail to create the DataFrame
@@ -166,7 +166,7 @@
 
     @property
     def version(self):
-        version = f"{self.con.server_version:0>6}"
+        version = f"{self.con.info.server_version:0>6}"
         major = int(version[:2])
         minor = int(version[2:4])
         patch = int(version[4:])
@@ -233,17 +233,17 @@
           year            int32
           month           int32
         """
-        import psycopg2
-        import psycopg2.extras
+        import psycopg
+        import psycopg.types.json
 
-        psycopg2.extras.register_default_json(loads=lambda x: x)
+        psycopg.types.json.set_json_loads(loads=lambda x: x)
 
-        self.con = psycopg2.connect(
+        self.con = psycopg.connect(
             host=host,
             port=port,
             user=user,
             password=password,
-            database=database,
+            dbname=database,
             options=(f"-csearch_path={schema}" * (schema is not None)) or None,
             **kwargs,
         )
@@ -252,7 +252,7 @@
 
     @util.experimental
     @classmethod
-    def from_connection(cls, con: psycopg2.extensions.connection) -> Backend:
+    def from_connection(cls, con: psycopg.Connection) -> Backend:
         """Create an Ibis client from an existing connection to a PostgreSQL database.
 
         Parameters
@@ -701,8 +701,9 @@
             yield result
 
     def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
-        import psycopg2
-        import psycopg2.extras
+        import psycopg
+        import psycopg.types
+        import psycopg.types.hstore
 
         with contextlib.suppress(AttributeError):
             query = query.sql(dialect=self.dialect)
@@ -711,14 +712,12 @@
         cursor = con.cursor()
 
         try:
-            # try to load hstore, uuid and ipaddress extensions
-            with contextlib.suppress(psycopg2.ProgrammingError):
-                psycopg2.extras.register_hstore(cursor)
-            with contextlib.suppress(psycopg2.ProgrammingError):
-                psycopg2.extras.register_uuid(conn_or_curs=cursor)
-            with contextlib.suppress(psycopg2.ProgrammingError):
-                psycopg2.extras.register_ipaddress(cursor)
-        except Exception:
+            # try to load hstore
+            psycopg.types.hstore.register_hstore(
+                psycopg.types.TypeInfo.fetch(con, "hstore"),
+                cursor,
+            )
+        except (psycopg.InternalError, psycopg.ProgrammingError, TypeError):
             cursor.close()
             raise
 

ibis/backends/postgres/tests/conftest.py

@@ -48,7 +48,7 @@
     supports_structs = False
     rounding_method = "half_to_even"
     service_name = "postgres"
-    deps = ("psycopg2",)
+    deps = ("psycopg",)
 
     driver_supports_multiple_statements = True
 

ibis/backends/postgres/tests/test_client.py

@@ -30,10 +30,10 @@
 import ibis.common.exceptions as com
 import ibis.expr.datatypes as dt
 import ibis.expr.types as ir
-from ibis.backends.tests.errors import PsycoPg2OperationalError
+from ibis.backends.tests.errors import PsycoPgOperationalError
 from ibis.util import gen_name
 
-pytest.importorskip("psycopg2")
+pytest.importorskip("psycopg")
 
 POSTGRES_TEST_DB = os.environ.get("IBIS_TEST_POSTGRES_DATABASE", "ibis_testing")
 IBIS_POSTGRES_HOST = os.environ.get("IBIS_TEST_POSTGRES_HOST", "localhost")
@@ -260,7 +260,7 @@
 
 def test_port():
     # check that we parse and use the port (and then of course fail cuz it's bogus)
-    with pytest.raises(PsycoPg2OperationalError):
+    with pytest.raises(PsycoPgOperationalError):
         ibis.connect("postgresql://postgres:postgres@localhost:1337/ibis_testing")
 
 
@@ -388,7 +388,7 @@
     quoted_pass = quote_plus(password)
     url = f"postgres://{IBIS_POSTGRES_USER}:{quoted_pass}@{IBIS_POSTGRES_HOST}:{IBIS_POSTGRES_PORT}/{POSTGRES_TEST_DB}"
     with pytest.raises(
-        PsycoPg2OperationalError,
+        PsycoPgOperationalError,
         match=f'password authentication failed for user "{IBIS_POSTGRES_USER}"',
     ):
         ibis.connect(url)

ibis/backends/postgres/tests/test_functions.py

@@ -16,7 +16,7 @@
 import ibis.expr.types as ir
 from ibis import literal as L
 
-pytest.importorskip("psycopg2")
+pytest.importorskip("psycopg")
 
 
 @pytest.mark.parametrize(
@@ -1195,7 +1195,7 @@
     )
     with con.begin() as c:
         c.execute(sql_string)
-        raw_data = [row[0][0] for row in c.fetchall()]
+        raw_data = [row[0] for row in c.fetchall()]
     expected = pd.Series(raw_data, name=name)
     tm.assert_series_equal(result, expected)
 
@@ -1212,7 +1212,7 @@
     )
     with con.begin() as c:
         c.execute(sql_string)
-        rows = [row[0] for (row,) in c.fetchall()]
+        rows = [row[0] for row in c.fetchall()]
     expected = pd.Series(rows, name=name)
     tm.assert_series_equal(result, expected)
 

ibis/backends/postgres/tests/test_postgis.py

@@ -7,7 +7,7 @@
 import pytest
 from numpy import testing
 
-pytest.importorskip("psycopg2")
+pytest.importorskip("psycopg")
 gpd = pytest.importorskip("geopandas")
 pytest.importorskip("shapely")
 

ibis/backends/postgres/tests/test_udf.py

@@ -12,7 +12,7 @@
 from ibis import udf
 from ibis.util import guid
 
-pytest.importorskip("psycopg2")
+pytest.importorskip("psycopg")
 
 
 @pytest.fixture(scope="session")

ibis/backends/risingwave/__init__.py

@@ -5,10 +5,9 @@
 from operator import itemgetter
 from typing import TYPE_CHECKING
 
-import psycopg2
+import psycopg
 import sqlglot as sg
 import sqlglot.expressions as sge
-from psycopg2 import extras
 
 import ibis
 import ibis.backends.sql.compilers as sc
@@ -110,12 +109,12 @@
           month           int32
         """
 
-        self.con = psycopg2.connect(
+        self.con = psycopg.connect(
             host=host,
             port=port,
             user=user,
             password=password,
-            database=database,
+            dbname=database,
             options=(f"-csearch_path={schema}" * (schema is not None)) or None,
         )
 
@@ -289,7 +288,7 @@
         )
         with self.begin() as cur:
             cur.execute(create_stmt_sql)
-            extras.execute_batch(cur, sql, data, 128)
+            cur.executemany(sql, data)
 
     def list_databases(
         self, *, like: str | None = None, catalog: str | None = None

ibis/backends/risingwave/tests/conftest.py

@@ -33,7 +33,7 @@
     supports_structs = False
     rounding_method = "half_to_even"
     service_name = "risingwave"
-    deps = ("psycopg2",)
+    deps = ("psycopg",)
 
     @property
     def test_files(self) -> Iterable[Path]:

ibis/backends/risingwave/tests/test_client.py

@@ -12,7 +12,7 @@
 import ibis.expr.types as ir
 from ibis.util import gen_name
 
-pytest.importorskip("psycopg2")
+pytest.importorskip("psycopg")
 
 RISINGWAVE_TEST_DB = os.environ.get("IBIS_TEST_RISINGWAVE_DATABASE", "dev")
 IBIS_RISINGWAVE_HOST = os.environ.get("IBIS_TEST_RISINGWAVE_HOST", "localhost")

ibis/backends/risingwave/tests/test_functions.py

@@ -14,7 +14,7 @@
 import ibis.expr.datatypes as dt
 from ibis import literal as L
 
-pytest.importorskip("psycopg2")
+pytest.importorskip("psycopg")
 
 
 @pytest.mark.parametrize(("value", "expected"), [(0, None), (5.5, 5.5)])

ibis/backends/tests/errors.py

@@ -113,23 +113,23 @@
     TrinoUserError = None
 
 try:
-    from psycopg2.errors import ArraySubscriptError as PsycoPg2ArraySubscriptError
-    from psycopg2.errors import DivisionByZero as PsycoPg2DivisionByZero
-    from psycopg2.errors import IndeterminateDatatype as PsycoPg2IndeterminateDatatype
-    from psycopg2.errors import InternalError_ as PsycoPg2InternalError
-    from psycopg2.errors import (
-        InvalidTextRepresentation as PsycoPg2InvalidTextRepresentation,
+    from psycopg.errors import ArraySubscriptError as PsycoPgArraySubscriptError
+    from psycopg.errors import DivisionByZero as PsycoPgDivisionByZero
+    from psycopg.errors import IndeterminateDatatype as PsycoPgIndeterminateDatatype
+    from psycopg.errors import InternalError_ as PsycoPgInternalError
+    from psycopg.errors import (
+        InvalidTextRepresentation as PsycoPgInvalidTextRepresentation,
     )
-    from psycopg2.errors import OperationalError as PsycoPg2OperationalError
-    from psycopg2.errors import ProgrammingError as PsycoPg2ProgrammingError
-    from psycopg2.errors import SyntaxError as PsycoPg2SyntaxError
-    from psycopg2.errors import UndefinedObject as PsycoPg2UndefinedObject
-except ImportError:
-    PsycoPg2SyntaxError = PsycoPg2IndeterminateDatatype = (
-        PsycoPg2InvalidTextRepresentation
-    ) = PsycoPg2DivisionByZero = PsycoPg2InternalError = PsycoPg2ProgrammingError = (
-        PsycoPg2OperationalError
-    ) = PsycoPg2UndefinedObject = PsycoPg2ArraySubscriptError = None
+    from psycopg.errors import OperationalError as PsycoPgOperationalError
+    from psycopg.errors import ProgrammingError as PsycoPgProgrammingError
+    from psycopg.errors import SyntaxError as PsycoPgSyntaxError
+    from psycopg.errors import UndefinedObject as PsycoPgUndefinedObject
+except ImportError:
+    PsycoPgSyntaxError = PsycoPgIndeterminateDatatype = (
+        PsycoPgInvalidTextRepresentation
+    ) = PsycoPgDivisionByZero = PsycoPgInternalError = PsycoPgProgrammingError = (
+        PsycoPgOperationalError
+    ) = PsycoPgUndefinedObject = PsycoPgArraySubscriptError = None
 
 try:
     from MySQLdb import NotSupportedError as MySQLNotSupportedError