Fixed #21803 -- Added support for post-commit callbacks

Made it possible to register and run callbacks after a database
transaction is committed with the `transaction.on_commit()` function.

This patch is heavily based on Carl Meyers django-transaction-hooks
<https://django-transaction-hooks.readthedocs.org/>. Thanks to
Aymeric Augustin, Carl Meyer, and Tim Graham for review and feedback.

2 files changed, 75 insertions, 15 deletions

diff --git a/django/db/backends/base/base.py b/django/db/backends/base/base.py
index 2390cd7940..96ca0ccfde 100644
--- a/django/db/backends/base/base.py
+++ b/django/db/backends/base/base.py
@@ -78,6 +78,15 @@ class BaseDatabaseWrapper(object):
         self.allow_thread_sharing = allow_thread_sharing
         self._thread_ident = thread.get_ident()
 
+        # A list of no-argument functions to run when the transaction commits.
+        # Each entry is an (sids, func) tuple, where sids is a set of the
+        # active savepoint IDs when this function was registered.
+        self.run_on_commit = []
+
+        # Should we run the on-commit hooks the next time set_autocommit(True)
+        # is called?
+        self.run_commit_hooks_on_set_autocommit_on = False
+
     @cached_property
     def timezone(self):
         """
@@ -163,6 +172,8 @@ class BaseDatabaseWrapper(object):
         self.init_connection_state()
         connection_created.send(sender=self.__class__, connection=self)
 
+        self.run_on_commit = []
+
     def check_settings(self):
         if self.settings_dict['TIME_ZONE'] is not None:
             if not settings.USE_TZ:
@@ -230,6 +241,7 @@ class BaseDatabaseWrapper(object):
         self._commit()
         # A successful commit means that the database connection works.
         self.errors_occurred = False
+        self.run_commit_hooks_on_set_autocommit_on = True
 
     def rollback(self):
         """
@@ -241,11 +253,15 @@ class BaseDatabaseWrapper(object):
         # A successful rollback means that the database connection works.
         self.errors_occurred = False
 
+        self.run_on_commit = []
+
     def close(self):
         """
         Closes the connection to the database.
         """
         self.validate_thread_sharing()
+        self.run_on_commit = []
+
         # Don't call validate_no_atomic_block() to avoid making it difficult
         # to get rid of a connection in an invalid state. The next connect()
         # will reset the transaction state anyway.
@@ -310,6 +326,11 @@ class BaseDatabaseWrapper(object):
         self.validate_thread_sharing()
         self._savepoint_rollback(sid)
 
+        # Remove any callbacks registered while this savepoint was active.
+        self.run_on_commit = [
+            (sids, func) for (sids, func) in self.run_on_commit if sid not in sids
+        ]
+
     def savepoint_commit(self, sid):
         """
         Releases a savepoint. Does nothing if savepoints are not supported.
@@ -343,15 +364,38 @@ class BaseDatabaseWrapper(object):
         self.ensure_connection()
         return self.autocommit
 
-    def set_autocommit(self, autocommit):
+    def set_autocommit(self, autocommit, force_begin_transaction_with_broken_autocommit=False):
         """
         Enable or disable autocommit.
+
+        The usual way to start a transaction is to turn autocommit off.
+        SQLite does not properly start a transaction when disabling
+        autocommit. To avoid this buggy behavior and to actually enter a new
+        transaction, an explcit BEGIN is required. Using
+        force_begin_transaction_with_broken_autocommit=True will issue an
+        explicit BEGIN with SQLite. This option will be ignored for other
+        backends.
         """
         self.validate_no_atomic_block()
         self.ensure_connection()
-        self._set_autocommit(autocommit)
+
+        start_transaction_under_autocommit = (
+            force_begin_transaction_with_broken_autocommit
+            and not autocommit
+            and self.features.autocommits_when_autocommit_is_off
+        )
+
+        if start_transaction_under_autocommit:
+            self._start_transaction_under_autocommit()
+        else:
+            self._set_autocommit(autocommit)
+
         self.autocommit = autocommit
 
+        if autocommit and self.run_commit_hooks_on_set_autocommit_on:
+            self.run_and_clear_commit_hooks()
+            self.run_commit_hooks_on_set_autocommit_on = False
+
     def get_rollback(self):
         """
         Get the "needs rollback" flag -- for *advanced use* only.
@@ -558,3 +602,23 @@ class BaseDatabaseWrapper(object):
             raise NotImplementedError(
                 'The SchemaEditorClass attribute of this database wrapper is still None')
         return self.SchemaEditorClass(self, *args, **kwargs)
+
+    def on_commit(self, func):
+        if self.in_atomic_block:
+            # Transaction in progress; save for execution on commit.
+            self.run_on_commit.append((set(self.savepoint_ids), func))
+        elif not self.get_autocommit():
+            raise TransactionManagementError('on_commit() cannot be used in manual transaction management')
+        else:
+            # No transaction in progress and in autocommit mode; execute
+            # immediately.
+            func()
+
+    def run_and_clear_commit_hooks(self):
+        self.validate_no_atomic_block()
+        try:
+            while self.run_on_commit:
+                sids, func = self.run_on_commit.pop(0)
+                func()
+        finally:
+            self.run_on_commit = []
diff --git a/django/db/transaction.py b/django/db/transaction.py
index d1388675d5..6c174bf2d3 100644
--- a/django/db/transaction.py
+++ b/django/db/transaction.py
@@ -103,6 +103,14 @@ def set_rollback(rollback, using=None):
     return get_connection(using).set_rollback(rollback)
 
 
+def on_commit(func, using=None):
+    """
+    Register `func` to be called when the current transaction is committed.
+    If the current transaction is rolled back, `func` will not be called.
+    """
+    get_connection(using).on_commit(func)
+
+
 #################################
 # Decorators / context managers #
 #################################
@@ -180,17 +188,7 @@ class Atomic(ContextDecorator):
             else:
                 connection.savepoint_ids.append(None)
         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.set_autocommit(False, force_begin_transaction_with_broken_autocommit=True)
             connection.in_atomic_block = True
 
     def __exit__(self, exc_type, exc_value, traceback):
@@ -272,8 +270,6 @@ class Atomic(ContextDecorator):
             if not connection.in_atomic_block:
                 if connection.closed_in_transaction:
                     connection.connection = None
-                elif connection.features.autocommits_when_autocommit_is_off:
-                    connection.autocommit = True
                 else:
                     connection.set_autocommit(True)
             # Outermost block exit when autocommit was disabled.