start-pack

myenv/lib/python3.12/site-packages/django/contrib/postgres/signals.py

@@ -0,0 +1,80 @@
+import functools
+
+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):
+    with connections[connection_alias].cursor() as cursor:
+        cursor.execute(
+            "SELECT oid, typarray FROM pg_type WHERE typname = %s", (type_name,)
+        )
+        oids = []
+        array_oids = []
+        for row in cursor:
+            oids.append(row[0])
+            array_oids.append(row[1])
+        return tuple(oids), tuple(array_oids)
+
+
+@functools.lru_cache
+def get_hstore_oids(connection_alias):
+    """Return hstore and hstore array OIDs."""
+    return get_type_oids(connection_alias, "hstore")
+
+
+@functools.lru_cache
+def get_citext_oids(connection_alias):
+    """Return citext and citext array OIDs."""
+    return get_type_oids(connection_alias, "citext")
+
+
+if is_psycopg3:
+    from psycopg.types import TypeInfo, 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)
+        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)