Fixed #33308 -- Added support for psycopg version 3.

Thanks Simon Charette, Tim Graham, and Adam Johnson for reviews.

Co-authored-by: Florian Apolloner <florian@apolloner.eu>
Co-authored-by: Mariusz Felisiak <felisiak.mariusz@gmail.com>

11 files changed, 54 insertions, 38 deletions

diff --git a/tests/backends/postgresql/tests.py b/tests/backends/postgresql/tests.py
index 41d445e6c7..7e1a2d000d 100644
--- a/tests/backends/postgresql/tests.py
+++ b/tests/backends/postgresql/tests.py
@@ -14,6 +14,11 @@ from django.db import (
 from django.db.backends.base.base import BaseDatabaseWrapper
 from django.test import TestCase, override_settings
 
+try:
+    from django.db.backends.postgresql.psycopg_any import is_psycopg3
+except ImportError:
+    is_psycopg3 = False
+
 
 @unittest.skipUnless(connection.vendor == "postgresql", "PostgreSQL tests")
 class Tests(TestCase):
@@ -228,7 +233,7 @@ class Tests(TestCase):
         # Since this is a django.test.TestCase, a transaction is in progress
         # and the isolation level isn't reported as 0. This test assumes that
         # PostgreSQL is configured with the default isolation level.
-        # Check the level on the psycopg2 connection, not the Django wrapper.
+        # Check the level on the psycopg connection, not the Django wrapper.
         self.assertIsNone(connection.connection.isolation_level)
 
         new_connection = connection.copy()
@@ -238,7 +243,7 @@ class Tests(TestCase):
         try:
             # Start a transaction so the isolation level isn't reported as 0.
             new_connection.set_autocommit(False)
-            # Check the level on the psycopg2 connection, not the Django wrapper.
+            # Check the level on the psycopg connection, not the Django wrapper.
             self.assertEqual(
                 new_connection.connection.isolation_level,
                 IsolationLevel.SERIALIZABLE,
@@ -252,7 +257,7 @@ class Tests(TestCase):
         new_connection.settings_dict["OPTIONS"]["isolation_level"] = -1
         msg = (
             "Invalid transaction isolation level -1 specified. Use one of the "
-            "IsolationLevel values."
+            "psycopg.IsolationLevel values."
         )
         with self.assertRaisesMessage(ImproperlyConfigured, msg):
             new_connection.ensure_connection()
@@ -268,7 +273,7 @@ class Tests(TestCase):
 
     def _select(self, val):
         with connection.cursor() as cursor:
-            cursor.execute("SELECT %s", (val,))
+            cursor.execute("SELECT %s::text[]", (val,))
             return cursor.fetchone()[0]
 
     def test_select_ascii_array(self):
@@ -308,17 +313,18 @@ class Tests(TestCase):
                     )
 
     def test_correct_extraction_psycopg_version(self):
-        from django.db.backends.postgresql.base import Database, psycopg2_version
+        from django.db.backends.postgresql.base import Database, psycopg_version
 
         with mock.patch.object(Database, "__version__", "4.2.1 (dt dec pq3 ext lo64)"):
-            self.assertEqual(psycopg2_version(), (4, 2, 1))
+            self.assertEqual(psycopg_version(), (4, 2, 1))
         with mock.patch.object(
             Database, "__version__", "4.2b0.dev1 (dt dec pq3 ext lo64)"
         ):
-            self.assertEqual(psycopg2_version(), (4, 2))
+            self.assertEqual(psycopg_version(), (4, 2))
 
     @override_settings(DEBUG=True)
-    def test_copy_cursors(self):
+    @unittest.skipIf(is_psycopg3, "psycopg2 specific test")
+    def test_copy_to_expert_cursors(self):
         out = StringIO()
         copy_expert_sql = "COPY django_session TO STDOUT (FORMAT CSV, HEADER)"
         with connection.cursor() as cursor:
@@ -329,6 +335,16 @@ class Tests(TestCase):
             [copy_expert_sql, "COPY django_session TO STDOUT"],
         )
 
+    @override_settings(DEBUG=True)
+    @unittest.skipUnless(is_psycopg3, "psycopg3 specific test")
+    def test_copy_cursors(self):
+        copy_sql = "COPY django_session TO STDOUT (FORMAT CSV, HEADER)"
+        with connection.cursor() as cursor:
+            with cursor.copy(copy_sql) as copy:
+                for row in copy:
+                    pass
+        self.assertEqual([q["sql"] for q in connection.queries], [copy_sql])
+
     def test_get_database_version(self):
         new_connection = connection.copy()
         new_connection.pg_version = 110009
diff --git a/tests/backends/tests.py b/tests/backends/tests.py
index c3cfa61fdb..5f11f91958 100644
--- a/tests/backends/tests.py
+++ b/tests/backends/tests.py
@@ -454,7 +454,7 @@ class BackendTestCase(TransactionTestCase):
         with connection.cursor() as cursor:
             self.assertIsInstance(cursor, CursorWrapper)
         # Both InterfaceError and ProgrammingError seem to be used when
