1from __future__ import annotations
2
3import collections
4import json
5import re
6from collections.abc import Generator, Iterable, Sequence
7from functools import cached_property
8from itertools import chain
9from typing import TYPE_CHECKING, Any, Protocol, cast
10
11from plain.postgres.constants import LOOKUP_SEP
12from plain.postgres.dialect import (
13 PK_DEFAULT_VALUE,
14 bulk_insert_sql,
15 distinct_sql,
16 explain_query_prefix,
17 for_update_sql,
18 limit_offset_sql,
19 on_conflict_suffix_sql,
20 quote_name,
21 return_insert_columns,
22)
23from plain.postgres.exceptions import EmptyResultSet, FieldError, FullResultSet
24from plain.postgres.expressions import (
25 BaseExpression,
26 F,
27 OrderBy,
28 Ref,
29 ResolvableExpression,
30 Value,
31)
32from plain.postgres.fields import DATABASE_DEFAULT
33from plain.postgres.fields.related import RelatedField
34from plain.postgres.functions import Cast, Random
35from plain.postgres.lookups import Lookup
36from plain.postgres.meta import Meta
37from plain.postgres.query_utils import select_related_descend
38from plain.postgres.sql.constants import (
39 CURSOR,
40 MULTI,
41 NO_RESULTS,
42 ORDER_DIR,
43 SINGLE,
44)
45from plain.postgres.sql.datastructures import Join
46from plain.postgres.sql.query import Query, get_order_dir
47from plain.postgres.transaction import TransactionManagementError
48from plain.utils.hashable import make_hashable
49from plain.utils.regex_helper import _lazy_re_compile
50
51if TYPE_CHECKING:
52 from plain.postgres.connection import DatabaseConnection
53 from plain.postgres.sql.query import AggregateQuery, InsertQuery
54
55# Type aliases for SQL compilation results
56SqlParams = tuple[Any, ...]
57SqlWithParams = tuple[str, SqlParams]
58
59
60class SQLCompilable(Protocol):
61 """Protocol for objects that can be compiled to SQL."""
62
63 def as_sql(
64 self, compiler: SQLCompiler, connection: DatabaseConnection
65 ) -> tuple[str, Sequence[Any]]:
66 """Return SQL string and parameters for this object."""
67 ...
68
69
70class PositionRef(Ref):
71 def __init__(self, ordinal: int, refs: str, source: Any):
72 self.ordinal = ordinal
73 super().__init__(refs, source)
74
75 def as_sql(
76 self, compiler: SQLCompiler, connection: DatabaseConnection
77 ) -> tuple[str, list[Any]]:
78 return str(self.ordinal), []
79
80
81def get_converters(
82 expressions: Iterable[Any], connection: DatabaseConnection
83) -> dict[int, tuple[list[Any], Any]]:
84 converters = {}
85 for i, expression in enumerate(expressions):
86 if expression:
87 field_converters = expression.get_db_converters(connection)
88 if field_converters:
89 converters[i] = (field_converters, expression)
90 return converters
91
92
93def apply_converters(
94 rows: Iterable, converters: dict, connection: DatabaseConnection
95) -> Generator[list]:
96 converters_list = list(converters.items())
97 for row in map(list, rows):
98 for pos, (convs, expression) in converters_list:
99 value = row[pos]
100 for converter in convs:
101 value = converter(value, expression, connection)
102 row[pos] = value
103 yield row
104
105
106class SQLCompiler:
107 # Multiline ordering SQL clause may appear from RawSQL.
108 ordering_parts = _lazy_re_compile(
109 r"^(.*)\s(?:ASC|DESC).*",
110 re.MULTILINE | re.DOTALL,
111 )
112
113 def __init__(
114 self, query: Query, connection: DatabaseConnection, elide_empty: bool = True
115 ):
116 self.query = query
117 self.connection = connection
118 # Some queries, e.g. coalesced aggregation, need to be executed even if
119 # they would return an empty result set.
120 self.elide_empty = elide_empty
121 self.quote_cache: dict[str, str] = {"*": "*"}
122 # The select, klass_info, and annotations are needed by QuerySet.iterator()
123 # these are set as a side-effect of executing the query. Note that we calculate
124 # separately a list of extra select columns needed for grammatical correctness
125 # of the query, but these columns are not included in self.select.
126 self.select: list[tuple[Any, SqlWithParams, str | None]] | None = None
127 self.annotation_col_map: dict[str, int] | None = None
128 self.klass_info: dict[str, Any] | None = None
129 self._meta_ordering: list[str] | None = None
130
131 def __repr__(self) -> str:
132 model_name = self.query.model.__qualname__ if self.query.model else "None"
133 return (
134 f"<{self.__class__.__qualname__} "
135 f"model={model_name} "
136 f"connection={self.connection!r}>"
137 )
138
139 def setup_query(self, with_col_aliases: bool = False) -> None:
140 if all(self.query.alias_refcount[a] == 0 for a in self.query.alias_map):
141 self.query.get_initial_alias()
142 self.select, self.klass_info, self.annotation_col_map = self.get_select(
143 with_col_aliases=with_col_aliases,
144 )
145 self.col_count = len(self.select)
146
147 def pre_sql_setup(
148 self, with_col_aliases: bool = False
149 ) -> tuple[list[Any], list[Any], list[SqlWithParams]] | None:
150 """
151 Do any necessary class setup immediately prior to producing SQL. This
152 is for things that can't necessarily be done in __init__ because we
153 might not have all the pieces in place at that time.
154 """
155 self.setup_query(with_col_aliases=with_col_aliases)
156 assert self.select is not None # Set by setup_query()
157 order_by = self.get_order_by()
158 self.where, self.having, self.qualify = self.query.where.split_having_qualify(
159 must_group_by=self.query.group_by is not None
160 )
161 extra_select = self.get_extra_select(order_by, self.select)
162 self.has_extra_select = bool(extra_select)
163 group_by = self.get_group_by(self.select + extra_select, order_by)
164 return extra_select, order_by, group_by
165
166 def get_group_by(
167 self, select: list[Any], order_by: list[Any]
168 ) -> list[SqlWithParams]:
169 """
170 Return a list of 2-tuples of form (sql, params).
171
172 The logic of what exactly the GROUP BY clause contains is hard
173 to describe in other words than "if it passes the test suite,
174 then it is correct".
175 """
176 # Some examples:
177 # SomeModel.query.annotate(Count('somecol'))
178 # GROUP BY: all fields of the model
179 #
180 # SomeModel.query.values('name').annotate(Count('somecol'))
181 # GROUP BY: name
182 #
183 # SomeModel.query.annotate(Count('somecol')).values('name')
184 # GROUP BY: all cols of the model
185 #
186 # SomeModel.query.values('name', 'id')
187 # .annotate(Count('somecol')).values('id')
188 # GROUP BY: name, id
189 #
190 # SomeModel.query.values('name').annotate(Count('somecol')).values('id')
191 # GROUP BY: name, id
192 #
193 # In fact, the self.query.group_by is the minimal set to GROUP BY. It
194 # can't be ever restricted to a smaller set, but additional columns in
195 # HAVING, ORDER BY, and SELECT clauses are added to it. Unfortunately
196 # the end result is that it is impossible to force the query to have
197 # a chosen GROUP BY clause - you can almost do this by using the form:
198 # .values(*wanted_cols).annotate(AnAggregate())
199 # but any later annotations, extra selects, values calls that
200 # refer some column outside of the wanted_cols, order_by, or even
201 # filter calls can alter the GROUP BY clause.
202
203 # The query.group_by is either None (no GROUP BY at all), True
204 # (group by select fields), or a list of expressions to be added
205 # to the group by.
206 if self.query.group_by is None:
207 return []
208 expressions = []
209 group_by_refs = set()
210 if self.query.group_by is not True:
211 # If the group by is set to a list (by .values() call most likely),
212 # then we need to add everything in it to the GROUP BY clause.
213 # Backwards compatibility hack for setting query.group_by. Remove
214 # when we have public API way of forcing the GROUP BY clause.
215 # Converts string references to expressions.
216 for expr in self.query.group_by:
217 if not hasattr(expr, "as_sql"):
218 expr = self.query.resolve_ref(expr)
219 if isinstance(expr, Ref):
220 if expr.refs not in group_by_refs:
221 group_by_refs.add(expr.refs)
222 expressions.append(expr.source)
223 else:
224 expressions.append(expr)
225 # Note that even if the group_by is set, it is only the minimal
226 # set to group by. So, we need to add cols in select, order_by, and
227 # having into the select in any case.
228 selected_expr_positions = {}
229 for ordinal, (expr, _, alias) in enumerate(select, start=1):
230 if alias:
231 selected_expr_positions[expr] = ordinal
232 # Skip members of the select clause that are already explicitly
233 # grouped against.
234 if alias in group_by_refs:
235 continue
236 expressions.extend(expr.get_group_by_cols())
237 if not self._meta_ordering:
238 for expr, (sql, params, is_ref) in order_by:
239 # Skip references to the SELECT clause, as all expressions in
240 # the SELECT clause are already part of the GROUP BY.
241 if not is_ref:
242 expressions.extend(expr.get_group_by_cols())
243 having_group_by = self.having.get_group_by_cols() if self.having else []
244 expressions.extend(having_group_by)
245 result = []
246 seen = set()
247 expressions = self.collapse_group_by(expressions, having_group_by)
248
249 for expr in expressions:
250 try:
251 sql, params = self.compile(expr)
252 except (EmptyResultSet, FullResultSet):
253 continue
254 # Use select index for GROUP BY when possible
255 if (position := selected_expr_positions.get(expr)) is not None:
256 sql, params = str(position), ()
257 else:
258 sql, params = expr.select_format(self, sql, params)
259 params_hash = make_hashable(params)
260 if (sql, params_hash) not in seen:
261 result.append((sql, params))
262 seen.add((sql, params_hash))
263 return result
264
265 def collapse_group_by(self, expressions: list[Any], having: list[Any]) -> list[Any]:
266 # Use group by functional dependence reduction:
267 # expressions can be reduced to the set of selected table
268 # primary keys as all other columns are functionally dependent on them.
269 # Filter out all expressions associated with a table's primary key
270 # present in the grouped columns. This is done by identifying all
271 # tables that have their primary key included in the grouped
272 # columns and removing non-primary key columns referring to them.
273 pks = {
274 expr
275 for expr in expressions
276 if hasattr(expr, "target") and expr.target.primary_key
277 }
278 aliases = {expr.alias for expr in pks}
279 return [
280 expr
281 for expr in expressions
282 if expr in pks
283 or expr in having
284 or getattr(expr, "alias", None) not in aliases
285 ]
286
287 def get_select(
288 self, with_col_aliases: bool = False
289 ) -> tuple[
290 list[tuple[Any, SqlWithParams, str | None]],
291 dict[str, Any] | None,
292 dict[str, int],
293 ]:
294 """
295 Return three values:
296 - a list of 3-tuples of (expression, (sql, params), alias)
297 - a klass_info structure,
298 - a dictionary of annotations
299
300 The (sql, params) is what the expression will produce, and alias is the
301 "AS alias" for the column (possibly None).
302
303 The klass_info structure contains the following information:
304 - The base model of the query.
305 - Which columns for that model are present in the query (by
306 position of the select clause).
307 - related_klass_infos: [f, klass_info] to descent into
308
309 The annotations is a dictionary of {'name': column position} values.
310 """
311 select = []
312 klass_info = None
313 annotations = {}
314 select_idx = 0
315 assert not (self.query.select and self.query.default_cols)
316 select_mask = self.query.get_select_mask()
317 if self.query.default_cols:
318 cols = self.get_default_columns(select_mask)
319 else:
320 # self.query.select is a special case. These columns never go to
321 # any model.
322 cols = self.query.select
323 if cols:
324 select_list = []
325 for col in cols:
326 select_list.append(select_idx)
327 select.append((col, None))
328 select_idx += 1
329 klass_info = {
330 "model": self.query.model,
331 "select_fields": select_list,
332 }
333 for alias, annotation in self.query.annotation_select.items():
334 annotations[alias] = select_idx
335 select.append((annotation, alias))
336 select_idx += 1
337
338 if self.query.select_related:
339 related_klass_infos = self.get_related_selections(select, select_mask)
340 if klass_info is not None:
341 klass_info["related_klass_infos"] = related_klass_infos
342
343 ret = []
344 col_idx = 1
345 for col, alias in select:
346 try:
347 sql, params = self.compile(col)
348 except EmptyResultSet:
349 empty_result_set_value = getattr(
350 col, "empty_result_set_value", NotImplemented
351 )
352 if empty_result_set_value is NotImplemented:
353 # Select a predicate that's always False.
354 sql, params = "0", ()
355 else:
356 sql, params = self.compile(Value(empty_result_set_value))
357 except FullResultSet:
358 sql, params = self.compile(Value(True))
359 else:
360 sql, params = col.select_format(self, sql, params)
361 if alias is None and with_col_aliases:
362 alias = f"col{col_idx}"
363 col_idx += 1
364 ret.append((col, (sql, params), alias))
365 return ret, klass_info, annotations # ty: ignore[invalid-return-type] (heterogeneous klass_info dict)
366
367 def _order_by_pairs(self) -> Generator[tuple[OrderBy, bool]]:
368 if not self.query.default_ordering or self.query.order_by:
369 ordering = self.query.order_by
370 elif (
371 self.query.model
372 and (options := self.query.model.model_options)
373 and options.ordering
374 ):
375 ordering = options.ordering
376 self._meta_ordering = list(ordering)
377 else:
378 ordering = []
379 if self.query.standard_ordering:
380 default_order, _ = ORDER_DIR["ASC"]
381 else:
382 default_order, _ = ORDER_DIR["DESC"]
383
384 selected_exprs = {}
385 if select := self.select:
386 for ordinal, (expr, _, alias) in enumerate(select, start=1):
387 pos_expr = PositionRef(ordinal, alias, expr) # ty: ignore[invalid-argument-type]
388 if alias:
389 selected_exprs[alias] = pos_expr
390 selected_exprs[expr] = pos_expr
391
392 for field in ordering:
393 if isinstance(field, ResolvableExpression):
394 # field is a BaseExpression (has asc/desc/copy methods)
395 field_expr = cast(BaseExpression, field)
396 if isinstance(field_expr, Value):
397 # output_field must be resolved for constants.
398 field_expr = Cast(field_expr, field_expr.output_field)
399 if not isinstance(field_expr, OrderBy):
400 field_expr = field_expr.asc()
401 if not self.query.standard_ordering:
402 field_expr = field_expr.copy()
403 field_expr.reverse_ordering()
404 field = field_expr
405 select_ref = selected_exprs.get(field.expression)
406 if select_ref or (
407 isinstance(field.expression, F)
408 and (select_ref := selected_exprs.get(field.expression.name))
409 ):
410 field = field.copy()
411 field.expression = select_ref
412 yield field, select_ref is not None
413 continue
414 if field == "?": # random
415 yield OrderBy(Random()), False
416 continue
417
418 col, order = get_order_dir(field, default_order)
419 descending = order == "DESC"
420
421 if select_ref := selected_exprs.get(col):
422 # Reference to expression in SELECT clause
423 yield (
424 OrderBy(
425 select_ref,
426 descending=descending,
427 ),
428 True,
429 )
430 continue
431 if col in self.query.annotations:
432 # References to an expression which is masked out of the SELECT
433 # clause.
434 expr = self.query.annotations[col]
435 if isinstance(expr, Value):
436 # output_field must be resolved for constants.
437 expr = Cast(expr, expr.output_field)
438 yield OrderBy(expr, descending=descending), False
439 continue
440
441 # 'col' is of the form 'field' or 'field1__field2' or
442 # '-field1__field2__field', etc.
443 assert self.query.model is not None, "Ordering by fields requires a model"
444 meta = self.query.model._model_meta
445 yield from self.find_ordering_name(
446 field,
447 meta,
448 default_order=default_order,
449 )
450
451 def get_order_by(self) -> list[tuple[Any, tuple[str, tuple, bool]]]:
452 """
453 Return a list of 2-tuples of the form (expr, (sql, params, is_ref)) for
454 the ORDER BY clause.
455
456 The order_by clause can alter the select clause (for example it can add
457 aliases to clauses that do not yet have one, or it can add totally new
458 select clauses).
459 """
460 result = []
461 seen = set()
462 for expr, is_ref in self._order_by_pairs():
463 resolved = expr.resolve_expression(self.query, allow_joins=True, reuse=None)
464 sql, params = self.compile(resolved)
465 # Don't add the same column twice, but the order direction is
466 # not taken into account so we strip it. When this entire method
467 # is refactored into expressions, then we can check each part as we
468 # generate it.
469 without_ordering = self.ordering_parts.search(sql)[1]
470 params_hash = make_hashable(params)
471 if (without_ordering, params_hash) in seen:
472 continue
473 seen.add((without_ordering, params_hash))
474 result.append((resolved, (sql, params, is_ref)))
475 return result
476
477 def get_extra_select(
478 self, order_by: list[Any], select: list[Any]
479 ) -> list[tuple[Any, SqlWithParams, None]]:
480 extra_select = []
481 if self.query.distinct and not self.query.distinct_fields:
482 select_sql = [t[1] for t in select]
483 for expr, (sql, params, is_ref) in order_by:
484 without_ordering = self.ordering_parts.search(sql)[1]
485 if not is_ref and (without_ordering, params) not in select_sql:
486 extra_select.append((expr, (without_ordering, params), None))
487 return extra_select
488
489 def quote_name_unless_alias(self, name: str) -> str:
490 """
491 A wrapper around quote_name() that doesn't quote aliases for table
492 names. This avoids problems with some SQL dialects that treat quoted
493 strings specially (e.g. PostgreSQL).
494 """
495 if name in self.quote_cache:
496 return self.quote_cache[name]
497 if (name in self.query.alias_map and name not in self.query.table_map) or (
498 self.query.external_aliases.get(name) and name not in self.query.table_map
499 ):
500 self.quote_cache[name] = name
501 return name
502 r = quote_name(name)
503 self.quote_cache[name] = r
504 return r
505
506 def compile(self, node: SQLCompilable) -> SqlWithParams:
507 sql, params = node.as_sql(self, self.connection)
508 return sql, tuple(params)
509
510 def get_qualify_sql(self) -> tuple[list[str], list[Any]]:
511 where_parts = []
512 if self.where:
513 where_parts.append(self.where)
514 if self.having:
515 where_parts.append(self.having)
516 inner_query = self.query.clone()
517 inner_query.subquery = True
518 inner_query.where = inner_query.where.__class__(where_parts)
519 # Augment the inner query with any window function references that
520 # might have been masked via values() and alias(). If any masked
521 # aliases are added they'll be masked again to avoid fetching
522 # the data in the `if qual_aliases` branch below.
523 select = {
524 expr: alias for expr, _, alias in self.get_select(with_col_aliases=True)[0]
525 }
526 select_aliases = set(select.values())
527 qual_aliases = set()
528 replacements = {}
529
530 def collect_replacements(expressions: list[Any]) -> None:
531 while expressions:
532 expr = expressions.pop()
533 if expr in replacements:
534 continue
535 elif select_alias := select.get(expr):
536 replacements[expr] = select_alias
537 elif isinstance(expr, Lookup):
538 expressions.extend(expr.get_source_expressions())
539 elif isinstance(expr, Ref):
540 if expr.refs not in select_aliases:
541 expressions.extend(expr.get_source_expressions())
542 else:
543 num_qual_alias = len(qual_aliases)
544 select_alias = f"qual{num_qual_alias}"
545 qual_aliases.add(select_alias)
546 inner_query.add_annotation(expr, select_alias)
547 replacements[expr] = select_alias
548
549 qualify = self.qualify
550 if qualify is None:
551 raise ValueError("QUALIFY clause expected but not provided")
552 collect_replacements(list(qualify.leaves()))
553 qualify = qualify.replace_expressions(
554 {expr: Ref(alias, expr) for expr, alias in replacements.items()}
555 )
556 self.qualify = qualify
557 order_by = []
558 for order_by_expr, *_ in self.get_order_by():
559 collect_replacements(order_by_expr.get_source_expressions())
560 order_by.append(
561 order_by_expr.replace_expressions(
562 {expr: Ref(alias, expr) for expr, alias in replacements.items()}
563 )
564 )
565 inner_query_compiler = inner_query.get_compiler(elide_empty=self.elide_empty)
566 inner_sql, inner_params = inner_query_compiler.as_sql(
567 # The limits must be applied to the outer query to avoid pruning
568 # results too eagerly.
569 with_limits=False,
570 # Force unique aliasing of selected columns to avoid collisions
571 # and make rhs predicates referencing easier.
572 with_col_aliases=True,
573 )
574 qualify_sql, qualify_params = self.compile(qualify)
575 result = [
576 "SELECT * FROM (",
577 inner_sql,
578 ")",
579 quote_name("qualify"),
580 "WHERE",
581 qualify_sql,
582 ]
583 if qual_aliases:
584 # If some select aliases were unmasked for filtering purposes they
585 # must be masked back.
586 cols = [quote_name(alias) for alias in select.values() if alias is not None]
587 result = [
588 "SELECT",
589 ", ".join(cols),
590 "FROM (",
591 *result,
592 ")",
593 quote_name("qualify_mask"),
594 ]
595 params = list(inner_params) + list(qualify_params)
596 # As the SQL spec is unclear on whether or not derived tables
597 # ordering must propagate it has to be explicitly repeated on the
598 # outer-most query to ensure it's preserved.
599 if order_by:
600 ordering_sqls = []
601 for ordering in order_by:
602 ordering_sql, ordering_params = self.compile(ordering)
603 ordering_sqls.append(ordering_sql)
604 params.extend(ordering_params)
605 result.extend(["ORDER BY", ", ".join(ordering_sqls)])
606 return result, params
607
608 def as_sql(
609 self, with_limits: bool = True, with_col_aliases: bool = False
610 ) -> SqlWithParams:
611 """
612 Create the SQL for this query. Return the SQL string and list of
613 parameters.
614
615 If 'with_limits' is False, any limit/offset information is not included
616 in the query.
617 """
618 refcounts_before = self.query.alias_refcount.copy()
619 try:
620 result = self.pre_sql_setup(with_col_aliases=with_col_aliases)
621 assert result is not None # SQLCompiler.pre_sql_setup always returns tuple
622 extra_select, order_by, group_by = result
623 assert self.select is not None # Set by pre_sql_setup()
624 for_update_part = None
625 # Is a LIMIT/OFFSET clause needed?
626 with_limit_offset = with_limits and self.query.is_sliced
627 if self.qualify:
628 result, params = self.get_qualify_sql()
629 order_by = None
630 else:
631 distinct_fields, distinct_params = self.get_distinct()
632 # This must come after 'select', 'ordering', and 'distinct'
633 # (see docstring of get_from_clause() for details).
634 from_, f_params = self.get_from_clause()
635 try:
636 where, w_params = (
637 self.compile(self.where) if self.where is not None else ("", [])
638 )
639 except EmptyResultSet:
640 if self.elide_empty:
641 raise
642 # Use a predicate that's always False.
643 where, w_params = "0 = 1", []
644 except FullResultSet:
645 where, w_params = "", []
646 try:
647 having, h_params = (
648 self.compile(self.having)
649 if self.having is not None
650 else ("", [])
651 )
652 except FullResultSet:
653 having, h_params = "", []
654 result = ["SELECT"]
655 params = []
656
657 if self.query.distinct:
658 distinct_result, distinct_params = distinct_sql(
659 distinct_fields,
660 distinct_params,
661 )
662 result += distinct_result
663 params += distinct_params
664
665 out_cols = []
666 for _, (s_sql, s_params), alias in self.select + extra_select:
667 if alias:
668 s_sql = f"{s_sql} AS {quote_name(alias)}"
669 params.extend(s_params)
670 out_cols.append(s_sql)
671
672 result += [", ".join(out_cols)]
673 if from_:
674 result += ["FROM", *from_]
675 params.extend(f_params)
676
677 if self.query.select_for_update:
678 if self.connection.get_autocommit():
679 raise TransactionManagementError(
680 "select_for_update cannot be used outside of a transaction."
681 )
682
683 for_update_part = for_update_sql(
684 nowait=self.query.select_for_update_nowait,
685 skip_locked=self.query.select_for_update_skip_locked,
686 of=tuple(self.get_select_for_update_of_arguments()),
687 no_key=self.query.select_for_no_key_update,
688 )
689
690 if where:
691 result.append(f"WHERE {where}")
692 params.extend(w_params)
693
694 grouping = []
695 for g_sql, g_params in group_by:
696 grouping.append(g_sql)
697 params.extend(g_params)
698 if grouping:
699 if distinct_fields:
700 raise NotImplementedError(
701 "annotate() + distinct(fields) is not implemented."
702 )
703 order_by = order_by or []
704 result.append("GROUP BY {}".format(", ".join(grouping)))
705 if self._meta_ordering:
706 order_by = None
707 if having:
708 result.append(f"HAVING {having}")
709 params.extend(h_params)
710
711 if self.query.explain_info:
712 result.insert(
713 0,
714 explain_query_prefix(
715 self.query.explain_info.format,
716 **self.query.explain_info.options,
717 ),
718 )
719
720 if order_by:
721 ordering = []
722 for _, (o_sql, o_params, _) in order_by:
723 ordering.append(o_sql)
724 params.extend(o_params)
725 result.append("ORDER BY {}".format(", ".join(ordering)))
726
727 if with_limit_offset:
728 result.append(
729 limit_offset_sql(self.query.low_mark, self.query.high_mark)
730 )
731
732 if for_update_part:
733 result.append(for_update_part)
734
735 if self.query.subquery and extra_select:
736 # If the query is used as a subquery, the extra selects would
737 # result in more columns than the left-hand side expression is
738 # expecting. This can happen when a subquery uses a combination
739 # of order_by() and distinct(), forcing the ordering expressions
740 # to be selected as well. Wrap the query in another subquery
741 # to exclude extraneous selects.
742 sub_selects = []
743 sub_params = []
744 for index, (select, _, alias) in enumerate(self.select, start=1):
745 if alias:
746 sub_selects.append(
747 "{}.{}".format(
748 quote_name("subquery"),
749 quote_name(alias),
750 )
751 )
752 else:
753 select_clone = select.relabeled_clone(
754 {select.alias: "subquery"}
755 )
756 subselect, subparams = select_clone.as_sql(
757 self, self.connection
758 )
759 sub_selects.append(subselect)
760 sub_params.extend(subparams)
761 return "SELECT {} FROM ({}) subquery".format(
762 ", ".join(sub_selects),
763 " ".join(result),
764 ), tuple(sub_params + params)
765
766 return " ".join(result), tuple(params)
767 finally:
768 # Finally do cleanup - get rid of the joins we created above.
769 self.query.reset_refcounts(refcounts_before)
770
771 def get_default_columns(
772 self,
773 select_mask: Any,
774 start_alias: str | None = None,
775 opts: Meta | None = None,
776 ) -> list[Any]:
777 """
778 Return Col expressions for every concrete field on the model. When
779 pulling in a related model (e.g. via select_related), the caller
780 passes ``opts`` and ``start_alias`` to traverse from that join.
781 """
782 result = []
783 if opts is None:
784 if self.query.model is None:
785 return result
786 opts = self.query.model._model_meta
787 start_alias = start_alias or self.query.get_initial_alias()
788
789 for field in opts.fields:
790 if select_mask and field not in select_mask:
791 continue
792 result.append(field.get_col(start_alias))
793 return result
794
795 def get_distinct(self) -> tuple[list[str], list]:
796 """
797 Return a quoted list of fields to use in DISTINCT ON part of the query.
798
799 This method can alter the tables in the query, and thus it must be
800 called before get_from_clause().
801 """
802 result = []
803 params = []
804 if not self.query.distinct_fields:
805 return result, params
806
807 if self.query.model is None:
808 return result, params
809 opts = self.query.model._model_meta
810
811 for name in self.query.distinct_fields:
812 parts = name.split(LOOKUP_SEP)
813 _, targets, alias, joins, path, _, transform_function = self._setup_joins(
814 parts, opts, None
815 )
816 targets, alias, _ = self.query.trim_joins(targets, joins, path)
817 for target in targets:
818 if name in self.query.annotation_select:
819 result.append(quote_name(name))
820 else:
821 r, p = self.compile(transform_function(target, alias))
822 result.append(r)
823 params.append(p)
824 return result, params
825
826 def find_ordering_name(
827 self,
828 name: str,
829 meta: Meta,
830 alias: str | None = None,
831 default_order: str = "ASC",
832 already_seen: set | None = None,
833 ) -> list[tuple[OrderBy, bool]]:
834 """
835 Return the table alias (the name might be ambiguous, the alias will
836 not be) and column name for ordering by the given 'name' parameter.
837 The 'name' is of the form 'field1__field2__...__fieldN'.
838 """
839 name, order = get_order_dir(name, default_order)
840 descending = order == "DESC"
841 pieces = name.split(LOOKUP_SEP)
842 (
843 field,
844 targets,
845 alias,
846 joins,
847 path,
848 meta,
849 transform_function,
850 ) = self._setup_joins(pieces, meta, alias)
851
852 # If we get to this point and the field is a relation to another model,
853 # append the default ordering for that model unless it is the
854 # attribute name of the field that is specified or
855 # there are transforms to process.
856 if (
857 isinstance(field, RelatedField)
858 and meta.model.model_options.ordering
859 and field.name != pieces[-1]
860 and not getattr(transform_function, "has_transforms", False)
861 ):
862 # Firstly, avoid infinite loops. Each join contributes its column
863 # pair to the signature; the base table (no join_col) contributes
864 # None. isinstance keeps `.join_col` greppable and type-checked.
865 already_seen = already_seen or set()
866 alias_map = self.query.alias_map
867 join_tuple = tuple(
868 alias_map[j].join_col if isinstance(alias_map[j], Join) else None
869 for j in joins
870 )
871 if join_tuple in already_seen:
872 raise FieldError("Infinite loop caused by ordering.")
873 already_seen.add(join_tuple)
874
875 results = []
876 for item in meta.model.model_options.ordering:
877 if isinstance(item, ResolvableExpression) and not isinstance(
878 item, OrderBy
879 ):
880 item_expr: BaseExpression = cast(BaseExpression, item)
881 item = item_expr.desc() if descending else item_expr.asc()
882 if isinstance(item, OrderBy):
883 results.append(
884 (item.prefix_references(f"{name}{LOOKUP_SEP}"), False)
885 )
886 continue
887 results.extend(
888 (expr.prefix_references(f"{name}{LOOKUP_SEP}"), is_ref)
889 for expr, is_ref in self.find_ordering_name(
890 item, meta, alias, order, already_seen
891 )
892 )
893 return results
894 targets, alias, _ = self.query.trim_joins(targets, joins, path)
895 return [
896 (OrderBy(transform_function(t, alias), descending=descending), False)
897 for t in targets
898 ]
899
900 def _setup_joins(
901 self, pieces: list[str], meta: Meta, alias: str | None
902 ) -> tuple[Any, Any, str, list, Any, Meta, Any]:
903 """
904 Helper method for get_order_by() and get_distinct().
905
906 get_ordering() and get_distinct() must produce same target columns on
907 same input, as the prefixes of get_ordering() and get_distinct() must
908 match. Executing SQL where this is not true is an error.
909 """
910 alias = alias or self.query.get_initial_alias()
911 assert alias is not None
912 field, targets, meta, joins, path, transform_function = self.query.setup_joins(
913 pieces, meta, alias
914 )
915 alias = joins[-1]
916 return field, targets, alias, joins, path, meta, transform_function
917
918 def get_from_clause(self) -> tuple[list[str], list]:
919 """
920 Return a list of strings that are joined together to go after the
921 "FROM" part of the query, as well as a list any extra parameters that
922 need to be included. Subclasses, can override this to create a
923 from-clause via a "select".
924
925 This should only be called after any SQL construction methods that
926 might change the tables that are needed. This means the select columns,
927 ordering, and distinct must be done first.
928 """
929 result = []
930 params = []
931 for alias in tuple(self.query.alias_map):
932 if not self.query.alias_refcount[alias]:
933 continue
934 from_clause = self.query.alias_map[alias]
935 clause_sql, clause_params = self.compile(from_clause)
936 result.append(clause_sql)
937 params.extend(clause_params)
938 return result, params
939
940 def get_related_selections(
941 self,
942 select: list[Any],
943 select_mask: Any,
944 opts: Meta | None = None,
945 root_alias: str | None = None,
946 cur_depth: int = 1,
947 requested: dict | None = None,
948 restricted: bool | None = None,
949 ) -> list[dict[str, Any]]:
950 """
951 Fill in the information needed for a select_related query. The current
952 depth is measured as the number of connections away from the root model
953 (for example, cur_depth=1 means we are looking at models with direct
954 connections to the root model).
955
956 Args:
957 opts: Meta for the model being queried (internal metadata)
958 """
959
960 related_klass_infos = []
961 if not restricted and cur_depth > self.query.max_depth:
962 # We've recursed far enough; bail out.
963 return related_klass_infos
964
965 if not opts:
966 assert self.query.model is not None, "select_related requires a model"
967 opts = self.query.model._model_meta
968 root_alias = self.query.get_initial_alias()
969
970 assert root_alias is not None # Must be provided or set above
971 assert opts is not None
972
973 def _get_field_choices() -> chain:
974 direct_choices = (
975 f.name for f in opts.fields if isinstance(f, RelatedField)
976 )
977 reverse_choices = (
978 f.field.related_query_name()
979 for f in opts.related_objects
980 if f.field.primary_key
981 )
982 return chain(direct_choices, reverse_choices)
983
984 # Setup for the case when only particular related fields should be
985 # included in the related selection.
986 fields_found = set()
987 if requested is None:
988 restricted = isinstance(self.query.select_related, dict)
989 if restricted:
990 requested = cast(dict, self.query.select_related)
991
992 def get_related_klass_infos(
993 klass_info: dict, related_klass_infos: list
994 ) -> None:
995 klass_info["related_klass_infos"] = related_klass_infos
996
997 for f in opts.fields:
998 fields_found.add(f.name)
999
1000 if restricted:
1001 assert requested is not None
1002 next = requested.get(f.name, {})
1003 # If a non-related field is used like a relation,
1004 # or if a single non-relational field is given.
1005 if not isinstance(f, RelatedField) and (next or f.name in requested):
1006 raise FieldError(
1007 "Non-relational field given in select_related: '{}'. "
1008 "Choices are: {}".format(
1009 f.name,
1010 ", ".join(_get_field_choices()) or "(none)",
1011 )
1012 )
1013 else:
1014 next = None
1015
1016 if not select_related_descend(f, restricted, requested, select_mask):
1017 continue
1018 related_select_mask = select_mask.get(f) or {}
1019 klass_info: dict[str, Any] = {
1020 "model": f.remote_field.model,
1021 "field": f,
1022 "reverse": False,
1023 "local_setter": f.set_cached_value,
1024 "remote_setter": f.remote_field.set_cached_value
1025 if f.primary_key
1026 else lambda x, y: None,
1027 }
1028 related_klass_infos.append(klass_info)
1029 select_fields = []
1030 _, _, _, joins, _, _ = self.query.setup_joins([f.name], opts, root_alias)
1031 alias = joins[-1]
1032 columns = self.get_default_columns(
1033 related_select_mask,
1034 start_alias=alias,
1035 opts=f.remote_field.model._model_meta,
1036 )
1037 for col in columns:
1038 select_fields.append(len(select))
1039 select.append((col, None))
1040 klass_info["select_fields"] = select_fields
1041 next_klass_infos = self.get_related_selections(
1042 select,
1043 related_select_mask,
1044 f.remote_field.model._model_meta,
1045 alias,
1046 cur_depth + 1,
1047 next,
1048 restricted,
1049 )
1050 get_related_klass_infos(klass_info, next_klass_infos)
1051
1052 if restricted:
1053 from plain.postgres.fields.reverse_related import ManyToManyRel
1054
1055 related_fields = [
1056 (o.field, o.related_model)
1057 for o in opts.related_objects
1058 if o.field.primary_key and not isinstance(o, ManyToManyRel)
1059 ]
1060 for related_field, model in related_fields:
1061 related_select_mask = select_mask.get(related_field) or {}
1062
1063 if not select_related_descend(
1064 related_field,
1065 restricted,
1066 requested,
1067 related_select_mask,
1068 reverse=True,
1069 ):
1070 continue
1071
1072 related_field_name = related_field.related_query_name()
1073 fields_found.add(related_field_name)
1074
1075 join_info = self.query.setup_joins(
1076 [related_field_name], opts, root_alias
1077 )
1078 alias = join_info.joins[-1]
1079 klass_info: dict[str, Any] = {
1080 "model": model,
1081 "field": related_field,
1082 "reverse": True,
1083 "local_setter": related_field.remote_field.set_cached_value,
1084 "remote_setter": related_field.set_cached_value,
1085 }
1086 related_klass_infos.append(klass_info)
1087 select_fields = []
1088 columns = self.get_default_columns(
1089 related_select_mask,
1090 start_alias=alias,
1091 opts=model._model_meta,
1092 )
1093 for col in columns:
1094 select_fields.append(len(select))
1095 select.append((col, None))
1096 klass_info["select_fields"] = select_fields
1097 assert requested is not None
1098 next = requested.get(related_field.related_query_name(), {})
1099 next_klass_infos = self.get_related_selections(
1100 select,
1101 related_select_mask,
1102 model._model_meta,
1103 alias,
1104 cur_depth + 1,
1105 next,
1106 restricted,
1107 )
1108 get_related_klass_infos(klass_info, next_klass_infos)
1109
1110 assert requested is not None
1111 fields_not_found = set(requested).difference(fields_found)
1112 if fields_not_found:
1113 invalid_fields = (f"'{s}'" for s in fields_not_found)
1114 raise FieldError(
1115 "Invalid field name(s) given in select_related: {}. "
1116 "Choices are: {}".format(
1117 ", ".join(invalid_fields),
1118 ", ".join(_get_field_choices()) or "(none)",
1119 )
1120 )
1121 return related_klass_infos
1122
1123 def get_select_for_update_of_arguments(self) -> list[str]:
1124 """
1125 Return a quoted list of arguments for the SELECT FOR UPDATE OF part of
1126 the query.
1127 """
1128
1129 def _get_first_selected_col_from_model(klass_info: dict) -> Any | None:
1130 """
1131 Find the first selected column whose target field belongs to this
1132 klass_info's model. Returns None when the model isn't represented
1133 in the select list — callers use that to skip locking the row.
1134 """
1135 assert self.select is not None
1136 model = klass_info["model"]
1137 for select_index in klass_info["select_fields"]:
1138 if self.select[select_index][0].target.model == model:
1139 return self.select[select_index][0]
1140 return None
1141
1142 def _get_field_choices(root_klass_info: dict[str, Any]) -> Generator[str]:
1143 """Yield all allowed field paths in breadth-first search order."""
1144 yield "self"
1145
1146 queue: collections.deque[tuple[list[str], dict[str, Any]]] = (
1147 collections.deque(
1148 ([], related_klass_info)
1149 for related_klass_info in root_klass_info.get(
1150 "related_klass_infos", []
1151 )
1152 )
1153 )
1154 while queue:
1155 parent_path, klass_info = queue.popleft()
1156 field = klass_info["field"]
1157 if klass_info["reverse"]:
1158 field = field.remote_field
1159 path = parent_path + [field.name]
1160 yield LOOKUP_SEP.join(path)
1161 queue.extend(
1162 (path, related_klass_info)
1163 for related_klass_info in klass_info.get("related_klass_infos", [])
1164 )
1165
1166 if not self.klass_info:
1167 return []
1168 result = []
1169 invalid_names = []
1170 for name in self.query.select_for_update_of:
1171 klass_info = self.klass_info
1172 if name == "self":
1173 col = _get_first_selected_col_from_model(klass_info)
1174 else:
1175 for part in name.split(LOOKUP_SEP):
1176 if klass_info is None:
1177 break
1178 klass_infos = (*klass_info.get("related_klass_infos", []),)
1179 for related_klass_info in klass_infos:
1180 field = related_klass_info["field"]
1181 if related_klass_info["reverse"]:
1182 field = field.remote_field
1183 if field.name == part:
1184 klass_info = related_klass_info
1185 break
1186 else:
1187 klass_info = None
1188 break
1189 if klass_info is None:
1190 invalid_names.append(name)
1191 continue
1192 col = _get_first_selected_col_from_model(klass_info)
1193 if col is not None:
1194 result.append(self.quote_name_unless_alias(col.alias))
1195 if invalid_names:
1196 raise FieldError(
1197 "Invalid field name(s) given in select_for_update(of=(...)): {}. "
1198 "Only relational fields followed in the query are allowed. "
1199 "Choices are: {}.".format(
1200 ", ".join(invalid_names),
1201 ", ".join(_get_field_choices(self.klass_info)),
1202 )
1203 )
1204 return result
1205
1206 def results_iter(
1207 self,
1208 results: Any = None,
1209 tuple_expected: bool = False,
1210 chunked_fetch: bool = False,
1211 ) -> Iterable[Any]:
1212 """Return an iterator over the results from executing this query."""
1213 if results is None:
1214 results = self.execute_sql(MULTI, chunked_fetch=chunked_fetch)
1215 assert self.select is not None # Set during query execution
1216 fields = [s[0] for s in self.select[0 : self.col_count]]
1217 converters = get_converters(fields, self.connection)
1218 rows = results
1219 if converters:
1220 rows = apply_converters(rows, converters, self.connection)
1221 if tuple_expected:
1222 rows = map(tuple, rows)
1223 return rows
1224
1225 def has_results(self) -> bool:
1226 """Check if the query returns any results."""
1227 return bool(self.execute_sql(SINGLE))
1228
1229 def execute_sql(
1230 self,
1231 result_type: str = MULTI,
1232 chunked_fetch: bool = False,
1233 ) -> Any:
1234 """
1235 Run the query against the database and return the result(s). The
1236 return value is a single data item if result_type is SINGLE, or a
1237 flat iterable of rows if the result_type is MULTI.
1238
1239 result_type is either MULTI (returns a list from fetchall(), or a
1240 streaming generator from cursor.stream() when chunked_fetch=True),
1241 SINGLE (only retrieve a single row), or None. In this last case, the
1242 cursor is returned if any query is executed, since it's used by
1243 subclasses such as InsertQuery). It's possible, however, that no query
1244 is needed, as the filters describe an empty set. In that case, None is
1245 returned, to avoid any unnecessary database interaction.
1246 """
1247 result_type = result_type or NO_RESULTS
1248 try:
1249 as_sql_result = self.as_sql()
1250 # SQLCompiler.as_sql returns SqlWithParams, subclasses may differ
1251 assert isinstance(as_sql_result, tuple)
1252 assert isinstance(as_sql_result[0], str)
1253 sql, params = as_sql_result
1254 if not sql:
1255 raise EmptyResultSet
1256 except EmptyResultSet:
1257 if result_type == MULTI:
1258 return iter([])
1259 else:
1260 return
1261 cursor = self.connection.cursor()
1262 if chunked_fetch:
1263 # Use psycopg3's cursor.stream() for server-side cursor iteration.
1264 result = cursor.stream(sql, params)
1265 if self.has_extra_select:
1266 col_count = self.col_count
1267 result = (r[:col_count] for r in result)
1268 return result
1269
1270 try:
1271 cursor.execute(sql, params)
1272 except Exception:
1273 cursor.close()
1274 raise
1275
1276 if result_type == CURSOR:
1277 # Give the caller the cursor to process and close.
1278 return cursor
1279 if result_type == SINGLE:
1280 try:
1281 val = cursor.fetchone()
1282 if val:
1283 return val[0 : self.col_count]
1284 return val
1285 finally:
1286 # done with the cursor
1287 cursor.close()
1288 if result_type == NO_RESULTS:
1289 cursor.close()
1290 return
1291
1292 try:
1293 rows = cursor.fetchall()
1294 finally:
1295 cursor.close()
1296 if self.has_extra_select:
1297 rows = [r[: self.col_count] for r in rows]
1298 return rows
1299
1300 def explain_query(self) -> Generator[str]:
1301 result = self.execute_sql()
1302 explain_info = self.query.explain_info
1303 # PostgreSQL may return tuples with integers and strings depending on
1304 # the EXPLAIN format. Flatten them out into strings.
1305 format_ = explain_info.format if explain_info is not None else None
1306 output_formatter = json.dumps if format_ and format_.lower() == "json" else str
1307 for row in result:
1308 if not isinstance(row, str):
1309 yield " ".join(output_formatter(c) for c in row)
1310 else:
1311 yield row
1312
1313
1314class SQLInsertCompiler(SQLCompiler):
1315 query: InsertQuery
1316 returning_fields: list | None = None
1317 returning_params: tuple = ()
1318
1319 def field_as_sql(self, field: Any, val: Any) -> tuple[str, list]:
1320 """
1321 Take a field and a value intended to be saved on that field, and
1322 return placeholder SQL and accompanying params. Check for raw values,
1323 expressions, and fields with get_placeholder() defined in that order.
1324
1325 When field is None, consider the value raw and use it as the
1326 placeholder, with no corresponding parameters returned.
1327 """
1328 if val is DATABASE_DEFAULT:
1329 # Emit the literal DEFAULT keyword so Postgres uses the column's
1330 # persistent DEFAULT (e.g. `gen_random_uuid()`). RETURNING then
1331 # populates the real value back onto the instance.
1332 sql, params = "DEFAULT", []
1333 elif field is None:
1334 # A field value of None means the value is raw.
1335 sql, params = val, []
1336 elif hasattr(val, "as_sql"):
1337 # This is an expression, let's compile it.
1338 sql, params_tuple = self.compile(val)
1339 params = list(params_tuple)
1340 elif hasattr(field, "get_placeholder"):
1341 # Some fields (e.g. geo fields) need special munging before
1342 # they can be inserted.
1343 sql, params = field.get_placeholder(val, self, self.connection), [val]
1344 else:
1345 # Return the common case for the placeholder
1346 sql, params = "%s", [val]
1347
1348 return sql, list(params) # Ensure params is a list
1349
1350 def prepare_value(self, field: Any, value: Any) -> Any:
1351 """
1352 Prepare a value to be used in a query by resolving it if it is an
1353 expression and otherwise calling the field's get_db_prep_save().
1354 """
1355 if value is DATABASE_DEFAULT:
1356 # Carry the sentinel through untouched — field_as_sql will emit
1357 # the literal DEFAULT keyword.
1358 return value
1359 if isinstance(value, ResolvableExpression):
1360 value = value.resolve_expression(
1361 self.query, allow_joins=False, for_save=True
1362 )
1363 # Don't allow values containing Col expressions. They refer to
1364 # existing columns on a row, but in the case of insert the row
1365 # doesn't exist yet.
1366 if value.contains_column_references:
1367 raise ValueError(
1368 f'Failed to insert expression "{value}" on {field}. F() expressions '
1369 "can only be used to update, not to insert."
1370 )
1371 if value.contains_aggregate:
1372 raise FieldError(
1373 "Aggregate functions are not allowed in this query "
1374 f"({field.name}={value!r})."
1375 )
1376 if value.contains_over_clause:
1377 raise FieldError(
1378 f"Window expressions are not allowed in this query ({field.name}={value!r})."
1379 )
1380 return field.get_db_prep_save(value, connection=self.connection)
1381
1382 def pre_save_val(self, field: Any, obj: Any) -> Any:
1383 """
1384 Get the given field's value off the given obj. pre_save() is used for
1385 things like update_now on DateTimeField.
1386 """
1387 return field.pre_save(obj, add=True)
1388
1389 def assemble_as_sql(
1390 self, fields: Sequence[Any], value_rows: list[list[Any]]
1391 ) -> tuple[Any, list[list[Any]]]:
1392 """
1393 Take a sequence of N fields and a sequence of M rows of values, and
1394 generate placeholder SQL and parameters for each field and value.
1395 Return a pair containing:
1396 * a sequence of M rows of N SQL placeholder strings, and
1397 * a sequence of M rows of corresponding parameter values.
1398
1399 Each placeholder string may contain any number of '%s' interpolation
1400 strings, and each parameter row will contain exactly as many params
1401 as the total number of '%s's in the corresponding placeholder row.
1402 """
1403 if not value_rows:
1404 return [], []
1405
1406 # list of (sql, [params]) tuples for each object to be saved
1407 # Shape: [n_objs][n_fields][2]
1408 rows_of_fields_as_sql = (
1409 (self.field_as_sql(field, v) for field, v in zip(fields, row))
1410 for row in value_rows
1411 )
1412
1413 # tuple like ([sqls], [[params]s]) for each object to be saved
1414 # Shape: [n_objs][2][n_fields]
1415 sql_and_param_pair_rows = (zip(*row) for row in rows_of_fields_as_sql)
1416
1417 # Extract separate lists for placeholders and params.
1418 # Each of these has shape [n_objs][n_fields]
1419 placeholder_rows, param_rows = zip(*sql_and_param_pair_rows)
1420
1421 # Params for each field are still lists, and need to be flattened.
1422 param_rows = [[p for ps in row for p in ps] for row in param_rows]
1423
1424 return placeholder_rows, param_rows
1425
1426 def as_sql( # ty: ignore[invalid-method-override] # Returns list for internal iteration in execute_sql
1427 self, with_limits: bool = True, with_col_aliases: bool = False
1428 ) -> list[SqlWithParams]:
1429 # We don't need quote_name_unless_alias() here, since these are all
1430 # going to be column names (so we can avoid the extra overhead).
1431 qn = quote_name
1432 assert self.query.model is not None, "INSERT requires a model"
1433 meta = self.query.model._model_meta
1434 options = self.query.model.model_options
1435 result = [f"INSERT INTO {qn(options.db_table)}"]
1436 if self.query.fields:
1437 fields = self.query.fields
1438 else:
1439 fields = [meta.get_forward_field("id")]
1440 result.append("({})".format(", ".join(qn(f.column) for f in fields)))
1441
1442 if self.query.fields:
1443 value_rows = [
1444 [
1445 self.prepare_value(field, self.pre_save_val(field, obj))
1446 for field in fields
1447 ]
1448 for obj in self.query.objs
1449 ]
1450 else:
1451 # An empty object.
1452 value_rows = [[PK_DEFAULT_VALUE] for _ in self.query.objs]
1453 fields = [None]
1454
1455 placeholder_rows, param_rows = self.assemble_as_sql(fields, value_rows)
1456
1457 conflict_suffix_sql = on_conflict_suffix_sql(
1458 fields, # ty: ignore[invalid-argument-type]
1459 self.query.on_conflict,
1460 (f.column for f in self.query.update_fields),
1461 (f.column for f in self.query.unique_fields),
1462 )
1463 if self.returning_fields:
1464 # Use RETURNING clause to get inserted values
1465 result.append(
1466 bulk_insert_sql(fields, placeholder_rows) # ty: ignore[invalid-argument-type]
1467 )
1468 params = param_rows
1469 if conflict_suffix_sql:
1470 result.append(conflict_suffix_sql)
1471 # Skip appending the RETURNING clause if it's an empty string.
1472 r_sql, self.returning_params = return_insert_columns(self.returning_fields)
1473 if r_sql:
1474 result.append(r_sql)
1475 params += [list(self.returning_params)]
1476 return [(" ".join(result), tuple(chain.from_iterable(params)))]
1477
1478 # Bulk insert without returning fields
1479 result.append(bulk_insert_sql(fields, placeholder_rows)) # ty: ignore[invalid-argument-type]
1480 if conflict_suffix_sql:
1481 result.append(conflict_suffix_sql)
1482 return [(" ".join(result), tuple(p for ps in param_rows for p in ps))]
1483
1484 def execute_sql( # ty: ignore[invalid-method-override]
1485 self, returning_fields: list | None = None
1486 ) -> list:
1487 assert self.query.model is not None, "INSERT execution requires a model"
1488 options = self.query.model.model_options
1489 self.returning_fields = returning_fields
1490 with self.connection.cursor() as cursor:
1491 for sql, params in self.as_sql():
1492 cursor.execute(sql, params)
1493 if not self.returning_fields:
1494 return []
1495 # Use RETURNING clause for both single and bulk inserts
1496 if len(self.query.objs) > 1:
1497 rows = cursor.fetchall()
1498 else:
1499 rows = [cursor.fetchone()]
1500 cols = [field.get_col(options.db_table) for field in self.returning_fields]
1501 converters = get_converters(cols, self.connection)
1502 if converters:
1503 rows = list(apply_converters(rows, converters, self.connection))
1504 return rows
1505
1506
1507class SQLDeleteCompiler(SQLCompiler):
1508 @cached_property
1509 def single_alias(self) -> bool:
1510 # Ensure base table is in aliases.
1511 self.query.get_initial_alias()
1512 return sum(self.query.alias_refcount[t] > 0 for t in self.query.alias_map) == 1
1513
1514 @classmethod
1515 def _expr_refs_base_model(cls, expr: Any, base_model: Any) -> bool:
1516 if isinstance(expr, Query):
1517 return expr.model == base_model
1518 if not hasattr(expr, "get_source_expressions"):
1519 return False
1520 return any(
1521 cls._expr_refs_base_model(source_expr, base_model)
1522 for source_expr in expr.get_source_expressions()
1523 )
1524
1525 @cached_property
1526 def contains_self_reference_subquery(self) -> bool:
1527 return any(
1528 self._expr_refs_base_model(expr, self.query.model)
1529 for expr in chain(
1530 self.query.annotations.values(), self.query.where.children
1531 )
1532 )
1533
1534 def _as_sql(self, query: Query) -> SqlWithParams:
1535 delete = f"DELETE FROM {self.quote_name_unless_alias(query.base_table)}" # ty: ignore[invalid-argument-type]
1536 try:
1537 where, params = self.compile(query.where)
1538 except FullResultSet:
1539 return delete, ()
1540 return f"{delete} WHERE {where}", tuple(params)
1541
1542 def as_sql(
1543 self, with_limits: bool = True, with_col_aliases: bool = False
1544 ) -> SqlWithParams:
1545 """
1546 Create the SQL for this query. Return the SQL string and list of
1547 parameters.
1548 """
1549 if self.single_alias and not self.contains_self_reference_subquery:
1550 return self._as_sql(self.query)
1551 innerq = self.query.clone()
1552 innerq.__class__ = Query
1553 innerq.clear_select_clause()
1554 assert self.query.model is not None, "DELETE requires a model"
1555 id_field = self.query.model._model_meta.get_forward_field("id")
1556 innerq.select = (id_field.get_col(self.query.get_initial_alias()),)
1557 outerq = Query(self.query.model)
1558 outerq.add_filter("id__in", innerq)
1559 return self._as_sql(outerq)
1560
1561
1562class SQLUpdateCompiler(SQLCompiler):
1563 def as_sql(
1564 self, with_limits: bool = True, with_col_aliases: bool = False
1565 ) -> SqlWithParams:
1566 """
1567 Create the SQL for this query. Return the SQL string and list of
1568 parameters.
1569 """
1570 self.pre_sql_setup()
1571 query_values = getattr(self.query, "values", None)
1572 if not query_values:
1573 return "", ()
1574 qn = self.quote_name_unless_alias
1575 values, update_params = [], []
1576 for field, val in query_values:
1577 if isinstance(val, ResolvableExpression):
1578 val = val.resolve_expression(
1579 self.query, allow_joins=False, for_save=True
1580 )
1581 if val.contains_aggregate:
1582 raise FieldError(
1583 "Aggregate functions are not allowed in this query "
1584 f"({field.name}={val!r})."
1585 )
1586 if val.contains_over_clause:
1587 raise FieldError(
1588 "Window expressions are not allowed in this query "
1589 f"({field.name}={val!r})."
1590 )
1591 elif hasattr(val, "prepare_database_save"):
1592 if isinstance(field, RelatedField):
1593 val = val.prepare_database_save(field)
1594 else:
1595 raise TypeError(
1596 f"Tried to update field {field} with a model instance, {val!r}. "
1597 f"Use a value compatible with {field.__class__.__name__}."
1598 )
1599 val = field.get_db_prep_save(val, connection=self.connection)
1600
1601 # Getting the placeholder for the field.
1602 if hasattr(field, "get_placeholder"):
1603 placeholder = field.get_placeholder(val, self, self.connection)
1604 else:
1605 placeholder = "%s"
1606 name = field.column
1607 if hasattr(val, "as_sql"):
1608 sql, params = self.compile(val)
1609 values.append(f"{qn(name)} = {placeholder % sql}")
1610 update_params.extend(params)
1611 elif val is not None:
1612 values.append(f"{qn(name)} = {placeholder}")
1613 update_params.append(val)
1614 else:
1615 values.append(f"{qn(name)} = NULL")
1616 table = self.query.base_table
1617 result = [
1618 f"UPDATE {qn(table)} SET", # ty: ignore[invalid-argument-type]
1619 ", ".join(values),
1620 ]
1621 try:
1622 where, params = self.compile(self.query.where)
1623 except FullResultSet:
1624 params = []
1625 else:
1626 result.append(f"WHERE {where}")
1627 return " ".join(result), tuple(update_params + list(params))
1628
1629 def execute_sql(self, result_type: str) -> int: # ty: ignore[invalid-method-override]
1630 """Execute the update and return the number of rows affected."""
1631 cursor = super().execute_sql(result_type)
1632 try:
1633 return cursor.rowcount if cursor else 0
1634 finally:
1635 if cursor:
1636 cursor.close()
1637
1638 def pre_sql_setup(
1639 self, with_col_aliases: bool = False
1640 ) -> tuple[list[Any], list[Any], list[SqlWithParams]] | None:
1641 """
1642 If the update depends on other tables (JOINs in the WHERE clause),
1643 rewrite the query so the current table is filtered by `id IN (subquery)`.
1644 """
1645 refcounts_before = self.query.alias_refcount.copy()
1646 # Ensure base table is in the query
1647 self.query.get_initial_alias()
1648 count = self.query.count_active_tables()
1649 if count == 1:
1650 return
1651 query = self.query.chain(klass=Query)
1652 query.select_related = False
1653 query.clear_ordering(force=True)
1654 query.select = ()
1655 query.add_fields(["id"])
1656 super().pre_sql_setup()
1657
1658 # Reset the where clause and drop the tables we no longer need (they
1659 # live in the sub-select now).
1660 self.query.clear_where()
1661 self.query.add_filter("id__in", query)
1662 self.query.reset_refcounts(refcounts_before)
1663
1664
1665class SQLAggregateCompiler(SQLCompiler):
1666 def as_sql(
1667 self, with_limits: bool = True, with_col_aliases: bool = False
1668 ) -> SqlWithParams:
1669 """
1670 Create the SQL for this query. Return the SQL string and list of
1671 parameters.
1672 """
1673 sql, params = [], []
1674 for annotation in self.query.annotation_select.values():
1675 ann_sql, ann_params = self.compile(annotation)
1676 ann_sql, ann_params = annotation.select_format(self, ann_sql, ann_params)
1677 sql.append(ann_sql)
1678 params.extend(ann_params)
1679 self.col_count = len(self.query.annotation_select)
1680 sql = ", ".join(sql)
1681 params = tuple(params)
1682
1683 inner_query = cast("AggregateQuery", self.query).inner_query
1684 inner_query_sql, inner_query_params = inner_query.get_compiler(
1685 elide_empty=self.elide_empty,
1686 ).as_sql(with_col_aliases=True)
1687 sql = f"SELECT {sql} FROM ({inner_query_sql}) subquery"
1688 params += inner_query_params
1689 return sql, params