Durable decorator

durability.py

import functools

from django.conf import settings
from django.db import transaction


def durable(func):
    """
    Decorator to ensure that a function is not being called from an atomic block.

    Usage:

        @durable
        def some_function(...):
            with transaction.atomic():
               ...   

    Code decorated like this is guaranteed to be *durable* - that is, not at risk of being rolled
    back due to an exception that happens after the use case has completed.

    This is achieved by enforcing that:

    1. The use case does not begin in the context of a currently open transaction.
    2. The use case does not leave work uncommitted.

    N.B. Only the default database connection is checked.

    Disabling this behaviour in tests
    ---------------------------------

    This behaviour doesn't play well with tests that use transactions for performance reasons.
    It can be disabled by setting DISABLE_DURABILITY_CHECKING to True in your Django settings.
    """

    @functools.wraps(func)
    def wrapper(*args, **kwargs):

        if not getattr(settings, "DISABLE_DURABILITY_CHECKING", False) and transaction.get_connection().in_atomic_block:
            raise exceptions.ProgrammingError(
                "A durable function must not be called within a database transaction."
            )

        return_value = func(*args, **kwargs)

        if not getattr(settings, "DISABLE_DURABILITY_CHECKING", False) and _db_may_have_uncommitted_work():
            # Clean up first, otherwise we may see errors later that will mask this one.
            transaction.rollback()
            raise exceptions.ProgrammingError("A durable function must not leave work uncommitted.")

        return return_value

    return wrapper


def _db_may_have_uncommitted_work():
    # Django doesn't seem to provide an API to tell this, but practically speaking there shouldn't
    # be any uncommitted work if we're not in an atomic block, and autocommit is turned on.
    return db.in_atomic_block() or not transaction.get_autocommit()