1"""
2Various data structures used in query construction.
3
4Factored out from plain.postgres.query to avoid making the main module very
5large and/or so that they can be used by other modules without getting into
6circular import difficulties.
7"""
8
9from __future__ import annotations
10
11import functools
12import inspect
13from collections.abc import Callable, Generator
14from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, Self, TypeGuard
15
16import psycopg
17from plain.logs import get_framework_logger
18from plain.postgres.constants import LOOKUP_SEP
19from plain.postgres.exceptions import FieldError
20from plain.utils import tree
21
22if TYPE_CHECKING:
23 from plain.postgres.base import Model
24 from plain.postgres.fields import Field
25 from plain.postgres.fields.related import ForeignKeyField, RelatedField
26 from plain.postgres.fields.reverse_related import ForeignKeyRel, ForeignObjectRel
27 from plain.postgres.lookups import Lookup, Transform
28 from plain.postgres.meta import Meta
29 from plain.postgres.sql.where import WhereNode
30
31logger = get_framework_logger()
32
33
34class PathInfo(NamedTuple):
35 """Information about a relation path when converting lookups (fk__somecol).
36
37 Describes the relation in Model terms (Meta and Fields for both
38 sides of the relation). The join_field is the field backing the relation.
39 """
40
41 from_meta: Meta
42 to_meta: Meta
43 target_field: Field
44 join_field: ForeignKeyField | ForeignKeyRel
45 m2m: bool
46 direct: bool
47
48
49def subclasses(cls: type) -> Generator[type]:
50 yield cls
51 for subclass in cls.__subclasses__():
52 yield from subclasses(subclass)
53
54
55class Q(tree.Node):
56 """
57 Encapsulate filters as objects that can then be combined logically (using
58 `&` and `|`).
59 """
60
61 # Connection types
62 AND = "AND"
63 OR = "OR"
64 default = AND
65 conditional = True
66
67 def __init__(
68 self,
69 *args: Any,
70 _connector: str | None = None,
71 _negated: bool = False,
72 **kwargs: Any,
73 ) -> None:
74 super().__init__(
75 children=[*args, *sorted(kwargs.items())],
76 connector=_connector,
77 negated=_negated,
78 )
79
80 def _combine(self, other: Any, conn: str) -> Q:
81 if getattr(other, "conditional", False) is False:
82 raise TypeError(other)
83 if not self:
84 return other.copy()
85 if not other and isinstance(other, Q):
86 return self.copy()
87
88 obj = self.create(connector=conn)
89 obj.add(self, conn)
90 obj.add(other, conn)
91 return obj
92
93 def __or__(self, other: Any) -> Q:
94 return self._combine(other, self.OR)
95
96 def __and__(self, other: Any) -> Q:
97 return self._combine(other, self.AND)
98
99 def __invert__(self) -> Q:
100 obj = self.copy()
101 obj.negate()
102 return obj
103
104 def resolve_expression(
105 self,
106 query: Any = None,
107 allow_joins: bool = True,
108 reuse: Any = None,
109 summarize: bool = False,
110 for_save: bool = False,
111 ) -> WhereNode:
112 # We must promote any new joins to left outer joins so that when Q is
113 # used as an expression, rows aren't filtered due to joins.
114 clause, joins = query._add_q(
115 self,
116 reuse,
117 allow_joins=allow_joins,
118 split_subq=False,
119 check_filterable=False,
120 summarize=summarize,
121 )
122 query.promote_joins(joins)
123 return clause
124
125 def flatten(self) -> Generator[Any]:
126 """
127 Recursively yield this Q object and all subexpressions, in depth-first
128 order.
129 """
130 yield self
131 for child in self.children:
132 if isinstance(child, tuple):
133 # Use the lookup.
134 child = child[1]
135 if hasattr(child, "flatten"):
136 yield from child.flatten()
137 else:
138 yield child
139
140 def check(self, against: dict[str, Any]) -> bool:
141 """
142 Do a database query to check if the expressions of the Q instance
143 matches against the expressions.
144 """
145 # Avoid circular imports.
146 from plain.postgres.expressions import ResolvableExpression, Value
147 from plain.postgres.fields import BooleanField
148 from plain.postgres.functions import Coalesce
149 from plain.postgres.sql import SINGLE, Query
150
151 query = Query(None)
152 for name, value in against.items():
153 if not isinstance(value, ResolvableExpression):
154 value = Value(value)
155 query.add_annotation(value, name, select=False)
156 query.add_annotation(Value(1), "_check")
157 # This will raise a FieldError if a field is missing in "against".
158 query.add_q(Q(Coalesce(self, True, output_field=BooleanField())))
159 compiler = query.get_compiler()
160 try:
161 return compiler.execute_sql(SINGLE) is not None
162 except psycopg.DatabaseError as e:
163 logger.warning(
164 "Got a database error calling check()",
165 extra={"expression": repr(self), "error": str(e)},
166 )
167 return True
168
169 def deconstruct(self) -> tuple[str, tuple[Any, ...], dict[str, Any]]:
170 path = f"{self.__class__.__module__}.{self.__class__.__name__}"
171 if path.startswith("plain.postgres.query_utils"):
172 path = path.replace("plain.postgres.query_utils", "plain.postgres")
173 args = tuple(self.children)
174 kwargs: dict[str, Any] = {}
175 if self.connector != self.default:
176 kwargs["_connector"] = self.connector
177 if self.negated:
178 kwargs["_negated"] = True
179 return path, args, kwargs
180
181
182class class_or_instance_method:
183 """
184 Hook used in RegisterLookupMixin to return partial functions depending on
185 the caller type (instance or class of models.Field).
186 """
187
188 def __init__(self, class_method: Any, instance_method: Any) -> None:
189 self.class_method = class_method
190 self.instance_method = instance_method
191
192 def __get__(self, instance: Any, owner: type) -> Any:
193 if instance is None:
194 return functools.partial(self.class_method, owner)
195 return functools.partial(self.instance_method, instance)
196
197
198class RegisterLookupMixin:
199 class_lookups: ClassVar[dict[str, type[Lookup | Transform]]]
200
201 def _get_lookup(self, lookup_name: str) -> type[Lookup | Transform] | None:
202 return self.get_lookups().get(lookup_name, None)
203
204 @functools.cache # noqa: B019 — keyed by class; classes live for the process
205 def get_class_lookups(cls: type[Self]) -> dict[str, type[Lookup | Transform]]:
206 class_lookups = [
207 parent.__dict__.get("class_lookups", {}) for parent in inspect.getmro(cls)
208 ]
209 return cls.merge_dicts(class_lookups)
210
211 def get_instance_lookups(self) -> dict[str, type[Lookup | Transform]]:
212 class_lookups = self.get_class_lookups()
213 if instance_lookups := getattr(self, "instance_lookups", None):
214 return {**class_lookups, **instance_lookups}
215 return class_lookups
216
217 get_lookups = class_or_instance_method(get_class_lookups, get_instance_lookups)
218 get_class_lookups: ClassVar[classmethod[Any, ..., Any]] = classmethod(
219 get_class_lookups
220 )
221
222 def get_lookup(self, lookup: str) -> type[Lookup] | None:
223 from plain.postgres.lookups import Lookup
224
225 found = self._get_lookup(lookup)
226 # output_field is a Field which inherits from RegisterLookupMixin
227 if found is None and (output_field := getattr(self, "output_field", None)):
228 return output_field.get_lookup(lookup)
229 if found is not None and not issubclass(found, Lookup):
230 return None
231 return found
232
233 def get_transform(self, name: str) -> Callable[..., Transform] | None:
234 from plain.postgres.lookups import Transform
235
236 found = self._get_lookup(name)
237 # output_field is a Field which inherits from RegisterLookupMixin
238 if found is None and (output_field := getattr(self, "output_field", None)):
239 return output_field.get_transform(name)
240 if found is not None and not issubclass(found, Transform):
241 return None
242 return found
243
244 @staticmethod
245 def merge_dicts(
246 dicts: list[dict[str, type[Lookup | Transform]]],
247 ) -> dict[str, type[Lookup | Transform]]:
248 """
249 Merge dicts in reverse to preference the order of the original list. e.g.,
250 merge_dicts([a, b]) will preference the keys in 'a' over those in 'b'.
251 """
252 merged: dict[str, type[Lookup | Transform]] = {}
253 for d in reversed(dicts):
254 merged.update(d)
255 return merged
256
257 @classmethod
258 def _clear_cached_class_lookups(cls: type[Self]) -> None:
259 for subclass in subclasses(cls):
260 if cached := getattr(subclass, "get_class_lookups", None):
261 cached.cache_clear()
262
263 def register_class_lookup(
264 cls: type[Self],
265 lookup: type[Lookup | Transform],
266 lookup_name: str | None = None,
267 ) -> type[Lookup | Transform]:
268 if lookup_name is None:
269 lookup_name = lookup.lookup_name
270 assert lookup_name is not None, "lookup_name must be set on the lookup class"
271 if "class_lookups" not in cls.__dict__:
272 cls.class_lookups = {}
273 cls.class_lookups[lookup_name] = lookup
274 cls._clear_cached_class_lookups()
275 return lookup
276
277 def register_instance_lookup(
278 self, lookup: type[Lookup | Transform], lookup_name: str | None = None
279 ) -> type[Lookup | Transform]:
280 if lookup_name is None:
281 lookup_name = lookup.lookup_name
282 if "instance_lookups" not in self.__dict__:
283 self.instance_lookups = {}
284 self.instance_lookups[lookup_name] = lookup
285 return lookup
286
287 register_lookup = class_or_instance_method(
288 register_class_lookup, register_instance_lookup
289 )
290 register_class_lookup: ClassVar[classmethod[Any, ..., Any]] = classmethod(
291 register_class_lookup
292 )
293
294 def _unregister_class_lookup(
295 cls: type[Self],
296 lookup: type[Lookup | Transform],
297 lookup_name: str | None = None,
298 ) -> None:
299 """
300 Remove given lookup from cls lookups. For use in tests only as it's
301 not thread-safe.
302 """
303 if lookup_name is None:
304 lookup_name = lookup.lookup_name
305 assert lookup_name is not None, "lookup_name must be set on the lookup class"
306 del cls.class_lookups[lookup_name]
307 cls._clear_cached_class_lookups()
308
309 def _unregister_instance_lookup(
310 self, lookup: type[Lookup | Transform], lookup_name: str | None = None
311 ) -> None:
312 """
313 Remove given lookup from instance lookups. For use in tests only as
314 it's not thread-safe.
315 """
316 if lookup_name is None:
317 lookup_name = lookup.lookup_name
318 del self.instance_lookups[lookup_name]
319
320 _unregister_lookup = class_or_instance_method(
321 _unregister_class_lookup, _unregister_instance_lookup
322 )
323 _unregister_class_lookup: ClassVar[classmethod[Any, ..., Any]] = classmethod(
324 _unregister_class_lookup
325 )
326
327
328def select_related_descend(
329 field: Any,
330 restricted: bool | None,
331 requested: dict[str, Any] | None,
332 select_mask: Any,
333 reverse: bool = False,
334) -> TypeGuard[RelatedField]:
335 """
336 Return True if this field should be used to descend deeper for
337 select_related() purposes. Used by both the query construction code
338 (compiler.get_related_selections()) and the model instance creation code
339 (compiler.klass_info).
340
341 Arguments:
342 * field - the field to be checked
343 * restricted - a boolean field, indicating if the field list has been
344 manually restricted using a requested clause)
345 * requested - The select_related() dictionary.
346 * select_mask - the dictionary of selected fields.
347 * reverse - boolean, True if we are checking a reverse select related
348 """
349 from plain.postgres.fields.related import RelatedField
350
351 if not isinstance(field, RelatedField):
352 return False
353 if restricted:
354 assert requested is not None, "requested must be provided when restricted=True"
355 if reverse and field.related_query_name() not in requested:
356 return False
357 if not reverse and field.name not in requested:
358 return False
359 if not restricted and field.allow_null:
360 return False
361 if (
362 restricted
363 and select_mask
364 and field.name in requested # ty: ignore[unsupported-operator]
365 and field not in select_mask
366 ):
367 raise FieldError(
368 f"Field {field.model.model_options.object_name}.{field.name} cannot be both "
369 "deferred and traversed using select_related at the same time."
370 )
371 return True
372
373
374def refs_expression(
375 lookup_parts: list[str], annotations: dict[str, Any]
376) -> tuple[str | None, tuple[str, ...]]:
377 """
378 Check if the lookup_parts contains references to the given annotations set.
379 Because the LOOKUP_SEP is contained in the default annotation names, check
380 each prefix of the lookup_parts for a match.
381 """
382 for n in range(1, len(lookup_parts) + 1):
383 level_n_lookup = LOOKUP_SEP.join(lookup_parts[0:n])
384 if annotations.get(level_n_lookup):
385 return level_n_lookup, tuple(lookup_parts[n:])
386 return None, ()
387
388
389def check_rel_lookup_compatibility(
390 model: type[Model], target_meta: Meta, field: Field | ForeignObjectRel
391) -> bool:
392 """
393 Check that model is compatible with target_meta — i.e. model matches
394 the target's model, or the field is a primary key whose model matches.
395 """
396
397 def check(meta: Meta) -> bool:
398 return model == meta.model
399
400 # Primary-key fields get a second chance: a queryset like
401 # `Model.query.filter(id__in=Model.query.all())` resolves `id__in` through
402 # the PK field, whose target meta is the remote model. Allow the match
403 # against the field's own model so the subquery (later reduced to
404 # `.values("id")`) is accepted.
405 return check(target_meta) or (
406 getattr(field, "primary_key", False) and check(field.model._model_meta)
407 )