Commit c1284c3d authored by Aymeric Augustin's avatar Aymeric Augustin
Browse files

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.
parent 88d5f321

django/db/backends/__init__.py

+9 −0
Original line number Diff line number Diff line
@@ -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.

django/db/transaction.py

+20 −0
Original line number Diff line number Diff line
@@ -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 #

#################################

docs/topics/db/transactions.txt

+21 −0
Original line number Diff line number Diff line
@@ -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

=======================

tests/transactions/tests.py

+24 −2
Original line number Diff line number Diff line 
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."""