Implemented an 'atomic' decorator and context manager.

Currently it only works in autocommit mode.

Based on @xact by Christophe Pettus.

7 files changed, 430 insertions, 23 deletions

diff --git a/AUTHORS b/AUTHORS
index 35a316d4c2..3c1cd81639 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -434,6 +434,7 @@ answer newbie questions, and generally made Django that much better:
     Andreas Pelme <andreas@pelme.se>
     permonik@mesias.brnonet.cz
     peter@mymart.com
+    Christophe Pettus <xof@thebuild.com>
     pgross@thoughtworks.com
     phaedo <http://phaedo.cx/>
     phil@produxion.net
diff --git a/django/db/backends/__init__.py b/django/db/backends/__init__.py
index 48190d3c62..818850bf43 100644
--- a/django/db/backends/__init__.py
+++ b/django/db/backends/__init__.py
@@ -50,6 +50,12 @@ class BaseDatabaseWrapper(object):
         # set somewhat aggressively, as the DBAPI doesn't make it easy to
         # deduce if the connection is in transaction or not.
         self._dirty = False
+        # Tracks if the connection is in a transaction managed by 'atomic'
+        self.in_atomic_block = False
+        # List of savepoints created by 'atomic'
+        self.savepoint_ids = []
+        # Hack to provide compatibility with legacy transaction management
+        self._atomic_forced_unmanaged = False
 
         # Connection termination related attributes
         self.close_at = None
@@ -148,7 +154,7 @@ class BaseDatabaseWrapper(object):
 
     def commit(self):
         """
-        Does the commit itself and resets the dirty flag.
+        Commits a transaction and resets the dirty flag.
         """
         self.validate_thread_sharing()
         self._commit()
@@ -156,7 +162,7 @@ class BaseDatabaseWrapper(object):
 
     def rollback(self):
         """
-        Does the rollback itself and resets the dirty flag.
+        Rolls back a transaction and resets the dirty flag.
         """
         self.validate_thread_sharing()
         self._rollback()
@@ -447,6 +453,12 @@ class BaseDatabaseWrapper(object):
             if must_close:
                 self.close()
 
+    def _start_transaction_under_autocommit(self):
+        """
+        Only required when autocommits_when_autocommit_is_off = True.
+        """
+        raise NotImplementedError
+
 
 class BaseDatabaseFeatures(object):
     allows_group_by_pk = False
@@ -549,6 +561,10 @@ class BaseDatabaseFeatures(object):
     # Support for the DISTINCT ON clause
     can_distinct_on_fields = False
 
+    # Does the backend decide to commit before SAVEPOINT statements
+    # when autocommit is disabled? http://bugs.python.org/issue8145#msg109965
+    autocommits_when_autocommit_is_off = False
+
     def __init__(self, connection):
         self.connection = connection
 
@@ -931,6 +947,9 @@ class BaseDatabaseOperations(object):
         return "BEGIN;"
 
     def end_transaction_sql(self, success=True):
+        """
+        Returns the SQL statement required to end a transaction.
+        """
         if not success:
             return "ROLLBACK;"
         return "COMMIT;"
diff --git a/django/db/backends/sqlite3/base.py b/django/db/backends/sqlite3/base.py
index f537860a53..f70c3872a8 100644
--- a/django/db/backends/sqlite3/base.py
+++ b/django/db/backends/sqlite3/base.py
@@ -99,6 +99,7 @@ class DatabaseFeatures(BaseDatabaseFeatures):
     supports_mixed_date_datetime_comparisons = False
     has_bulk_insert = True
     can_combine_inserts_with_and_without_auto_increment_pk = False
+    autocommits_when_autocommit_is_off = True
 
     @cached_property
     def uses_savepoints(self):
@@ -360,10 +361,12 @@ class DatabaseWrapper(BaseDatabaseWrapper):
             BaseDatabaseWrapper.close(self)
 
     def _savepoint_allowed(self):
-        # When 'isolation_level' is None, Django doesn't provide a way to
-        # create a transaction (yet) so savepoints can't be created. When it
-        # isn't, sqlite3 commits before each savepoint -- it's a bug.
-        return False
+        # When 'isolation_level' is not None, sqlite3 commits before each
+        # savepoint; it's a bug. When it is None, savepoints don't make sense
+        # because autocommit is enabled. The only exception is inside atomic
+        # blocks. To work around that bug, on SQLite, atomic starts a
+        # transaction explicitly rather than simply disable autocommit.
+        return self.in_atomic_block
 
     def _set_autocommit(self, autocommit):
         if autocommit:
@@ -413,6 +416,14 @@ class DatabaseWrapper(BaseDatabaseWrapper):
     def is_usable(self):
         return True
 
