1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980 |
- 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)
|