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