Fixed #20571 -- Added an API to control connection.needs_rollback.

This is useful:
- to force a rollback on the exit of an atomic block without having to
  raise and catch an exception;
- to prevent a rollback after handling an exception manually.

1 files changed, 24 insertions, 2 deletions

diff --git a/tests/transactions/tests.py b/tests/transactions/tests.py
index 24b7615d6f..756fa40abd 100644
--- a/tests/transactions/tests.py
+++ b/tests/transactions/tests.py
@@ -1,9 +1,8 @@
 from __future__ import absolute_import
 
 import sys
-import warnings
 
-from django.db import connection, transaction, IntegrityError
+from django.db import connection, transaction, DatabaseError, IntegrityError
 from django.test import TransactionTestCase, skipUnlessDBFeature
 from django.test.utils import IgnorePendingDeprecationWarningsMixin
 from django.utils import six
@@ -188,6 +187,29 @@ class AtomicTests(TransactionTestCase):
                 raise Exception("Oops, that's his first name")
         self.assertQuerysetEqual(Reporter.objects.all(), [])
 
+    def test_force_rollback(self):
+        with transaction.atomic():
+            Reporter.objects.create(first_name="Tintin")
+            # atomic block shouldn't rollback, but force it.
+            self.assertFalse(transaction.get_rollback())
+            transaction.set_rollback(True)
+        self.assertQuerysetEqual(Reporter.objects.all(), [])
+
+    def test_prevent_rollback(self):
+        with transaction.atomic():
+            Reporter.objects.create(first_name="Tintin")
+            sid = transaction.savepoint()
+            # trigger a database error inside an inner atomic without savepoint
+            with self.assertRaises(DatabaseError):
+                with transaction.atomic(savepoint=False):
+                    connection.cursor().execute(
+                            "SELECT no_such_col FROM transactions_reporter")
+            transaction.savepoint_rollback(sid)
+            # atomic block should rollback, but prevent it, as we just did it.
+            self.assertTrue(transaction.get_rollback())
+            transaction.set_rollback(False)
+        self.assertQuerysetEqual(Reporter.objects.all(), ['<Reporter: Tintin>'])
+
 
 class AtomicInsideTransactionTests(AtomicTests):
     """All basic tests for atomic should also pass within an existing transaction."""