From c1284c3d3c6131a9d0ded9601ae0feb9a2e81a65 Mon Sep 17 00:00:00 2001
From: Aymeric Augustin <aymeric.augustin@m4x.org>
Date: Thu, 27 Jun 2013 22:19:54 +0200
Subject: [PATCH] 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.
---
 django/db/backends/__init__.py  |  9 +++++++++
 django/db/transaction.py        | 20 ++++++++++++++++++++
 docs/topics/db/transactions.txt | 21 +++++++++++++++++++++
 tests/transactions/tests.py     | 26 ++++++++++++++++++++++++--
 4 files changed, 74 insertions(+), 2 deletions(-)

diff --git a/django/db/backends/__init__.py b/django/db/backends/__init__.py
index fa3cc5ac02..1a74232704 100644
--- a/django/db/backends/__init__.py
+++ b/django/db/backends/__init__.py
@@ -330,6 +330,15 @@ class BaseDatabaseWrapper(object):
         self._set_autocommit(autocommit)
         self.autocommit = autocommit
 
+    def set_rollback(self, rollback):
+        """
+        Set or unset the "needs rollback" flag -- for *advanced use* only.
+        """
+        if not self.in_atomic_block:
+            raise TransactionManagementError(
+                "needs_rollback doesn't work outside of an 'atomic' block.")
+        self.needs_rollback = rollback
+
     def validate_no_atomic_block(self):
         """
         Raise an error if an atomic block is active.
diff --git a/django/db/transaction.py b/django/db/transaction.py
index f770f2efa7..95b9ae165e 100644
--- a/django/db/transaction.py
+++ b/django/db/transaction.py
@@ -171,6 +171,26 @@ def clean_savepoints(using=None):
     """
     get_connection(using).clean_savepoints()
 
+def get_rollback(using=None):
+    """
+    Gets the "needs rollback" flag -- for *advanced use* only.
+    """
+    return get_connection(using).needs_rollback
+
+def set_rollback(rollback, using=None):
+    """
+    Sets or unsets the "needs rollback" flag -- for *advanced use* only.
+
+    When `rollback` is `True`, it triggers 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`, it prevents 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)
+
 #################################
 # Decorators / context managers #
 #################################
diff --git a/docs/topics/db/transactions.txt b/docs/topics/db/transactions.txt
index e9a626f56b..903579cc38 100644
--- a/docs/topics/db/transactions.txt
+++ b/docs/topics/db/transactions.txt
@@ -389,6 +389,27 @@ The following example demonstrates the use of savepoints::
             transaction.savepoint_rollback(sid)
             # open transaction now contains only a.save()
 
+.. versionadded:: 1.6
+
+Savepoints may be used to recover from a database error by performing a partial
+rollback. If you're doing this inside an :func:`atomic` block, the entire block
+will still be rolled back, because it doesn't know you've handled the situation
+at a lower level! To prevent this, you can control the rollback behavior with
+the following functions.
+
+.. function:: get_rollback(using=None)
+
+.. function:: set_rollback(rollback, using=None)
+
+Setting the rollback flag to ``True`` forces a rollback when exiting the
+innermost atomic block. This may be useful to trigger a rollback without
+raising an exception.
+
+Setting it to ``False`` prevents such a rollback. Before doing that, make sure
+you've rolled back the transaction to a known-good savepoint within the current
+atomic block! Otherwise you're breaking atomicity and data corruption may
+occur.
+
 Database-specific notes
 =======================
 
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."""