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>

1 files changed, 92 insertions, 21 deletions

diff --git a/django/db/backends/postgresql/psycopg_any.py b/django/db/backends/postgresql/psycopg_any.py
index e9bb84f313..579104dead 100644
--- a/django/db/backends/postgresql/psycopg_any.py
+++ b/django/db/backends/postgresql/psycopg_any.py
@@ -1,31 +1,102 @@
-from enum import IntEnum
+import ipaddress
+from functools import lru_cache
 
-from psycopg2 import errors, extensions, sql  # NOQA
-from psycopg2.extras import DateRange, DateTimeRange, DateTimeTZRange, Inet  # NOQA
-from psycopg2.extras import Json as Jsonb  # NOQA
-from psycopg2.extras import NumericRange, Range  # NOQA
+try:
+    from psycopg import ClientCursor, IsolationLevel, adapt, adapters, errors, sql
+    from psycopg.postgres import types
+    from psycopg.types.datetime import TimestamptzLoader
+    from psycopg.types.json import Jsonb
+    from psycopg.types.range import Range, RangeDumper
+    from psycopg.types.string import TextLoader
 
-RANGE_TYPES = (DateRange, DateTimeRange, DateTimeTZRange, NumericRange)
+    Inet = ipaddress.ip_address
 
+    DateRange = DateTimeRange = DateTimeTZRange = NumericRange = Range
+    RANGE_TYPES = (Range,)
 
-class IsolationLevel(IntEnum):
-    READ_UNCOMMITTED = extensions.ISOLATION_LEVEL_READ_UNCOMMITTED
-    READ_COMMITTED = extensions.ISOLATION_LEVEL_READ_COMMITTED
-    REPEATABLE_READ = extensions.ISOLATION_LEVEL_REPEATABLE_READ
-    SERIALIZABLE = extensions.ISOLATION_LEVEL_SERIALIZABLE
+    TSRANGE_OID = types["tsrange"].oid
+    TSTZRANGE_OID = types["tstzrange"].oid
 
+    def mogrify(sql, params, connection):
+        return ClientCursor(connection.connection).mogrify(sql, params)
 
-def _quote(value, connection=None):
-    adapted = extensions.adapt(value)
-    if hasattr(adapted, "encoding"):
-        adapted.encoding = "utf8"
-    # getquoted() returns a quoted bytestring of the adapted value.
-    return adapted.getquoted().decode()
+    # Adapters.
+    class BaseTzLoader(TimestamptzLoader):
+        """
+        Load a PostgreSQL timestamptz using the a specific timezone.
+        The timezone can be None too, in which case it will be chopped.
+        """
 
+        timezone = None
 
-sql.quote = _quote
+        def load(self, data):
+            res = super().load(data)
+            return res.replace(tzinfo=self.timezone)
 
+    def register_tzloader(tz, context):
+        class SpecificTzLoader(BaseTzLoader):
+            timezone = tz
 
-def mogrify(sql, params, connection):
-    with connection.cursor() as cursor:
-        return cursor.mogrify(sql, params).decode()
+        context.adapters.register_loader("timestamptz", SpecificTzLoader)
+
+    class DjangoRangeDumper(RangeDumper):
+        """A Range dumper customized for Django."""
+
+        def upgrade(self, obj, format):
+            # Dump ranges containing naive datetimes as tstzrange, because
+            # Django doesn't use tz-aware ones.
+            dumper = super().upgrade(obj, format)
+            if dumper is not self and dumper.oid == TSRANGE_OID:
+                dumper.oid = TSTZRANGE_OID
+            return dumper
+
+    @lru_cache
+    def get_adapters_template(use_tz, timezone):
+        # Create at adapters map extending the base one.
+        ctx = adapt.AdaptersMap(adapters)
+        # Register a no-op dumper to avoid a round trip from psycopg version 3
+        # decode to json.dumps() to json.loads(), when using a custom decoder
+        # in JSONField.
+        ctx.register_loader("jsonb", TextLoader)
+        # Don't convert automatically from PostgreSQL network types to Python
+        # ipaddress.
+        ctx.register_loader("inet", TextLoader)
+        ctx.register_loader("cidr", TextLoader)
+        ctx.register_dumper(Range, DjangoRangeDumper)
+        # Register a timestamptz loader configured on self.timezone.
+        # This, however, can be overridden by create_cursor.
+        register_tzloader(timezone, ctx)
+        return ctx
+
+    is_psycopg3 = True
+
+except ImportError:
+    from enum import IntEnum
+
+    from psycopg2 import errors, extensions, sql  # NOQA
+    from psycopg2.extras import DateRange, DateTimeRange, DateTimeTZRange, Inet  # NOQA
+    from psycopg2.extras import Json as Jsonb  # NOQA
+    from psycopg2.extras import NumericRange, Range  # NOQA
+
+    RANGE_TYPES = (DateRange, DateTimeRange, DateTimeTZRange, NumericRange)
+
+    class IsolationLevel(IntEnum):
+        READ_UNCOMMITTED = extensions.ISOLATION_LEVEL_READ_UNCOMMITTED
+        READ_COMMITTED = extensions.ISOLATION_LEVEL_READ_COMMITTED
+        REPEATABLE_READ = extensions.ISOLATION_LEVEL_REPEATABLE_READ
+        SERIALIZABLE = extensions.ISOLATION_LEVEL_SERIALIZABLE
+
+    def _quote(value, connection=None):
+        adapted = extensions.adapt(value)
+        if hasattr(adapted, "encoding"):
+            adapted.encoding = "utf8"
+        # getquoted() returns a quoted bytestring of the adapted value.
+        return adapted.getquoted().decode()
+
+    sql.quote = _quote
+
+    def mogrify(sql, params, connection):
+        with connection.cursor() as cursor:
+            return cursor.mogrify(sql, params).decode()
+
+    is_psycopg3 = False