Hide keyboard shortcuts

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. 

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 copy 

17import logging 

18 

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 

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 one. 

114 

115 This is equivalent to using the ``+`` operator to 

116 produce a new :class:`.BakedQuery` with modifications. 

117 

118 """ 

119 return self._clone().add_criteria(fn, *args) 

120 

121 def for_session(self, session): 

122 """Return a :class:`.Result` object for this :class:`.BakedQuery`. 

123 

124 This is equivalent to calling the :class:`.BakedQuery` as a 

125 Python callable, e.g. ``result = my_baked_query(session)``. 

126 

127 """ 

128 return Result(self, session) 

129 

130 def __call__(self, session): 

131 return self.for_session(session) 

132 

133 def spoil(self, full=False): 

134 """Cancel any query caching that will occur on this BakedQuery object. 

135 

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. 

139 

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. 

143 

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. 

151 

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 

159 

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`. 

164 

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. 

169 

170 """ 

171 return self._cache_key + (session._query_cls,) 

172 

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 

179 

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. 

183 

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. 

187 

188 """ 

189 

190 key = () 

191 

192 if not cache_path: 

193 cache_path = effective_path 

194 

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 

206 

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 ) 

214 

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) 

223 

224 def _bake(self, session): 

225 query = self._as_query(session) 

226 

227 context = query._compile_context() 

228 

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) 

246 

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 

252 

253 return context 

254 

255 def to_query(self, query_or_session): 

256 """Return the :class:`_query.Query` object for use as a subquery. 

257 

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:: 

262 

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) 

266 

267 main_bq = self.bakery(lambda s: s.query(Address)) 

268 main_bq += lambda q: q.filter( 

269 sub_bq.to_query(q).exists()) 

270 

271 In the case where the subquery is used in the first callable against 

272 a :class:`.Session`, the :class:`.Session` is also accepted:: 

273 

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) 

277 

278 main_bq = self.bakery( 

279 lambda s: s.query(Address.id, sub_bq.to_query(q).as_scalar()) 

280 ) 

281 

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. 

285 

286 

287 .. versionadded:: 1.3 

288 

289 

290 """ 

291 

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) 

306 

307 def _as_query(self, session): 

308 query = self.steps[0](session) 

309 

310 for step in self.steps[1:]: 

311 query = step(query) 

312 return query 

313 

314 def _bake_subquery_loaders(self, session, context): 

315 """convert subquery eager loaders in the cache into baked queries. 

316 

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. 

321 

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] 

332 

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. 

339 

340 """ 

341 if "baked_queries" not in context.attributes: 

342 return 

343 

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) 

353 

354 

355class Result(object): 

356 """Invokes a :class:`.BakedQuery` against a :class:`.Session`. 

357 

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. 

361 

362 """ 

363 

364 __slots__ = "bq", "session", "_params", "_post_criteria" 

365 

366 def __init__(self, bq, session): 

367 self.bq = bq 

368 self.session = session 

369 self._params = {} 

370 self._post_criteria = [] 

371 

372 def params(self, *args, **kw): 

373 """Specify parameters to be replaced into the string SQL statement.""" 

374 

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 

384 

385 def _using_post_criteria(self, fns): 

386 if fns: 

387 self._post_criteria.extend(fns) 

388 return self 

389 

390 def with_post_criteria(self, fn): 

391 """Add a criteria function that will be applied post-cache. 

392 

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) 

398 

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. 

405 

406 .. versionadded:: 1.2 

407 

408 

409 """ 

410 return self._using_post_criteria([fn]) 

411 

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 

417 

418 def __str__(self): 

419 return str(self._as_query()) 

420 

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()) 

425 

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) 

429 

430 context = copy.copy(baked_context) 

431 context.session = self.session 

432 context.attributes = context.attributes.copy() 

433 

434 bq._unbake_subquery_loaders( 

435 self.session, context, self._params, self._post_criteria 

436 ) 

437 

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) 

444 

445 return q._execute_and_instances(context) 

446 

447 def count(self): 

448 """return the 'count'. 

449 

450 Equivalent to :meth:`_query.Query.count`. 

451 

452 Note this uses a subquery to ensure an accurate count regardless 

453 of the structure of the original statement. 

454 

455 .. versionadded:: 1.1.6 

456 

457 """ 

458 

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() 

462 

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. 

467 

468 Equivalent to :meth:`_query.Query.scalar`. 

469 

470 .. versionadded:: 1.1.6 

471 

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 

480 

481 def first(self): 

482 """Return the first row. 

483 

484 Equivalent to :meth:`_query.Query.first`. 

485 

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 

497 

498 def one(self): 

499 """Return exactly one result or raise an exception. 

500 

501 Equivalent to :meth:`_query.Query.one`. 

502 

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 

517 

518 def one_or_none(self): 

519 """Return one or zero results, or raise an exception for multiple 

520 rows. 

521 

522 Equivalent to :meth:`_query.Query.one_or_none`. 

523 

524 .. versionadded:: 1.0.9 

525 

526 """ 

527 ret = list(self) 

528 

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 ) 

538 

539 def all(self): 

540 """Return all rows. 

541 

542 Equivalent to :meth:`_query.Query.all`. 

543 

544 """ 

545 return list(self) 

546 

547 def get(self, ident): 

548 """Retrieve an object based on identity. 

549 

550 Equivalent to :meth:`_query.Query.get`. 

551 

552 """ 

553 

554 query = self.bq.steps[0](self.session) 

555 return query._get_impl(ident, self._load_on_pk_identity) 

556 

557 def _load_on_pk_identity(self, query, primary_key_identity): 

558 """Load the given primary key identity from the database.""" 

559 

560 mapper = query._mapper_zero() 

561 

562 _get_clause, _get_params = mapper._get_clause 

563 

564 def setup(query): 

565 _lcl_get_clause = _get_clause 

566 q = query._clone() 

567 q._get_condition() 

568 q._order_by = None 

569 

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 ) 

585 

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 

591 

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 

596 

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,) 

602 

603 bq = bq.with_criteria( 

604 setup, tuple(elem is None for elem in primary_key_identity) 

605 ) 

606 

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 ) 

615 

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 

624 

625 

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. 

631 

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. 

634 

635 """ 

636 pass 

637 

638 

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. 

644 

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. 

649 

650 """ 

651 raise NotImplementedError( 

652 "Baked lazy loading is now the default implementation" 

653 ) 

654 

655 

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. 

660 

661 """ 

662 return loadopt.set_relationship_strategy(attr, {"lazy": "baked_select"}) 

663 

664 

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 ) 

675 

676 

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 ) 

687 

688 

689baked_lazyload = baked_lazyload._unbound_fn 

690baked_lazyload_all = baked_lazyload_all._unbound_all_fn 

691 

692bakery = BakedQuery.bakery