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, 48 insertions, 29 deletions

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
-
-    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
-        )
+if is_psycopg3:
+    from psycopg.types import TypeInfo, hstore
 
-    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)
+    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)
+        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)