add transaction hooks

Author: timgraham

django_mongodb_backend/base.py

@@ -1,4 +1,5 @@
 import contextlib
+import logging
 import os
 
 from django.core.exceptions import ImproperlyConfigured
@@ -46,6 +47,9 @@ def wrapper(self, *args, **kwargs):
     return wrapper
 
 
+logger = logging.getLogger("django.db.backends.base")
+
+
 class DatabaseWrapper(BaseDatabaseWrapper):
     data_types = {
         "AutoField": "int",
@@ -319,3 +323,24 @@ def _end_session(self):
         # Private API, specific to this backend.
         self.session.end_session()
         self.session = None
+
+    def on_commit(self, func, robust=False):
+        if not callable(func):
+            raise TypeError("on_commit()'s callback must be a callable.")
+        if self.in_atomic_block_mongo:
+            # Transaction in progress; save for execution on commit.
+            self.run_on_commit.append((set(self.savepoint_ids), func, robust))
+        else:
+            # No transaction in progress and in autocommit mode; execute
+            # immediately.
+            if robust:
+                try:
+                    func()
+                except Exception as e:
+                    logger.exception(
+                        "Error calling %s in on_commit() (%s).",
+                        func.__qualname__,
+                        e,
+                    )
+            else:
+                func()

django_mongodb_backend/transaction.py

@@ -115,7 +115,8 @@ def __exit__(self, exc_type, exc_value, traceback):
         finally:
             # Outermost block exit when autocommit was enabled.
             if not connection.in_atomic_block_mongo:
-                pass
+                if connection.run_commit_hooks_on_set_autocommit_on:
+                    connection.run_and_clear_commit_hooks()
             # connection.set_autocommit(True)
             # Outermost block exit when autocommit was disabled.
             elif not connection.commit_on_exit:

tests/transaction_hooks_/__init__.py

@@ -0,0 +1,5 @@
+from django.db import models
+
+
+class Thing(models.Model):
+    num = models.IntegerField()

tests/transaction_hooks_/models.py

@@ -0,0 +1,214 @@
+from django.db import connection
+from django.test import TransactionTestCase, skipUnlessDBFeature
+
+from django_mongodb_backend import transaction
+
+from .models import Thing
+
+
+class ForcedError(Exception):
+    pass
+
+
+@skipUnlessDBFeature("_supports_transactions")
+class TestConnectionOnCommit(TransactionTestCase):
+    """
+    Tests for transaction.on_commit().
+
+    Creation/checking of database objects in parallel with callback tracking is
+    to verify that the behavior of the two match in all tested cases.
+    """
+
+    available_apps = ["transaction_hooks_"]
+
+    def setUp(self):
+        self.notified = []
+
+    def notify(self, id_):
+        if id_ == "error":
+            raise ForcedError()
+        self.notified.append(id_)
+
+    def do(self, num):
+        """Create a Thing instance and notify about it."""
+        Thing.objects.create(num=num)
+        transaction.on_commit(lambda: self.notify(num))
+
+    def assertDone(self, nums):
+        self.assertNotified(nums)
+        self.assertEqual(sorted(t.num for t in Thing.objects.all()), sorted(nums))
+
+    def assertNotified(self, nums):
+        self.assertEqual(self.notified, nums)
+
+    def test_executes_immediately_if_no_transaction(self):
+        self.do(1)
+        self.assertDone([1])
+
+    def test_robust_if_no_transaction(self):
+        def robust_callback():
+            raise ForcedError("robust callback")
+
+        with self.assertLogs("django.db.backends.base", "ERROR") as cm:
+            transaction.on_commit(robust_callback, robust=True)
+            self.do(1)
+
+        self.assertDone([1])
+        log_record = cm.records[0]
+        self.assertEqual(
+            log_record.getMessage(),
+            "Error calling TestConnectionOnCommit.test_robust_if_no_transaction."
+            "<locals>.robust_callback in on_commit() (robust callback).",
+        )
+        self.assertIsNotNone(log_record.exc_info)
+        raised_exception = log_record.exc_info[1]
+        self.assertIsInstance(raised_exception, ForcedError)
+        self.assertEqual(str(raised_exception), "robust callback")
+
+    def test_robust_transaction(self):
+        def robust_callback():
+            raise ForcedError("robust callback")
+
+        with self.assertLogs("django.db.backends", "ERROR") as cm, transaction.atomic():
+            transaction.on_commit(robust_callback, robust=True)
+            self.do(1)
+
+        self.assertDone([1])
+        log_record = cm.records[0]
+        self.assertEqual(
+            log_record.getMessage(),
+            "Error calling TestConnectionOnCommit.test_robust_transaction.<locals>."
+            "robust_callback in on_commit() during transaction (robust callback).",
+        )
+        self.assertIsNotNone(log_record.exc_info)
+        raised_exception = log_record.exc_info[1]
+        self.assertIsInstance(raised_exception, ForcedError)
+        self.assertEqual(str(raised_exception), "robust callback")
+
+    def test_delays_execution_until_after_transaction_commit(self):
+        with transaction.atomic():
+            self.do(1)
+            self.assertNotified([])
+        self.assertDone([1])
+
+    def test_does_not_execute_if_transaction_rolled_back(self):
+        try:
+            with transaction.atomic():
+                self.do(1)
+                raise ForcedError()
+        except ForcedError:
+            pass
+
+        self.assertDone([])
+
+    def test_executes_only_after_final_transaction_committed(self):
+        with transaction.atomic():
+            with transaction.atomic():
+                self.do(1)
+                self.assertNotified([])
+            self.assertNotified([])
+        self.assertDone([1])
+
+    def test_no_hooks_run_from_failed_transaction(self):
+        """If outer transaction fails, no hooks from within it run."""
+        try:
+            with transaction.atomic():
+                with transaction.atomic():
+                    self.do(1)
+                raise ForcedError()
+        except ForcedError:
+            pass
+
+        self.assertDone([])
+
+    def test_runs_hooks_in_order_registered(self):
+        with transaction.atomic():
+            self.do(1)
+            with transaction.atomic():
+                self.do(2)
+            self.do(3)
+
+        self.assertDone([1, 2, 3])
+
+    def test_hooks_cleared_after_successful_commit(self):
+        with transaction.atomic():
+            self.do(1)
+        with transaction.atomic():
+            self.do(2)
+
+        self.assertDone([1, 2])  # not [1, 1, 2]
+
+    def test_hooks_cleared_after_rollback(self):
+        try:
+            with transaction.atomic():
+                self.do(1)
+                raise ForcedError()
+        except ForcedError:
+            pass
+
+        with transaction.atomic():
+            self.do(2)
+
+        self.assertDone([2])
+
+    @skipUnlessDBFeature("test_db_allows_multiple_connections")
+    def test_hooks_cleared_on_reconnect(self):
+        with transaction.atomic():
+            self.do(1)
+            connection.close_pool()
+
+        connection.connect()
+
+        with transaction.atomic():
+            self.do(2)
+
+        self.assertDone([2])
+
+    def test_error_in_hook_doesnt_prevent_clearing_hooks(self):
+        try:
+            with transaction.atomic():
+                transaction.on_commit(lambda: self.notify("error"))
+        except ForcedError:
+            pass
+
+        with transaction.atomic():
+            self.do(1)
+
+        self.assertDone([1])
+
+    def test_db_query_in_hook(self):
+        with transaction.atomic():
+            Thing.objects.create(num=1)
+            transaction.on_commit(lambda: [self.notify(t.num) for t in Thing.objects.all()])
+
+        self.assertDone([1])
+
+    def test_transaction_in_hook(self):
+        def on_commit():
+            with transaction.atomic():
+                t = Thing.objects.create(num=1)
+                self.notify(t.num)
+
+        with transaction.atomic():
+            transaction.on_commit(on_commit)
+
+        self.assertDone([1])
+
+    def test_hook_in_hook(self):
+        def on_commit(i, add_hook):
+            with transaction.atomic():
+                if add_hook:
+                    transaction.on_commit(lambda: on_commit(i + 10, False))
+                t = Thing.objects.create(num=i)
+                self.notify(t.num)
+
+        with transaction.atomic():
+            transaction.on_commit(lambda: on_commit(1, True))
+            transaction.on_commit(lambda: on_commit(2, True))
+
+        self.assertDone([1, 11, 2, 12])
+
+    def test_raises_exception_non_callable(self):
+        msg = "on_commit()'s callback must be a callable."
+        with self.assertRaisesMessage(TypeError, msg):
+            transaction.on_commit(None)

tests/transaction_hooks_/tests.py

@@ -1,12 +1,3 @@
-"""
-Transactions
-
-Django handles transactions in three different ways. The default is to commit
-each transaction upon a write, but you can decorate a function to get
-commit-on-success behavior. Alternatively, you can manage the transaction
-manually.
-"""
-
 from django.db import models