81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
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)
|