You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

signals.py 2.2KB

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