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)