-        # accessing closed cursor (psycopg2 has InterfaceError, rest seem
+        # accessing closed cursor (psycopg has InterfaceError, rest seem
         # to use ProgrammingError).
         with self.assertRaises(connection.features.closed_cursor_error_class):
             # cursor should be closed, so no queries should be possible.
@@ -462,12 +462,12 @@ class BackendTestCase(TransactionTestCase):
 
     @unittest.skipUnless(
         connection.vendor == "postgresql",
-        "Psycopg2 specific cursor.closed attribute needed",
+        "Psycopg specific cursor.closed attribute needed",
     )
     def test_cursor_contextmanager_closing(self):
         # There isn't a generic way to test that cursors are closed, but
-        # psycopg2 offers us a way to check that by closed attribute.
-        # So, run only on psycopg2 for that reason.
+        # psycopg offers us a way to check that by closed attribute.
+        # So, run only on psycopg for that reason.
         with connection.cursor() as cursor:
             self.assertIsInstance(cursor, CursorWrapper)
         self.assertTrue(cursor.closed)
diff --git a/tests/db_functions/datetime/test_extract_trunc.py b/tests/db_functions/datetime/test_extract_trunc.py
index 2d38234981..b2327931f0 100644
--- a/tests/db_functions/datetime/test_extract_trunc.py
+++ b/tests/db_functions/datetime/test_extract_trunc.py
@@ -245,7 +245,7 @@ class DateFunctionTests(TestCase):
         self.create_model(start_datetime, end_datetime)
         self.create_model(end_datetime, start_datetime)
 
