progress

Author: timgraham

django_mongodb_backend/base.py

@@ -2,7 +2,9 @@
 import os
 
 from django.core.exceptions import ImproperlyConfigured
+from django.db import DEFAULT_DB_ALIAS
 from django.db.backends.base.base import BaseDatabaseWrapper
+from django.db.backends.utils import debug_transaction
 from django.utils.asyncio import async_unsafe
 from django.utils.functional import cached_property
 from pymongo.collection import Collection
@@ -32,6 +34,17 @@ def __exit__(self, exception_type, exception_value, exception_traceback):
         pass
 
 
+def requires_transaction_support(func):
+    """Make a method a no-op if transactions aren't supported."""
+
+    def wrapper(self, *args, **kwargs):
+        if not self.features._supports_transactions:
+            return
+        func(self, *args, **kwargs)
+
+    return wrapper
+
+
 class DatabaseWrapper(BaseDatabaseWrapper):
     data_types = {
         "AutoField": "int",
@@ -142,6 +155,10 @@ def _isnull_operator(a, b):
     ops_class = DatabaseOperations
     validation_class = DatabaseValidation
 
+    def __init__(self, settings_dict, alias=DEFAULT_DB_ALIAS):
+        super().__init__(settings_dict, alias=alias)
+        self.session = None
+
     def get_collection(self, name, **kwargs):
         collection = Collection(self.database, name, **kwargs)
         if self.queries_logged:
@@ -212,6 +229,10 @@ def close(self):
 
     def close_pool(self):
         """Close the MongoClient."""
+        # Clear commit hooks and session.
+        self.run_on_commit = []
+        if self.session:
+            self._end_session()
         connection = self.connection
         if connection is None:
             return
@@ -227,6 +248,56 @@ def close_pool(self):
     def cursor(self):
         return Cursor()
 
+    @requires_transaction_support
+    def validate_no_broken_transaction(self):
+        super().validate_no_broken_transaction()
+
     def get_database_version(self):
         """Return a tuple of the database's version."""
         return tuple(self.connection.server_info()["versionArray"])
+
+    @requires_transaction_support
+    def _start_transaction(self, autocommit, force_begin_transaction_with_broken_autocommit=False):
+        # Besides @transaction.atomic() (which uses
+        # _start_transaction_under_autocommit(), disabling autocommit is
+        # another way to start a transaction.
+        #        if not autocommit:
+        #        self._start_transaction()
+        #   def _start_transaction(self):
+        # Private API, specific to this backend.
+        if self.session is None:
+            self.session = self.connection.start_session()
+            with debug_transaction(self, "session.start_transaction()"):
+                self.session.start_transaction()
+
+    @requires_transaction_support
+    def _commit_transaction(self):
+        self.validate_thread_sharing()
+        self.validate_no_atomic_block()
+        if self.session:
+            with debug_transaction(self, "session.commit_transaction()"):
+                self.session.commit_transaction()
+            self._end_session()
+        # A successful commit means that the database connection works.
+        self.errors_occurred = False
+        self.run_commit_hooks_on_set_autocommit_on = True
+
+    @async_unsafe
+    @requires_transaction_support
+    def _rollback_transaction(self):
+        """Roll back a MongoDB transaction and reset the dirty flag."""
+        self.validate_thread_sharing()
+        self.validate_no_atomic_block()
+        if self.session:
+            with debug_transaction(self, "session.abort_transaction()"):
+                self.session.abort_transaction()
+            self._end_session()
+        # A successful rollback means that the database connection works.
+        self.errors_occurred = False
+        self.needs_rollback = False
+        self.run_on_commit = []
+
+    def _end_session(self):
+        # Private API, specific to this backend.
+        self.session.end_session()
+        self.session = None

django_mongodb_backend/compiler.py

@@ -678,7 +678,10 @@ def execute_sql(self, returning_fields=None):
     @wrap_database_errors
     def insert(self, docs, returning_fields=None):
         """Store a list of documents using field columns as element names."""
-        inserted_ids = self.collection.insert_many(docs).inserted_ids
+        self.connection.validate_no_broken_transaction()
+        inserted_ids = self.collection.insert_many(
+            docs, session=self.connection.session
+        ).inserted_ids
         return [(x,) for x in inserted_ids] if returning_fields else []
 
     @cached_property
@@ -759,7 +762,10 @@ def execute_sql(self, result_type):
 
     @wrap_database_errors
     def update(self, criteria, pipeline):
-        return self.collection.update_many(criteria, pipeline).matched_count
+        self.connection.validate_no_broken_transaction()
+        return self.collection.update_many(
+            criteria, pipeline, session=self.connection.session
+        ).matched_count
 
     def check_query(self):
         super().check_query()

django_mongodb_backend/query.py

@@ -61,17 +61,23 @@ def __repr__(self):
     @wrap_database_errors
     def delete(self):
         """Execute a delete query."""
+        self.compiler.connection.validate_no_broken_transaction()
         if self.compiler.subqueries:
             raise NotSupportedError("Cannot use QuerySet.delete() when a subquery is required.")
-        return self.compiler.collection.delete_many(self.match_mql).deleted_count
+        return self.compiler.collection.delete_many(
+            self.match_mql, session=self.compiler.connection.session
+        ).deleted_count
 
     @wrap_database_errors
     def get_cursor(self):
         """
         Return a pymongo CommandCursor that can be iterated on to give the
         results of the query.
         """
-        return self.compiler.collection.aggregate(self.get_pipeline())
+        self.compiler.connection.validate_no_broken_transaction()
+        return self.compiler.collection.aggregate(
+            self.get_pipeline(), session=self.compiler.connection.session
+        )
 
     def get_pipeline(self):
         pipeline = []

django_mongodb_backend/queryset.py

@@ -35,7 +35,7 @@ def __init__(self, pipeline, using, model):
     def _execute_query(self):
         connection = connections[self.using]
         collection = connection.get_collection(self.model._meta.db_table)
-        self.cursor = collection.aggregate(self.pipeline)
+        self.cursor = collection.aggregate(self.pipeline, session=connection.session)
 
     def __str__(self):
         return str(self.pipeline)

django_mongodb_backend/transaction.py

@@ -23,16 +23,6 @@ def get_connection(using=None):
     return connections[using]
 
 
-def get_autocommit(using=None):
-    """Get the autocommit status of the connection."""
-    return get_connection(using).get_autocommit()
-
-
-def set_autocommit(autocommit, using=None):
-    """Set the autocommit status of the connection."""
-    return get_connection(using).set_autocommit(autocommit)
-
-
 def commit(using=None):
     """Commit a transaction."""
     get_connection(using).commit()
@@ -141,7 +131,7 @@ def __enter__(self):
             and not connection.atomic_blocks[-1]._from_testcase
         ):
             raise RuntimeError(
-                "A durable atomic block cannot be nested within another " "atomic block."
+                "A durable atomic block cannot be nested within another atomic block."
             )
         if not connection.in_atomic_block:
             # Reset state when entering an outermost atomic block.
