Update 2025-04-24_11:44:19
This commit is contained in:
oib
570
venv/lib/python3.11/site-packages/sqlalchemy/ext/baked.py
Normal file
570
venv/lib/python3.11/site-packages/sqlalchemy/ext/baked.py
Normal file
@ -0,0 +1,570 @@
|
||||
# ext/baked.py
|
||||
# Copyright (C) 2005-2025 the SQLAlchemy authors and contributors
|
||||
# <see AUTHORS file>
|
||||
#
|
||||
# This module is part of SQLAlchemy and is released under
|
||||
# the MIT License: https://www.opensource.org/licenses/mit-license.php
|
||||
# mypy: ignore-errors
|
||||
|
||||
|
||||
"""Baked query extension.
|
||||
|
||||
Provides a creational pattern for the :class:`.query.Query` object which
|
||||
allows the fully constructed object, Core select statement, and string
|
||||
compiled result to be fully cached.
|
||||
|
||||
|
||||
"""
|
||||
|
||||
import collections.abc as collections_abc
|
||||
import logging
|
||||
|
||||
from .. import exc as sa_exc
|
||||
from .. import util
|
||||
from ..orm import exc as orm_exc
|
||||
from ..orm.query import Query
|
||||
from ..orm.session import Session
|
||||
from ..sql import func
|
||||
from ..sql import literal_column
|
||||
from ..sql import util as sql_util
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Bakery:
|
||||
"""Callable which returns a :class:`.BakedQuery`.
|
||||
|
||||
This object is returned by the class method
|
||||
:meth:`.BakedQuery.bakery`. It exists as an object
|
||||
so that the "cache" can be easily inspected.
|
||||
|
||||
.. versionadded:: 1.2
|
||||
|
||||
|
||||
"""
|
||||
|
||||
__slots__ = "cls", "cache"
|
||||
|
||||
def __init__(self, cls_, cache):
|
||||
self.cls = cls_
|
||||
self.cache = cache
|
||||
|
||||
def __call__(self, initial_fn, *args):
|
||||
return self.cls(self.cache, initial_fn, args)
|
||||
|
||||
|
||||
class BakedQuery:
|
||||
"""A builder object for :class:`.query.Query` objects."""
|
||||
|
||||
__slots__ = "steps", "_bakery", "_cache_key", "_spoiled"
|
||||
|
||||
def __init__(self, bakery, initial_fn, args=()):
|
||||
self._cache_key = ()
|
||||
self._update_cache_key(initial_fn, args)
|
||||
self.steps = [initial_fn]
|
||||
self._spoiled = False
|
||||
self._bakery = bakery
|
||||
|
||||
@classmethod
|
||||
def bakery(cls, size=200, _size_alert=None):
|
||||
"""Construct a new bakery.
|
||||
|
||||
:return: an instance of :class:`.Bakery`
|
||||
|
||||
"""
|
||||
|
||||
return Bakery(cls, util.LRUCache(size, size_alert=_size_alert))
|
||||
|
||||
def _clone(self):
|
||||
b1 = BakedQuery.__new__(BakedQuery)
|
||||
b1._cache_key = self._cache_key
|
||||
b1.steps = list(self.steps)
|
||||
b1._bakery = self._bakery
|
||||
b1._spoiled = self._spoiled
|
||||
return b1
|
||||
|
||||
def _update_cache_key(self, fn, args=()):
|
||||
self._cache_key += (fn.__code__,) + args
|
||||
|
||||
def __iadd__(self, other):
|
||||
if isinstance(other, tuple):
|
||||
self.add_criteria(*other)
|
||||
else:
|
||||
self.add_criteria(other)
|
||||
return self
|
||||
|
||||
def __add__(self, other):
|
||||
if isinstance(other, tuple):
|
||||
return self.with_criteria(*other)
|
||||
else:
|
||||
return self.with_criteria(other)
|
||||
|
||||
def add_criteria(self, fn, *args):
|
||||
"""Add a criteria function to this :class:`.BakedQuery`.
|
||||
|
||||
This is equivalent to using the ``+=`` operator to
|
||||
modify a :class:`.BakedQuery` in-place.
|
||||
|
||||
"""
|
||||
self._update_cache_key(fn, args)
|
||||
self.steps.append(fn)
|
||||
return self
|
||||
|
||||
def with_criteria(self, fn, *args):
|
||||
"""Add a criteria function to a :class:`.BakedQuery` cloned from this
|
||||
one.
|
||||
|
||||
This is equivalent to using the ``+`` operator to
|
||||
produce a new :class:`.BakedQuery` with modifications.
|
||||
|
||||
"""
|
||||
return self._clone().add_criteria(fn, *args)
|
||||
|
||||
def for_session(self, session):
|
||||
"""Return a :class:`_baked.Result` object for this
|
||||
:class:`.BakedQuery`.
|
||||
|
||||
This is equivalent to calling the :class:`.BakedQuery` as a
|
||||
Python callable, e.g. ``result = my_baked_query(session)``.
|
||||
|
||||
"""
|
||||
return Result(self, session)
|
||||
|
||||
def __call__(self, session):
|
||||
return self.for_session(session)
|
||||
|
||||
def spoil(self, full=False):
|
||||
"""Cancel any query caching that will occur on this BakedQuery object.
|
||||
|
||||
The BakedQuery can continue to be used normally, however additional
|
||||
creational functions will not be cached; they will be called
|
||||
on every invocation.
|
||||
|
||||
This is to support the case where a particular step in constructing
|
||||
a baked query disqualifies the query from being cacheable, such
|
||||
as a variant that relies upon some uncacheable value.
|
||||
|
||||
:param full: if False, only functions added to this
|
||||
:class:`.BakedQuery` object subsequent to the spoil step will be
|
||||
non-cached; the state of the :class:`.BakedQuery` up until
|
||||
this point will be pulled from the cache. If True, then the
|
||||
entire :class:`_query.Query` object is built from scratch each
|
||||
time, with all creational functions being called on each
|
||||
invocation.
|
||||
|
||||
"""
|
||||
if not full and not self._spoiled:
|
||||
_spoil_point = self._clone()
|
||||
_spoil_point._cache_key += ("_query_only",)
|
||||
self.steps = [_spoil_point._retrieve_baked_query]
|
||||
self._spoiled = True
|
||||
return self
|
||||
|
||||
def _effective_key(self, session):
|
||||
"""Return the key that actually goes into the cache dictionary for
|
||||
this :class:`.BakedQuery`, taking into account the given
|
||||
:class:`.Session`.
|
||||
|
||||
This basically means we also will include the session's query_class,
|
||||
as the actual :class:`_query.Query` object is part of what's cached
|
||||
and needs to match the type of :class:`_query.Query` that a later
|
||||
session will want to use.
|
||||
|
||||
"""
|
||||
return self._cache_key + (session._query_cls,)
|
||||
|
||||
def _with_lazyload_options(self, options, effective_path, cache_path=None):
|
||||
"""Cloning version of _add_lazyload_options."""
|
||||
q = self._clone()
|
||||
q._add_lazyload_options(options, effective_path, cache_path=cache_path)
|
||||
return q
|
||||
|
||||
def _add_lazyload_options(self, options, effective_path, cache_path=None):
|
||||
"""Used by per-state lazy loaders to add options to the
|
||||
"lazy load" query from a parent query.
|
||||
|
||||
Creates a cache key based on given load path and query options;
|
||||
if a repeatable cache key cannot be generated, the query is
|
||||
"spoiled" so that it won't use caching.
|
||||
|
||||
"""
|
||||
|
||||
key = ()
|
||||
|
||||
if not cache_path:
|
||||
cache_path = effective_path
|
||||
|
||||
for opt in options:
|
||||
if opt._is_legacy_option or opt._is_compile_state:
|
||||
ck = opt._generate_cache_key()
|
||||
if ck is None:
|
||||
self.spoil(full=True)
|
||||
else:
|
||||
assert not ck[1], (
|
||||
"loader options with variable bound parameters "
|
||||
"not supported with baked queries. Please "
|
||||
"use new-style select() statements for cached "
|
||||
"ORM queries."
|
||||
)
|
||||
key += ck[0]
|
||||
|
||||
self.add_criteria(
|
||||
lambda q: q._with_current_path(effective_path).options(*options),
|
||||
cache_path.path,
|
||||
key,
|
||||
)
|
||||
|
||||
def _retrieve_baked_query(self, session):
|
||||
query = self._bakery.get(self._effective_key(session), None)
|
||||
if query is None:
|
||||
query = self._as_query(session)
|
||||
self._bakery[self._effective_key(session)] = query.with_session(
|
||||
None
|
||||
)
|
||||
return query.with_session(session)
|
||||
|
||||
def _bake(self, session):
|
||||
query = self._as_query(session)
|
||||
query.session = None
|
||||
|
||||
# in 1.4, this is where before_compile() event is
|
||||
# invoked
|
||||
statement = query._statement_20()
|
||||
|
||||
# if the query is not safe to cache, we still do everything as though
|
||||
# we did cache it, since the receiver of _bake() assumes subqueryload
|
||||
# context was set up, etc.
|
||||
#
|
||||
# note also we want to cache the statement itself because this
|
||||
# allows the statement itself to hold onto its cache key that is
|
||||
# used by the Connection, which in itself is more expensive to
|
||||
# generate than what BakedQuery was able to provide in 1.3 and prior
|
||||
|
||||
if statement._compile_options._bake_ok:
|
||||
self._bakery[self._effective_key(session)] = (
|
||||
query,
|
||||
statement,
|
||||
)
|
||||
|
||||
return query, statement
|
||||
|
||||
def to_query(self, query_or_session):
|
||||
"""Return the :class:`_query.Query` object for use as a subquery.
|
||||
|
||||
This method should be used within the lambda callable being used
|
||||
to generate a step of an enclosing :class:`.BakedQuery`. The
|
||||
parameter should normally be the :class:`_query.Query` object that
|
||||
is passed to the lambda::
|
||||
|
||||
sub_bq = self.bakery(lambda s: s.query(User.name))
|
||||
sub_bq += lambda q: q.filter(User.id == Address.user_id).correlate(Address)
|
||||
|
||||
main_bq = self.bakery(lambda s: s.query(Address))
|
||||
main_bq += lambda q: q.filter(sub_bq.to_query(q).exists())
|
||||
|
||||
In the case where the subquery is used in the first callable against
|
||||
a :class:`.Session`, the :class:`.Session` is also accepted::
|
||||
|
||||
sub_bq = self.bakery(lambda s: s.query(User.name))
|
||||
sub_bq += lambda q: q.filter(User.id == Address.user_id).correlate(Address)
|
||||
|
||||
main_bq = self.bakery(
|
||||
lambda s: s.query(Address.id, sub_bq.to_query(q).scalar_subquery())
|
||||
)
|
||||
|
||||
:param query_or_session: a :class:`_query.Query` object or a class
|
||||
:class:`.Session` object, that is assumed to be within the context
|
||||
of an enclosing :class:`.BakedQuery` callable.
|
||||
|
||||
|
||||
.. versionadded:: 1.3
|
||||
|
||||
|
||||
""" # noqa: E501
|
||||
|
||||
if isinstance(query_or_session, Session):
|
||||
session = query_or_session
|
||||
elif isinstance(query_or_session, Query):
|
||||
session = query_or_session.session
|
||||
if session is None:
|
||||
raise sa_exc.ArgumentError(
|
||||
"Given Query needs to be associated with a Session"
|
||||
)
|
||||
else:
|
||||
raise TypeError(
|
||||
"Query or Session object expected, got %r."
|
||||
% type(query_or_session)
|
||||
)
|
||||
return self._as_query(session)
|
||||
|
||||
def _as_query(self, session):
|
||||
query = self.steps[0](session)
|
||||
|
||||
for step in self.steps[1:]:
|
||||
query = step(query)
|
||||
|
||||
return query
|
||||
|
||||
|
||||
class Result:
|
||||
"""Invokes a :class:`.BakedQuery` against a :class:`.Session`.
|
||||
|
||||
The :class:`_baked.Result` object is where the actual :class:`.query.Query`
|
||||
object gets created, or retrieved from the cache,
|
||||
against a target :class:`.Session`, and is then invoked for results.
|
||||
|
||||
"""
|
||||
|
||||
__slots__ = "bq", "session", "_params", "_post_criteria"
|
||||
|
||||
def __init__(self, bq, session):
|
||||
self.bq = bq
|
||||
self.session = session
|
||||
self._params = {}
|
||||
self._post_criteria = []
|
||||
|
||||
def params(self, *args, **kw):
|
||||
"""Specify parameters to be replaced into the string SQL statement."""
|
||||
|
||||
if len(args) == 1:
|
||||
kw.update(args[0])
|
||||
elif len(args) > 0:
|
||||
raise sa_exc.ArgumentError(
|
||||
"params() takes zero or one positional argument, "
|
||||
"which is a dictionary."
|
||||
)
|
||||
self._params.update(kw)
|
||||
return self
|
||||
|
||||
def _using_post_criteria(self, fns):
|
||||
if fns:
|
||||
self._post_criteria.extend(fns)
|
||||
return self
|
||||
|
||||
def with_post_criteria(self, fn):
|
||||
"""Add a criteria function that will be applied post-cache.
|
||||
|
||||
This adds a function that will be run against the
|
||||
:class:`_query.Query` object after it is retrieved from the
|
||||
cache. This currently includes **only** the
|
||||
:meth:`_query.Query.params` and :meth:`_query.Query.execution_options`
|
||||
methods.
|
||||
|
||||
.. warning:: :meth:`_baked.Result.with_post_criteria`
|
||||
functions are applied
|
||||
to the :class:`_query.Query`
|
||||
object **after** the query's SQL statement
|
||||
object has been retrieved from the cache. Only
|
||||
:meth:`_query.Query.params` and
|
||||
:meth:`_query.Query.execution_options`
|
||||
methods should be used.
|
||||
|
||||
|
||||
.. versionadded:: 1.2
|
||||
|
||||
|
||||
"""
|
||||
return self._using_post_criteria([fn])
|
||||
|
||||
def _as_query(self):
|
||||
q = self.bq._as_query(self.session).params(self._params)
|
||||
for fn in self._post_criteria:
|
||||
q = fn(q)
|
||||
return q
|
||||
|
||||
def __str__(self):
|
||||
return str(self._as_query())
|
||||
|
||||
def __iter__(self):
|
||||
return self._iter().__iter__()
|
||||
|
||||
def _iter(self):
|
||||
bq = self.bq
|
||||
|
||||
if not self.session.enable_baked_queries or bq._spoiled:
|
||||
return self._as_query()._iter()
|
||||
|
||||
query, statement = bq._bakery.get(
|
||||
bq._effective_key(self.session), (None, None)
|
||||
)
|
||||
if query is None:
|
||||
query, statement = bq._bake(self.session)
|
||||
|
||||
if self._params:
|
||||
q = query.params(self._params)
|
||||
else:
|
||||
q = query
|
||||
for fn in self._post_criteria:
|
||||
q = fn(q)
|
||||
|
||||
params = q._params
|
||||
execution_options = dict(q._execution_options)
|
||||
execution_options.update(
|
||||
{
|
||||
"_sa_orm_load_options": q.load_options,
|
||||
"compiled_cache": bq._bakery,
|
||||
}
|
||||
)
|
||||
|
||||
result = self.session.execute(
|
||||
statement, params, execution_options=execution_options
|
||||
)
|
||||
if result._attributes.get("is_single_entity", False):
|
||||
result = result.scalars()
|
||||
|
||||
if result._attributes.get("filtered", False):
|
||||
result = result.unique()
|
||||
|
||||
return result
|
||||
|
||||
def count(self):
|
||||
"""return the 'count'.
|
||||
|
||||
Equivalent to :meth:`_query.Query.count`.
|
||||
|
||||
Note this uses a subquery to ensure an accurate count regardless
|
||||
of the structure of the original statement.
|
||||
|
||||
"""
|
||||
|
||||
col = func.count(literal_column("*"))
|
||||
bq = self.bq.with_criteria(lambda q: q._legacy_from_self(col))
|
||||
return bq.for_session(self.session).params(self._params).scalar()
|
||||
|
||||
def scalar(self):
|
||||
"""Return the first element of the first result or None
|
||||
if no rows present. If multiple rows are returned,
|
||||
raises MultipleResultsFound.
|
||||
|
||||
Equivalent to :meth:`_query.Query.scalar`.
|
||||
|
||||
"""
|
||||
try:
|
||||
ret = self.one()
|
||||
if not isinstance(ret, collections_abc.Sequence):
|
||||
return ret
|
||||
return ret[0]
|
||||
except orm_exc.NoResultFound:
|
||||
return None
|
||||
|
||||
def first(self):
|
||||
"""Return the first row.
|
||||
|
||||
Equivalent to :meth:`_query.Query.first`.
|
||||
|
||||
"""
|
||||
|
||||
bq = self.bq.with_criteria(lambda q: q.slice(0, 1))
|
||||
return (
|
||||
bq.for_session(self.session)
|
||||
.params(self._params)
|
||||
._using_post_criteria(self._post_criteria)
|
||||
._iter()
|
||||
.first()
|
||||
)
|
||||
|
||||
def one(self):
|
||||
"""Return exactly one result or raise an exception.
|
||||
|
||||
Equivalent to :meth:`_query.Query.one`.
|
||||
|
||||
"""
|
||||
return self._iter().one()
|
||||
|
||||
def one_or_none(self):
|
||||
"""Return one or zero results, or raise an exception for multiple
|
||||
rows.
|
||||
|
||||
Equivalent to :meth:`_query.Query.one_or_none`.
|
||||
|
||||
"""
|
||||
return self._iter().one_or_none()
|
||||
|
||||
def all(self):
|
||||
"""Return all rows.
|
||||
|
||||
Equivalent to :meth:`_query.Query.all`.
|
||||
|
||||
"""
|
||||
return self._iter().all()
|
||||
|
||||
def get(self, ident):
|
||||
"""Retrieve an object based on identity.
|
||||
|
||||
Equivalent to :meth:`_query.Query.get`.
|
||||
|
||||
"""
|
||||
|
||||
query = self.bq.steps[0](self.session)
|
||||
return query._get_impl(ident, self._load_on_pk_identity)
|
||||
|
||||
def _load_on_pk_identity(self, session, query, primary_key_identity, **kw):
|
||||
"""Load the given primary key identity from the database."""
|
||||
|
||||
mapper = query._raw_columns[0]._annotations["parententity"]
|
||||
|
||||
_get_clause, _get_params = mapper._get_clause
|
||||
|
||||
def setup(query):
|
||||
_lcl_get_clause = _get_clause
|
||||
q = query._clone()
|
||||
q._get_condition()
|
||||
q._order_by = None
|
||||
|
||||
# None present in ident - turn those comparisons
|
||||
# into "IS NULL"
|
||||
if None in primary_key_identity:
|
||||
nones = {
|
||||
_get_params[col].key
|
||||
for col, value in zip(
|
||||
mapper.primary_key, primary_key_identity
|
||||
)
|
||||
if value is None
|
||||
}
|
||||
_lcl_get_clause = sql_util.adapt_criterion_to_null(
|
||||
_lcl_get_clause, nones
|
||||
)
|
||||
|
||||
# TODO: can mapper._get_clause be pre-adapted?
|
||||
q._where_criteria = (
|
||||
sql_util._deep_annotate(_lcl_get_clause, {"_orm_adapt": True}),
|
||||
)
|
||||
|
||||
for fn in self._post_criteria:
|
||||
q = fn(q)
|
||||
return q
|
||||
|
||||
# cache the query against a key that includes
|
||||
# which positions in the primary key are NULL
|
||||
# (remember, we can map to an OUTER JOIN)
|
||||
bq = self.bq
|
||||
|
||||
# add the clause we got from mapper._get_clause to the cache
|
||||
# key so that if a race causes multiple calls to _get_clause,
|
||||
# we've cached on ours
|
||||
bq = bq._clone()
|
||||
bq._cache_key += (_get_clause,)
|
||||
|
||||
bq = bq.with_criteria(
|
||||
setup, tuple(elem is None for elem in primary_key_identity)
|
||||
)
|
||||
|
||||
params = {
|
||||
_get_params[primary_key].key: id_val
|
||||
for id_val, primary_key in zip(
|
||||
primary_key_identity, mapper.primary_key
|
||||
)
|
||||
}
|
||||
|
||||
result = list(bq.for_session(self.session).params(**params))
|
||||
l = len(result)
|
||||
if l > 1:
|
||||
raise orm_exc.MultipleResultsFound()
|
||||
elif l:
|
||||
return result[0]
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
bakery = BakedQuery.bakery
|
||||
Reference in New Issue
Block a user