codekingpro/portable-devtools
115k
1# ext/baked.py
2# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of SQLAlchemy and is released under
6# the MIT License: https://www.opensource.org/licenses/mit-license.php
7# mypy: ignore-errors
8
9
10"""Baked query extension.
11
12Provides a creational pattern for the :class:`.query.Query` object which
13allows the fully constructed object, Core select statement, and string
14compiled result to be fully cached.
15
16
17"""
18
19import collections.abc as collections_abc
20import logging
21
22from .. import exc as sa_exc
23from .. import util
24from ..orm import exc as orm_exc
25from ..orm.query import Query
26from ..orm.session import Session
27from ..sql import func
28from ..sql import literal_column
29from ..sql import util as sql_util
30
31
32log = logging.getLogger(__name__)
33
34
35class Bakery:
36 """Callable which returns a :class:`.BakedQuery`.
37
38 This object is returned by the class method
39 :meth:`.BakedQuery.bakery`. It exists as an object
40 so that the "cache" can be easily inspected.
41
42 .. versionadded:: 1.2
43
44
45 """
46
47 __slots__ = "cls", "cache"
48
49 def __init__(self, cls_, cache):
50 self.cls = cls_
51 self.cache = cache
52
53 def __call__(self, initial_fn, *args):
54 return self.cls(self.cache, initial_fn, args)
55
56
57class BakedQuery:
58 """A builder object for :class:`.query.Query` objects."""
59
60 __slots__ = "steps", "_bakery", "_cache_key", "_spoiled"
61
62 def __init__(self, bakery, initial_fn, args=()):
63 self._cache_key = ()
64 self._update_cache_key(initial_fn, args)
65 self.steps = [initial_fn]
66 self._spoiled = False
67 self._bakery = bakery
68
69 @classmethod
70 def bakery(cls, size=200, _size_alert=None):
71 """Construct a new bakery.
72
73 :return: an instance of :class:`.Bakery`
74
75 """
76
77 return Bakery(cls, util.LRUCache(size, size_alert=_size_alert))
78
79 def _clone(self):
80 b1 = BakedQuery.__new__(BakedQuery)
81 b1._cache_key = self._cache_key
82 b1.steps = list(self.steps)
83 b1._bakery = self._bakery
84 b1._spoiled = self._spoiled
85 return b1
86
87 def _update_cache_key(self, fn, args=()):
88 self._cache_key += (fn.__code__,) + args
89
90 def __iadd__(self, other):
91 if isinstance(other, tuple):
92 self.add_criteria(*other)
93 else:
94 self.add_criteria(other)
95 return self
96
97 def __add__(self, other):
98 if isinstance(other, tuple):
99 return self.with_criteria(*other)
100 else:
101 return self.with_criteria(other)
102
103 def add_criteria(self, fn, *args):
104 """Add a criteria function to this :class:`.BakedQuery`.
105
106 This is equivalent to using the ``+=`` operator to
107 modify a :class:`.BakedQuery` in-place.
108
109 """
110 self._update_cache_key(fn, args)
111 self.steps.append(fn)
112 return self
113
114 def with_criteria(self, fn, *args):
115 """Add a criteria function to a :class:`.BakedQuery` cloned from this
116 one.
117
118 This is equivalent to using the ``+`` operator to
119 produce a new :class:`.BakedQuery` with modifications.
120
121 """
122 return self._clone().add_criteria(fn, *args)
123
124 def for_session(self, session):
125 """Return a :class:`_baked.Result` object for this
126 :class:`.BakedQuery`.
127
128 This is equivalent to calling the :class:`.BakedQuery` as a
129 Python callable, e.g. ``result = my_baked_query(session)``.
130
131 """
132 return Result(self, session)
133
134 def __call__(self, session):
135 return self.for_session(session)
136
137 def spoil(self, full=False):
138 """Cancel any query caching that will occur on this BakedQuery object.
139
140 The BakedQuery can continue to be used normally, however additional
141 creational functions will not be cached; they will be called
142 on every invocation.
143
144 This is to support the case where a particular step in constructing
145 a baked query disqualifies the query from being cacheable, such
146 as a variant that relies upon some uncacheable value.
147
148 :param full: if False, only functions added to this
149 :class:`.BakedQuery` object subsequent to the spoil step will be
150 non-cached; the state of the :class:`.BakedQuery` up until
151 this point will be pulled from the cache. If True, then the
152 entire :class:`_query.Query` object is built from scratch each
153 time, with all creational functions being called on each
154 invocation.
155
156 """
157 if not full and not self._spoiled:
158 _spoil_point = self._clone()
159 _spoil_point._cache_key += ("_query_only",)
160 self.steps = [_spoil_point._retrieve_baked_query]
161 self._spoiled = True
162 return self
163
164 def _effective_key(self, session):
165 """Return the key that actually goes into the cache dictionary for
166 this :class:`.BakedQuery`, taking into account the given
167 :class:`.Session`.
168
169 This basically means we also will include the session's query_class,
170 as the actual :class:`_query.Query` object is part of what's cached
171 and needs to match the type of :class:`_query.Query` that a later
172 session will want to use.
173
174 """
175 return self._cache_key + (session._query_cls,)
176
177 def _with_lazyload_options(self, options, effective_path, cache_path=None):
178 """Cloning version of _add_lazyload_options."""
179 q = self._clone()
180 q._add_lazyload_options(options, effective_path, cache_path=cache_path)
181 return q
182
183 def _add_lazyload_options(self, options, effective_path, cache_path=None):
184 """Used by per-state lazy loaders to add options to the
185 "lazy load" query from a parent query.
186
187 Creates a cache key based on given load path and query options;
188 if a repeatable cache key cannot be generated, the query is
189 "spoiled" so that it won't use caching.
190
191 """
192
193 key = ()
194
195 if not cache_path:
196 cache_path = effective_path
197
198 for opt in options:
199 if opt._is_legacy_option or opt._is_compile_state:
200 ck = opt._generate_cache_key()
201 if ck is None:
202 self.spoil(full=True)
203 else:
204 assert not ck[1], (
205 "loader options with variable bound parameters "
206 "not supported with baked queries. Please "
207 "use new-style select() statements for cached "
208 "ORM queries."
209 )
210 key += ck[0]
211
212 self.add_criteria(
213 lambda q: q._with_current_path(effective_path).options(*options),
214 cache_path.path,
215 key,
216 )
217
218 def _retrieve_baked_query(self, session):
219 query = self._bakery.get(self._effective_key(session), None)
220 if query is None:
221 query = self._as_query(session)
222 self._bakery[self._effective_key(session)] = query.with_session(
223 None
224 )
225 return query.with_session(session)
226
227 def _bake(self, session):
228 query = self._as_query(session)
229 query.session = None
230
231 # in 1.4, this is where before_compile() event is
232 # invoked
233 statement = query._statement_20()
234
235 # if the query is not safe to cache, we still do everything as though
236 # we did cache it, since the receiver of _bake() assumes subqueryload
237 # context was set up, etc.
238 #
239 # note also we want to cache the statement itself because this
240 # allows the statement itself to hold onto its cache key that is
241 # used by the Connection, which in itself is more expensive to
242 # generate than what BakedQuery was able to provide in 1.3 and prior
243
244 if statement._compile_options._bake_ok:
245 self._bakery[self._effective_key(session)] = (
246 query,
247 statement,
248 )
249
250 return query, statement
251
252 def to_query(self, query_or_session):
253 """Return the :class:`_query.Query` object for use as a subquery.
254
255 This method should be used within the lambda callable being used
256 to generate a step of an enclosing :class:`.BakedQuery`. The
257 parameter should normally be the :class:`_query.Query` object that
258 is passed to the lambda::
259
260 sub_bq = self.bakery(lambda s: s.query(User.name))
261 sub_bq += lambda q: q.filter(User.id == Address.user_id).correlate(Address)
262
263 main_bq = self.bakery(lambda s: s.query(Address))
264 main_bq += lambda q: q.filter(sub_bq.to_query(q).exists())
265
266 In the case where the subquery is used in the first callable against
267 a :class:`.Session`, the :class:`.Session` is also accepted::
268
269 sub_bq = self.bakery(lambda s: s.query(User.name))
270 sub_bq += lambda q: q.filter(User.id == Address.user_id).correlate(Address)
271
272 main_bq = self.bakery(
273 lambda s: s.query(Address.id, sub_bq.to_query(q).scalar_subquery())
274 )
275
276 :param query_or_session: a :class:`_query.Query` object or a class
277 :class:`.Session` object, that is assumed to be within the context
278 of an enclosing :class:`.BakedQuery` callable.
279
280
281 .. versionadded:: 1.3
282
283
284 """ # noqa: E501
285
286 if isinstance(query_or_session, Session):
287 session = query_or_session
288 elif isinstance(query_or_session, Query):
289 session = query_or_session.session
290 if session is None:
291 raise sa_exc.ArgumentError(
292 "Given Query needs to be associated with a Session"
293 )
294 else:
295 raise TypeError(
296 "Query or Session object expected, got %r."
297 % type(query_or_session)
298 )
299 return self._as_query(session)
300
301 def _as_query(self, session):
302 query = self.steps[0](session)
303
304 for step in self.steps[1:]:
305 query = step(query)
306
307 return query
308
309
310class Result:
311 """Invokes a :class:`.BakedQuery` against a :class:`.Session`.
312
313 The :class:`_baked.Result` object is where the actual :class:`.query.Query`
314 object gets created, or retrieved from the cache,
315 against a target :class:`.Session`, and is then invoked for results.
316
317 """
318
319 __slots__ = "bq", "session", "_params", "_post_criteria"
320
321 def __init__(self, bq, session):
322 self.bq = bq
323 self.session = session
324 self._params = {}
325 self._post_criteria = []
326
327 def params(self, *args, **kw):
328 """Specify parameters to be replaced into the string SQL statement."""
329
330 if len(args) == 1:
331 kw.update(args[0])
332 elif len(args) > 0:
333 raise sa_exc.ArgumentError(
334 "params() takes zero or one positional argument, "
335 "which is a dictionary."
336 )
337 self._params.update(kw)
338 return self
339
340 def _using_post_criteria(self, fns):
341 if fns:
342 self._post_criteria.extend(fns)
343 return self
344
345 def with_post_criteria(self, fn):
346 """Add a criteria function that will be applied post-cache.
347
348 This adds a function that will be run against the
349 :class:`_query.Query` object after it is retrieved from the
350 cache. This currently includes **only** the
351 :meth:`_query.Query.params` and :meth:`_query.Query.execution_options`
352 methods.
353
354 .. warning:: :meth:`_baked.Result.with_post_criteria`
355 functions are applied
356 to the :class:`_query.Query`
357 object **after** the query's SQL statement
358 object has been retrieved from the cache. Only
359 :meth:`_query.Query.params` and
360 :meth:`_query.Query.execution_options`
361 methods should be used.
362
363
364 .. versionadded:: 1.2
365
366
367 """
368 return self._using_post_criteria([fn])
369
370 def _as_query(self):
371 q = self.bq._as_query(self.session).params(self._params)
372 for fn in self._post_criteria:
373 q = fn(q)
374 return q
375
376 def __str__(self):
377 return str(self._as_query())
378
379 def __iter__(self):
380 return self._iter().__iter__()
381
382 def _iter(self):
383 bq = self.bq
384
385 if not self.session.enable_baked_queries or bq._spoiled:
386 return self._as_query()._iter()
387
388 query, statement = bq._bakery.get(
389 bq._effective_key(self.session), (None, None)
390 )
391 if query is None:
392 query, statement = bq._bake(self.session)
393
394 if self._params:
395 q = query.params(self._params)
396 else:
397 q = query
398 for fn in self._post_criteria:
399 q = fn(q)
400
401 params = q._params
402 execution_options = dict(q._execution_options)
403 execution_options.update(
404 {
405 "_sa_orm_load_options": q.load_options,
406 "compiled_cache": bq._bakery,
407 }
408 )
409
410 result = self.session.execute(
411 statement, params, execution_options=execution_options
412 )
413 if result._attributes.get("is_single_entity", False):
414 result = result.scalars()
415
416 if result._attributes.get("filtered", False):
417 result = result.unique()
418
419 return result
420
421 def count(self):
422 """return the 'count'.
423
424 Equivalent to :meth:`_query.Query.count`.
425
426 Note this uses a subquery to ensure an accurate count regardless
427 of the structure of the original statement.
428
429 """
430
431 col = func.count(literal_column("*"))
432 bq = self.bq.with_criteria(lambda q: q._legacy_from_self(col))
433 return bq.for_session(self.session).params(self._params).scalar()
434
435 def scalar(self):
436 """Return the first element of the first result or None
437 if no rows present. If multiple rows are returned,
438 raises MultipleResultsFound.
439
440 Equivalent to :meth:`_query.Query.scalar`.
441
442 """
443 try:
444 ret = self.one()
445 if not isinstance(ret, collections_abc.Sequence):
446 return ret
447 return ret[0]
448 except orm_exc.NoResultFound:
449 return None
450
451 def first(self):
452 """Return the first row.
453
454 Equivalent to :meth:`_query.Query.first`.
455
456 """
457
458 bq = self.bq.with_criteria(lambda q: q.slice(0, 1))
459 return (
460 bq.for_session(self.session)
461 .params(self._params)
462 ._using_post_criteria(self._post_criteria)
463 ._iter()
464 .first()
465 )
466
467 def one(self):
468 """Return exactly one result or raise an exception.
469
470 Equivalent to :meth:`_query.Query.one`.
471
472 """
473 return self._iter().one()
474
475 def one_or_none(self):
476 """Return one or zero results, or raise an exception for multiple
477 rows.
478
479 Equivalent to :meth:`_query.Query.one_or_none`.
480
481 """
482 return self._iter().one_or_none()
483
484 def all(self):
485 """Return all rows.
486
487 Equivalent to :meth:`_query.Query.all`.
488
489 """
490 return self._iter().all()
491
492 def get(self, ident):
493 """Retrieve an object based on identity.
494
495 Equivalent to :meth:`_query.Query.get`.
496
497 """
498
499 query = self.bq.steps[0](self.session)
500 return query._get_impl(ident, self._load_on_pk_identity)
501
502 def _load_on_pk_identity(self, session, query, primary_key_identity, **kw):
503 """Load the given primary key identity from the database."""
504
505 mapper = query._raw_columns[0]._annotations["parententity"]
506
507 _get_clause, _get_params = mapper._get_clause
508
509 def setup(query):
510 _lcl_get_clause = _get_clause
511 q = query._clone()
512 q._get_condition()
513 q._order_by = None
514
515 # None present in ident - turn those comparisons
516 # into "IS NULL"
517 if None in primary_key_identity:
518 nones = {
519 _get_params[col].key
520 for col, value in zip(
521 mapper.primary_key, primary_key_identity
522 )
523 if value is None
524 }
525 _lcl_get_clause = sql_util.adapt_criterion_to_null(
526 _lcl_get_clause, nones
527 )
528
529 # TODO: can mapper._get_clause be pre-adapted?
530 q._where_criteria = (
531 sql_util._deep_annotate(_lcl_get_clause, {"_orm_adapt": True}),
532 )
533
534 for fn in self._post_criteria:
535 q = fn(q)
536 return q
537
538 # cache the query against a key that includes
539 # which positions in the primary key are NULL
540 # (remember, we can map to an OUTER JOIN)
541 bq = self.bq
542
543 # add the clause we got from mapper._get_clause to the cache
544 # key so that if a race causes multiple calls to _get_clause,
545 # we've cached on ours
546 bq = bq._clone()
547 bq._cache_key += (_get_clause,)
548
549 bq = bq.with_criteria(
550 setup, tuple(elem is None for elem in primary_key_identity)
551 )
552
553 params = {
554 _get_params[primary_key].key: id_val
555 for id_val, primary_key in zip(
556 primary_key_identity, mapper.primary_key
557 )
558 }
559
560 result = list(bq.for_session(self.session).params(**params))
561 l = len(result)
562 if l > 1:
563 raise orm_exc.MultipleResultsFound()
564 elif l:
565 return result[0]
566 else:
567 return None
568
569
570bakery = BakedQuery.bakery
571