123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339 |
- from contextlib import ContextDecorator, contextmanager
-
- from django.db import (
- DEFAULT_DB_ALIAS,
- DatabaseError,
- Error,
- ProgrammingError,
- connections,
- )
-
-
- class TransactionManagementError(ProgrammingError):
- """Transaction management is used improperly."""
-
- pass
-
-
- def get_connection(using=None):
- """
- Get a database connection by name, or the default database connection
- if no name is provided. This is a private API.
- """
- if using is None:
- using = DEFAULT_DB_ALIAS
- return connections[using]
-
-
- def get_autocommit(using=None):
- """Get the autocommit status of the connection."""
- return get_connection(using).get_autocommit()
-
-
- def set_autocommit(autocommit, using=None):
- """Set the autocommit status of the connection."""
- return get_connection(using).set_autocommit(autocommit)
-
-
- def commit(using=None):
- """Commit a transaction."""
- get_connection(using).commit()
-
-
- def rollback(using=None):
- """Roll back a transaction."""
- get_connection(using).rollback()
-
-
- def savepoint(using=None):
- """
- Create a savepoint (if supported and required by the backend) inside the
- current transaction. Return an identifier for the savepoint that will be
- used for the subsequent rollback or commit.
- """
- return get_connection(using).savepoint()
-
-
- def savepoint_rollback(sid, using=None):
- """
- Roll back the most recent savepoint (if one exists). Do nothing if
- savepoints are not supported.
- """
- get_connection(using).savepoint_rollback(sid)
-
-
- def savepoint_commit(sid, using=None):
- """
- Commit the most recent savepoint (if one exists). Do nothing if
- savepoints are not supported.
- """
- get_connection(using).savepoint_commit(sid)
-
-
- def clean_savepoints(using=None):
- """
- Reset the counter used to generate unique savepoint ids in this thread.
- """
- get_connection(using).clean_savepoints()
-
-
- def get_rollback(using=None):
- """Get the "needs rollback" flag -- for *advanced use* only."""
- return get_connection(using).get_rollback()
-
-
- def set_rollback(rollback, using=None):
- """
- Set or unset the "needs rollback" flag -- for *advanced use* only.
-
- When `rollback` is `True`, trigger a rollback when exiting the innermost
- enclosing atomic block that has `savepoint=True` (that's the default). Use
- this to force a rollback without raising an exception.
-
- When `rollback` is `False`, prevent such a rollback. Use this only after
- rolling back to a known-good state! Otherwise, you break the atomic block
- and data corruption may occur.
- """
- return get_connection(using).set_rollback(rollback)
-
-
- @contextmanager
- def mark_for_rollback_on_error(using=None):
- """
- Internal low-level utility to mark a transaction as "needs rollback" when
- an exception is raised while not enforcing the enclosed block to be in a
- transaction. This is needed by Model.save() and friends to avoid starting a
- transaction when in autocommit mode and a single query is executed.
-
- It's equivalent to:
-
- connection = get_connection(using)
- if connection.get_autocommit():
- yield
- else:
- with transaction.atomic(using=using, savepoint=False):
- yield
-
- but it uses low-level utilities to avoid performance overhead.
- """
- try:
- yield
- except Exception:
- connection = get_connection(using)
- if connection.in_atomic_block:
- connection.needs_rollback = True
- raise
-
-
- def on_commit(func, using=None):
- """
- Register `func` to be called when the current transaction is committed.
- If the current transaction is rolled back, `func` will not be called.
- """
- get_connection(using).on_commit(func)
-
-
- #################################
- # Decorators / context managers #
- #################################
-
-
- class Atomic(ContextDecorator):
- """
- Guarantee the atomic execution of a given block.
-
- An instance can be used either as a decorator or as a context manager.
-
- When it's used as a decorator, __call__ wraps the execution of the
- decorated function in the instance itself, used as a context manager.
-
- When it's used as a context manager, __enter__ creates a transaction or a
- savepoint, depending on whether a transaction is already in progress, and
- __exit__ commits the transaction or releases the savepoint on normal exit,
- and rolls back the transaction or to the savepoint on exceptions.
-
- It's possible to disable the creation of savepoints if the goal is to
- ensure that some code runs within a transaction without creating overhead.
-
- A stack of savepoints identifiers is maintained as an attribute of the
- connection. None denotes the absence of a savepoint.
-
- This allows reentrancy even if the same AtomicWrapper is reused. For
- example, it's possible to define `oa = atomic('other')` and use `@oa` or
- `with oa:` multiple times.
-
- Since database connections are thread-local, this is thread-safe.
-
- An atomic block can be tagged as durable. In this case, raise a
- RuntimeError if it's nested within another atomic block. This guarantees
- that database changes in a durable block are committed to the database when
- the block exists without error.
-
- This is a private API.
- """
-
- def __init__(self, using, savepoint, durable):
- self.using = using
- self.savepoint = savepoint
- self.durable = durable
- self._from_testcase = False
-
- def __enter__(self):
- connection = get_connection(self.using)
-
- if (
- self.durable
- and connection.atomic_blocks
- and not connection.atomic_blocks[-1]._from_testcase
- ):
- raise RuntimeError(
- "A durable atomic block cannot be nested within another "
- "atomic block."
- )
- if not connection.in_atomic_block:
- # Reset state when entering an outermost atomic block.
- connection.commit_on_exit = True
- connection.needs_rollback = False
- if not connection.get_autocommit():
- # Pretend we're already in an atomic block to bypass the code
- # that disables autocommit to enter a transaction, and make a
- # note to deal with this case in __exit__.
- connection.in_atomic_block = True
- connection.commit_on_exit = False
-
- if connection.in_atomic_block:
- # We're already in a transaction; create a savepoint, unless we
- # were told not to or we're already waiting for a rollback. The
- # second condition avoids creating useless savepoints and prevents
- # overwriting needs_rollback until the rollback is performed.
- if self.savepoint and not connection.needs_rollback:
- sid = connection.savepoint()
- connection.savepoint_ids.append(sid)
- else:
- connection.savepoint_ids.append(None)
- else:
- connection.set_autocommit(
- False, force_begin_transaction_with_broken_autocommit=True
- )
- connection.in_atomic_block = True
-
- if connection.in_atomic_block:
- connection.atomic_blocks.append(self)
-
- def __exit__(self, exc_type, exc_value, traceback):
- connection = get_connection(self.using)
-
- if connection.in_atomic_block:
- connection.atomic_blocks.pop()
-
- if connection.savepoint_ids:
- sid = connection.savepoint_ids.pop()
- else:
- # Prematurely unset this flag to allow using commit or rollback.
- connection.in_atomic_block = False
-
- try:
- if connection.closed_in_transaction:
- # The database will perform a rollback by itself.
- # Wait until we exit the outermost block.
- pass
-
- elif exc_type is None and not connection.needs_rollback:
- if connection.in_atomic_block:
- # Release savepoint if there is one
- if sid is not None:
- try:
- connection.savepoint_commit(sid)
- except DatabaseError:
- try:
- connection.savepoint_rollback(sid)
- # The savepoint won't be reused. Release it to
- # minimize overhead for the database server.
- connection.savepoint_commit(sid)
- except Error:
- # If rolling back to a savepoint fails, mark for
- # rollback at a higher level and avoid shadowing
- # the original exception.
- connection.needs_rollback = True
- raise
- else:
- # Commit transaction
- try:
- connection.commit()
- except DatabaseError:
- try:
- connection.rollback()
- except Error:
- # An error during rollback means that something
- # went wrong with the connection. Drop it.
- connection.close()
- raise
- else:
- # This flag will be set to True again if there isn't a savepoint
- # allowing to perform the rollback at this level.
- connection.needs_rollback = False
- if connection.in_atomic_block:
- # Roll back to savepoint if there is one, mark for rollback
- # otherwise.
- if sid is None:
- connection.needs_rollback = True
- else:
- try:
- connection.savepoint_rollback(sid)
- # The savepoint won't be reused. Release it to
- # minimize overhead for the database server.
- connection.savepoint_commit(sid)
- except Error:
- # If rolling back to a savepoint fails, mark for
- # rollback at a higher level and avoid shadowing
- # the original exception.
- connection.needs_rollback = True
- else:
- # Roll back transaction
- try:
- connection.rollback()
- except Error:
- # An error during rollback means that something
- # went wrong with the connection. Drop it.
- connection.close()
-
- finally:
- # Outermost block exit when autocommit was enabled.
- if not connection.in_atomic_block:
- if connection.closed_in_transaction:
- connection.connection = None
- else:
- connection.set_autocommit(True)
- # Outermost block exit when autocommit was disabled.
- elif not connection.savepoint_ids and not connection.commit_on_exit:
- if connection.closed_in_transaction:
- connection.connection = None
- else:
- connection.in_atomic_block = False
-
-
- def atomic(using=None, savepoint=True, durable=False):
- # Bare decorator: @atomic -- although the first argument is called
- # `using`, it's actually the function being decorated.
- if callable(using):
- return Atomic(DEFAULT_DB_ALIAS, savepoint, durable)(using)
- # Decorator: @atomic(...) or context manager: with atomic(...): ...
- else:
- return Atomic(using, savepoint, durable)
-
-
- def _non_atomic_requests(view, using):
- try:
- view._non_atomic_requests.add(using)
- except AttributeError:
- view._non_atomic_requests = {using}
- return view
-
-
- def non_atomic_requests(using=None):
- if callable(using):
- return _non_atomic_requests(using, DEFAULT_DB_ALIAS)
- else:
- if using is None:
- using = DEFAULT_DB_ALIAS
- return lambda view: _non_atomic_requests(view, using)
|