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>

5 files changed, 74 insertions, 31 deletions

diff --git a/django/contrib/postgres/fields/array.py b/django/contrib/postgres/fields/array.py
index eaff032465..8477dd9fff 100644
--- a/django/contrib/postgres/fields/array.py
+++ b/django/contrib/postgres/fields/array.py
@@ -237,7 +237,7 @@ class ArrayField(CheckFieldDefaultMixin, Field):
 
 class ArrayRHSMixin:
     def __init__(self, lhs, rhs):
-        # Don't wrap arrays that contains only None values, psycopg2 doesn't
+        # Don't wrap arrays that contains only None values, psycopg doesn't
         # allow this.
         if isinstance(rhs, (tuple, list)) and any(self._rhs_not_none_values(rhs)):
             expressions = []
diff --git a/django/contrib/postgres/fields/ranges.py b/django/contrib/postgres/fields/ranges.py
index d5c438dbdc..fbb6012660 100644
--- a/django/contrib/postgres/fields/ranges.py
+++ b/django/contrib/postgres/fields/ranges.py
@@ -9,6 +9,7 @@ from django.db.backends.postgresql.psycopg_any import (
     NumericRange,
     Range,
 )
+from django.db.models.functions import Cast
 from django.db.models.lookups import PostgresOperatorLookup
 
 from .utils import AttributeSetter
@@ -208,7 +209,14 @@ class DateRangeField(RangeField):
         return "daterange"
 
 
-RangeField.register_lookup(lookups.DataContains)
+class RangeContains(lookups.DataContains):
+    def get_prep_lookup(self):
+        if not isinstance(self.rhs, (list, tuple, Range)):
+            return Cast(self.rhs, self.lhs.field.base_field)
+        return super().get_prep_lookup()
+
+
+RangeField.register_lookup(RangeContains)
 RangeField.register_lookup(lookups.ContainedBy)
 RangeField.register_lookup(lookups.Overlap)
 
diff --git a/django/contrib/postgres/operations.py b/django/contrib/postgres/operations.py
index 9dbd491773..5ac396bedf 100644
--- a/django/contrib/postgres/operations.py
+++ b/django/contrib/postgres/operations.py
@@ -35,6 +35,10 @@ class CreateExtension(Operation):
         # installed, otherwise a subsequent data migration would use the same
         # connection.
         register_type_handlers(schema_editor.connection)
+        if hasattr(schema_editor.connection, "register_geometry_adapters"):
+            schema_editor.connection.register_geometry_adapters(
+                schema_editor.connection.connection, True
+            )
 
     def database_backwards(self, app_label, schema_editor, from_state, to_state):
         if not router.allow_migrate(schema_editor.connection.alias, app_label):
diff --git a/django/contrib/postgres/search.py b/django/contrib/postgres/search.py
index 05c8f72f6f..4e370aa167 100644
--- a/django/contrib/postgres/search.py
+++ b/django/contrib/postgres/search.py
@@ -39,6 +39,11 @@ class SearchQueryField(Field):
         return "tsquery"
 
 
+class _Float4Field(Field):
+    def db_type(self, connection):
+        return "float4"
+
+
 class SearchConfig(Expression):
     def __init__(self, config):
         super().__init__()
@@ -138,7 +143,11 @@ class SearchVector(SearchVectorCombinable, Func):
         if clone.weight:
             weight_sql, extra_params = compiler.compile(clone.weight)
             sql = "setweight({}, {})".format(sql, weight_sql)
-        return sql, config_params + params + extra_params
+
+        # These parameters must be bound on the client side because we may
+        # want to create an index on this expression.
+        sql = connection.ops.compose_sql(sql, config_params + params + extra_params)
+        return sql, []
 
 
 class CombinedSearchVector(SearchVectorCombinable, CombinedExpression):
@@ -244,6 +253,8 @@ class SearchRank(Func):
         normalization=None,
         cover_density=False,
     ):
+        from .fields.array import ArrayField
+
         if not hasattr(vector, "resolve_expression"):
             vector = SearchVector(vector)
         if not hasattr(query, "resolve_expression"):