+    def _start_transaction_under_autocommit(self):
+        """
+        Start a transaction explicitly in autocommit mode.
+
+        Staying in autocommit mode works around a bug of sqlite3 that breaks
+        savepoints when autocommit is disabled.
+        """
+        self.cursor().execute("BEGIN")
 
 FORMAT_QMARK_REGEX = re.compile(r'(?<!%)%s')
 
diff --git a/django/db/transaction.py b/django/db/transaction.py
index 49b67f4122..8126c18a70 100644
--- a/django/db/transaction.py
+++ b/django/db/transaction.py
@@ -16,7 +16,7 @@ import warnings
 
 from functools import wraps
 
-from django.db import connections, DEFAULT_DB_ALIAS
+from django.db import connections, DatabaseError, DEFAULT_DB_ALIAS
 
 
 class TransactionManagementError(Exception):
@@ -134,13 +134,13 @@ def rollback_unless_managed(using=None):
 
 def commit(using=None):
     """
-    Does the commit itself and resets the dirty flag.
+    Commits a transaction and resets the dirty flag.
     """
     get_connection(using).commit()
 
 def rollback(using=None):
     """
-    This function does the rollback itself and resets the dirty flag.
+    Rolls back a transaction and resets the dirty flag.
     """
     get_connection(using).rollback()
 
@@ -166,9 +166,151 @@ def savepoint_commit(sid, using=None):
     """
     get_connection(using).savepoint_commit(sid)
 
-##############
-# DECORATORS #
-##############
+
+#################################
+# Decorators / context managers #
+#################################
+
+class Atomic(object):
+    """
+    This class guarantees the atomic execution of a given block.
+
+    An instance can be used either as a decorator or as a context manager.
+
+    When it's used as a decorator, __call__ wraps the execution of the
+    decorated function in the instance itself, used as a context manager.
+
+    When it's used as a context manager, __enter__ creates a transaction or a
+    savepoint, depending on whether a transaction is already in progress, and
+    __exit__ commits the transaction or releases the savepoint on normal exit,
+    and rolls back the transaction or to the savepoint on exceptions.
+
+    A stack of savepoints identifiers is maintained as an attribute of the
+    connection. None denotes a plain transaction.
+
+    This allows reentrancy even if the same AtomicWrapper is reused. For
+    example, it's possible to define `oa = @atomic('other')` and use `@ao` or
+    `with oa:` multiple times.
+
+    Since database connections are thread-local, this is thread-safe.
+    """
+
+    def __init__(self, using):
+        self.using = using
+
+    def _legacy_enter_transaction_management(self, connection):
+        if not connection.in_atomic_block:
+            if connection.transaction_state and connection.transaction_state[-1]:
+                connection._atomic_forced_unmanaged = True
+                connection.enter_transaction_management(managed=False)
+            else:
+                connection._atomic_forced_unmanaged = False
+
+    def _legacy_leave_transaction_management(self, connection):
+        if not connection.in_atomic_block and connection._atomic_forced_unmanaged:
+            connection.leave_transaction_management()
+
+    def __enter__(self):
+        connection = get_connection(self.using)
+
+        # Ensure we have a connection to the database before testing
+        # autocommit status.
+        connection.ensure_connection()
+
+        # Remove this when the legacy transaction management goes away.
+        self._legacy_enter_transaction_management(connection)
+
+        if not connection.in_atomic_block and not connection.autocommit:
+            raise TransactionManagementError(
+                "'atomic' cannot be used when autocommit is disabled.")
+
+        if connection.in_atomic_block:
+            # We're already in a transaction; create a savepoint.
+            sid = connection.savepoint()
+            connection.savepoint_ids.append(sid)
+        else:
+            # We aren't in a transaction yet; create one.
+            # The usual way to start a transaction is to turn autocommit off.
+            # However, some database adapters (namely sqlite3) don't handle
+            # transactions and savepoints properly when autocommit is off.
+            # In such cases, start an explicit transaction instead, which has
+            # the side-effect of disabling autocommit.
+            if connection.features.autocommits_when_autocommit_is_off:
+                connection._start_transaction_under_autocommit()
+                connection.autocommit = False
+            else:
+                connection.set_autocommit(False)
+            connection.in_atomic_block = True
+            connection.savepoint_ids.append(None)
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        connection = get_connection(self.using)
+        sid = connection.savepoint_ids.pop()
+        if exc_value is None:
+            if sid is None:
+                # Commit transaction
+                connection.in_atomic_block = False
+                try:
+                    connection.commit()
+                except DatabaseError:
+                    connection.rollback()
+                    # Remove this when the legacy transaction management goes away.
+                    self._legacy_leave_transaction_management(connection)
+                    raise
+                finally:
+                    if connection.features.autocommits_when_autocommit_is_off:
+                        connection.autocommit = True
+                    else:
+                        connection.set_autocommit(True)
+            else:
+                # Release savepoint
+                try:
+                    connection.savepoint_commit(sid)
+                except DatabaseError:
+                    connection.savepoint_rollback(sid)
+                    # Remove this when the legacy transaction management goes away.
+                    self._legacy_leave_transaction_management(connection)
+                    raise
+        else:
+            if sid is None:
+                # Roll back transaction
+                connection.in_atomic_block = False
+                try:
+                    connection.rollback()
+                finally:
+                    if connection.features.autocommits_when_autocommit_is_off:
+                        connection.autocommit = True
+                    else:
+                        connection.set_autocommit(True)
+            else:
+                # Roll back to savepoint
+                connection.savepoint_rollback(sid)
+
+        # Remove this when the legacy transaction management goes away.
+        self._legacy_leave_transaction_management(connection)
+
+
+    def __call__(self, func):
+        @wraps(func)
+        def inner(*args, **kwargs):
+            with self:
+                return func(*args, **kwargs)
+        return inner
+
+
+def atomic(using=None):
+    # 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)(using)
+    # Decorator: @atomic(...) or context manager: with atomic(...): ...
+    else:
+        return Atomic(using)
+
+
+############################################
+# Deprecated decorators / context managers #
+############################################
 
 class Transaction(object):
     """
