lambdas.py 44 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317
  1. # sql/lambdas.py
  2. # Copyright (C) 2005-2024 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. import inspect
  8. import itertools
  9. import operator
  10. import sys
  11. import threading
  12. import types
  13. import weakref
  14. from . import coercions
  15. from . import elements
  16. from . import roles
  17. from . import schema
  18. from . import traversals
  19. from . import visitors
  20. from .base import _clone
  21. from .base import Options
  22. from .operators import ColumnOperators
  23. from .. import exc
  24. from .. import inspection
  25. from .. import util
  26. from ..util import collections_abc
  27. from ..util import compat
  28. _closure_per_cache_key = util.LRUCache(1000)
  29. class LambdaOptions(Options):
  30. enable_tracking = True
  31. track_closure_variables = True
  32. track_on = None
  33. global_track_bound_values = True
  34. track_bound_values = True
  35. lambda_cache = None
  36. def lambda_stmt(
  37. lmb,
  38. enable_tracking=True,
  39. track_closure_variables=True,
  40. track_on=None,
  41. global_track_bound_values=True,
  42. track_bound_values=True,
  43. lambda_cache=None,
  44. ):
  45. """Produce a SQL statement that is cached as a lambda.
  46. The Python code object within the lambda is scanned for both Python
  47. literals that will become bound parameters as well as closure variables
  48. that refer to Core or ORM constructs that may vary. The lambda itself
  49. will be invoked only once per particular set of constructs detected.
  50. E.g.::
  51. from sqlalchemy import lambda_stmt
  52. stmt = lambda_stmt(lambda: table.select())
  53. stmt += lambda s: s.where(table.c.id == 5)
  54. result = connection.execute(stmt)
  55. The object returned is an instance of :class:`_sql.StatementLambdaElement`.
  56. .. versionadded:: 1.4
  57. :param lmb: a Python function, typically a lambda, which takes no arguments
  58. and returns a SQL expression construct
  59. :param enable_tracking: when False, all scanning of the given lambda for
  60. changes in closure variables or bound parameters is disabled. Use for
  61. a lambda that produces the identical results in all cases with no
  62. parameterization.
  63. :param track_closure_variables: when False, changes in closure variables
  64. within the lambda will not be scanned. Use for a lambda where the
  65. state of its closure variables will never change the SQL structure
  66. returned by the lambda.
  67. :param track_bound_values: when False, bound parameter tracking will
  68. be disabled for the given lambda. Use for a lambda that either does
  69. not produce any bound values, or where the initial bound values never
  70. change.
  71. :param global_track_bound_values: when False, bound parameter tracking
  72. will be disabled for the entire statement including additional links
  73. added via the :meth:`_sql.StatementLambdaElement.add_criteria` method.
  74. :param lambda_cache: a dictionary or other mapping-like object where
  75. information about the lambda's Python code as well as the tracked closure
  76. variables in the lambda itself will be stored. Defaults
  77. to a global LRU cache. This cache is independent of the "compiled_cache"
  78. used by the :class:`_engine.Connection` object.
  79. .. seealso::
  80. :ref:`engine_lambda_caching`
  81. """
  82. return StatementLambdaElement(
  83. lmb,
  84. roles.StatementRole,
  85. LambdaOptions(
  86. enable_tracking=enable_tracking,
  87. track_on=track_on,
  88. track_closure_variables=track_closure_variables,
  89. global_track_bound_values=global_track_bound_values,
  90. track_bound_values=track_bound_values,
  91. lambda_cache=lambda_cache,
  92. ),
  93. )
  94. class LambdaElement(elements.ClauseElement):
  95. """A SQL construct where the state is stored as an un-invoked lambda.
  96. The :class:`_sql.LambdaElement` is produced transparently whenever
  97. passing lambda expressions into SQL constructs, such as::
  98. stmt = select(table).where(lambda: table.c.col == parameter)
  99. The :class:`_sql.LambdaElement` is the base of the
  100. :class:`_sql.StatementLambdaElement` which represents a full statement
  101. within a lambda.
  102. .. versionadded:: 1.4
  103. .. seealso::
  104. :ref:`engine_lambda_caching`
  105. """
  106. __visit_name__ = "lambda_element"
  107. _is_lambda_element = True
  108. _traverse_internals = [
  109. ("_resolved", visitors.InternalTraversal.dp_clauseelement)
  110. ]
  111. _transforms = ()
  112. parent_lambda = None
  113. def __repr__(self):
  114. return "%s(%r)" % (self.__class__.__name__, self.fn.__code__)
  115. def __init__(
  116. self, fn, role, opts=LambdaOptions, apply_propagate_attrs=None
  117. ):
  118. self.fn = fn
  119. self.role = role
  120. self.tracker_key = (fn.__code__,)
  121. self.opts = opts
  122. if apply_propagate_attrs is None and (role is roles.StatementRole):
  123. apply_propagate_attrs = self
  124. rec = self._retrieve_tracker_rec(fn, apply_propagate_attrs, opts)
  125. if apply_propagate_attrs is not None:
  126. propagate_attrs = rec.propagate_attrs
  127. if propagate_attrs:
  128. apply_propagate_attrs._propagate_attrs = propagate_attrs
  129. def _retrieve_tracker_rec(self, fn, apply_propagate_attrs, opts):
  130. lambda_cache = opts.lambda_cache
  131. if lambda_cache is None:
  132. lambda_cache = _closure_per_cache_key
  133. tracker_key = self.tracker_key
  134. fn = self.fn
  135. closure = fn.__closure__
  136. tracker = AnalyzedCode.get(
  137. fn,
  138. self,
  139. opts,
  140. )
  141. self._resolved_bindparams = bindparams = []
  142. if self.parent_lambda is not None:
  143. parent_closure_cache_key = self.parent_lambda.closure_cache_key
  144. else:
  145. parent_closure_cache_key = ()
  146. if parent_closure_cache_key is not traversals.NO_CACHE:
  147. anon_map = traversals.anon_map()
  148. cache_key = tuple(
  149. [
  150. getter(closure, opts, anon_map, bindparams)
  151. for getter in tracker.closure_trackers
  152. ]
  153. )
  154. if traversals.NO_CACHE not in anon_map:
  155. cache_key = parent_closure_cache_key + cache_key
  156. self.closure_cache_key = cache_key
  157. try:
  158. rec = lambda_cache[tracker_key + cache_key]
  159. except KeyError:
  160. rec = None
  161. else:
  162. cache_key = traversals.NO_CACHE
  163. rec = None
  164. else:
  165. cache_key = traversals.NO_CACHE
  166. rec = None
  167. self.closure_cache_key = cache_key
  168. if rec is None:
  169. if cache_key is not traversals.NO_CACHE:
  170. with AnalyzedCode._generation_mutex:
  171. key = tracker_key + cache_key
  172. if key not in lambda_cache:
  173. rec = AnalyzedFunction(
  174. tracker, self, apply_propagate_attrs, fn
  175. )
  176. rec.closure_bindparams = bindparams
  177. lambda_cache[key] = rec
  178. else:
  179. rec = lambda_cache[key]
  180. else:
  181. rec = NonAnalyzedFunction(self._invoke_user_fn(fn))
  182. else:
  183. bindparams[:] = [
  184. orig_bind._with_value(new_bind.value, maintain_key=True)
  185. for orig_bind, new_bind in zip(
  186. rec.closure_bindparams, bindparams
  187. )
  188. ]
  189. self._rec = rec
  190. if cache_key is not traversals.NO_CACHE:
  191. if self.parent_lambda is not None:
  192. bindparams[:0] = self.parent_lambda._resolved_bindparams
  193. lambda_element = self
  194. while lambda_element is not None:
  195. rec = lambda_element._rec
  196. if rec.bindparam_trackers:
  197. tracker_instrumented_fn = rec.tracker_instrumented_fn
  198. for tracker in rec.bindparam_trackers:
  199. tracker(
  200. lambda_element.fn,
  201. tracker_instrumented_fn,
  202. bindparams,
  203. )
  204. lambda_element = lambda_element.parent_lambda
  205. return rec
  206. def __getattr__(self, key):
  207. return getattr(self._rec.expected_expr, key)
  208. @property
  209. def _is_sequence(self):
  210. return self._rec.is_sequence
  211. @property
  212. def _select_iterable(self):
  213. if self._is_sequence:
  214. return itertools.chain.from_iterable(
  215. [element._select_iterable for element in self._resolved]
  216. )
  217. else:
  218. return self._resolved._select_iterable
  219. @property
  220. def _from_objects(self):
  221. if self._is_sequence:
  222. return itertools.chain.from_iterable(
  223. [element._from_objects for element in self._resolved]
  224. )
  225. else:
  226. return self._resolved._from_objects
  227. def _param_dict(self):
  228. return {b.key: b.value for b in self._resolved_bindparams}
  229. def _setup_binds_for_tracked_expr(self, expr):
  230. bindparam_lookup = {b.key: b for b in self._resolved_bindparams}
  231. def replace(thing):
  232. if isinstance(thing, elements.BindParameter):
  233. if thing.key in bindparam_lookup:
  234. bind = bindparam_lookup[thing.key]
  235. if thing.expanding:
  236. bind.expanding = True
  237. bind.expand_op = thing.expand_op
  238. bind.type = thing.type
  239. return bind
  240. if self._rec.is_sequence:
  241. expr = [
  242. visitors.replacement_traverse(sub_expr, {}, replace)
  243. for sub_expr in expr
  244. ]
  245. elif getattr(expr, "is_clause_element", False):
  246. expr = visitors.replacement_traverse(expr, {}, replace)
  247. return expr
  248. def _copy_internals(
  249. self, clone=_clone, deferred_copy_internals=None, **kw
  250. ):
  251. # TODO: this needs A LOT of tests
  252. self._resolved = clone(
  253. self._resolved,
  254. deferred_copy_internals=deferred_copy_internals,
  255. **kw
  256. )
  257. @util.memoized_property
  258. def _resolved(self):
  259. expr = self._rec.expected_expr
  260. if self._resolved_bindparams:
  261. expr = self._setup_binds_for_tracked_expr(expr)
  262. return expr
  263. def _gen_cache_key(self, anon_map, bindparams):
  264. if self.closure_cache_key is traversals.NO_CACHE:
  265. anon_map[traversals.NO_CACHE] = True
  266. return None
  267. cache_key = (
  268. self.fn.__code__,
  269. self.__class__,
  270. ) + self.closure_cache_key
  271. parent = self.parent_lambda
  272. while parent is not None:
  273. cache_key = (
  274. (parent.fn.__code__,) + parent.closure_cache_key + cache_key
  275. )
  276. parent = parent.parent_lambda
  277. if self._resolved_bindparams:
  278. bindparams.extend(self._resolved_bindparams)
  279. return cache_key
  280. def _invoke_user_fn(self, fn, *arg):
  281. return fn()
  282. class DeferredLambdaElement(LambdaElement):
  283. """A LambdaElement where the lambda accepts arguments and is
  284. invoked within the compile phase with special context.
  285. This lambda doesn't normally produce its real SQL expression outside of the
  286. compile phase. It is passed a fixed set of initial arguments
  287. so that it can generate a sample expression.
  288. """
  289. def __init__(self, fn, role, opts=LambdaOptions, lambda_args=()):
  290. self.lambda_args = lambda_args
  291. super(DeferredLambdaElement, self).__init__(fn, role, opts)
  292. def _invoke_user_fn(self, fn, *arg):
  293. return fn(*self.lambda_args)
  294. def _resolve_with_args(self, *lambda_args):
  295. tracker_fn = self._rec.tracker_instrumented_fn
  296. expr = tracker_fn(*lambda_args)
  297. expr = coercions.expect(self.role, expr)
  298. expr = self._setup_binds_for_tracked_expr(expr)
  299. # this validation is getting very close, but not quite, to achieving
  300. # #5767. The problem is if the base lambda uses an unnamed column
  301. # as is very common with mixins, the parameter name is different
  302. # and it produces a false positive; that is, for the documented case
  303. # that is exactly what people will be doing, it doesn't work, so
  304. # I'm not really sure how to handle this right now.
  305. # expected_binds = [
  306. # b._orig_key
  307. # for b in self._rec.expr._generate_cache_key()[1]
  308. # if b.required
  309. # ]
  310. # got_binds = [
  311. # b._orig_key for b in expr._generate_cache_key()[1] if b.required
  312. # ]
  313. # if expected_binds != got_binds:
  314. # raise exc.InvalidRequestError(
  315. # "Lambda callable at %s produced a different set of bound "
  316. # "parameters than its original run: %s"
  317. # % (self.fn.__code__, ", ".join(got_binds))
  318. # )
  319. # TODO: TEST TEST TEST, this is very out there
  320. for deferred_copy_internals in self._transforms:
  321. expr = deferred_copy_internals(expr)
  322. return expr
  323. def _copy_internals(
  324. self, clone=_clone, deferred_copy_internals=None, **kw
  325. ):
  326. super(DeferredLambdaElement, self)._copy_internals(
  327. clone=clone,
  328. deferred_copy_internals=deferred_copy_internals, # **kw
  329. opts=kw,
  330. )
  331. # TODO: A LOT A LOT of tests. for _resolve_with_args, we don't know
  332. # our expression yet. so hold onto the replacement
  333. if deferred_copy_internals:
  334. self._transforms += (deferred_copy_internals,)
  335. class StatementLambdaElement(roles.AllowsLambdaRole, LambdaElement):
  336. """Represent a composable SQL statement as a :class:`_sql.LambdaElement`.
  337. The :class:`_sql.StatementLambdaElement` is constructed using the
  338. :func:`_sql.lambda_stmt` function::
  339. from sqlalchemy import lambda_stmt
  340. stmt = lambda_stmt(lambda: select(table))
  341. Once constructed, additional criteria can be built onto the statement
  342. by adding subsequent lambdas, which accept the existing statement
  343. object as a single parameter::
  344. stmt += lambda s: s.where(table.c.col == parameter)
  345. .. versionadded:: 1.4
  346. .. seealso::
  347. :ref:`engine_lambda_caching`
  348. """
  349. def __add__(self, other):
  350. return self.add_criteria(other)
  351. def add_criteria(
  352. self,
  353. other,
  354. enable_tracking=True,
  355. track_on=None,
  356. track_closure_variables=True,
  357. track_bound_values=True,
  358. ):
  359. """Add new criteria to this :class:`_sql.StatementLambdaElement`.
  360. E.g.::
  361. >>> def my_stmt(parameter):
  362. ... stmt = lambda_stmt(
  363. ... lambda: select(table.c.x, table.c.y),
  364. ... )
  365. ... stmt = stmt.add_criteria(
  366. ... lambda: table.c.x > parameter
  367. ... )
  368. ... return stmt
  369. The :meth:`_sql.StatementLambdaElement.add_criteria` method is
  370. equivalent to using the Python addition operator to add a new
  371. lambda, except that additional arguments may be added including
  372. ``track_closure_values`` and ``track_on``::
  373. >>> def my_stmt(self, foo):
  374. ... stmt = lambda_stmt(
  375. ... lambda: select(func.max(foo.x, foo.y)),
  376. ... track_closure_variables=False
  377. ... )
  378. ... stmt = stmt.add_criteria(
  379. ... lambda: self.where_criteria,
  380. ... track_on=[self]
  381. ... )
  382. ... return stmt
  383. See :func:`_sql.lambda_stmt` for a description of the parameters
  384. accepted.
  385. """
  386. opts = self.opts + dict(
  387. enable_tracking=enable_tracking,
  388. track_closure_variables=track_closure_variables,
  389. global_track_bound_values=self.opts.global_track_bound_values,
  390. track_on=track_on,
  391. track_bound_values=track_bound_values,
  392. )
  393. return LinkedLambdaElement(other, parent_lambda=self, opts=opts)
  394. def _execute_on_connection(
  395. self, connection, multiparams, params, execution_options
  396. ):
  397. if self._rec.expected_expr.supports_execution:
  398. return connection._execute_clauseelement(
  399. self, multiparams, params, execution_options
  400. )
  401. else:
  402. raise exc.ObjectNotExecutableError(self)
  403. @property
  404. def _with_options(self):
  405. return self._rec.expected_expr._with_options
  406. @property
  407. def _effective_plugin_target(self):
  408. return self._rec.expected_expr._effective_plugin_target
  409. @property
  410. def _execution_options(self):
  411. return self._rec.expected_expr._execution_options
  412. def spoil(self):
  413. """Return a new :class:`.StatementLambdaElement` that will run
  414. all lambdas unconditionally each time.
  415. """
  416. return NullLambdaStatement(self.fn())
  417. class NullLambdaStatement(roles.AllowsLambdaRole, elements.ClauseElement):
  418. """Provides the :class:`.StatementLambdaElement` API but does not
  419. cache or analyze lambdas.
  420. the lambdas are instead invoked immediately.
  421. The intended use is to isolate issues that may arise when using
  422. lambda statements.
  423. """
  424. __visit_name__ = "lambda_element"
  425. _is_lambda_element = True
  426. _traverse_internals = [
  427. ("_resolved", visitors.InternalTraversal.dp_clauseelement)
  428. ]
  429. def __init__(self, statement):
  430. self._resolved = statement
  431. self._propagate_attrs = statement._propagate_attrs
  432. def __getattr__(self, key):
  433. return getattr(self._resolved, key)
  434. def __add__(self, other):
  435. statement = other(self._resolved)
  436. return NullLambdaStatement(statement)
  437. def add_criteria(self, other, **kw):
  438. statement = other(self._resolved)
  439. return NullLambdaStatement(statement)
  440. def _execute_on_connection(
  441. self, connection, multiparams, params, execution_options
  442. ):
  443. if self._resolved.supports_execution:
  444. return connection._execute_clauseelement(
  445. self, multiparams, params, execution_options
  446. )
  447. else:
  448. raise exc.ObjectNotExecutableError(self)
  449. class LinkedLambdaElement(StatementLambdaElement):
  450. """Represent subsequent links of a :class:`.StatementLambdaElement`."""
  451. role = None
  452. def __init__(self, fn, parent_lambda, opts):
  453. self.opts = opts
  454. self.fn = fn
  455. self.parent_lambda = parent_lambda
  456. self.tracker_key = parent_lambda.tracker_key + (fn.__code__,)
  457. self._retrieve_tracker_rec(fn, self, opts)
  458. self._propagate_attrs = parent_lambda._propagate_attrs
  459. def _invoke_user_fn(self, fn, *arg):
  460. return fn(self.parent_lambda._resolved)
  461. class AnalyzedCode(object):
  462. __slots__ = (
  463. "track_closure_variables",
  464. "track_bound_values",
  465. "bindparam_trackers",
  466. "closure_trackers",
  467. "build_py_wrappers",
  468. )
  469. _fns = weakref.WeakKeyDictionary()
  470. _generation_mutex = threading.RLock()
  471. @classmethod
  472. def get(cls, fn, lambda_element, lambda_kw, **kw):
  473. try:
  474. # TODO: validate kw haven't changed?
  475. return cls._fns[fn.__code__]
  476. except KeyError:
  477. pass
  478. with cls._generation_mutex:
  479. # check for other thread already created object
  480. if fn.__code__ in cls._fns:
  481. return cls._fns[fn.__code__]
  482. cls._fns[fn.__code__] = analyzed = AnalyzedCode(
  483. fn, lambda_element, lambda_kw, **kw
  484. )
  485. return analyzed
  486. def __init__(self, fn, lambda_element, opts):
  487. if inspect.ismethod(fn):
  488. raise exc.ArgumentError(
  489. "Method %s may not be passed as a SQL expression" % fn
  490. )
  491. closure = fn.__closure__
  492. self.track_bound_values = (
  493. opts.track_bound_values and opts.global_track_bound_values
  494. )
  495. enable_tracking = opts.enable_tracking
  496. track_on = opts.track_on
  497. track_closure_variables = opts.track_closure_variables
  498. self.track_closure_variables = track_closure_variables and not track_on
  499. # a list of callables generated from _bound_parameter_getter_*
  500. # functions. Each of these uses a PyWrapper object to retrieve
  501. # a parameter value
  502. self.bindparam_trackers = []
  503. # a list of callables generated from _cache_key_getter_* functions
  504. # these callables work to generate a cache key for the lambda
  505. # based on what's inside its closure variables.
  506. self.closure_trackers = []
  507. self.build_py_wrappers = []
  508. if enable_tracking:
  509. if track_on:
  510. self._init_track_on(track_on)
  511. self._init_globals(fn)
  512. if closure:
  513. self._init_closure(fn)
  514. self._setup_additional_closure_trackers(fn, lambda_element, opts)
  515. def _init_track_on(self, track_on):
  516. self.closure_trackers.extend(
  517. self._cache_key_getter_track_on(idx, elem)
  518. for idx, elem in enumerate(track_on)
  519. )
  520. def _init_globals(self, fn):
  521. build_py_wrappers = self.build_py_wrappers
  522. bindparam_trackers = self.bindparam_trackers
  523. track_bound_values = self.track_bound_values
  524. for name in fn.__code__.co_names:
  525. if name not in fn.__globals__:
  526. continue
  527. _bound_value = self._roll_down_to_literal(fn.__globals__[name])
  528. if coercions._deep_is_literal(_bound_value):
  529. build_py_wrappers.append((name, None))
  530. if track_bound_values:
  531. bindparam_trackers.append(
  532. self._bound_parameter_getter_func_globals(name)
  533. )
  534. def _init_closure(self, fn):
  535. build_py_wrappers = self.build_py_wrappers
  536. closure = fn.__closure__
  537. track_bound_values = self.track_bound_values
  538. track_closure_variables = self.track_closure_variables
  539. bindparam_trackers = self.bindparam_trackers
  540. closure_trackers = self.closure_trackers
  541. for closure_index, (fv, cell) in enumerate(
  542. zip(fn.__code__.co_freevars, closure)
  543. ):
  544. _bound_value = self._roll_down_to_literal(cell.cell_contents)
  545. if coercions._deep_is_literal(_bound_value):
  546. build_py_wrappers.append((fv, closure_index))
  547. if track_bound_values:
  548. bindparam_trackers.append(
  549. self._bound_parameter_getter_func_closure(
  550. fv, closure_index
  551. )
  552. )
  553. else:
  554. # for normal cell contents, add them to a list that
  555. # we can compare later when we get new lambdas. if
  556. # any identities have changed, then we will
  557. # recalculate the whole lambda and run it again.
  558. if track_closure_variables:
  559. closure_trackers.append(
  560. self._cache_key_getter_closure_variable(
  561. fn, fv, closure_index, cell.cell_contents
  562. )
  563. )
  564. def _setup_additional_closure_trackers(self, fn, lambda_element, opts):
  565. # an additional step is to actually run the function, then
  566. # go through the PyWrapper objects that were set up to catch a bound
  567. # parameter. then if they *didn't* make a param, oh they're another
  568. # object in the closure we have to track for our cache key. so
  569. # create trackers to catch those.
  570. analyzed_function = AnalyzedFunction(
  571. self,
  572. lambda_element,
  573. None,
  574. fn,
  575. )
  576. closure_trackers = self.closure_trackers
  577. for pywrapper in analyzed_function.closure_pywrappers:
  578. if not pywrapper._sa__has_param:
  579. closure_trackers.append(
  580. self._cache_key_getter_tracked_literal(fn, pywrapper)
  581. )
  582. @classmethod
  583. def _roll_down_to_literal(cls, element):
  584. is_clause_element = hasattr(element, "__clause_element__")
  585. if is_clause_element:
  586. while not isinstance(
  587. element, (elements.ClauseElement, schema.SchemaItem, type)
  588. ):
  589. try:
  590. element = element.__clause_element__()
  591. except AttributeError:
  592. break
  593. if not is_clause_element:
  594. insp = inspection.inspect(element, raiseerr=False)
  595. if insp is not None:
  596. try:
  597. return insp.__clause_element__()
  598. except AttributeError:
  599. return insp
  600. # TODO: should we coerce consts None/True/False here?
  601. return element
  602. else:
  603. return element
  604. def _bound_parameter_getter_func_globals(self, name):
  605. """Return a getter that will extend a list of bound parameters
  606. with new entries from the ``__globals__`` collection of a particular
  607. lambda.
  608. """
  609. def extract_parameter_value(
  610. current_fn, tracker_instrumented_fn, result
  611. ):
  612. wrapper = tracker_instrumented_fn.__globals__[name]
  613. object.__getattribute__(wrapper, "_extract_bound_parameters")(
  614. current_fn.__globals__[name], result
  615. )
  616. return extract_parameter_value
  617. def _bound_parameter_getter_func_closure(self, name, closure_index):
  618. """Return a getter that will extend a list of bound parameters
  619. with new entries from the ``__closure__`` collection of a particular
  620. lambda.
  621. """
  622. def extract_parameter_value(
  623. current_fn, tracker_instrumented_fn, result
  624. ):
  625. wrapper = tracker_instrumented_fn.__closure__[
  626. closure_index
  627. ].cell_contents
  628. object.__getattribute__(wrapper, "_extract_bound_parameters")(
  629. current_fn.__closure__[closure_index].cell_contents, result
  630. )
  631. return extract_parameter_value
  632. def _cache_key_getter_track_on(self, idx, elem):
  633. """Return a getter that will extend a cache key with new entries
  634. from the "track_on" parameter passed to a :class:`.LambdaElement`.
  635. """
  636. if isinstance(elem, tuple):
  637. # tuple must contain hascachekey elements
  638. def get(closure, opts, anon_map, bindparams):
  639. return tuple(
  640. tup_elem._gen_cache_key(anon_map, bindparams)
  641. for tup_elem in opts.track_on[idx]
  642. )
  643. elif isinstance(elem, traversals.HasCacheKey):
  644. def get(closure, opts, anon_map, bindparams):
  645. return opts.track_on[idx]._gen_cache_key(anon_map, bindparams)
  646. else:
  647. def get(closure, opts, anon_map, bindparams):
  648. return opts.track_on[idx]
  649. return get
  650. def _cache_key_getter_closure_variable(
  651. self,
  652. fn,
  653. variable_name,
  654. idx,
  655. cell_contents,
  656. use_clause_element=False,
  657. use_inspect=False,
  658. ):
  659. """Return a getter that will extend a cache key with new entries
  660. from the ``__closure__`` collection of a particular lambda.
  661. """
  662. if isinstance(cell_contents, traversals.HasCacheKey):
  663. def get(closure, opts, anon_map, bindparams):
  664. obj = closure[idx].cell_contents
  665. if use_inspect:
  666. obj = inspection.inspect(obj)
  667. elif use_clause_element:
  668. while hasattr(obj, "__clause_element__"):
  669. if not getattr(obj, "is_clause_element", False):
  670. obj = obj.__clause_element__()
  671. return obj._gen_cache_key(anon_map, bindparams)
  672. elif isinstance(cell_contents, types.FunctionType):
  673. def get(closure, opts, anon_map, bindparams):
  674. return closure[idx].cell_contents.__code__
  675. elif isinstance(cell_contents, collections_abc.Sequence):
  676. def get(closure, opts, anon_map, bindparams):
  677. contents = closure[idx].cell_contents
  678. try:
  679. return tuple(
  680. elem._gen_cache_key(anon_map, bindparams)
  681. for elem in contents
  682. )
  683. except AttributeError as ae:
  684. self._raise_for_uncacheable_closure_variable(
  685. variable_name, fn, from_=ae
  686. )
  687. else:
  688. # if the object is a mapped class or aliased class, or some
  689. # other object in the ORM realm of things like that, imitate
  690. # the logic used in coercions.expect() to roll it down to the
  691. # SQL element
  692. element = cell_contents
  693. is_clause_element = False
  694. while hasattr(element, "__clause_element__"):
  695. is_clause_element = True
  696. if not getattr(element, "is_clause_element", False):
  697. element = element.__clause_element__()
  698. else:
  699. break
  700. if not is_clause_element:
  701. insp = inspection.inspect(element, raiseerr=False)
  702. if insp is not None:
  703. return self._cache_key_getter_closure_variable(
  704. fn, variable_name, idx, insp, use_inspect=True
  705. )
  706. else:
  707. return self._cache_key_getter_closure_variable(
  708. fn, variable_name, idx, element, use_clause_element=True
  709. )
  710. self._raise_for_uncacheable_closure_variable(variable_name, fn)
  711. return get
  712. def _raise_for_uncacheable_closure_variable(
  713. self, variable_name, fn, from_=None
  714. ):
  715. util.raise_(
  716. exc.InvalidRequestError(
  717. "Closure variable named '%s' inside of lambda callable %s "
  718. "does not refer to a cacheable SQL element, and also does not "
  719. "appear to be serving as a SQL literal bound value based on "
  720. "the default "
  721. "SQL expression returned by the function. This variable "
  722. "needs to remain outside the scope of a SQL-generating lambda "
  723. "so that a proper cache key may be generated from the "
  724. "lambda's state. Evaluate this variable outside of the "
  725. "lambda, set track_on=[<elements>] to explicitly select "
  726. "closure elements to track, or set "
  727. "track_closure_variables=False to exclude "
  728. "closure variables from being part of the cache key."
  729. % (variable_name, fn.__code__),
  730. ),
  731. from_=from_,
  732. )
  733. def _cache_key_getter_tracked_literal(self, fn, pytracker):
  734. """Return a getter that will extend a cache key with new entries
  735. from the ``__closure__`` collection of a particular lambda.
  736. this getter differs from _cache_key_getter_closure_variable
  737. in that these are detected after the function is run, and PyWrapper
  738. objects have recorded that a particular literal value is in fact
  739. not being interpreted as a bound parameter.
  740. """
  741. elem = pytracker._sa__to_evaluate
  742. closure_index = pytracker._sa__closure_index
  743. variable_name = pytracker._sa__name
  744. return self._cache_key_getter_closure_variable(
  745. fn, variable_name, closure_index, elem
  746. )
  747. class NonAnalyzedFunction(object):
  748. __slots__ = ("expr",)
  749. closure_bindparams = None
  750. bindparam_trackers = None
  751. def __init__(self, expr):
  752. self.expr = expr
  753. @property
  754. def expected_expr(self):
  755. return self.expr
  756. class AnalyzedFunction(object):
  757. __slots__ = (
  758. "analyzed_code",
  759. "fn",
  760. "closure_pywrappers",
  761. "tracker_instrumented_fn",
  762. "expr",
  763. "bindparam_trackers",
  764. "expected_expr",
  765. "is_sequence",
  766. "propagate_attrs",
  767. "closure_bindparams",
  768. )
  769. def __init__(
  770. self,
  771. analyzed_code,
  772. lambda_element,
  773. apply_propagate_attrs,
  774. fn,
  775. ):
  776. self.analyzed_code = analyzed_code
  777. self.fn = fn
  778. self.bindparam_trackers = analyzed_code.bindparam_trackers
  779. self._instrument_and_run_function(lambda_element)
  780. self._coerce_expression(lambda_element, apply_propagate_attrs)
  781. def _instrument_and_run_function(self, lambda_element):
  782. analyzed_code = self.analyzed_code
  783. fn = self.fn
  784. self.closure_pywrappers = closure_pywrappers = []
  785. build_py_wrappers = analyzed_code.build_py_wrappers
  786. if not build_py_wrappers:
  787. self.tracker_instrumented_fn = tracker_instrumented_fn = fn
  788. self.expr = lambda_element._invoke_user_fn(tracker_instrumented_fn)
  789. else:
  790. track_closure_variables = analyzed_code.track_closure_variables
  791. closure = fn.__closure__
  792. # will form the __closure__ of the function when we rebuild it
  793. if closure:
  794. new_closure = {
  795. fv: cell.cell_contents
  796. for fv, cell in zip(fn.__code__.co_freevars, closure)
  797. }
  798. else:
  799. new_closure = {}
  800. # will form the __globals__ of the function when we rebuild it
  801. new_globals = fn.__globals__.copy()
  802. for name, closure_index in build_py_wrappers:
  803. if closure_index is not None:
  804. value = closure[closure_index].cell_contents
  805. new_closure[name] = bind = PyWrapper(
  806. fn,
  807. name,
  808. value,
  809. closure_index=closure_index,
  810. track_bound_values=(
  811. self.analyzed_code.track_bound_values
  812. ),
  813. )
  814. if track_closure_variables:
  815. closure_pywrappers.append(bind)
  816. else:
  817. value = fn.__globals__[name]
  818. new_globals[name] = bind = PyWrapper(fn, name, value)
  819. # rewrite the original fn. things that look like they will
  820. # become bound parameters are wrapped in a PyWrapper.
  821. self.tracker_instrumented_fn = (
  822. tracker_instrumented_fn
  823. ) = self._rewrite_code_obj(
  824. fn,
  825. [new_closure[name] for name in fn.__code__.co_freevars],
  826. new_globals,
  827. )
  828. # now invoke the function. This will give us a new SQL
  829. # expression, but all the places that there would be a bound
  830. # parameter, the PyWrapper in its place will give us a bind
  831. # with a predictable name we can match up later.
  832. # additionally, each PyWrapper will log that it did in fact
  833. # create a parameter, otherwise, it's some kind of Python
  834. # object in the closure and we want to track that, to make
  835. # sure it doesn't change to something else, or if it does,
  836. # that we create a different tracked function with that
  837. # variable.
  838. self.expr = lambda_element._invoke_user_fn(tracker_instrumented_fn)
  839. def _coerce_expression(self, lambda_element, apply_propagate_attrs):
  840. """Run the tracker-generated expression through coercion rules.
  841. After the user-defined lambda has been invoked to produce a statement
  842. for re-use, run it through coercion rules to both check that it's the
  843. correct type of object and also to coerce it to its useful form.
  844. """
  845. parent_lambda = lambda_element.parent_lambda
  846. expr = self.expr
  847. if parent_lambda is None:
  848. if isinstance(expr, collections_abc.Sequence):
  849. self.expected_expr = [
  850. coercions.expect(
  851. lambda_element.role,
  852. sub_expr,
  853. apply_propagate_attrs=apply_propagate_attrs,
  854. )
  855. for sub_expr in expr
  856. ]
  857. self.is_sequence = True
  858. else:
  859. self.expected_expr = coercions.expect(
  860. lambda_element.role,
  861. expr,
  862. apply_propagate_attrs=apply_propagate_attrs,
  863. )
  864. self.is_sequence = False
  865. else:
  866. self.expected_expr = expr
  867. self.is_sequence = False
  868. if apply_propagate_attrs is not None:
  869. self.propagate_attrs = apply_propagate_attrs._propagate_attrs
  870. else:
  871. self.propagate_attrs = util.EMPTY_DICT
  872. def _rewrite_code_obj(self, f, cell_values, globals_):
  873. """Return a copy of f, with a new closure and new globals
  874. yes it works in pypy :P
  875. """
  876. argrange = range(len(cell_values))
  877. code = "def make_cells():\n"
  878. if cell_values:
  879. code += " (%s) = (%s)\n" % (
  880. ", ".join("i%d" % i for i in argrange),
  881. ", ".join("o%d" % i for i in argrange),
  882. )
  883. code += " def closure():\n"
  884. code += " return %s\n" % ", ".join("i%d" % i for i in argrange)
  885. code += " return closure.__closure__"
  886. vars_ = {"o%d" % i: cell_values[i] for i in argrange}
  887. compat.exec_(code, vars_, vars_)
  888. closure = vars_["make_cells"]()
  889. func = type(f)(
  890. f.__code__, globals_, f.__name__, f.__defaults__, closure
  891. )
  892. if sys.version_info >= (3,):
  893. func.__annotations__ = f.__annotations__
  894. func.__kwdefaults__ = f.__kwdefaults__
  895. func.__doc__ = f.__doc__
  896. func.__module__ = f.__module__
  897. return func
  898. class PyWrapper(ColumnOperators):
  899. """A wrapper object that is injected into the ``__globals__`` and
  900. ``__closure__`` of a Python function.
  901. When the function is instrumented with :class:`.PyWrapper` objects, it is
  902. then invoked just once in order to set up the wrappers. We look through
  903. all the :class:`.PyWrapper` objects we made to find the ones that generated
  904. a :class:`.BindParameter` object, e.g. the expression system interpreted
  905. something as a literal. Those positions in the globals/closure are then
  906. ones that we will look at, each time a new lambda comes in that refers to
  907. the same ``__code__`` object. In this way, we keep a single version of
  908. the SQL expression that this lambda produced, without calling upon the
  909. Python function that created it more than once, unless its other closure
  910. variables have changed. The expression is then transformed to have the
  911. new bound values embedded into it.
  912. """
  913. def __init__(
  914. self,
  915. fn,
  916. name,
  917. to_evaluate,
  918. closure_index=None,
  919. getter=None,
  920. track_bound_values=True,
  921. ):
  922. self.fn = fn
  923. self._name = name
  924. self._to_evaluate = to_evaluate
  925. self._param = None
  926. self._has_param = False
  927. self._bind_paths = {}
  928. self._getter = getter
  929. self._closure_index = closure_index
  930. self.track_bound_values = track_bound_values
  931. def __call__(self, *arg, **kw):
  932. elem = object.__getattribute__(self, "_to_evaluate")
  933. value = elem(*arg, **kw)
  934. if (
  935. self._sa_track_bound_values
  936. and coercions._deep_is_literal(value)
  937. and not isinstance(
  938. # TODO: coverage where an ORM option or similar is here
  939. value,
  940. traversals.HasCacheKey,
  941. )
  942. ):
  943. name = object.__getattribute__(self, "_name")
  944. raise exc.InvalidRequestError(
  945. "Can't invoke Python callable %s() inside of lambda "
  946. "expression argument at %s; lambda SQL constructs should "
  947. "not invoke functions from closure variables to produce "
  948. "literal values since the "
  949. "lambda SQL system normally extracts bound values without "
  950. "actually "
  951. "invoking the lambda or any functions within it. Call the "
  952. "function outside of the "
  953. "lambda and assign to a local variable that is used in the "
  954. "lambda as a closure variable, or set "
  955. "track_bound_values=False if the return value of this "
  956. "function is used in some other way other than a SQL bound "
  957. "value." % (name, self._sa_fn.__code__)
  958. )
  959. else:
  960. return value
  961. def operate(self, op, *other, **kwargs):
  962. elem = object.__getattribute__(self, "_py_wrapper_literal")()
  963. return op(elem, *other, **kwargs)
  964. def reverse_operate(self, op, other, **kwargs):
  965. elem = object.__getattribute__(self, "_py_wrapper_literal")()
  966. return op(other, elem, **kwargs)
  967. def _extract_bound_parameters(self, starting_point, result_list):
  968. param = object.__getattribute__(self, "_param")
  969. if param is not None:
  970. param = param._with_value(starting_point, maintain_key=True)
  971. result_list.append(param)
  972. for pywrapper in object.__getattribute__(self, "_bind_paths").values():
  973. getter = object.__getattribute__(pywrapper, "_getter")
  974. element = getter(starting_point)
  975. pywrapper._sa__extract_bound_parameters(element, result_list)
  976. def _py_wrapper_literal(self, expr=None, operator=None, **kw):
  977. param = object.__getattribute__(self, "_param")
  978. to_evaluate = object.__getattribute__(self, "_to_evaluate")
  979. if param is None:
  980. name = object.__getattribute__(self, "_name")
  981. self._param = param = elements.BindParameter(
  982. name,
  983. required=False,
  984. unique=True,
  985. _compared_to_operator=operator,
  986. _compared_to_type=expr.type if expr is not None else None,
  987. )
  988. self._has_param = True
  989. return param._with_value(to_evaluate, maintain_key=True)
  990. def __bool__(self):
  991. to_evaluate = object.__getattribute__(self, "_to_evaluate")
  992. return bool(to_evaluate)
  993. def __nonzero__(self):
  994. to_evaluate = object.__getattribute__(self, "_to_evaluate")
  995. return bool(to_evaluate)
  996. def __getattribute__(self, key):
  997. if key.startswith("_sa_"):
  998. return object.__getattribute__(self, key[4:])
  999. elif key in (
  1000. "__clause_element__",
  1001. "operate",
  1002. "reverse_operate",
  1003. "_py_wrapper_literal",
  1004. "__class__",
  1005. "__dict__",
  1006. ):
  1007. return object.__getattribute__(self, key)
  1008. if key.startswith("__"):
  1009. elem = object.__getattribute__(self, "_to_evaluate")
  1010. return getattr(elem, key)
  1011. else:
  1012. return self._sa__add_getter(key, operator.attrgetter)
  1013. def __iter__(self):
  1014. elem = object.__getattribute__(self, "_to_evaluate")
  1015. return iter(elem)
  1016. def __getitem__(self, key):
  1017. elem = object.__getattribute__(self, "_to_evaluate")
  1018. if not hasattr(elem, "__getitem__"):
  1019. raise AttributeError("__getitem__")
  1020. if isinstance(key, PyWrapper):
  1021. # TODO: coverage
  1022. raise exc.InvalidRequestError(
  1023. "Dictionary keys / list indexes inside of a cached "
  1024. "lambda must be Python literals only"
  1025. )
  1026. return self._sa__add_getter(key, operator.itemgetter)
  1027. def _add_getter(self, key, getter_fn):
  1028. bind_paths = object.__getattribute__(self, "_bind_paths")
  1029. bind_path_key = (key, getter_fn)
  1030. if bind_path_key in bind_paths:
  1031. return bind_paths[bind_path_key]
  1032. getter = getter_fn(key)
  1033. elem = object.__getattribute__(self, "_to_evaluate")
  1034. value = getter(elem)
  1035. rolled_down_value = AnalyzedCode._roll_down_to_literal(value)
  1036. if coercions._deep_is_literal(rolled_down_value):
  1037. wrapper = PyWrapper(self._sa_fn, key, value, getter=getter)
  1038. bind_paths[bind_path_key] = wrapper
  1039. return wrapper
  1040. else:
  1041. return value
  1042. @inspection._inspects(LambdaElement)
  1043. def insp(lmb):
  1044. return inspection.inspect(lmb._resolved)