operations.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415
  1. import json
  2. from functools import lru_cache, partial
  3. from django.conf import settings
  4. from django.db.backends.base.operations import BaseDatabaseOperations
  5. from django.db.backends.postgresql.psycopg_any import (
  6. Inet,
  7. Jsonb,
  8. errors,
  9. is_psycopg3,
  10. mogrify,
  11. )
  12. from django.db.backends.utils import split_tzname_delta
  13. from django.db.models.constants import OnConflict
  14. from django.db.models.functions import Cast
  15. from django.utils.regex_helper import _lazy_re_compile
  16. @lru_cache
  17. def get_json_dumps(encoder):
  18. if encoder is None:
  19. return json.dumps
  20. return partial(json.dumps, cls=encoder)
  21. class DatabaseOperations(BaseDatabaseOperations):
  22. cast_char_field_without_max_length = "varchar"
  23. explain_prefix = "EXPLAIN"
  24. explain_options = frozenset(
  25. [
  26. "ANALYZE",
  27. "BUFFERS",
  28. "COSTS",
  29. "SETTINGS",
  30. "SUMMARY",
  31. "TIMING",
  32. "VERBOSE",
  33. "WAL",
  34. ]
  35. )
  36. cast_data_types = {
  37. "AutoField": "integer",
  38. "BigAutoField": "bigint",
  39. "SmallAutoField": "smallint",
  40. }
  41. if is_psycopg3:
  42. from psycopg.types import numeric
  43. integerfield_type_map = {
  44. "SmallIntegerField": numeric.Int2,
  45. "IntegerField": numeric.Int4,
  46. "BigIntegerField": numeric.Int8,
  47. "PositiveSmallIntegerField": numeric.Int2,
  48. "PositiveIntegerField": numeric.Int4,
  49. "PositiveBigIntegerField": numeric.Int8,
  50. }
  51. def unification_cast_sql(self, output_field):
  52. internal_type = output_field.get_internal_type()
  53. if internal_type in (
  54. "GenericIPAddressField",
  55. "IPAddressField",
  56. "TimeField",
  57. "UUIDField",
  58. ):
  59. # PostgreSQL will resolve a union as type 'text' if input types are
  60. # 'unknown'.
  61. # https://www.postgresql.org/docs/current/typeconv-union-case.html
  62. # These fields cannot be implicitly cast back in the default
  63. # PostgreSQL configuration so we need to explicitly cast them.
  64. # We must also remove components of the type within brackets:
  65. # varchar(255) -> varchar.
  66. return (
  67. "CAST(%%s AS %s)" % output_field.db_type(self.connection).split("(")[0]
  68. )
  69. return "%s"
  70. # EXTRACT format cannot be passed in parameters.
  71. _extract_format_re = _lazy_re_compile(r"[A-Z_]+")
  72. def date_extract_sql(self, lookup_type, sql, params):
  73. # https://www.postgresql.org/docs/current/functions-datetime.html#FUNCTIONS-DATETIME-EXTRACT
  74. if lookup_type == "week_day":
  75. # For consistency across backends, we return Sunday=1, Saturday=7.
  76. return f"EXTRACT(DOW FROM {sql}) + 1", params
  77. elif lookup_type == "iso_week_day":
  78. return f"EXTRACT(ISODOW FROM {sql})", params
  79. elif lookup_type == "iso_year":
  80. return f"EXTRACT(ISOYEAR FROM {sql})", params
  81. lookup_type = lookup_type.upper()
  82. if not self._extract_format_re.fullmatch(lookup_type):
  83. raise ValueError(f"Invalid lookup type: {lookup_type!r}")
  84. return f"EXTRACT({lookup_type} FROM {sql})", params
  85. def date_trunc_sql(self, lookup_type, sql, params, tzname=None):
  86. sql, params = self._convert_sql_to_tz(sql, params, tzname)
  87. # https://www.postgresql.org/docs/current/functions-datetime.html#FUNCTIONS-DATETIME-TRUNC
  88. return f"DATE_TRUNC(%s, {sql})", (lookup_type, *params)
  89. def _prepare_tzname_delta(self, tzname):
  90. tzname, sign, offset = split_tzname_delta(tzname)
  91. if offset:
  92. sign = "-" if sign == "+" else "+"
  93. return f"{tzname}{sign}{offset}"
  94. return tzname
  95. def _convert_sql_to_tz(self, sql, params, tzname):
  96. if tzname and settings.USE_TZ:
  97. tzname_param = self._prepare_tzname_delta(tzname)
  98. return f"{sql} AT TIME ZONE %s", (*params, tzname_param)
  99. return sql, params
  100. def datetime_cast_date_sql(self, sql, params, tzname):
  101. sql, params = self._convert_sql_to_tz(sql, params, tzname)
  102. return f"({sql})::date", params
  103. def datetime_cast_time_sql(self, sql, params, tzname):
  104. sql, params = self._convert_sql_to_tz(sql, params, tzname)
  105. return f"({sql})::time", params
  106. def datetime_extract_sql(self, lookup_type, sql, params, tzname):
  107. sql, params = self._convert_sql_to_tz(sql, params, tzname)
  108. if lookup_type == "second":
  109. # Truncate fractional seconds.
  110. return f"EXTRACT(SECOND FROM DATE_TRUNC(%s, {sql}))", ("second", *params)
  111. return self.date_extract_sql(lookup_type, sql, params)
  112. def datetime_trunc_sql(self, lookup_type, sql, params, tzname):
  113. sql, params = self._convert_sql_to_tz(sql, params, tzname)
  114. # https://www.postgresql.org/docs/current/functions-datetime.html#FUNCTIONS-DATETIME-TRUNC
  115. return f"DATE_TRUNC(%s, {sql})", (lookup_type, *params)
  116. def time_extract_sql(self, lookup_type, sql, params):
  117. if lookup_type == "second":
  118. # Truncate fractional seconds.
  119. return f"EXTRACT(SECOND FROM DATE_TRUNC(%s, {sql}))", ("second", *params)
  120. return self.date_extract_sql(lookup_type, sql, params)
  121. def time_trunc_sql(self, lookup_type, sql, params, tzname=None):
  122. sql, params = self._convert_sql_to_tz(sql, params, tzname)
  123. return f"DATE_TRUNC(%s, {sql})::time", (lookup_type, *params)
  124. def deferrable_sql(self):
  125. return " DEFERRABLE INITIALLY DEFERRED"
  126. def fetch_returned_insert_rows(self, cursor):
  127. """
  128. Given a cursor object that has just performed an INSERT...RETURNING
  129. statement into a table, return the tuple of returned data.
  130. """
  131. return cursor.fetchall()
  132. def lookup_cast(self, lookup_type, internal_type=None):
  133. lookup = "%s"
  134. # Cast text lookups to text to allow things like filter(x__contains=4)
  135. if lookup_type in (
  136. "iexact",
  137. "contains",
  138. "icontains",
  139. "startswith",
  140. "istartswith",
  141. "endswith",
  142. "iendswith",
  143. "regex",
  144. "iregex",
  145. ):
  146. if internal_type in ("IPAddressField", "GenericIPAddressField"):
  147. lookup = "HOST(%s)"
  148. # RemovedInDjango51Warning.
  149. elif internal_type in ("CICharField", "CIEmailField", "CITextField"):
  150. lookup = "%s::citext"
  151. else:
  152. lookup = "%s::text"
  153. # Use UPPER(x) for case-insensitive lookups; it's faster.
  154. if lookup_type in ("iexact", "icontains", "istartswith", "iendswith"):
  155. lookup = "UPPER(%s)" % lookup
  156. return lookup
  157. def no_limit_value(self):
  158. return None
  159. def prepare_sql_script(self, sql):
  160. return [sql]
  161. def quote_name(self, name):
  162. if name.startswith('"') and name.endswith('"'):
  163. return name # Quoting once is enough.
  164. return '"%s"' % name
  165. def compose_sql(self, sql, params):
  166. return mogrify(sql, params, self.connection)
  167. def set_time_zone_sql(self):
  168. return "SELECT set_config('TimeZone', %s, false)"
  169. def sql_flush(self, style, tables, *, reset_sequences=False, allow_cascade=False):
  170. if not tables:
  171. return []
  172. # Perform a single SQL 'TRUNCATE x, y, z...;' statement. It allows us
  173. # to truncate tables referenced by a foreign key in any other table.
  174. sql_parts = [
  175. style.SQL_KEYWORD("TRUNCATE"),
  176. ", ".join(style.SQL_FIELD(self.quote_name(table)) for table in tables),
  177. ]
  178. if reset_sequences:
  179. sql_parts.append(style.SQL_KEYWORD("RESTART IDENTITY"))
  180. if allow_cascade:
  181. sql_parts.append(style.SQL_KEYWORD("CASCADE"))
  182. return ["%s;" % " ".join(sql_parts)]
  183. def sequence_reset_by_name_sql(self, style, sequences):
  184. # 'ALTER SEQUENCE sequence_name RESTART WITH 1;'... style SQL statements
  185. # to reset sequence indices
  186. sql = []
  187. for sequence_info in sequences:
  188. table_name = sequence_info["table"]
  189. # 'id' will be the case if it's an m2m using an autogenerated
  190. # intermediate table (see BaseDatabaseIntrospection.sequence_list).
  191. column_name = sequence_info["column"] or "id"
  192. sql.append(
  193. "%s setval(pg_get_serial_sequence('%s','%s'), 1, false);"
  194. % (
  195. style.SQL_KEYWORD("SELECT"),
  196. style.SQL_TABLE(self.quote_name(table_name)),
  197. style.SQL_FIELD(column_name),
  198. )
  199. )
  200. return sql
  201. def tablespace_sql(self, tablespace, inline=False):
  202. if inline:
  203. return "USING INDEX TABLESPACE %s" % self.quote_name(tablespace)
  204. else:
  205. return "TABLESPACE %s" % self.quote_name(tablespace)
  206. def sequence_reset_sql(self, style, model_list):
  207. from django.db import models
  208. output = []
  209. qn = self.quote_name
  210. for model in model_list:
  211. # Use `coalesce` to set the sequence for each model to the max pk
  212. # value if there are records, or 1 if there are none. Set the
  213. # `is_called` property (the third argument to `setval`) to true if
  214. # there are records (as the max pk value is already in use),
  215. # otherwise set it to false. Use pg_get_serial_sequence to get the
  216. # underlying sequence name from the table name and column name.
  217. for f in model._meta.local_fields:
  218. if isinstance(f, models.AutoField):
  219. output.append(
  220. "%s setval(pg_get_serial_sequence('%s','%s'), "
  221. "coalesce(max(%s), 1), max(%s) %s null) %s %s;"
  222. % (
  223. style.SQL_KEYWORD("SELECT"),
  224. style.SQL_TABLE(qn(model._meta.db_table)),
  225. style.SQL_FIELD(f.column),
  226. style.SQL_FIELD(qn(f.column)),
  227. style.SQL_FIELD(qn(f.column)),
  228. style.SQL_KEYWORD("IS NOT"),
  229. style.SQL_KEYWORD("FROM"),
  230. style.SQL_TABLE(qn(model._meta.db_table)),
  231. )
  232. )
  233. # Only one AutoField is allowed per model, so don't bother
  234. # continuing.
  235. break
  236. return output
  237. def prep_for_iexact_query(self, x):
  238. return x
  239. def max_name_length(self):
  240. """
  241. Return the maximum length of an identifier.
  242. The maximum length of an identifier is 63 by default, but can be
  243. changed by recompiling PostgreSQL after editing the NAMEDATALEN
  244. macro in src/include/pg_config_manual.h.
  245. This implementation returns 63, but can be overridden by a custom
  246. database backend that inherits most of its behavior from this one.
  247. """
  248. return 63
  249. def distinct_sql(self, fields, params):
  250. if fields:
  251. params = [param for param_list in params for param in param_list]
  252. return (["DISTINCT ON (%s)" % ", ".join(fields)], params)
  253. else:
  254. return ["DISTINCT"], []
  255. if is_psycopg3:
  256. def last_executed_query(self, cursor, sql, params):
  257. try:
  258. return self.compose_sql(sql, params)
  259. except errors.DataError:
  260. return None
  261. else:
  262. def last_executed_query(self, cursor, sql, params):
  263. # https://www.psycopg.org/docs/cursor.html#cursor.query
  264. # The query attribute is a Psycopg extension to the DB API 2.0.
  265. if cursor.query is not None:
  266. return cursor.query.decode()
  267. return None
  268. def return_insert_columns(self, fields):
  269. if not fields:
  270. return "", ()
  271. columns = [
  272. "%s.%s"
  273. % (
  274. self.quote_name(field.model._meta.db_table),
  275. self.quote_name(field.column),
  276. )
  277. for field in fields
  278. ]
  279. return "RETURNING %s" % ", ".join(columns), ()
  280. def bulk_insert_sql(self, fields, placeholder_rows):
  281. placeholder_rows_sql = (", ".join(row) for row in placeholder_rows)
  282. values_sql = ", ".join("(%s)" % sql for sql in placeholder_rows_sql)
  283. return "VALUES " + values_sql
  284. if is_psycopg3:
  285. def adapt_integerfield_value(self, value, internal_type):
  286. if value is None or hasattr(value, "resolve_expression"):
  287. return value
  288. return self.integerfield_type_map[internal_type](value)
  289. def adapt_datefield_value(self, value):
  290. return value
  291. def adapt_datetimefield_value(self, value):
  292. return value
  293. def adapt_timefield_value(self, value):
  294. return value
  295. def adapt_decimalfield_value(self, value, max_digits=None, decimal_places=None):
  296. return value
  297. def adapt_ipaddressfield_value(self, value):
  298. if value:
  299. return Inet(value)
  300. return None
  301. def adapt_json_value(self, value, encoder):
  302. return Jsonb(value, dumps=get_json_dumps(encoder))
  303. def subtract_temporals(self, internal_type, lhs, rhs):
  304. if internal_type == "DateField":
  305. lhs_sql, lhs_params = lhs
  306. rhs_sql, rhs_params = rhs
  307. params = (*lhs_params, *rhs_params)
  308. return "(interval '1 day' * (%s - %s))" % (lhs_sql, rhs_sql), params
  309. return super().subtract_temporals(internal_type, lhs, rhs)
  310. def explain_query_prefix(self, format=None, **options):
  311. extra = {}
  312. # Normalize options.
  313. if options:
  314. options = {
  315. name.upper(): "true" if value else "false"
  316. for name, value in options.items()
  317. }
  318. for valid_option in self.explain_options:
  319. value = options.pop(valid_option, None)
  320. if value is not None:
  321. extra[valid_option] = value
  322. prefix = super().explain_query_prefix(format, **options)
  323. if format:
  324. extra["FORMAT"] = format
  325. if extra:
  326. prefix += " (%s)" % ", ".join("%s %s" % i for i in extra.items())
  327. return prefix
  328. def on_conflict_suffix_sql(self, fields, on_conflict, update_fields, unique_fields):
  329. if on_conflict == OnConflict.IGNORE:
  330. return "ON CONFLICT DO NOTHING"
  331. if on_conflict == OnConflict.UPDATE:
  332. return "ON CONFLICT(%s) DO UPDATE SET %s" % (
  333. ", ".join(map(self.quote_name, unique_fields)),
  334. ", ".join(
  335. [
  336. f"{field} = EXCLUDED.{field}"
  337. for field in map(self.quote_name, update_fields)
  338. ]
  339. ),
  340. )
  341. return super().on_conflict_suffix_sql(
  342. fields,
  343. on_conflict,
  344. update_fields,
  345. unique_fields,
  346. )
  347. def prepare_join_on_clause(self, lhs_table, lhs_field, rhs_table, rhs_field):
  348. lhs_expr, rhs_expr = super().prepare_join_on_clause(
  349. lhs_table, lhs_field, rhs_table, rhs_field
  350. )
  351. if lhs_field.db_type(self.connection) != rhs_field.db_type(self.connection):
  352. rhs_expr = Cast(rhs_expr, lhs_field)
  353. return lhs_expr, rhs_expr