@@ -279,7 +421,8 @@ def commit_on_success_unless_managed(using=None):
     """
     Transitory API to preserve backwards-compatibility while refactoring.
     """
-    if get_autocommit(using):
+    connection = get_connection(using)
+    if connection.autocommit and not connection.in_atomic_block:
         return commit_on_success(using)
     else:
         def entering(using):
diff --git a/docs/topics/db/transactions.txt b/docs/topics/db/transactions.txt
index e2c8e4e3f5..2a4cd306c6 100644
--- a/docs/topics/db/transactions.txt
+++ b/docs/topics/db/transactions.txt
@@ -1,13 +1,16 @@
-==============================
-Managing database transactions
-==============================
+=====================
+Database transactions
+=====================
 
 .. module:: django.db.transaction
 
 Django gives you a few ways to control how database transactions are managed.
 
+Managing database transactions
+==============================
+
 Django's default transaction behavior
-=====================================
+-------------------------------------
 
 Django's default behavior is to run in autocommit mode. Each query is
 immediately committed to the database. :ref:`See below for details
@@ -24,7 +27,7 @@ immediately committed to the database. :ref:`See below for details
     behavior <transactions-changes-from-1.5>`.
 
 Tying transactions to HTTP requests
-===================================
+-----------------------------------
 
 The recommended way to handle transactions in Web requests is to tie them to
 the request and response phases via Django's ``TransactionMiddleware``.
@@ -63,6 +66,85 @@ connection internally.
     multiple databases and want transaction control over databases other than
     "default", you will need to write your own transaction middleware.
 
+Controlling transactions explicitly
+-----------------------------------
+
+.. versionadded:: 1.6
+
+Django provides a single API to control database transactions.
+
+.. function:: atomic(using=None)
+
+    This function creates an atomic block for writes to the database.
+    (Atomicity is the defining property of database transactions.)
+
+    When the block completes successfully, the changes are committed to the
+    database. When it raises an exception, the changes are rolled back.
+
+    ``atomic`` can be nested. In this case, when an inner block completes
+    successfully, its effects can still be rolled back if an exception is
+    raised in the outer block at a later point.
+
+    ``atomic`` takes a ``using`` argument which should be the name of a
+    database. If this argument isn't provided, Django uses the ``"default"``
+    database.
+
+    ``atomic`` is usable both as a decorator::
+
+        from django.db import transaction
+
+        @transaction.atomic
+        def viewfunc(request):
+            # This code executes inside a transaction.
+            do_stuff()
+
+    and as a context manager::
+
+        from django.db import transaction
+
+        def viewfunc(request):
+            # This code executes in autocommit mode (Django's default).
+            do_stuff()
+
+            with transaction.atomic():
+                # This code executes inside a transaction.
+                do_more_stuff()
+
+    Wrapping ``atomic`` in a try/except block allows for natural handling of
+    integrity errors::
+
+        from django.db import IntegrityError, transaction
+
+        @transaction.atomic
+        def viewfunc(request):
+            do_stuff()
+
+            try:
+                with transaction.atomic():
+                    do_stuff_that_could_fail()
+            except IntegrityError:
+                handle_exception()
+
+            do_more_stuff()
+
+    In this example, even if ``do_stuff_that_could_fail()`` causes a database
+    error by breaking an integrity constraint, you can execute queries in
+    ``do_more_stuff()``, and the changes from ``do_stuff()`` are still there.
+
+    In order to guarantee atomicity, ``atomic`` disables some APIs. Attempting
+    to commit, roll back, or change the autocommit state of the database
+    connection within an ``atomic`` block will raise an exception.
+
+    ``atomic`` can only be used in autocommit mode. It will raise an exception
+    if autocommit is turned off.
+
+    Under the hood, Django's transaction management code:
+
+    - opens a transaction when entering the outermost ``atomic`` block;
+    - creates a savepoint when entering an inner ``atomic`` block;
+    - releases or rolls back to the savepoint when exiting an inner block;
+    - commits or rolls back the transaction when exiting the outermost block.
+
 .. _transaction-management-functions:
 
 Controlling transaction management in views
@@ -325,9 +407,8 @@ When autocommit is enabled, savepoints don't make sense. When it's disabled,
 commits before any statement other than ``SELECT``, ``INSERT``, ``UPDATE``,
 ``DELETE`` and ``REPLACE``.)
 
-As a consequence, savepoints are only usable if you start a transaction
-manually while in autocommit mode, and Django doesn't provide an API to
-achieve that.
+As a consequence, savepoints are only usable inside a transaction ie. inside
+an :func:`atomic` block.
 
 Transactions in MySQL
 ---------------------
diff --git a/tests/transactions/models.py b/tests/transactions/models.py
index 0f8d6b16ec..6c2bcfd23f 100644
--- a/tests/transactions/models.py
+++ b/tests/transactions/models.py
@@ -22,4 +22,4 @@ class Reporter(models.Model):
         ordering = ('first_name', 'last_name')
 
     def __str__(self):
-        return "%s %s" % (self.first_name, self.last_name)
+        return ("%s %s" % (self.first_name, self.last_name)).strip()
diff --git a/tests/transactions/tests.py b/tests/transactions/tests.py
index a1edf53fcb..14252dd6dc 100644
--- a/tests/transactions/tests.py
+++ b/tests/transactions/tests.py
@@ -1,11 +1,163 @@
 from __future__ import absolute_import
 
+import sys
+
 from django.db import connection, transaction, IntegrityError
-from django.test import TransactionTestCase, skipUnlessDBFeature
+from django.test import TestCase, TransactionTestCase, skipUnlessDBFeature
+from django.utils import six
+from django.utils.unittest import skipUnless
 
 from .models import Reporter
 
 
+@skipUnless(connection.features.uses_savepoints,
+        "'atomic' requires transactions and savepoints.")
+class AtomicTests(TransactionTestCase):
+    """
+    Tests for the atomic decorator and context manager.
+
+    The tests make assertions on internal attributes because there isn't a
+    robust way to ask the database for its current transaction state.
+
+    Since the decorator syntax is converted into a context manager (see the
+    implementation), there are only a few basic tests with the decorator
+    syntax and the bulk of the tests use the context manager syntax.
+    """
+
+    def test_decorator_syntax_commit(self):
+        @transaction.atomic
+        def make_reporter():
+            Reporter.objects.create(first_name="Tintin")
+        make_reporter()
+        self.assertQuerysetEqual(Reporter.objects.all(), ['<Reporter: Tintin>'])
+
+    def test_decorator_syntax_rollback(self):
+        @transaction.atomic
+        def make_reporter():
+            Reporter.objects.create(first_name="Haddock")
+            raise Exception("Oops, that's his last name")
+        with six.assertRaisesRegex(self, Exception, "Oops"):
+            make_reporter()
+        self.assertQuerysetEqual(Reporter.objects.all(), [])
+
+    def test_alternate_decorator_syntax_commit(self):
+        @transaction.atomic()
+        def make_reporter():
+            Reporter.objects.create(first_name="Tintin")
+        make_reporter()
+        self.assertQuerysetEqual(Reporter.objects.all(), ['<Reporter: Tintin>'])
+
+    def test_alternate_decorator_syntax_rollback(self):
+        @transaction.atomic()
+        def make_reporter():
+            Reporter.objects.create(first_name="Haddock")
+            raise Exception("Oops, that's his last name")
+        with six.assertRaisesRegex(self, Exception, "Oops"):
+            make_reporter()
+        self.assertQuerysetEqual(Reporter.objects.all(), [])
+
+    def test_commit(self):
+        with transaction.atomic():
+            Reporter.objects.create(first_name="Tintin")
+        self.assertQuerysetEqual(Reporter.objects.all(), ['<Reporter: Tintin>'])
+
+    def test_rollback(self):
+        with six.assertRaisesRegex(self, Exception, "Oops"):
+            with transaction.atomic():
+                Reporter.objects.create(first_name="Haddock")
+                raise Exception("Oops, that's his last name")
+        self.assertQuerysetEqual(Reporter.objects.all(), [])
+
+    def test_nested_commit_commit(self):
+        with transaction.atomic():
+            Reporter.objects.create(first_name="Tintin")
+            with transaction.atomic():
+                Reporter.objects.create(first_name="Archibald", last_name="Haddock")
+        self.assertQuerysetEqual(Reporter.objects.all(),
+                ['<Reporter: Archibald Haddock>', '<Reporter: Tintin>'])
+
+    def test_nested_commit_rollback(self):
+        with transaction.atomic():
+            Reporter.objects.create(first_name="Tintin")
+            with six.assertRaisesRegex(self, Exception, "Oops"):
+                with transaction.atomic():
+                    Reporter.objects.create(first_name="Haddock")
+                    raise Exception("Oops, that's his last name")
+        self.assertQuerysetEqual(Reporter.objects.all(), ['<Reporter: Tintin>'])
+
+    def test_nested_rollback_commit(self):
+        with six.assertRaisesRegex(self, Exception, "Oops"):
+            with transaction.atomic():
+                Reporter.objects.create(last_name="Tintin")
+                with transaction.atomic():
+                    Reporter.objects.create(last_name="Haddock")
+                raise Exception("Oops, that's his first name")
+        self.assertQuerysetEqual(Reporter.objects.all(), [])
+
+    def test_nested_rollback_rollback(self):
+        with six.assertRaisesRegex(self, Exception, "Oops"):
+            with transaction.atomic():
+                Reporter.objects.create(last_name="Tintin")
+                with six.assertRaisesRegex(self, Exception, "Oops"):
+                    with transaction.atomic():
+                        Reporter.objects.create(first_name="Haddock")
+                    raise Exception("Oops, that's his last name")
+                raise Exception("Oops, that's his first name")
+        self.assertQuerysetEqual(Reporter.objects.all(), [])
+
+    def test_reuse_commit_commit(self):
+        atomic = transaction.atomic()
+        with atomic:
+            Reporter.objects.create(first_name="Tintin")
+            with atomic:
+                Reporter.objects.create(first_name="Archibald", last_name="Haddock")
+        self.assertQuerysetEqual(Reporter.objects.all(),
+                ['<Reporter: Archibald Haddock>', '<Reporter: Tintin>'])
+
+    def test_reuse_commit_rollback(self):
+        atomic = transaction.atomic()
+        with atomic:
+            Reporter.objects.create(first_name="Tintin")
+            with six.assertRaisesRegex(self, Exception, "Oops"):
+                with atomic:
+                    Reporter.objects.create(first_name="Haddock")
+                    raise Exception("Oops, that's his last name")
+        self.assertQuerysetEqual(Reporter.objects.all(), ['<Reporter: Tintin>'])
+
+    def test_reuse_rollback_commit(self):
+        atomic = transaction.atomic()
+        with six.assertRaisesRegex(self, Exception, "Oops"):
+            with atomic:
+                Reporter.objects.create(last_name="Tintin")
+                with atomic:
+                    Reporter.objects.create(last_name="Haddock")
+                raise Exception("Oops, that's his first name")
+        self.assertQuerysetEqual(Reporter.objects.all(), [])
+
+    def test_reuse_rollback_rollback(self):
+        atomic = transaction.atomic()
+        with six.assertRaisesRegex(self, Exception, "Oops"):
+            with atomic:
+                Reporter.objects.create(last_name="Tintin")
+                with six.assertRaisesRegex(self, Exception, "Oops"):
+                    with atomic:
+                        Reporter.objects.create(first_name="Haddock")
+                    raise Exception("Oops, that's his last name")
+                raise Exception("Oops, that's his first name")
+        self.assertQuerysetEqual(Reporter.objects.all(), [])
+
+
+class AtomicInsideTransactionTests(AtomicTests):
+    """All basic tests for atomic should also pass within an existing transaction."""
+
+    def setUp(self):
+        self.atomic = transaction.atomic()
+        self.atomic.__enter__()
+
+    def tearDown(self):
+        self.atomic.__exit__(*sys.exc_info())
+
+
 class TransactionTests(TransactionTestCase):
     def create_a_reporter_then_fail(self, first, last):
         a = Reporter(first_name=first, last_name=last)