signals.py 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980
  1. import functools
  2. from django.db import connections
  3. from django.db.backends.base.base import NO_DB_ALIAS
  4. from django.db.backends.postgresql.psycopg_any import is_psycopg3
  5. def get_type_oids(connection_alias, type_name):
  6. with connections[connection_alias].cursor() as cursor:
  7. cursor.execute(
  8. "SELECT oid, typarray FROM pg_type WHERE typname = %s", (type_name,)
  9. )
  10. oids = []
  11. array_oids = []
  12. for row in cursor:
  13. oids.append(row[0])
  14. array_oids.append(row[1])
  15. return tuple(oids), tuple(array_oids)
  16. @functools.lru_cache
  17. def get_hstore_oids(connection_alias):
  18. """Return hstore and hstore array OIDs."""
  19. return get_type_oids(connection_alias, "hstore")
  20. @functools.lru_cache
  21. def get_citext_oids(connection_alias):
  22. """Return citext and citext array OIDs."""
  23. return get_type_oids(connection_alias, "citext")
  24. if is_psycopg3:
  25. from psycopg.types import TypeInfo, hstore
  26. def register_type_handlers(connection, **kwargs):
  27. if connection.vendor != "postgresql" or connection.alias == NO_DB_ALIAS:
  28. return
  29. oids, array_oids = get_hstore_oids(connection.alias)
  30. for oid, array_oid in zip(oids, array_oids):
  31. ti = TypeInfo("hstore", oid, array_oid)
  32. hstore.register_hstore(ti, connection.connection)
  33. _, citext_oids = get_citext_oids(connection.alias)
  34. for array_oid in citext_oids:
  35. ti = TypeInfo("citext", 0, array_oid)
  36. ti.register(connection.connection)
  37. else:
  38. import psycopg2
  39. from psycopg2.extras import register_hstore
  40. def register_type_handlers(connection, **kwargs):
  41. if connection.vendor != "postgresql" or connection.alias == NO_DB_ALIAS:
  42. return
  43. oids, array_oids = get_hstore_oids(connection.alias)
  44. # Don't register handlers when hstore is not available on the database.
  45. #
  46. # If someone tries to create an hstore field it will error there. This is
  47. # necessary as someone may be using PSQL without extensions installed but
  48. # be using other features of contrib.postgres.
  49. #
  50. # This is also needed in order to create the connection in order to install
  51. # the hstore extension.
  52. if oids:
  53. register_hstore(
  54. connection.connection, globally=True, oid=oids, array_oid=array_oids
  55. )
  56. oids, citext_oids = get_citext_oids(connection.alias)
  57. # Don't register handlers when citext is not available on the database.
  58. #
  59. # The same comments in the above call to register_hstore() also apply here.
  60. if oids:
  61. array_type = psycopg2.extensions.new_array_type(
  62. citext_oids, "citext[]", psycopg2.STRING
  63. )
  64. psycopg2.extensions.register_type(array_type, None)