1"""
2Code to manage the creation and SQL rendering of 'where' constraints.
3"""
4
5from functools import cached_property
6from typing import TYPE_CHECKING, Any
7
8from plain.postgres.exceptions import EmptyResultSet, FullResultSet
9from plain.postgres.expressions import ResolvableExpression
10from plain.utils import tree
11
12if TYPE_CHECKING:
13 from plain.postgres.connection import DatabaseConnection
14 from plain.postgres.lookups import Lookup
15 from plain.postgres.sql.compiler import SQLCompiler
16
17# Connection types
18AND = "AND"
19OR = "OR"
20
21
22class WhereNode(tree.Node):
23 """
24 An SQL WHERE clause.
25
26 The class is tied to the Query class that created it (in order to create
27 the correct SQL).
28
29 A child is usually an expression producing boolean values. Most likely the
30 expression is a Lookup instance.
31
32 However, a child could also be any class with as_sql() and either
33 relabeled_clone() method or relabel_aliases() and clone() methods and
34 contains_aggregate attribute.
35 """
36
37 default = AND
38 resolved = False
39 conditional = True
40
41 def split_having_qualify(
42 self, negated: bool = False, must_group_by: bool = False
43 ) -> tuple[WhereNode | None, WhereNode | None, WhereNode | None]:
44 """
45 Return three possibly None nodes: one for those parts of self that
46 should be included in the WHERE clause, one for those parts of self
47 that must be included in the HAVING clause, and one for those parts
48 that refer to window functions.
49 """
50 if not self.contains_aggregate and not self.contains_over_clause:
51 return self, None, None
52 in_negated = negated ^ self.negated
53 # Whether or not children must be connected in the same filtering
54 # clause (WHERE > HAVING > QUALIFY) to maintain logical semantic.
55 must_remain_connected = (in_negated and self.connector == AND) or (
56 not in_negated and self.connector == OR
57 )
58 if (
59 must_remain_connected
60 and self.contains_aggregate
61 and not self.contains_over_clause
62 ):
63 # It's must cheaper to short-circuit and stash everything in the
64 # HAVING clause than split children if possible.
65 return None, self, None
66 where_parts = []
67 having_parts = []
68 qualify_parts = []
69 for c in self.children:
70 if hasattr(c, "split_having_qualify"):
71 where_part, having_part, qualify_part = c.split_having_qualify(
72 in_negated, must_group_by
73 )
74 if where_part is not None:
75 where_parts.append(where_part)
76 if having_part is not None:
77 having_parts.append(having_part)
78 if qualify_part is not None:
79 qualify_parts.append(qualify_part)
80 elif c.contains_over_clause:
81 qualify_parts.append(c)
82 elif c.contains_aggregate:
83 having_parts.append(c)
84 else:
85 where_parts.append(c)
86 if must_remain_connected and qualify_parts:
87 # Disjunctive heterogeneous predicates can be pushed down to
88 # qualify as long as no conditional aggregation is involved.
89 if not where_parts or (where_parts and not must_group_by):
90 return None, None, self
91 elif where_parts:
92 # In theory this should only be enforced when dealing with
93 # where_parts containing predicates against multi-valued
94 # relationships that could affect aggregation results but this
95 # is complex to infer properly.
96 raise NotImplementedError(
97 "Heterogeneous disjunctive predicates against window functions are "
98 "not implemented when performing conditional aggregation."
99 )
100 where_node = (
101 self.create(where_parts, self.connector, self.negated)
102 if where_parts
103 else None
104 )
105 having_node = (
106 self.create(having_parts, self.connector, self.negated)
107 if having_parts
108 else None
109 )
110 qualify_node = (
111 self.create(qualify_parts, self.connector, self.negated)
112 if qualify_parts
113 else None
114 )
115 return where_node, having_node, qualify_node
116
117 def as_sql(
118 self, compiler: SQLCompiler, connection: DatabaseConnection
119 ) -> tuple[str, list[Any]]:
120 """
121 Return the SQL version of the where clause and the value to be
122 substituted in. Return '', [] if this node matches everything,
123 None, [] if this node is empty, and raise EmptyResultSet if this
124 node can't match anything.
125 """
126 result = []
127 result_params = []
128 if self.connector == AND:
129 full_needed, empty_needed = len(self.children), 1
130 else:
131 full_needed, empty_needed = 1, len(self.children)
132
133 for child in self.children:
134 try:
135 sql, params = compiler.compile(child)
136 except EmptyResultSet:
137 empty_needed -= 1
138 except FullResultSet:
139 full_needed -= 1
140 else:
141 if sql:
142 result.append(sql)
143 result_params.extend(params)
144 else:
145 full_needed -= 1
146 # Check if this node matches nothing or everything.
147 # First check the amount of full nodes and empty nodes
148 # to make this node empty/full.
149 # Now, check if this node is full/empty using the
150 # counts.
151 if empty_needed == 0:
152 if self.negated:
153 raise FullResultSet
154 else:
155 raise EmptyResultSet
156 if full_needed == 0:
157 if self.negated:
158 raise EmptyResultSet
159 else:
160 raise FullResultSet
161 conn = f" {self.connector} "
162 sql_string = conn.join(result)
163 if not sql_string:
164 raise FullResultSet
165 if self.negated:
166 sql_string = f"NOT ({sql_string})"
167 elif len(result) > 1 or self.resolved:
168 sql_string = f"({sql_string})"
169 return sql_string, result_params
170
171 def get_group_by_cols(self) -> list[Any]:
172 cols = []
173 for child in self.children:
174 cols.extend(child.get_group_by_cols())
175 return cols
176
177 def get_source_expressions(self) -> list[Any]:
178 return self.children[:]
179
180 def set_source_expressions(self, children: list[Any]) -> None:
181 assert len(children) == len(self.children)
182 self.children = children
183
184 def relabel_aliases(self, change_map: dict[str, str]) -> None:
185 """
186 Relabel the alias values of any children. 'change_map' is a dictionary
187 mapping old (current) alias values to the new values.
188 """
189 for pos, child in enumerate(self.children):
190 if hasattr(child, "relabel_aliases"):
191 # For example another WhereNode
192 child.relabel_aliases(change_map)
193 elif hasattr(child, "relabeled_clone"):
194 self.children[pos] = child.relabeled_clone(change_map)
195
196 def clone(self) -> WhereNode:
197 clone = self.create(connector=self.connector, negated=self.negated)
198 for child in self.children:
199 if hasattr(child, "clone"):
200 child = child.clone()
201 clone.children.append(child)
202 return clone
203
204 def relabeled_clone(self, change_map: dict[str, str]) -> WhereNode:
205 clone = self.clone()
206 clone.relabel_aliases(change_map)
207 return clone
208
209 def replace_expressions(self, replacements: dict[Any, Any]) -> WhereNode:
210 if replacement := replacements.get(self):
211 return replacement
212 clone = self.create(connector=self.connector, negated=self.negated)
213 for child in self.children:
214 clone.children.append(child.replace_expressions(replacements))
215 return clone
216
217 def get_refs(self) -> set[Any]:
218 refs = set()
219 for child in self.children:
220 refs |= child.get_refs()
221 return refs
222
223 @classmethod
224 def _contains_aggregate(cls, obj: Any) -> bool:
225 if isinstance(obj, tree.Node):
226 return any(cls._contains_aggregate(c) for c in obj.children)
227 return obj.contains_aggregate
228
229 @cached_property
230 def contains_aggregate(self) -> bool:
231 return self._contains_aggregate(self)
232
233 @classmethod
234 def _contains_over_clause(cls, obj: Any) -> bool:
235 if isinstance(obj, tree.Node):
236 return any(cls._contains_over_clause(c) for c in obj.children)
237 return obj.contains_over_clause
238
239 @cached_property
240 def contains_over_clause(self) -> bool:
241 return self._contains_over_clause(self)
242
243 @property
244 def is_summary(self) -> bool:
245 return any(child.is_summary for child in self.children)
246
247 @staticmethod
248 def _resolve_leaf(expr: Any, query: Any, *args: Any, **kwargs: Any) -> Any:
249 if isinstance(expr, ResolvableExpression):
250 expr = expr.resolve_expression(query, *args, **kwargs)
251 return expr
252
253 @classmethod
254 def _resolve_node(cls, node: Any, query: Any, *args: Any, **kwargs: Any) -> None:
255 if hasattr(node, "children"):
256 for child in node.children:
257 cls._resolve_node(child, query, *args, **kwargs)
258 if hasattr(node, "lhs"):
259 node.lhs = cls._resolve_leaf(node.lhs, query, *args, **kwargs)
260 if hasattr(node, "rhs"):
261 node.rhs = cls._resolve_leaf(node.rhs, query, *args, **kwargs)
262
263 def resolve_expression(self, *args: Any, **kwargs: Any) -> WhereNode:
264 clone = self.clone()
265 clone._resolve_node(clone, *args, **kwargs)
266 clone.resolved = True
267 return clone
268
269 @cached_property
270 def output_field(self) -> Any:
271 from plain.postgres.fields import BooleanField
272
273 return BooleanField()
274
275 @property
276 def _output_field_or_none(self) -> Any:
277 return self.output_field
278
279 def select_format(
280 self, compiler: SQLCompiler, sql: str, params: list[Any]
281 ) -> tuple[str, list[Any]]:
282 # Boolean expressions work directly in SELECT
283 return sql, params
284
285 def get_db_converters(self, connection: DatabaseConnection) -> list[Any]:
286 return self.output_field.get_db_converters(connection)
287
288 def get_lookup(self, lookup: str) -> type[Lookup] | None:
289 return self.output_field.get_lookup(lookup)
290
291 def leaves(self) -> Any:
292 for child in self.children:
293 if isinstance(child, WhereNode):
294 yield from child.leaves()
295 else:
296 yield child
297
298
299class NothingNode:
300 """A node that matches nothing."""
301
302 contains_aggregate = False
303 contains_over_clause = False
304
305 def as_sql(
306 self,
307 compiler: SQLCompiler | None = None,
308 connection: DatabaseConnection | None = None,
309 ) -> tuple[str, list[Any]]:
310 raise EmptyResultSet