@@ -158,7 +148,9 @@ def __enter__(self):
             # We're already in a transaction
             pass
         else:
-            connection.set_autocommit(False, force_begin_transaction_with_broken_autocommit=True)
+            connection._start_transaction(
+                False, force_begin_transaction_with_broken_autocommit=True
+            )
             connection.in_atomic_block = True
 
         if connection.in_atomic_block:
@@ -173,22 +165,22 @@ def __exit__(self, exc_type, exc_value, traceback):
         # Prematurely unset this flag to allow using commit or rollback.
         connection._in_atomic_block = False
         try:
-            if connection._closed_in_transaction:
+            if connection.closed_in_transaction:
                 # The database will perform a rollback by itself.
                 # Wait until we exit the outermost block.
                 pass
 
-            elif exc_type is None and not connection._needs_rollback:
+            elif exc_type is None and not connection.needs_rollback:
                 if connection._in_atomic_block:
                     # Release savepoint if there is one
                     pass
                 else:
                     # Commit transaction
                     try:
-                        connection._commit()
+                        connection._commit_transaction()
                     except DatabaseError:
                         try:
-                            connection._rollback()
+                            connection._rollback_transaction()
                         except Error:
                             # An error during rollback means that something
                             # went wrong with the connection. Drop it.
@@ -204,7 +196,7 @@ def __exit__(self, exc_type, exc_value, traceback):
                 else:
                     # Roll back transaction
                     try:
