util.py 8.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253
  1. # engine/util.py
  2. # Copyright (C) 2005-2022 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. from .. import exc
  8. from .. import util
  9. from ..util import collections_abc
  10. from ..util import immutabledict
  11. def connection_memoize(key):
  12. """Decorator, memoize a function in a connection.info stash.
  13. Only applicable to functions which take no arguments other than a
  14. connection. The memo will be stored in ``connection.info[key]``.
  15. """
  16. @util.decorator
  17. def decorated(fn, self, connection):
  18. connection = connection.connect()
  19. try:
  20. return connection.info[key]
  21. except KeyError:
  22. connection.info[key] = val = fn(self, connection)
  23. return val
  24. return decorated
  25. _no_tuple = ()
  26. _no_kw = util.immutabledict()
  27. def _distill_params(connection, multiparams, params):
  28. r"""Given arguments from the calling form \*multiparams, \**params,
  29. return a list of bind parameter structures, usually a list of
  30. dictionaries.
  31. In the case of 'raw' execution which accepts positional parameters,
  32. it may be a list of tuples or lists.
  33. """
  34. if not multiparams:
  35. if params:
  36. connection._warn_for_legacy_exec_format()
  37. return [params]
  38. else:
  39. return []
  40. elif len(multiparams) == 1:
  41. zero = multiparams[0]
  42. if isinstance(zero, (list, tuple)):
  43. if (
  44. not zero
  45. or hasattr(zero[0], "__iter__")
  46. and not hasattr(zero[0], "strip")
  47. ):
  48. # execute(stmt, [{}, {}, {}, ...])
  49. # execute(stmt, [(), (), (), ...])
  50. return zero
  51. else:
  52. # this is used by exec_driver_sql only, so a deprecation
  53. # warning would already be coming from passing a plain
  54. # textual statement with positional parameters to
  55. # execute().
  56. # execute(stmt, ("value", "value"))
  57. return [zero]
  58. elif hasattr(zero, "keys"):
  59. # execute(stmt, {"key":"value"})
  60. return [zero]
  61. else:
  62. connection._warn_for_legacy_exec_format()
  63. # execute(stmt, "value")
  64. return [[zero]]
  65. else:
  66. connection._warn_for_legacy_exec_format()
  67. if hasattr(multiparams[0], "__iter__") and not hasattr(
  68. multiparams[0], "strip"
  69. ):
  70. return multiparams
  71. else:
  72. return [multiparams]
  73. def _distill_cursor_params(connection, multiparams, params):
  74. """_distill_params without any warnings. more appropriate for
  75. "cursor" params that can include tuple arguments, lists of tuples,
  76. etc.
  77. """
  78. if not multiparams:
  79. if params:
  80. return [params]
  81. else:
  82. return []
  83. elif len(multiparams) == 1:
  84. zero = multiparams[0]
  85. if isinstance(zero, (list, tuple)):
  86. if (
  87. not zero
  88. or hasattr(zero[0], "__iter__")
  89. and not hasattr(zero[0], "strip")
  90. ):
  91. # execute(stmt, [{}, {}, {}, ...])
  92. # execute(stmt, [(), (), (), ...])
  93. return zero
  94. else:
  95. # this is used by exec_driver_sql only, so a deprecation
  96. # warning would already be coming from passing a plain
  97. # textual statement with positional parameters to
  98. # execute().
  99. # execute(stmt, ("value", "value"))
  100. return [zero]
  101. elif hasattr(zero, "keys"):
  102. # execute(stmt, {"key":"value"})
  103. return [zero]
  104. else:
  105. # execute(stmt, "value")
  106. return [[zero]]
  107. else:
  108. if hasattr(multiparams[0], "__iter__") and not hasattr(
  109. multiparams[0], "strip"
  110. ):
  111. return multiparams
  112. else:
  113. return [multiparams]
  114. def _distill_params_20(params):
  115. if params is None:
  116. return _no_tuple, _no_kw
  117. elif isinstance(params, list):
  118. # collections_abc.MutableSequence): # avoid abc.__instancecheck__
  119. if params and not isinstance(
  120. params[0], (collections_abc.Mapping, tuple)
  121. ):
  122. raise exc.ArgumentError(
  123. "List argument must consist only of tuples or dictionaries"
  124. )
  125. return (params,), _no_kw
  126. elif isinstance(
  127. params,
  128. (tuple, dict, immutabledict),
  129. # only do abc.__instancecheck__ for Mapping after we've checked
  130. # for plain dictionaries and would otherwise raise
  131. ) or isinstance(params, collections_abc.Mapping):
  132. return (params,), _no_kw
  133. else:
  134. raise exc.ArgumentError("mapping or sequence expected for parameters")
  135. class TransactionalContext(object):
  136. """Apply Python context manager behavior to transaction objects.
  137. Performs validation to ensure the subject of the transaction is not
  138. used if the transaction were ended prematurely.
  139. """
  140. _trans_subject = None
  141. def _transaction_is_active(self):
  142. raise NotImplementedError()
  143. def _transaction_is_closed(self):
  144. raise NotImplementedError()
  145. def _rollback_can_be_called(self):
  146. """indicates the object is in a state that is known to be acceptable
  147. for rollback() to be called.
  148. This does not necessarily mean rollback() will succeed or not raise
  149. an error, just that there is currently no state detected that indicates
  150. rollback() would fail or emit warnings.
  151. It also does not mean that there's a transaction in progress, as
  152. it is usually safe to call rollback() even if no transaction is
  153. present.
  154. .. versionadded:: 1.4.28
  155. """
  156. raise NotImplementedError()
  157. def _get_subject(self):
  158. raise NotImplementedError()
  159. @classmethod
  160. def _trans_ctx_check(cls, subject):
  161. trans_context = subject._trans_context_manager
  162. if trans_context:
  163. if not trans_context._transaction_is_active():
  164. raise exc.InvalidRequestError(
  165. "Can't operate on closed transaction inside context "
  166. "manager. Please complete the context manager "
  167. "before emitting further commands."
  168. )
  169. def __enter__(self):
  170. subject = self._get_subject()
  171. # none for outer transaction, may be non-None for nested
  172. # savepoint, legacy nesting cases
  173. trans_context = subject._trans_context_manager
  174. self._outer_trans_ctx = trans_context
  175. self._trans_subject = subject
  176. subject._trans_context_manager = self
  177. return self
  178. def __exit__(self, type_, value, traceback):
  179. subject = self._trans_subject
  180. # simplistically we could assume that
  181. # "subject._trans_context_manager is self". However, any calling
  182. # code that is manipulating __exit__ directly would break this
  183. # assumption. alembic context manager
  184. # is an example of partial use that just calls __exit__ and
  185. # not __enter__ at the moment. it's safe to assume this is being done
  186. # in the wild also
  187. out_of_band_exit = (
  188. subject is None or subject._trans_context_manager is not self
  189. )
  190. if type_ is None and self._transaction_is_active():
  191. try:
  192. self.commit()
  193. except:
  194. with util.safe_reraise():
  195. if self._rollback_can_be_called():
  196. self.rollback()
  197. finally:
  198. if not out_of_band_exit:
  199. subject._trans_context_manager = self._outer_trans_ctx
  200. self._trans_subject = self._outer_trans_ctx = None
  201. else:
  202. try:
  203. if not self._transaction_is_active():
  204. if not self._transaction_is_closed():
  205. self.close()
  206. else:
  207. if self._rollback_can_be_called():
  208. self.rollback()
  209. finally:
  210. if not out_of_band_exit:
  211. subject._trans_context_manager = self._outer_trans_ctx
  212. self._trans_subject = self._outer_trans_ctx = None