Coverage for /home/martinb/.local/share/virtualenvs/camcops/lib/python3.6/site-packages/sqlalchemy/ext/baked.py : 55%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# sqlalchemy/ext/baked.py
2# Copyright (C) 2005-2020 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: http://www.opensource.org/licenses/mit-license.php
7"""Baked query extension.
9Provides a creational pattern for the :class:`.query.Query` object which
10allows the fully constructed object, Core select statement, and string
11compiled result to be fully cached.
14"""
16import copy
17import logging
19from .. import exc as sa_exc
20from .. import util
21from ..orm import exc as orm_exc
22from ..orm import strategy_options
23from ..orm.query import Query
24from ..orm.session import Session
25from ..sql import func
26from ..sql import literal_column
27from ..sql import util as sql_util
30log = logging.getLogger(__name__)
33class Bakery(object):
34 """Callable which returns a :class:`.BakedQuery`.
36 This object is returned by the class method
37 :meth:`.BakedQuery.bakery`. It exists as an object
38 so that the "cache" can be easily inspected.
40 .. versionadded:: 1.2
43 """
45 __slots__ = "cls", "cache"
47 def __init__(self, cls_, cache):
48 self.cls = cls_
49 self.cache = cache
51 def __call__(self, initial_fn, *args):
52 return self.cls(self.cache, initial_fn, args)
55class BakedQuery(object):
56 """A builder object for :class:`.query.Query` objects."""
58 __slots__ = "steps", "_bakery", "_cache_key", "_spoiled"
60 def __init__(self, bakery, initial_fn, args=()):
61 self._cache_key = ()
62 self._update_cache_key(initial_fn, args)
63 self.steps = [initial_fn]
64 self._spoiled = False
65 self._bakery = bakery
67 @classmethod
68 def bakery(cls, size=200, _size_alert=None):
69 """Construct a new bakery.
71 :return: an instance of :class:`.Bakery`
73 """
75 return Bakery(cls, util.LRUCache(size, size_alert=_size_alert))
77 def _clone(self):
78 b1 = BakedQuery.__new__(BakedQuery)
79 b1._cache_key = self._cache_key
80 b1.steps = list(self.steps)
81 b1._bakery = self._bakery
82 b1._spoiled = self._spoiled
83 return b1
85 def _update_cache_key(self, fn, args=()):
86 self._cache_key += (fn.__code__,) + args
88 def __iadd__(self, other):
89 if isinstance(other, tuple):
90 self.add_criteria(*other)
91 else:
92 self.add_criteria(other)
93 return self
95 def __add__(self, other):
96 if isinstance(other, tuple):
97 return self.with_criteria(*other)
98 else:
99 return self.with_criteria(other)
101 def add_criteria(self, fn, *args):
102 """Add a criteria function to this :class:`.BakedQuery`.
104 This is equivalent to using the ``+=`` operator to
105 modify a :class:`.BakedQuery` in-place.
107 """
108 self._update_cache_key(fn, args)
109 self.steps.append(fn)
110 return self
112 def with_criteria(self, fn, *args):
113 """Add a criteria function to a :class:`.BakedQuery` cloned from this one.
115 This is equivalent to using the ``+`` operator to
116 produce a new :class:`.BakedQuery` with modifications.
118 """
119 return self._clone().add_criteria(fn, *args)
121 def for_session(self, session):
122 """Return a :class:`.Result` object for this :class:`.BakedQuery`.
124 This is equivalent to calling the :class:`.BakedQuery` as a
125 Python callable, e.g. ``result = my_baked_query(session)``.
127 """
128 return Result(self, session)
130 def __call__(self, session):
131 return self.for_session(session)
133 def spoil(self, full=False):
134 """Cancel any query caching that will occur on this BakedQuery object.
136 The BakedQuery can continue to be used normally, however additional
137 creational functions will not be cached; they will be called
138 on every invocation.
140 This is to support the case where a particular step in constructing
141 a baked query disqualifies the query from being cacheable, such
142 as a variant that relies upon some uncacheable value.
144 :param full: if False, only functions added to this
145 :class:`.BakedQuery` object subsequent to the spoil step will be
146 non-cached; the state of the :class:`.BakedQuery` up until
147 this point will be pulled from the cache. If True, then the
148 entire :class:`_query.Query` object is built from scratch each
149 time, with all creational functions being called on each
150 invocation.
152 """
153 if not full and not self._spoiled:
154 _spoil_point = self._clone()
155 _spoil_point._cache_key += ("_query_only",)
156 self.steps = [_spoil_point._retrieve_baked_query]
157 self._spoiled = True
158 return self
160 def _effective_key(self, session):
161 """Return the key that actually goes into the cache dictionary for
162 this :class:`.BakedQuery`, taking into account the given
163 :class:`.Session`.
165 This basically means we also will include the session's query_class,
166 as the actual :class:`_query.Query` object is part of what's cached
167 and needs to match the type of :class:`_query.Query` that a later
168 session will want to use.
170 """
171 return self._cache_key + (session._query_cls,)
173 def _with_lazyload_options(self, options, effective_path, cache_path=None):
174 """Cloning version of _add_lazyload_options.
175 """
176 q = self._clone()
177 q._add_lazyload_options(options, effective_path, cache_path=cache_path)
178 return q
180 def _add_lazyload_options(self, options, effective_path, cache_path=None):
181 """Used by per-state lazy loaders to add options to the
182 "lazy load" query from a parent query.
184 Creates a cache key based on given load path and query options;
185 if a repeatable cache key cannot be generated, the query is
186 "spoiled" so that it won't use caching.
188 """
190 key = ()
192 if not cache_path:
193 cache_path = effective_path
195 if cache_path.path[0].is_aliased_class:
196 # paths that are against an AliasedClass are unsafe to cache
197 # with since the AliasedClass is an ad-hoc object.
198 self.spoil(full=True)
199 else:
200 for opt in options:
201 cache_key = opt._generate_cache_key(cache_path)
202 if cache_key is False:
203 self.spoil(full=True)
204 elif cache_key is not None:
205 key += cache_key
207 self.add_criteria(
208 lambda q: q._with_current_path(
209 effective_path
210 )._conditional_options(*options),
211 cache_path.path,
212 key,
213 )
215 def _retrieve_baked_query(self, session):
216 query = self._bakery.get(self._effective_key(session), None)
217 if query is None:
218 query = self._as_query(session)
219 self._bakery[self._effective_key(session)] = query.with_session(
220 None
221 )
222 return query.with_session(session)
224 def _bake(self, session):
225 query = self._as_query(session)
227 context = query._compile_context()
229 self._bake_subquery_loaders(session, context)
230 context.session = None
231 context.query = query = context.query.with_session(None)
232 query._execution_options = query._execution_options.union(
233 {"compiled_cache": self._bakery}
234 )
235 # we'll be holding onto the query for some of its state,
236 # so delete some compilation-use-only attributes that can take up
237 # space
238 for attr in (
239 "_correlate",
240 "_from_obj",
241 "_mapper_adapter_map",
242 "_joinpath",
243 "_joinpoint",
244 ):
245 query.__dict__.pop(attr, None)
247 # if the query is not safe to cache, we still do everything as though
248 # we did cache it, since the receiver of _bake() assumes subqueryload
249 # context was set up, etc.
250 if context.query._bake_ok:
251 self._bakery[self._effective_key(session)] = context
253 return context
255 def to_query(self, query_or_session):
256 """Return the :class:`_query.Query` object for use as a subquery.
258 This method should be used within the lambda callable being used
259 to generate a step of an enclosing :class:`.BakedQuery`. The
260 parameter should normally be the :class:`_query.Query` object that
261 is passed to the lambda::
263 sub_bq = self.bakery(lambda s: s.query(User.name))
264 sub_bq += lambda q: q.filter(
265 User.id == Address.user_id).correlate(Address)
267 main_bq = self.bakery(lambda s: s.query(Address))
268 main_bq += lambda q: q.filter(
269 sub_bq.to_query(q).exists())
271 In the case where the subquery is used in the first callable against
272 a :class:`.Session`, the :class:`.Session` is also accepted::
274 sub_bq = self.bakery(lambda s: s.query(User.name))
275 sub_bq += lambda q: q.filter(
276 User.id == Address.user_id).correlate(Address)
278 main_bq = self.bakery(
279 lambda s: s.query(Address.id, sub_bq.to_query(q).as_scalar())
280 )
282 :param query_or_session: a :class:`_query.Query` object or a class
283 :class:`.Session` object, that is assumed to be within the context
284 of an enclosing :class:`.BakedQuery` callable.
287 .. versionadded:: 1.3
290 """
292 if isinstance(query_or_session, Session):
293 session = query_or_session
294 elif isinstance(query_or_session, Query):
295 session = query_or_session.session
296 if session is None:
297 raise sa_exc.ArgumentError(
298 "Given Query needs to be associated with a Session"
299 )
300 else:
301 raise TypeError(
302 "Query or Session object expected, got %r."
303 % type(query_or_session)
304 )
305 return self._as_query(session)
307 def _as_query(self, session):
308 query = self.steps[0](session)
310 for step in self.steps[1:]:
311 query = step(query)
312 return query
314 def _bake_subquery_loaders(self, session, context):
315 """convert subquery eager loaders in the cache into baked queries.
317 For subquery eager loading to work, all we need here is that the
318 Query point to the correct session when it is run. However, since
319 we are "baking" anyway, we may as well also turn the query into
320 a "baked" query so that we save on performance too.
322 """
323 context.attributes["baked_queries"] = baked_queries = []
324 for k, v in list(context.attributes.items()):
325 if isinstance(v, Query):
326 if "subquery" in k:
327 bk = BakedQuery(self._bakery, lambda *args: v)
328 bk._cache_key = self._cache_key + k
329 bk._bake(session)
330 baked_queries.append((k, bk._cache_key, v))
331 del context.attributes[k]
333 def _unbake_subquery_loaders(
334 self, session, context, params, post_criteria
335 ):
336 """Retrieve subquery eager loaders stored by _bake_subquery_loaders
337 and turn them back into Result objects that will iterate just
338 like a Query object.
340 """
341 if "baked_queries" not in context.attributes:
342 return
344 for k, cache_key, query in context.attributes["baked_queries"]:
345 bk = BakedQuery(
346 self._bakery, lambda sess, q=query: q.with_session(sess)
347 )
348 bk._cache_key = cache_key
349 q = bk.for_session(session)
350 for fn in post_criteria:
351 q = q.with_post_criteria(fn)
352 context.attributes[k] = q.params(**params)
355class Result(object):
356 """Invokes a :class:`.BakedQuery` against a :class:`.Session`.
358 The :class:`.Result` object is where the actual :class:`.query.Query`
359 object gets created, or retrieved from the cache,
360 against a target :class:`.Session`, and is then invoked for results.
362 """
364 __slots__ = "bq", "session", "_params", "_post_criteria"
366 def __init__(self, bq, session):
367 self.bq = bq
368 self.session = session
369 self._params = {}
370 self._post_criteria = []
372 def params(self, *args, **kw):
373 """Specify parameters to be replaced into the string SQL statement."""
375 if len(args) == 1:
376 kw.update(args[0])
377 elif len(args) > 0:
378 raise sa_exc.ArgumentError(
379 "params() takes zero or one positional argument, "
380 "which is a dictionary."
381 )
382 self._params.update(kw)
383 return self
385 def _using_post_criteria(self, fns):
386 if fns:
387 self._post_criteria.extend(fns)
388 return self
390 def with_post_criteria(self, fn):
391 """Add a criteria function that will be applied post-cache.
393 This adds a function that will be run against the
394 :class:`_query.Query` object after it is retrieved from the
395 cache. Functions here can be used to alter the query in ways
396 that **do not affect the SQL output**, such as execution options
397 and shard identifiers (when using a shard-enabled query object)
399 .. warning:: :meth:`.Result.with_post_criteria` functions are applied
400 to the :class:`_query.Query`
401 object **after** the query's SQL statement
402 object has been retrieved from the cache. Any operations here
403 which intend to modify the SQL should ensure that
404 :meth:`.BakedQuery.spoil` was called first.
406 .. versionadded:: 1.2
409 """
410 return self._using_post_criteria([fn])
412 def _as_query(self):
413 q = self.bq._as_query(self.session).params(self._params)
414 for fn in self._post_criteria:
415 q = fn(q)
416 return q
418 def __str__(self):
419 return str(self._as_query())
421 def __iter__(self):
422 bq = self.bq
423 if not self.session.enable_baked_queries or bq._spoiled:
424 return iter(self._as_query())
426 baked_context = bq._bakery.get(bq._effective_key(self.session), None)
427 if baked_context is None:
428 baked_context = bq._bake(self.session)
430 context = copy.copy(baked_context)
431 context.session = self.session
432 context.attributes = context.attributes.copy()
434 bq._unbake_subquery_loaders(
435 self.session, context, self._params, self._post_criteria
436 )
438 context.statement.use_labels = True
439 if context.autoflush and not context.populate_existing:
440 self.session._autoflush()
441 q = context.query.params(self._params).with_session(self.session)
442 for fn in self._post_criteria:
443 q = fn(q)
445 return q._execute_and_instances(context)
447 def count(self):
448 """return the 'count'.
450 Equivalent to :meth:`_query.Query.count`.
452 Note this uses a subquery to ensure an accurate count regardless
453 of the structure of the original statement.
455 .. versionadded:: 1.1.6
457 """
459 col = func.count(literal_column("*"))
460 bq = self.bq.with_criteria(lambda q: q.from_self(col))
461 return bq.for_session(self.session).params(self._params).scalar()
463 def scalar(self):
464 """Return the first element of the first result or None
465 if no rows present. If multiple rows are returned,
466 raises MultipleResultsFound.
468 Equivalent to :meth:`_query.Query.scalar`.
470 .. versionadded:: 1.1.6
472 """
473 try:
474 ret = self.one()
475 if not isinstance(ret, tuple):
476 return ret
477 return ret[0]
478 except orm_exc.NoResultFound:
479 return None
481 def first(self):
482 """Return the first row.
484 Equivalent to :meth:`_query.Query.first`.
486 """
487 bq = self.bq.with_criteria(lambda q: q.slice(0, 1))
488 ret = list(
489 bq.for_session(self.session)
490 .params(self._params)
491 ._using_post_criteria(self._post_criteria)
492 )
493 if len(ret) > 0:
494 return ret[0]
495 else:
496 return None
498 def one(self):
499 """Return exactly one result or raise an exception.
501 Equivalent to :meth:`_query.Query.one`.
503 """
504 try:
505 ret = self.one_or_none()
506 except orm_exc.MultipleResultsFound as err:
507 util.raise_(
508 orm_exc.MultipleResultsFound(
509 "Multiple rows were found for one()"
510 ),
511 replace_context=err,
512 )
513 else:
514 if ret is None:
515 raise orm_exc.NoResultFound("No row was found for one()")
516 return ret
518 def one_or_none(self):
519 """Return one or zero results, or raise an exception for multiple
520 rows.
522 Equivalent to :meth:`_query.Query.one_or_none`.
524 .. versionadded:: 1.0.9
526 """
527 ret = list(self)
529 l = len(ret)
530 if l == 1:
531 return ret[0]
532 elif l == 0:
533 return None
534 else:
535 raise orm_exc.MultipleResultsFound(
536 "Multiple rows were found for one_or_none()"
537 )
539 def all(self):
540 """Return all rows.
542 Equivalent to :meth:`_query.Query.all`.
544 """
545 return list(self)
547 def get(self, ident):
548 """Retrieve an object based on identity.
550 Equivalent to :meth:`_query.Query.get`.
552 """
554 query = self.bq.steps[0](self.session)
555 return query._get_impl(ident, self._load_on_pk_identity)
557 def _load_on_pk_identity(self, query, primary_key_identity):
558 """Load the given primary key identity from the database."""
560 mapper = query._mapper_zero()
562 _get_clause, _get_params = mapper._get_clause
564 def setup(query):
565 _lcl_get_clause = _get_clause
566 q = query._clone()
567 q._get_condition()
568 q._order_by = None
570 # None present in ident - turn those comparisons
571 # into "IS NULL"
572 if None in primary_key_identity:
573 nones = set(
574 [
575 _get_params[col].key
576 for col, value in zip(
577 mapper.primary_key, primary_key_identity
578 )
579 if value is None
580 ]
581 )
582 _lcl_get_clause = sql_util.adapt_criterion_to_null(
583 _lcl_get_clause, nones
584 )
586 _lcl_get_clause = q._adapt_clause(_lcl_get_clause, True, False)
587 q._criterion = _lcl_get_clause
588 for fn in self._post_criteria:
589 q = fn(q)
590 return q
592 # cache the query against a key that includes
593 # which positions in the primary key are NULL
594 # (remember, we can map to an OUTER JOIN)
595 bq = self.bq
597 # add the clause we got from mapper._get_clause to the cache
598 # key so that if a race causes multiple calls to _get_clause,
599 # we've cached on ours
600 bq = bq._clone()
601 bq._cache_key += (_get_clause,)
603 bq = bq.with_criteria(
604 setup, tuple(elem is None for elem in primary_key_identity)
605 )
607 params = dict(
608 [
609 (_get_params[primary_key].key, id_val)
610 for id_val, primary_key in zip(
611 primary_key_identity, mapper.primary_key
612 )
613 ]
614 )
616 result = list(bq.for_session(self.session).params(**params))
617 l = len(result)
618 if l > 1:
619 raise orm_exc.MultipleResultsFound()
620 elif l:
621 return result[0]
622 else:
623 return None
626@util.deprecated(
627 "1.2", "Baked lazy loading is now the default implementation."
628)
629def bake_lazy_loaders():
630 """Enable the use of baked queries for all lazyloaders systemwide.
632 The "baked" implementation of lazy loading is now the sole implementation
633 for the base lazy loader; this method has no effect except for a warning.
635 """
636 pass
639@util.deprecated(
640 "1.2", "Baked lazy loading is now the default implementation."
641)
642def unbake_lazy_loaders():
643 """Disable the use of baked queries for all lazyloaders systemwide.
645 This method now raises NotImplementedError() as the "baked" implementation
646 is the only lazy load implementation. The
647 :paramref:`_orm.relationship.bake_queries` flag may be used to disable
648 the caching of queries on a per-relationship basis.
650 """
651 raise NotImplementedError(
652 "Baked lazy loading is now the default implementation"
653 )
656@strategy_options.loader_option()
657def baked_lazyload(loadopt, attr):
658 """Indicate that the given attribute should be loaded using "lazy"
659 loading with a "baked" query used in the load.
661 """
662 return loadopt.set_relationship_strategy(attr, {"lazy": "baked_select"})
665@baked_lazyload._add_unbound_fn
666@util.deprecated(
667 "1.2",
668 "Baked lazy loading is now the default "
669 "implementation for lazy loading.",
670)
671def baked_lazyload(*keys):
672 return strategy_options._UnboundLoad._from_keys(
673 strategy_options._UnboundLoad.baked_lazyload, keys, False, {}
674 )
677@baked_lazyload._add_unbound_all_fn
678@util.deprecated(
679 "1.2",
680 "Baked lazy loading is now the default "
681 "implementation for lazy loading.",
682)
683def baked_lazyload_all(*keys):
684 return strategy_options._UnboundLoad._from_keys(
685 strategy_options._UnboundLoad.baked_lazyload, keys, True, {}
686 )
689baked_lazyload = baked_lazyload._unbound_fn
690baked_lazyload_all = baked_lazyload_all._unbound_all_fn
692bakery = BakedQuery.bakery