-                        connection.rollback()
+                        connection._rollback_transaction()
                     except Error:
                         # An error during rollback means that something
                         # went wrong with the connection. Drop it.
@@ -214,8 +206,8 @@ def __exit__(self, exc_type, exc_value, traceback):
             if not connection.in_atomic_block:
                 if connection.closed_in_transaction:
                     connection.connection = None
-                else:
-                    connection.set_autocommit(True)
+                # else:
+                #     connection.set_autocommit(True)
             # Outermost block exit when autocommit was disabled.
             elif not connection.commit_on_exit:
                 if connection.closed_in_transaction:

tests/transactions_/tests.py

@@ -103,40 +103,6 @@ def test_nested_rollback_rollback(self):
             raise Exception("Oops, that's his first name")
         self.assertSequenceEqual(Reporter.objects.all(), [])
 
-    def test_merged_commit_commit(self):
-        with transaction.atomic():
-            reporter1 = Reporter.objects.create(first_name="Tintin")
-            with transaction.atomic(savepoint=False):
-                reporter2 = Reporter.objects.create(first_name="Archibald", last_name="Haddock")
-        self.assertSequenceEqual(Reporter.objects.all(), [reporter2, reporter1])
-
-    def test_merged_commit_rollback(self):
-        with transaction.atomic():
-            Reporter.objects.create(first_name="Tintin")
-            with self.assertRaisesMessage(Exception, "Oops"), transaction.atomic(savepoint=False):
-                Reporter.objects.create(first_name="Haddock")
-                raise Exception("Oops, that's his last name")
-        # Writes in the outer block are rolled back too.
-        self.assertSequenceEqual(Reporter.objects.all(), [])
-
-    def test_merged_rollback_commit(self):
-        with self.assertRaisesMessage(Exception, "Oops"), transaction.atomic():
-            Reporter.objects.create(last_name="Tintin")
-            with transaction.atomic(savepoint=False):
-                Reporter.objects.create(last_name="Haddock")
-            raise Exception("Oops, that's his first name")
-        self.assertSequenceEqual(Reporter.objects.all(), [])
-
-    def test_merged_rollback_rollback(self):
-        with self.assertRaisesMessage(Exception, "Oops"), transaction.atomic():
-            Reporter.objects.create(last_name="Tintin")
-            with self.assertRaisesMessage(Exception, "Oops"):
-                with transaction.atomic(savepoint=False):
-                    Reporter.objects.create(first_name="Haddock")
-                raise Exception("Oops, that's his last name")
-            raise Exception("Oops, that's his first name")
-        self.assertSequenceEqual(Reporter.objects.all(), [])
-
     def test_reuse_commit_commit(self):
         atomic = transaction.atomic()
         with atomic:
@@ -212,8 +178,11 @@ class AtomicErrorsTests(TransactionTestCase):
 
     def test_atomic_prevents_setting_autocommit(self):
         autocommit = transaction.get_autocommit()
-        with transaction.atomic(), self.assertRaisesMessage(
-            transaction.TransactionManagementError, self.forbidden_atomic_msg
+        with (
+            transaction.atomic(),
+            self.assertRaisesMessage(
+                transaction.TransactionManagementError, self.forbidden_atomic_msg
+            ),
         ):
             transaction.set_autocommit(not autocommit)
         # Make sure autocommit wasn't changed.