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": ["/psycopg2/", "/risingwave/"]
     },
     {
       "addLabels": ["snowflake"],

conda/environment-arm64-flink.yml

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

conda/environment.yml

@@ -27,7 +27,8 @@ dependencies:
   - pip
   - uv>=0.4.29
   - polars >=1,<2
-  - psycopg2 >=2.8.4
+  - 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 @@ def _from_url(self, url: ParseResult, **kwargs):
         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 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
 
         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 @@ def begin(self):
         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 @@ def _fetch_from_cursor(self, cursor, schema: sch.Schema) -> pd.DataFrame:
 
     @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 @@ def do_connect(
           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 @@ def do_connect(
 
     @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 @@ def _safe_raw_sql(self, *args, **kwargs):
             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,16 +712,16 @@ def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
         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):
             cursor.close()
             raise
+        except TypeError:
+            pass
 
         try:
             cursor.execute(query, **kwargs)

ibis/backends/postgres/tests/conftest.py

@@ -48,7 +48,7 @@ class TestConf(ServiceBackendTest):
     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_kwargs_passthrough_in_connect():
 
 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 @@ def test_password_with_bracket():
     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 @@ def test_string_to_binary_cast(con):
     )
     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 @@ def test_string_to_binary_round_trip(con):
     )
     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

@@ -10,7 +10,6 @@
 import psycopg2
 import sqlglot as sg
 import sqlglot.expressions as sge
-from pandas.api.types import is_float_dtype
 from psycopg2 import extras
 
 import ibis
@@ -106,48 +105,6 @@ def _from_url(self, url: ParseResult, **kwargs):
 
         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(
-                f"{self.name} cannot yet reliably handle `null` typed columns; "
-                f"got null typed columns: {null_columns}"
-            )
-
-        name = op.name
-        quoted = self.compiler.quoted
-        create_stmt = sg.exp.Create(
-            kind="TABLE",
-            this=sg.exp.Schema(
-                this=sg.to_identifier(name, quoted=quoted),
-                expressions=schema.to_sqlglot(self.dialect),
-            ),
-            properties=sg.exp.Properties(expressions=[sge.TemporaryProperty()]),
-        )
-        create_stmt_sql = create_stmt.sql(self.dialect)
-
-        df = op.data.to_frame()
-        # nan gets compiled into 'NaN'::float which throws errors in non-float columns
-        # In order to hold NaN values, pandas automatically converts integer columns
-        # to float columns if there are NaN values in them. Therefore, we need to convert
-        # them to their original dtypes (that support pd.NA) to figure out which columns
-        # are actually non-float, then fill the NaN values in those columns with None.
-        convert_df = df.convert_dtypes()
-        for col in convert_df.columns:
-            if not is_float_dtype(convert_df[col]):
-                df[col] = df[col].replace(float("nan"), None)
-
-        data = df.itertuples(index=False)
-        sql = self._build_insert_template(
-            name, schema=schema, columns=True, placeholder="%s"
-        )
-
-        with self.begin() as cur:
-            cur.execute(create_stmt_sql)
-            execute_batch(cur, sql, data, 128)
-
     @contextlib.contextmanager
     def begin(self):
         con = self.con

ibis/backends/tests/errors.py

@@ -131,6 +131,25 @@
         PsycoPg2OperationalError
     ) = PsycoPg2UndefinedObject = PsycoPg2ArraySubscriptError = None
 
+try:
+    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 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
     from MySQLdb import OperationalError as MySQLOperationalError

ibis/backends/tests/test_array.py

@@ -27,6 +27,7 @@
     PsycoPg2InternalError,
     PsycoPg2ProgrammingError,
     PsycoPg2SyntaxError,
+    PsycoPgSyntaxError,
     Py4JJavaError,
     PyAthenaDatabaseError,
     PyAthenaOperationalError,
@@ -1094,7 +1095,7 @@ def test_array_intersect(con, data):
 
 
 @builtin_array
-@pytest.mark.notimpl(["postgres"], raises=PsycoPg2SyntaxError)
+@pytest.mark.notimpl(["postgres"], raises=PsycoPgSyntaxError)
 @pytest.mark.notimpl(["risingwave"], raises=PsycoPg2InternalError)
 @pytest.mark.notimpl(
     ["trino"], reason="inserting maps into structs doesn't work", raises=TrinoUserError

ibis/backends/tests/test_numeric.py

@@ -22,8 +22,8 @@
     ImpalaHiveServer2Error,
     MySQLOperationalError,
     OracleDatabaseError,
-    PsycoPg2DivisionByZero,
     PsycoPg2InternalError,
+    PsycoPgDivisionByZero,
     Py4JError,
     Py4JJavaError,
     PyAthenaOperationalError,
@@ -1323,7 +1323,7 @@ def test_floating_mod(backend, alltypes, df):
 )
 @pytest.mark.notyet(["mssql"], raises=PyODBCDataError)
 @pytest.mark.notyet(["snowflake"], raises=SnowflakeProgrammingError)
-@pytest.mark.notyet(["postgres"], raises=PsycoPg2DivisionByZero)
+@pytest.mark.notyet(["postgres"], raises=PsycoPgDivisionByZero)
 @pytest.mark.notimpl(["exasol"], raises=ExaQueryError)
 @pytest.mark.xfail_version(duckdb=["duckdb<1.1"])
 def test_divide_by_zero(backend, alltypes, df, column, denominator):