Fixed #32220 -- Added durable argument to transaction.atomic().

5 files changed, 139 insertions, 23 deletions

diff --git a/django/db/transaction.py b/django/db/transaction.py
index 508a10c924..6d39e4a573 100644
--- a/django/db/transaction.py
+++ b/django/db/transaction.py
@@ -158,16 +158,30 @@ class Atomic(ContextDecorator):
 
     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.
     """
+    # This private flag is provided only to disable the durability checks in
+    # TestCase.
+    _ensure_durability = True
 
-    def __init__(self, using, savepoint):
+    def __init__(self, using, savepoint, durable):
         self.using = using
         self.savepoint = savepoint
+        self.durable = durable
 
     def __enter__(self):
         connection = get_connection(self.using)
 
+        if self.durable and self._ensure_durability and connection.in_atomic_block:
+            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
@@ -282,14 +296,14 @@ class Atomic(ContextDecorator):
                     connection.in_atomic_block = False
 
 
-def atomic(using=None, savepoint=True):
+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)(using)
+        return Atomic(DEFAULT_DB_ALIAS, savepoint, durable)(using)
     # Decorator: @atomic(...) or context manager: with atomic(...): ...
     else:
-        return Atomic(using, savepoint)
+        return Atomic(using, savepoint, durable)
 
 
 def _non_atomic_requests(view, using):
diff --git a/django/test/testcases.py b/django/test/testcases.py
index 716eef5e42..114ca85012 100644
--- a/django/test/testcases.py
+++ b/django/test/testcases.py
@@ -1181,29 +1181,37 @@ class TestCase(TransactionTestCase):
         super().setUpClass()
         if not cls._databases_support_transactions():
             return
-        cls.cls_atomics = cls._enter_atomics()
-
-        if cls.fixtures:
-            for db_name in cls._databases_names(include_mirrors=False):
-                try:
-                    call_command('loaddata', *cls.fixtures, **{'verbosity': 0, 'database': db_name})
-                except Exception:
-                    cls._rollback_atomics(cls.cls_atomics)
-                    cls._remove_databases_failures()
-                    raise
-        pre_attrs = cls.__dict__.copy()
+        # Disable the durability check to allow testing durable atomic blocks
+        # in a transaction for performance reasons.
+        transaction.Atomic._ensure_durability = False
         try:
-            cls.setUpTestData()
+            cls.cls_atomics = cls._enter_atomics()
+
+            if cls.fixtures:
+                for db_name in cls._databases_names(include_mirrors=False):
+                    try:
+                        call_command('loaddata', *cls.fixtures, **{'verbosity': 0, 'database': db_name})
+                    except Exception:
+                        cls._rollback_atomics(cls.cls_atomics)
+                        cls._remove_databases_failures()
+                        raise
+            pre_attrs = cls.__dict__.copy()
+            try:
+                cls.setUpTestData()
+            except Exception:
+                cls._rollback_atomics(cls.cls_atomics)
+                cls._remove_databases_failures()
+                raise
+            for name, value in cls.__dict__.items():
+                if value is not pre_attrs.get(name):
+                    setattr(cls, name, TestData(name, value))
         except Exception:
-            cls._rollback_atomics(cls.cls_atomics)
-            cls._remove_databases_failures()
+            transaction.Atomic._ensure_durability = True
             raise
-        for name, value in cls.__dict__.items():
-            if value is not pre_attrs.get(name):
-                setattr(cls, name, TestData(name, value))
 
     @classmethod
     def tearDownClass(cls):
+        transaction.Atomic._ensure_durability = True
         if cls._databases_support_transactions():
             cls._rollback_atomics(cls.cls_atomics)
             for conn in connections.all():
diff --git a/docs/releases/3.2.txt b/docs/releases/3.2.txt
index 222ef6c870..1209b357e2 100644
--- a/docs/releases/3.2.txt
+++ b/docs/releases/3.2.txt
@@ -356,6 +356,11 @@ Models
   allow using transforms. See :ref:`using-transforms-in-expressions` for
   details.
 
+* The new ``durable`` argument for :func:`~django.db.transaction.atomic`
+  guarantees that changes made in the atomic block will be committed if the
+  block exits without errors. A nested atomic block marked as durable will
+  raise a ``RuntimeError``.
+
 Pagination
 ~~~~~~~~~~
 
diff --git a/docs/topics/db/transactions.txt b/docs/topics/db/transactions.txt
index 996dd7534d..bdfb99cdfd 100644
--- a/docs/topics/db/transactions.txt
+++ b/docs/topics/db/transactions.txt
@@ -93,7 +93,7 @@ Controlling transactions explicitly
 
 Django provides a single API to control database transactions.
 
-.. function:: atomic(using=None, savepoint=True)
+.. function:: atomic(using=None, savepoint=True, durable=False)
 
     Atomicity is the defining property of database transactions. ``atomic``
     allows us to create a block of code within which the atomicity on the
@@ -105,6 +105,12 @@ Django provides a single API to control database transactions.
     completes successfully, its effects can still be rolled back if an
     exception is raised in the outer block at a later point.
 
+    It is sometimes useful to ensure an ``atomic`` block is always the
+    outermost ``atomic`` block, ensuring that any database changes are
+    committed when the block is exited without errors. This is known as
+    durability and can be achieved by setting ``durable=True``. If the
+    ``atomic`` block is nested within another it raises a ``RuntimeError``.
+
     ``atomic`` is usable both as a :py:term:`decorator`::
 
         from django.db import transaction
@@ -232,6 +238,16 @@ Django provides a single API to control database transactions.
     is especially important if you're using :func:`atomic` in long-running
     processes, outside of Django's request / response cycle.
 
+.. warning::
+
+    :class:`django.test.TestCase` disables the durability check to allow
+    testing durable atomic blocks in a transaction for performance reasons. Use
+    :class:`django.test.TransactionTestCase` for testing durability.
+
+.. versionchanged:: 3.2
+
+    The ``durable`` argument was added.
+
 Autocommit
 ==========
 
diff --git a/tests/transactions/tests.py b/tests/transactions/tests.py
index a4d64222be..dc163be6c5 100644
--- a/tests/transactions/tests.py
+++ b/tests/transactions/tests.py
@@ -8,7 +8,7 @@ from django.db import (
     transaction,
 )
 from django.test import (
-    TransactionTestCase, skipIfDBFeature, skipUnlessDBFeature,
+    TestCase, TransactionTestCase, skipIfDBFeature, skipUnlessDBFeature,
 )
 
 from .models import Reporter
@@ -498,3 +498,76 @@ class NonAutocommitTests(TransactionTestCase):
         finally:
             transaction.rollback()
             transaction.set_autocommit(True)
+
+
+class DurableTests(TransactionTestCase):
+    available_apps = ['transactions']
+
+    def test_commit(self):
+        with transaction.atomic(durable=True):
+            reporter = Reporter.objects.create(first_name='Tintin')
+        self.assertEqual(Reporter.objects.get(), reporter)
+
+    def test_nested_outer_durable(self):
+        with transaction.atomic(durable=True):
+            reporter1 = Reporter.objects.create(first_name='Tintin')
+            with transaction.atomic():
+                reporter2 = Reporter.objects.create(
+                    first_name='Archibald',
+                    last_name='Haddock',
+                )
+        self.assertSequenceEqual(Reporter.objects.all(), [reporter2, reporter1])
+
+    def test_nested_both_durable(self):
+        msg = 'A durable atomic block cannot be nested within another atomic block.'
+        with transaction.atomic(durable=True):
+            with self.assertRaisesMessage(RuntimeError, msg):
+                with transaction.atomic(durable=True):
+                    pass
+
+    def test_nested_inner_durable(self):
+        msg = 'A durable atomic block cannot be nested within another atomic block.'
+        with transaction.atomic():
+            with self.assertRaisesMessage(RuntimeError, msg):
+                with transaction.atomic(durable=True):
+                    pass
+
+
+class DisableDurabiltityCheckTests(TestCase):
+    """
+    TestCase runs all tests in a transaction by default. Code using
+    durable=True would always fail when run from TestCase. This would mean
+    these tests would be forced to use the slower TransactionTestCase even when
+    not testing durability. For this reason, TestCase disables the durability
+    check.
+    """
+    available_apps = ['transactions']
+
+    def test_commit(self):
+        with transaction.atomic(durable=True):
+            reporter = Reporter.objects.create(first_name='Tintin')
+        self.assertEqual(Reporter.objects.get(), reporter)
+
+    def test_nested_outer_durable(self):
+        with transaction.atomic(durable=True):
+            reporter1 = Reporter.objects.create(first_name='Tintin')
+            with transaction.atomic():
+                reporter2 = Reporter.objects.create(
+                    first_name='Archibald',
+                    last_name='Haddock',
+                )
+        self.assertSequenceEqual(Reporter.objects.all(), [reporter2, reporter1])
+
+    def test_nested_both_durable(self):
+        with transaction.atomic(durable=True):
+            # Error is not raised.
+            with transaction.atomic(durable=True):
+                reporter = Reporter.objects.create(first_name='Tintin')
+        self.assertEqual(Reporter.objects.get(), reporter)
+
+    def test_nested_inner_durable(self):
+        with transaction.atomic():
+            # Error is not raised.
+            with transaction.atomic(durable=True):
+                reporter = Reporter.objects.create(first_name='Tintin')
+        self.assertEqual(Reporter.objects.get(), reporter)