Use psycopg rather than psycopg2 for Postgres

Author: judahrand

.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,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 @@ 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,13 +712,11 @@ 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)
+            # try to load hstore
+            with contextlib.suppress(TypeError):
+                type_info = psycopg.types.TypeInfo.fetch(con, "hstore")
+            with contextlib.suppress(psycopg.ProgrammingError, TypeError):
+                psycopg.types.hstore.register_hstore(type_info, cursor)
         except Exception:
             cursor.close()
             raise

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(
@@ -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/tests/errors.py

@@ -112,6 +112,25 @@
 except ImportError:
     TrinoUserError = 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 psycopg2.errors import ArraySubscriptError as PsycoPg2ArraySubscriptError
     from psycopg2.errors import DivisionByZero as PsycoPg2DivisionByZero

ibis/backends/tests/test_array.py

@@ -22,11 +22,13 @@
     GoogleBadRequest,
     MySQLOperationalError,
     PolarsComputeError,
-    PsycoPg2ArraySubscriptError,
     PsycoPg2IndeterminateDatatype,
     PsycoPg2InternalError,
     PsycoPg2ProgrammingError,
-    PsycoPg2SyntaxError,
+    PsycoPgIndeterminateDatatype,
+    PsycoPgInternalError,
+    PsycoPgInvalidTextRepresentation,
+    PsycoPgSyntaxError,
     Py4JJavaError,
     PyAthenaDatabaseError,
     PyAthenaOperationalError,
@@ -1094,7 +1096,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
@@ -1114,7 +1116,7 @@ def test_unnest_struct(con):
 
 
 @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
@@ -1205,7 +1207,7 @@ def test_zip_null(con, fn):
 
 
 @builtin_array
-@pytest.mark.notimpl(["postgres"], raises=PsycoPg2SyntaxError)
+@pytest.mark.notimpl(["postgres"], raises=PsycoPgSyntaxError)
 @pytest.mark.notimpl(["risingwave"], raises=PsycoPg2ProgrammingError)
 @pytest.mark.notimpl(["datafusion"], raises=Exception, reason="not yet supported")
 @pytest.mark.notimpl(
@@ -1276,8 +1278,17 @@ def flatten_data():
     ["bigquery"], reason="BigQuery doesn't support arrays of arrays", raises=TypeError
 )
 @pytest.mark.notyet(
-    ["postgres", "risingwave"],
+    ["postgres"],
     reason="Postgres doesn't truly support arrays of arrays",
+    raises=(
+        com.OperationNotDefinedError,
+        PsycoPgIndeterminateDatatype,
+        PsycoPgInternalError,
+    ),
+)
+@pytest.mark.notyet(
+    ["risingwave"],
+    reason="Risingwave doesn't truly support arrays of arrays",
     raises=(
         com.OperationNotDefinedError,
         PsycoPg2IndeterminateDatatype,
@@ -1769,7 +1780,7 @@ def test_table_unnest_column_expr(backend):
 )
 @pytest.mark.notimpl(["trino"], raises=TrinoUserError)
 @pytest.mark.notimpl(["athena"], raises=PyAthenaOperationalError)
-@pytest.mark.notimpl(["postgres"], raises=PsycoPg2SyntaxError)
+@pytest.mark.notimpl(["postgres"], raises=PsycoPgSyntaxError)
 @pytest.mark.notimpl(["risingwave"], raises=PsycoPg2ProgrammingError)
 @pytest.mark.notyet(
     ["risingwave"], raises=PsycoPg2InternalError, reason="not supported in risingwave"
@@ -1887,7 +1898,7 @@ def test_array_agg_bool(con, data, agg, baseline_func):
 
 @pytest.mark.notyet(
     ["postgres"],
-    raises=PsycoPg2ArraySubscriptError,
+    raises=PsycoPgInvalidTextRepresentation,
     reason="all dimensions must match in size",
 )
 @pytest.mark.notimpl(["risingwave", "flink"], raises=com.OperationNotDefinedError)

ibis/backends/tests/test_client.py

@@ -32,7 +32,7 @@
     ImpalaHiveServer2Error,
     OracleDatabaseError,
     PsycoPg2InternalError,
-    PsycoPg2UndefinedObject,
+    PsycoPgUndefinedObject,
     Py4JJavaError,
     PyAthenaDatabaseError,
     PyODBCProgrammingError,
@@ -725,7 +725,7 @@ def test_list_database_contents(con):
 @pytest.mark.notyet(["databricks"], raises=DatabricksServerOperationError)
 @pytest.mark.notyet(["bigquery"], raises=com.UnsupportedBackendType)
 @pytest.mark.notyet(
-    ["postgres"], raises=PsycoPg2UndefinedObject, reason="no unsigned int types"
+    ["postgres"], raises=PsycoPgUndefinedObject, reason="no unsigned int types"
 )
 @pytest.mark.notyet(
     ["oracle"], raises=OracleDatabaseError, reason="no unsigned int types"

ibis/backends/tests/test_generic.py

@@ -25,7 +25,7 @@
     OracleDatabaseError,
     PolarsInvalidOperationError,
     PsycoPg2InternalError,
-    PsycoPg2SyntaxError,
+    PsycoPgSyntaxError,
     Py4JJavaError,
     PyAthenaDatabaseError,
     PyAthenaOperationalError,
@@ -1739,7 +1739,7 @@ def hash_256(col):
                 pytest.mark.notimpl(["flink"], raises=Py4JJavaError),
                 pytest.mark.notimpl(["druid"], raises=PyDruidProgrammingError),
                 pytest.mark.notimpl(["oracle"], raises=OracleDatabaseError),
-                pytest.mark.notimpl(["postgres"], raises=PsycoPg2SyntaxError),
+                pytest.mark.notimpl(["postgres"], raises=PsycoPgSyntaxError),
                 pytest.mark.notimpl(["risingwave"], raises=PsycoPg2InternalError),
                 pytest.mark.notimpl(["snowflake"], raises=AssertionError),
                 pytest.mark.never(

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):