@@ -252,6 +263,7 @@ class SearchRank(Func):
         if weights is not None:
             if not hasattr(weights, "resolve_expression"):
                 weights = Value(weights)
+            weights = Cast(weights, ArrayField(_Float4Field()))
             expressions = (weights,) + expressions
         if normalization is not None:
             if not hasattr(normalization, "resolve_expression"):
diff --git a/django/contrib/postgres/signals.py b/django/contrib/postgres/signals.py
index 5c6ca3687a..a3816d3d30 100644
--- a/django/contrib/postgres/signals.py
+++ b/django/contrib/postgres/signals.py
@@ -1,10 +1,8 @@
 import functools
 
-import psycopg2
-from psycopg2.extras import register_hstore
-
 from django.db import connections
 from django.db.backends.base.base import NO_DB_ALIAS
+from django.db.backends.postgresql.psycopg_any import is_psycopg3
 
 
 def get_type_oids(connection_alias, type_name):
@@ -32,30 +30,51 @@ def get_citext_oids(connection_alias):
     return get_type_oids(connection_alias, "citext")
 
 
-def register_type_handlers(connection, **kwargs):
-    if connection.vendor != "postgresql" or connection.alias == NO_DB_ALIAS:
-        return
+if is_psycopg3:
+    from psycopg.types import TypeInfo, hstore
 
-    oids, array_oids = get_hstore_oids(connection.alias)
-    # Don't register handlers when hstore is not available on the database.
-    #
-    # If someone tries to create an hstore field it will error there. This is
-    # necessary as someone may be using PSQL without extensions installed but
-    # be using other features of contrib.postgres.
-    #
-    # This is also needed in order to create the connection in order to install
-    # the hstore extension.
-    if oids:
-        register_hstore(
-            connection.connection, globally=True, oid=oids, array_oid=array_oids
-        )
+    def register_type_handlers(connection, **kwargs):
+        if connection.vendor != "postgresql" or connection.alias == NO_DB_ALIAS:
+            return
 
-    oids, citext_oids = get_citext_oids(connection.alias)
-    # Don't register handlers when citext is not available on the database.
-    #
-    # The same comments in the above call to register_hstore() also apply here.
-    if oids:
-        array_type = psycopg2.extensions.new_array_type(
-            citext_oids, "citext[]", psycopg2.STRING
-        )
-        psycopg2.extensions.register_type(array_type, None)
+        oids, array_oids = get_hstore_oids(connection.alias)
+        for oid, array_oid in zip(oids, array_oids):
+            ti = TypeInfo("hstore", oid, array_oid)
+            hstore.register_hstore(ti, connection.connection)
+
+        _, citext_oids = get_citext_oids(connection.alias)
+        for array_oid in citext_oids:
+            ti = TypeInfo("citext", 0, array_oid)
+            ti.register(connection.connection)
+
+else:
+    import psycopg2
+    from psycopg2.extras import register_hstore
+
+    def register_type_handlers(connection, **kwargs):
+        if connection.vendor != "postgresql" or connection.alias == NO_DB_ALIAS:
+            return
+
+        oids, array_oids = get_hstore_oids(connection.alias)
+        # Don't register handlers when hstore is not available on the database.
+        #
+        # If someone tries to create an hstore field it will error there. This is
+        # necessary as someone may be using PSQL without extensions installed but
+        # be using other features of contrib.postgres.
+        #
+        # This is also needed in order to create the connection in order to install
+        # the hstore extension.
+        if oids:
+            register_hstore(
+                connection.connection, globally=True, oid=oids, array_oid=array_oids
+            )
+
+        oids, citext_oids = get_citext_oids(connection.alias)
+        # Don't register handlers when citext is not available on the database.
+        #
+        # The same comments in the above call to register_hstore() also apply here.
+        if oids:
+            array_type = psycopg2.extensions.new_array_type(
+                citext_oids, "citext[]", psycopg2.STRING
+            )
+            psycopg2.extensions.register_type(array_type, None)