You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

transaction.py 11 KiB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300
  1. from django.db import (
  2. DEFAULT_DB_ALIAS, DatabaseError, Error, ProgrammingError, connections,
  3. )
  4. from django.utils.decorators import ContextDecorator
  5. class TransactionManagementError(ProgrammingError):
  6. """
  7. This exception is thrown when transaction management is used improperly.
  8. """
  9. pass
  10. def get_connection(using=None):
  11. """
  12. Get a database connection by name, or the default database connection
  13. if no name is provided. This is a private API.
  14. """
  15. if using is None:
  16. using = DEFAULT_DB_ALIAS
  17. return connections[using]
  18. def get_autocommit(using=None):
  19. """
  20. Get the autocommit status of the connection.
  21. """
  22. return get_connection(using).get_autocommit()
  23. def set_autocommit(autocommit, using=None):
  24. """
  25. Set the autocommit status of the connection.
  26. """
  27. return get_connection(using).set_autocommit(autocommit)
  28. def commit(using=None):
  29. """
  30. Commits a transaction.
  31. """
  32. get_connection(using).commit()
  33. def rollback(using=None):
  34. """
  35. Rolls back a transaction.
  36. """
  37. get_connection(using).rollback()
  38. def savepoint(using=None):
  39. """
  40. Creates a savepoint (if supported and required by the backend) inside the
  41. current transaction. Returns an identifier for the savepoint that will be
  42. used for the subsequent rollback or commit.
  43. """
  44. return get_connection(using).savepoint()
  45. def savepoint_rollback(sid, using=None):
  46. """
  47. Rolls back the most recent savepoint (if one exists). Does nothing if
  48. savepoints are not supported.
  49. """
  50. get_connection(using).savepoint_rollback(sid)
  51. def savepoint_commit(sid, using=None):
  52. """
  53. Commits the most recent savepoint (if one exists). Does nothing if
  54. savepoints are not supported.
  55. """
  56. get_connection(using).savepoint_commit(sid)
  57. def clean_savepoints(using=None):
  58. """
  59. Resets the counter used to generate unique savepoint ids in this thread.
  60. """
  61. get_connection(using).clean_savepoints()
  62. def get_rollback(using=None):
  63. """
  64. Gets the "needs rollback" flag -- for *advanced use* only.
  65. """
  66. return get_connection(using).get_rollback()
  67. def set_rollback(rollback, using=None):
  68. """
  69. Sets or unsets the "needs rollback" flag -- for *advanced use* only.
  70. When `rollback` is `True`, it triggers a rollback when exiting the
  71. innermost enclosing atomic block that has `savepoint=True` (that's the
  72. default). Use this to force a rollback without raising an exception.
  73. When `rollback` is `False`, it prevents such a rollback. Use this only
  74. after rolling back to a known-good state! Otherwise, you break the atomic
  75. block and data corruption may occur.
  76. """
  77. return get_connection(using).set_rollback(rollback)
  78. def on_commit(func, using=None):
  79. """
  80. Register `func` to be called when the current transaction is committed.
  81. If the current transaction is rolled back, `func` will not be called.
  82. """
  83. get_connection(using).on_commit(func)
  84. #################################
  85. # Decorators / context managers #
  86. #################################
  87. class Atomic(ContextDecorator):
  88. """
  89. This class guarantees the atomic execution of a given block.
  90. An instance can be used either as a decorator or as a context manager.
  91. When it's used as a decorator, __call__ wraps the execution of the
  92. decorated function in the instance itself, used as a context manager.
  93. When it's used as a context manager, __enter__ creates a transaction or a
  94. savepoint, depending on whether a transaction is already in progress, and
  95. __exit__ commits the transaction or releases the savepoint on normal exit,
  96. and rolls back the transaction or to the savepoint on exceptions.
  97. It's possible to disable the creation of savepoints if the goal is to
  98. ensure that some code runs within a transaction without creating overhead.
  99. A stack of savepoints identifiers is maintained as an attribute of the
  100. connection. None denotes the absence of a savepoint.
  101. This allows reentrancy even if the same AtomicWrapper is reused. For
  102. example, it's possible to define `oa = @atomic('other')` and use `@oa` or
  103. `with oa:` multiple times.
  104. Since database connections are thread-local, this is thread-safe.
  105. This is a private API.
  106. """
  107. def __init__(self, using, savepoint):
  108. self.using = using
  109. self.savepoint = savepoint
  110. def __enter__(self):
  111. connection = get_connection(self.using)
  112. if not connection.in_atomic_block:
  113. # Reset state when entering an outermost atomic block.
  114. connection.commit_on_exit = True
  115. connection.needs_rollback = False
  116. if not connection.get_autocommit():
  117. # Some database adapters (namely sqlite3) don't handle
  118. # transactions and savepoints properly when autocommit is off.
  119. # Turning autocommit back on isn't an option; it would trigger
  120. # a premature commit. Give up if that happens.
  121. if connection.features.autocommits_when_autocommit_is_off:
  122. raise TransactionManagementError(
  123. "Your database backend doesn't behave properly when "
  124. "autocommit is off. Turn it on before using 'atomic'.")
  125. # Pretend we're already in an atomic block to bypass the code
  126. # that disables autocommit to enter a transaction, and make a
  127. # note to deal with this case in __exit__.
  128. connection.in_atomic_block = True
  129. connection.commit_on_exit = False
  130. if connection.in_atomic_block:
  131. # We're already in a transaction; create a savepoint, unless we
  132. # were told not to or we're already waiting for a rollback. The
  133. # second condition avoids creating useless savepoints and prevents
  134. # overwriting needs_rollback until the rollback is performed.
  135. if self.savepoint and not connection.needs_rollback:
  136. sid = connection.savepoint()
  137. connection.savepoint_ids.append(sid)
  138. else:
  139. connection.savepoint_ids.append(None)
  140. else:
  141. connection.set_autocommit(False, force_begin_transaction_with_broken_autocommit=True)
  142. connection.in_atomic_block = True
  143. def __exit__(self, exc_type, exc_value, traceback):
  144. connection = get_connection(self.using)
  145. if connection.savepoint_ids:
  146. sid = connection.savepoint_ids.pop()
  147. else:
  148. # Prematurely unset this flag to allow using commit or rollback.
  149. connection.in_atomic_block = False
  150. try:
  151. if connection.closed_in_transaction:
  152. # The database will perform a rollback by itself.
  153. # Wait until we exit the outermost block.
  154. pass
  155. elif exc_type is None and not connection.needs_rollback:
  156. if connection.in_atomic_block:
  157. # Release savepoint if there is one
  158. if sid is not None:
  159. try:
  160. connection.savepoint_commit(sid)
  161. except DatabaseError:
  162. try:
  163. connection.savepoint_rollback(sid)
  164. # The savepoint won't be reused. Release it to
  165. # minimize overhead for the database server.
  166. connection.savepoint_commit(sid)
  167. except Error:
  168. # If rolling back to a savepoint fails, mark for
  169. # rollback at a higher level and avoid shadowing
  170. # the original exception.
  171. connection.needs_rollback = True
  172. raise
  173. else:
  174. # Commit transaction
  175. try:
  176. connection.commit()
  177. except DatabaseError:
  178. try:
  179. connection.rollback()
  180. except Error:
  181. # An error during rollback means that something
  182. # went wrong with the connection. Drop it.
  183. connection.close()
  184. raise
  185. else:
  186. # This flag will be set to True again if there isn't a savepoint
  187. # allowing to perform the rollback at this level.
  188. connection.needs_rollback = False
  189. if connection.in_atomic_block:
  190. # Roll back to savepoint if there is one, mark for rollback
  191. # otherwise.
  192. if sid is None:
  193. connection.needs_rollback = True
  194. else:
  195. try:
  196. connection.savepoint_rollback(sid)
  197. # The savepoint won't be reused. Release it to
  198. # minimize overhead for the database server.
  199. connection.savepoint_commit(sid)
  200. except Error:
  201. # If rolling back to a savepoint fails, mark for
  202. # rollback at a higher level and avoid shadowing
  203. # the original exception.
  204. connection.needs_rollback = True
  205. else:
  206. # Roll back transaction
  207. try:
  208. connection.rollback()
  209. except Error:
  210. # An error during rollback means that something
  211. # went wrong with the connection. Drop it.
  212. connection.close()
  213. finally:
  214. # Outermost block exit when autocommit was enabled.
  215. if not connection.in_atomic_block:
  216. if connection.closed_in_transaction:
  217. connection.connection = None
  218. else:
  219. connection.set_autocommit(True)
  220. # Outermost block exit when autocommit was disabled.
  221. elif not connection.savepoint_ids and not connection.commit_on_exit:
  222. if connection.closed_in_transaction:
  223. connection.connection = None
  224. else:
  225. connection.in_atomic_block = False
  226. def atomic(using=None, savepoint=True):
  227. # Bare decorator: @atomic -- although the first argument is called
  228. # `using`, it's actually the function being decorated.
  229. if callable(using):
  230. return Atomic(DEFAULT_DB_ALIAS, savepoint)(using)
  231. # Decorator: @atomic(...) or context manager: with atomic(...): ...
  232. else:
  233. return Atomic(using, savepoint)
  234. def _non_atomic_requests(view, using):
  235. try:
  236. view._non_atomic_requests.add(using)
  237. except AttributeError:
  238. view._non_atomic_requests = {using}
  239. return view
  240. def non_atomic_requests(using=None):
  241. if callable(using):
  242. return _non_atomic_requests(using, DEFAULT_DB_ALIAS)
  243. else:
  244. if using is None:
  245. using = DEFAULT_DB_ALIAS
  246. return lambda view: _non_atomic_requests(view, using)