-        with self.assertRaises((DataError, OperationalError, ValueError)):
+        with self.assertRaises((OperationalError, ValueError)):
             DTModel.objects.filter(
                 start_datetime__year=Extract(
                     "start_datetime", "day' FROM start_datetime)) OR 1=1;--"
diff --git a/tests/db_utils/tests.py b/tests/db_utils/tests.py
index 9c0ec905cc..a2d9cc7b5e 100644
--- a/tests/db_utils/tests.py
+++ b/tests/db_utils/tests.py
@@ -62,14 +62,20 @@ class ConnectionHandlerTests(SimpleTestCase):
 class DatabaseErrorWrapperTests(TestCase):
     @unittest.skipUnless(connection.vendor == "postgresql", "PostgreSQL test")
     def test_reraising_backend_specific_database_exception(self):
+        from django.db.backends.postgresql.psycopg_any import is_psycopg3
+
         with connection.cursor() as cursor:
             msg = 'table "X" does not exist'
             with self.assertRaisesMessage(ProgrammingError, msg) as cm:
                 cursor.execute('DROP TABLE "X"')
         self.assertNotEqual(type(cm.exception), type(cm.exception.__cause__))
         self.assertIsNotNone(cm.exception.__cause__)
-        self.assertIsNotNone(cm.exception.__cause__.pgcode)
-        self.assertIsNotNone(cm.exception.__cause__.pgerror)
+        if is_psycopg3:
+            self.assertIsNotNone(cm.exception.__cause__.diag.sqlstate)
+            self.assertIsNotNone(cm.exception.__cause__.diag.message_primary)
+        else:
+            self.assertIsNotNone(cm.exception.__cause__.pgcode)
+            self.assertIsNotNone(cm.exception.__cause__.pgerror)
 
 
 class LoadBackendTests(SimpleTestCase):
diff --git a/tests/fixtures/tests.py b/tests/fixtures/tests.py
index 9eb2740c90..deac1c2d77 100644
--- a/tests/fixtures/tests.py
+++ b/tests/fixtures/tests.py
@@ -916,15 +916,11 @@ class FixtureLoadingTests(DumpDataAssertMixin, TestCase):
         with self.assertRaisesMessage(IntegrityError, msg):
             management.call_command("loaddata", "invalid.json", verbosity=0)
 
-    @unittest.skipUnless(
-        connection.vendor == "postgresql", "psycopg2 prohibits null characters in data."
-    )
+    @skipUnlessDBFeature("prohibits_null_characters_in_text_exception")
     def test_loaddata_null_characters_on_postgresql(self):
-        msg = (
-            "Could not load fixtures.Article(pk=2): "
-            "A string literal cannot contain NUL (0x00) characters."
-        )
-        with self.assertRaisesMessage(ValueError, msg):
+        error, msg = connection.features.prohibits_null_characters_in_text_exception
+        msg = f"Could not load fixtures.Article(pk=2): {msg}"
+        with self.assertRaisesMessage(error, msg):
             management.call_command("loaddata", "null_character_in_field_value.json")
 
     def test_loaddata_app_option(self):
diff --git a/tests/gis_tests/tests.py b/tests/gis_tests/tests.py
index d1c93592a8..9da2b4df99 100644
--- a/tests/gis_tests/tests.py
+++ b/tests/gis_tests/tests.py
@@ -36,7 +36,7 @@ if HAS_POSTGRES:
                 raise NotImplementedError("This function was not expected to be called")
 
 
-@unittest.skipUnless(HAS_POSTGRES, "The psycopg2 driver is needed for these tests")
+@unittest.skipUnless(HAS_POSTGRES, "The psycopg driver is needed for these tests")
 class TestPostGISVersionCheck(unittest.TestCase):
     """
     The PostGIS version check parses correctly the version numbers
diff --git a/tests/model_fields/test_jsonfield.py b/tests/model_fields/test_jsonfield.py
index 05816817ef..60357d87b2 100644
--- a/tests/model_fields/test_jsonfield.py
+++ b/tests/model_fields/test_jsonfield.py
@@ -1007,7 +1007,7 @@ class TestQuerying(TestCase):
                 False,
             )
         self.assertIn(
-            """."value" -> 'test'' = ''"a"'') OR 1 = 1 OR (''d') = '"x"' """,
+            """."value" -> 'test'' = ''"a"'') OR 1 = 1 OR (''d') = '"x"'""",
             queries[0]["sql"],
         )
 
diff --git a/tests/postgres_tests/test_apps.py b/tests/postgres_tests/test_apps.py
index 7c4cc38183..d9fb962251 100644
--- a/tests/postgres_tests/test_apps.py
+++ b/tests/postgres_tests/test_apps.py
@@ -19,6 +19,7 @@ try:
         DateTimeRange,
         DateTimeTZRange,
         NumericRange,
+        is_psycopg3,
     )
 except ImportError:
     pass
@@ -59,6 +60,7 @@ class PostgresConfigTests(TestCase):
                         MigrationWriter.serialize(field)
 
         assertNotSerializable()
+        import_name = "psycopg.types.range" if is_psycopg3 else "psycopg2.extras"
         with self.modify_settings(INSTALLED_APPS={"append": "django.contrib.postgres"}):
             for default, test_field in tests:
                 with self.subTest(default=default):
@@ -68,16 +70,12 @@ class PostgresConfigTests(TestCase):
                         imports,
                         {
                             "import django.contrib.postgres.fields.ranges",
-                            "import psycopg2.extras",
+                            f"import {import_name}",
                         },
                     )
                     self.assertIn(
-                        "%s.%s(default=psycopg2.extras.%r)"
-                        % (
-                            field.__module__,
-                            field.__class__.__name__,
-                            default,
-                        ),
+                        f"{field.__module__}.{field.__class__.__name__}"
+                        f"(default={import_name}.{default!r})",
                         serialized_field,
                     )
         assertNotSerializable()
diff --git a/tests/postgres_tests/test_array.py b/tests/postgres_tests/test_array.py
index a3c26fddae..86e9d00b41 100644
--- a/tests/postgres_tests/test_array.py
+++ b/tests/postgres_tests/test_array.py
@@ -317,7 +317,7 @@ class TestQuerying(PostgreSQLTestCase):
     def test_in_including_F_object(self):
         # This test asserts that Array objects passed to filters can be
         # constructed to contain F objects. This currently doesn't work as the
-        # psycopg2 mogrify method that generates the ARRAY() syntax is
+        # psycopg mogrify method that generates the ARRAY() syntax is
         # expecting literals, not column references (#27095).
         self.assertSequenceEqual(
             NullableIntegerArrayModel.objects.filter(field__in=[[models.F("id")]]),
diff --git a/tests/requirements/postgres.txt b/tests/requirements/postgres.txt
index f0288c8b74..726a08b3e4 100644
--- a/tests/requirements/postgres.txt
+++ b/tests/requirements/postgres.txt
@@ -1 +1 @@
-psycopg2>=2.8.4
+psycopg[binary]>=3.1
diff --git a/tests/schema/test_logging.py b/tests/schema/test_logging.py
index 2821e5f406..9c7069c874 100644
--- a/tests/schema/test_logging.py
+++ b/tests/schema/test_logging.py
@@ -9,9 +9,9 @@ class SchemaLoggerTests(TestCase):
         params = [42, 1337]
         with self.assertLogs("django.db.backends.schema", "DEBUG") as cm:
             editor.execute(sql, params)
+        if connection.features.schema_editor_uses_clientside_param_binding:
+            sql = "SELECT * FROM foo WHERE id in (42, 1337)"
+            params = None
         self.assertEqual(cm.records[0].sql, sql)
         self.assertEqual(cm.records[0].params, params)
-        self.assertEqual(
-            cm.records[0].getMessage(),
-            "SELECT * FROM foo WHERE id in (%s, %s); (params [42, 1337])",
-        )
+        self.assertEqual(cm.records[0].getMessage(), f"{sql}; (params {params})")