base.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326
  1. """
  2. PostgreSQL database backend for Django.
  3. Requires psycopg 2: https://www.psycopg.org/
  4. """
  5. import asyncio
  6. import threading
  7. import warnings
  8. from django.conf import settings
  9. from django.core.exceptions import ImproperlyConfigured
  10. from django.db import connections
  11. from django.db.backends.base.base import BaseDatabaseWrapper
  12. from django.db.backends.utils import (
  13. CursorDebugWrapper as BaseCursorDebugWrapper,
  14. )
  15. from django.db.utils import DatabaseError as WrappedDatabaseError
  16. from django.utils.asyncio import async_unsafe
  17. from django.utils.functional import cached_property
  18. from django.utils.safestring import SafeString
  19. from django.utils.version import get_version_tuple
  20. try:
  21. import psycopg2 as Database
  22. import psycopg2.extensions
  23. import psycopg2.extras
  24. except ImportError as e:
  25. raise ImproperlyConfigured("Error loading psycopg2 module: %s" % e)
  26. def psycopg2_version():
  27. version = psycopg2.__version__.split(' ', 1)[0]
  28. return get_version_tuple(version)
  29. PSYCOPG2_VERSION = psycopg2_version()
  30. if PSYCOPG2_VERSION < (2, 5, 4):
  31. raise ImproperlyConfigured("psycopg2_version 2.5.4 or newer is required; you have %s" % psycopg2.__version__)
  32. # Some of these import psycopg2, so import them after checking if it's installed.
  33. from .client import DatabaseClient # NOQA isort:skip
  34. from .creation import DatabaseCreation # NOQA isort:skip
  35. from .features import DatabaseFeatures # NOQA isort:skip
  36. from .introspection import DatabaseIntrospection # NOQA isort:skip
  37. from .operations import DatabaseOperations # NOQA isort:skip
  38. from .schema import DatabaseSchemaEditor # NOQA isort:skip
  39. from .utils import utc_tzinfo_factory # NOQA isort:skip
  40. psycopg2.extensions.register_adapter(SafeString, psycopg2.extensions.QuotedString)
  41. psycopg2.extras.register_uuid()
  42. # Register support for inet[] manually so we don't have to handle the Inet()
  43. # object on load all the time.
  44. INETARRAY_OID = 1041
  45. INETARRAY = psycopg2.extensions.new_array_type(
  46. (INETARRAY_OID,),
  47. 'INETARRAY',
  48. psycopg2.extensions.UNICODE,
  49. )
  50. psycopg2.extensions.register_type(INETARRAY)
  51. class DatabaseWrapper(BaseDatabaseWrapper):
  52. vendor = 'postgresql'
  53. display_name = 'PostgreSQL'
  54. # This dictionary maps Field objects to their associated PostgreSQL column
  55. # types, as strings. Column-type strings can contain format strings; they'll
  56. # be interpolated against the values of Field.__dict__ before being output.
  57. # If a column type is set to None, it won't be included in the output.
  58. data_types = {
  59. 'AutoField': 'serial',
  60. 'BigAutoField': 'bigserial',
  61. 'BinaryField': 'bytea',
  62. 'BooleanField': 'boolean',
  63. 'CharField': 'varchar(%(max_length)s)',
  64. 'DateField': 'date',
  65. 'DateTimeField': 'timestamp with time zone',
  66. 'DecimalField': 'numeric(%(max_digits)s, %(decimal_places)s)',
  67. 'DurationField': 'interval',
  68. 'FileField': 'varchar(%(max_length)s)',
  69. 'FilePathField': 'varchar(%(max_length)s)',
  70. 'FloatField': 'double precision',
  71. 'IntegerField': 'integer',
  72. 'BigIntegerField': 'bigint',
  73. 'IPAddressField': 'inet',
  74. 'GenericIPAddressField': 'inet',
  75. 'NullBooleanField': 'boolean',
  76. 'OneToOneField': 'integer',
  77. 'PositiveIntegerField': 'integer',
  78. 'PositiveSmallIntegerField': 'smallint',
  79. 'SlugField': 'varchar(%(max_length)s)',
  80. 'SmallAutoField': 'smallserial',
  81. 'SmallIntegerField': 'smallint',
  82. 'TextField': 'text',
  83. 'TimeField': 'time',
  84. 'UUIDField': 'uuid',
  85. }
  86. data_type_check_constraints = {
  87. 'PositiveIntegerField': '"%(column)s" >= 0',
  88. 'PositiveSmallIntegerField': '"%(column)s" >= 0',
  89. }
  90. operators = {
  91. 'exact': '= %s',
  92. 'iexact': '= UPPER(%s)',
  93. 'contains': 'LIKE %s',
  94. 'icontains': 'LIKE UPPER(%s)',
  95. 'regex': '~ %s',
  96. 'iregex': '~* %s',
  97. 'gt': '> %s',
  98. 'gte': '>= %s',
  99. 'lt': '< %s',
  100. 'lte': '<= %s',
  101. 'startswith': 'LIKE %s',
  102. 'endswith': 'LIKE %s',
  103. 'istartswith': 'LIKE UPPER(%s)',
  104. 'iendswith': 'LIKE UPPER(%s)',
  105. }
  106. # The patterns below are used to generate SQL pattern lookup clauses when
  107. # the right-hand side of the lookup isn't a raw string (it might be an expression
  108. # or the result of a bilateral transformation).
  109. # In those cases, special characters for LIKE operators (e.g. \, *, _) should be
  110. # escaped on database side.
  111. #
  112. # Note: we use str.format() here for readability as '%' is used as a wildcard for
  113. # the LIKE operator.
  114. pattern_esc = r"REPLACE(REPLACE(REPLACE({}, E'\\', E'\\\\'), E'%%', E'\\%%'), E'_', E'\\_')"
  115. pattern_ops = {
  116. 'contains': "LIKE '%%' || {} || '%%'",
  117. 'icontains': "LIKE '%%' || UPPER({}) || '%%'",
  118. 'startswith': "LIKE {} || '%%'",
  119. 'istartswith': "LIKE UPPER({}) || '%%'",
  120. 'endswith': "LIKE '%%' || {}",
  121. 'iendswith': "LIKE '%%' || UPPER({})",
  122. }
  123. Database = Database
  124. SchemaEditorClass = DatabaseSchemaEditor
  125. # Classes instantiated in __init__().
  126. client_class = DatabaseClient
  127. creation_class = DatabaseCreation
  128. features_class = DatabaseFeatures
  129. introspection_class = DatabaseIntrospection
  130. ops_class = DatabaseOperations
  131. # PostgreSQL backend-specific attributes.
  132. _named_cursor_idx = 0
  133. def get_connection_params(self):
  134. settings_dict = self.settings_dict
  135. # None may be used to connect to the default 'postgres' db
  136. if settings_dict['NAME'] == '':
  137. raise ImproperlyConfigured(
  138. "settings.DATABASES is improperly configured. "
  139. "Please supply the NAME value.")
  140. if len(settings_dict['NAME'] or '') > self.ops.max_name_length():
  141. raise ImproperlyConfigured(
  142. "The database name '%s' (%d characters) is longer than "
  143. "PostgreSQL's limit of %d characters. Supply a shorter NAME "
  144. "in settings.DATABASES." % (
  145. settings_dict['NAME'],
  146. len(settings_dict['NAME']),
  147. self.ops.max_name_length(),
  148. )
  149. )
  150. conn_params = {
  151. 'database': settings_dict['NAME'] or 'postgres',
  152. **settings_dict['OPTIONS'],
  153. }
  154. conn_params.pop('isolation_level', None)
  155. if settings_dict['USER']:
  156. conn_params['user'] = settings_dict['USER']
  157. if settings_dict['PASSWORD']:
  158. conn_params['password'] = settings_dict['PASSWORD']
  159. if settings_dict['HOST']:
  160. conn_params['host'] = settings_dict['HOST']
  161. if settings_dict['PORT']:
  162. conn_params['port'] = settings_dict['PORT']
  163. return conn_params
  164. @async_unsafe
  165. def get_new_connection(self, conn_params):
  166. connection = Database.connect(**conn_params)
  167. # self.isolation_level must be set:
  168. # - after connecting to the database in order to obtain the database's
  169. # default when no value is explicitly specified in options.
  170. # - before calling _set_autocommit() because if autocommit is on, that
  171. # will set connection.isolation_level to ISOLATION_LEVEL_AUTOCOMMIT.
  172. options = self.settings_dict['OPTIONS']
  173. try:
  174. self.isolation_level = options['isolation_level']
  175. except KeyError:
  176. self.isolation_level = connection.isolation_level
  177. else:
  178. # Set the isolation level to the value from OPTIONS.
  179. if self.isolation_level != connection.isolation_level:
  180. connection.set_session(isolation_level=self.isolation_level)
  181. return connection
  182. def ensure_timezone(self):
  183. if self.connection is None:
  184. return False
  185. conn_timezone_name = self.connection.get_parameter_status('TimeZone')
  186. timezone_name = self.timezone_name
  187. if timezone_name and conn_timezone_name != timezone_name:
  188. with self.connection.cursor() as cursor:
  189. cursor.execute(self.ops.set_time_zone_sql(), [timezone_name])
  190. return True
  191. return False
  192. def init_connection_state(self):
  193. self.connection.set_client_encoding('UTF8')
  194. timezone_changed = self.ensure_timezone()
  195. if timezone_changed:
  196. # Commit after setting the time zone (see #17062)
  197. if not self.get_autocommit():
  198. self.connection.commit()
  199. @async_unsafe
  200. def create_cursor(self, name=None):
  201. if name:
  202. # In autocommit mode, the cursor will be used outside of a
  203. # transaction, hence use a holdable cursor.
  204. cursor = self.connection.cursor(name, scrollable=False, withhold=self.connection.autocommit)
  205. else:
  206. cursor = self.connection.cursor()
  207. cursor.tzinfo_factory = utc_tzinfo_factory if settings.USE_TZ else None
  208. return cursor
  209. @async_unsafe
  210. def chunked_cursor(self):
  211. self._named_cursor_idx += 1
  212. # Get the current async task
  213. # Note that right now this is behind @async_unsafe, so this is
  214. # unreachable, but in future we'll start loosening this restriction.
  215. # For now, it's here so that every use of "threading" is
  216. # also async-compatible.
  217. try:
  218. if hasattr(asyncio, 'current_task'):
  219. # Python 3.7 and up
  220. current_task = asyncio.current_task()
  221. else:
  222. # Python 3.6
  223. current_task = asyncio.Task.current_task()
  224. except RuntimeError:
  225. current_task = None
  226. # Current task can be none even if the current_task call didn't error
  227. if current_task:
  228. task_ident = str(id(current_task))
  229. else:
  230. task_ident = 'sync'
  231. # Use that and the thread ident to get a unique name
  232. return self._cursor(
  233. name='_django_curs_%d_%s_%d' % (
  234. # Avoid reusing name in other threads / tasks
  235. threading.current_thread().ident,
  236. task_ident,
  237. self._named_cursor_idx,
  238. )
  239. )
  240. def _set_autocommit(self, autocommit):
  241. with self.wrap_database_errors:
  242. self.connection.autocommit = autocommit
  243. def check_constraints(self, table_names=None):
  244. """
  245. Check constraints by setting them to immediate. Return them to deferred
  246. afterward.
  247. """
  248. self.cursor().execute('SET CONSTRAINTS ALL IMMEDIATE')
  249. self.cursor().execute('SET CONSTRAINTS ALL DEFERRED')
  250. def is_usable(self):
  251. try:
  252. # Use a psycopg cursor directly, bypassing Django's utilities.
  253. self.connection.cursor().execute("SELECT 1")
  254. except Database.Error:
  255. return False
  256. else:
  257. return True
  258. @property
  259. def _nodb_connection(self):
  260. nodb_connection = super()._nodb_connection
  261. try:
  262. nodb_connection.ensure_connection()
  263. except (Database.DatabaseError, WrappedDatabaseError):
  264. warnings.warn(
  265. "Normally Django will use a connection to the 'postgres' database "
  266. "to avoid running initialization queries against the production "
  267. "database when it's not needed (for example, when running tests). "
  268. "Django was unable to create a connection to the 'postgres' database "
  269. "and will use the first PostgreSQL database instead.",
  270. RuntimeWarning
  271. )
  272. for connection in connections.all():
  273. if connection.vendor == 'postgresql' and connection.settings_dict['NAME'] != 'postgres':
  274. return self.__class__(
  275. {**self.settings_dict, 'NAME': connection.settings_dict['NAME']},
  276. alias=self.alias,
  277. )
  278. return nodb_connection
  279. @cached_property
  280. def pg_version(self):
  281. with self.temporary_connection():
  282. return self.connection.server_version
  283. def make_debug_cursor(self, cursor):
  284. return CursorDebugWrapper(cursor, self)
  285. class CursorDebugWrapper(BaseCursorDebugWrapper):
  286. def copy_expert(self, sql, file, *args):
  287. with self.debug_sql(sql):
  288. return self.cursor.copy_expert(sql, file, *args)
  289. def copy_to(self, file, table, *args, **kwargs):
  290. with self.debug_sql(sql='COPY %s TO STDOUT' % table):
  291. return self.cursor.copy_to(file, table, *args, **kwargs)