INTPYTHON-355 Add transaction.atomic()

Author: timgraham

django_mongodb_backend/base.py

@@ -1,8 +1,11 @@
 import contextlib
+import logging
 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 +35,9 @@ def __exit__(self, exception_type, exception_value, exception_traceback):
         pass
 
 
+logger = logging.getLogger("django.db.backends.base")
+
+
 class DatabaseWrapper(BaseDatabaseWrapper):
     data_types = {
         "AutoField": "int",
@@ -142,6 +148,17 @@ 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
+        # Tracks whether the connection is in a transaction managed by
+        # django_mongodb_backend.transaction.atomic. `in_atomic_block` isn't
+        # used in case Django's atomic() (used internally in Django) is called
+        # within this package's atomic().
+        self.in_atomic_block_mongo = False
+        # Current number of nested 'atomic' calls.
+        self.nested_atomics = 0
+
     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
@@ -230,3 +251,57 @@ def cursor(self):
     def get_database_version(self):
         """Return a tuple of the database's version."""
         return tuple(self.connection.server_info()["versionArray"])
+
+    ## Transaction API for django_mongodb_backend.transaction.atomic()
+    @async_unsafe
+    def start_transaction_mongo(self):
+        if self.session is None:
+            self.session = self.connection.start_session()
+            with debug_transaction(self, "session.start_transaction()"):
+                self.session.start_transaction()
+
+    @async_unsafe
+    def commit_mongo(self):
+        if self.session:
+            with debug_transaction(self, "session.commit_transaction()"):
+                self.session.commit_transaction()
+            self._end_session()
+        self.run_and_clear_commit_hooks()
+
+    @async_unsafe
+    def rollback_mongo(self):
+        if self.session:
+            with debug_transaction(self, "session.abort_transaction()"):
+                self.session.abort_transaction()
+            self._end_session()
+        self.run_on_commit = []
+
+    def _end_session(self):
+        self.session.end_session()
+        self.session = None
+
+    def on_commit(self, func, robust=False):
+        """
+        Copied from BaseDatabaseWrapper.on_commit() except that it checks
+        in_atomic_block_mongo instead of in_atomic_block.
+        """
+        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.
+            # The first item in the tuple (an empty list) is normally the
+            # savepoint IDs, which isn't applicable on MongoDB.
+            self.run_on_commit.append(([], func, robust))
+        else:
+            # No transaction in progress; 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/compiler.py

@@ -696,7 +696,9 @@ 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
+        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
@@ -777,7 +779,9 @@ def execute_sql(self, result_type):
 
     @wrap_database_errors
     def update(self, criteria, pipeline):
-        return self.collection.update_many(criteria, pipeline).matched_count
+        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

@@ -63,15 +63,19 @@ def delete(self):
         """Execute a delete query."""
         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())
+        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

@@ -0,0 +1,61 @@
+from contextlib import ContextDecorator
+
+from django.db import DEFAULT_DB_ALIAS, DatabaseError
+from django.db.transaction import get_connection, on_commit
+
+__all__ = [
+    "atomic",
+    "on_commit",  # convenience alias
+]
+
+
+class Atomic(ContextDecorator):
+    """
+    Guarantee the atomic execution of a given block.
+
+    Simplified from django.db.transaction.
+    """
+
+    def __init__(self, using):
+        self.using = using
+
+    def __enter__(self):
+        connection = get_connection(self.using)
+        if connection.in_atomic_block_mongo:
+            # Track the number of nested atomic() calls.
+            connection.nested_atomics += 1
+        else:
+            # Start a transaction for the outermost atomic().
+            connection.start_transaction_mongo()
+            connection.in_atomic_block_mongo = True
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        connection = get_connection(self.using)
+        if connection.nested_atomics:
+            # Exiting inner atomic.
+            connection.nested_atomics -= 1
+        else:
+            # Reset flag when exiting outer atomic.
+            connection.in_atomic_block_mongo = False
+        if exc_type is None:
+            # atomic() exited without an error.
+            if not connection.in_atomic_block_mongo:
+                # Commit transaction if outer atomic().
+                try:
+                    connection.commit_mongo()
+                except DatabaseError:
+                    connection.rollback_mongo()
+        else:
+            # atomic() exited with an error.
+            if not connection.in_atomic_block_mongo:
+                # Rollback transaction if outer atomic().
+                connection.rollback_mongo()
+
+
+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(...): ...
+    return Atomic(using)

docs/source/releases/5.2.x.rst

@@ -13,6 +13,7 @@ New features
 - Added subquery support for :class:`~.fields.EmbeddedModelArrayField`.
 - Added the ``options`` parameter to
   :func:`~django_mongodb_backend.utils.parse_uri`.
+- Added support for :ref:`database transactions <transactions>`.
 - Added :class:`~.fields.PolymorphicEmbeddedModelField` and
   :class:`~.fields.PolymorphicEmbeddedModelArrayField` for storing a model
   instance or list of model instances that may be of more than one model class.

docs/source/topics/index.rst

@@ -9,4 +9,5 @@ know:
    :maxdepth: 2
 
    embedded-models
+   transactions
    known-issues

docs/source/topics/known-issues.rst

@@ -80,11 +80,12 @@ Database functions
 Transaction management
 ======================
 
-Query execution uses Django and MongoDB's default behavior of autocommit mode.
-Each query is immediately committed to the database.
+By default, query execution uses Django and MongoDB's default behavior of autocommit
+mode. Each query is immediately committed to the database.
 
 Django's :doc:`transaction management APIs <django:topics/db/transactions>`
-are not supported.
+are not supported. Instead, this package provides its own :doc:`transaction APIs
+</topics/transactions>`.
 
 Database introspection
 ======================

docs/source/topics/transactions.rst

@@ -0,0 +1,142 @@
+============
+Transactions
+============
+
+.. versionadded:: 5.2.0b2
+
+.. module:: django_mongodb_backend.transaction
+
+MongoDB supports :doc:`transactions <manual:core/transactions>` if it's
+configured as a :doc:`replica set <manual:replication>` or a :doc:`sharded
+cluster <manual:sharding>`.
+
+Because MongoDB transactions have some limitations and are not meant to be used
+as freely as SQL transactions, :doc:`Django's transactions APIs
+<django:topics/db/transactions>`, including most notably
+:func:`django.db.transaction.atomic`, function as no-ops.
+
+Instead, Django MongoDB Backend provides its own
+:func:`django_mongodb_backend.transaction.atomic` function.
+
+Outside of a transaction, query execution uses Django and MongoDB's default
+behavior of autocommit mode. Each query is immediately committed to the
+database.
+
+Controlling transactions
+========================
+
+.. function:: atomic(using=None)
+
+    Atomicity is the defining property of database transactions. ``atomic``
+    allows creating a block of code within which the atomicity on the database
+    is guaranteed. If the block of code is successfully completed, the changes
+    are committed to the database. If there is an exception, the changes are
+    rolled back.
+
+    ``atomic`` is usable both as a :py:term:`decorator`::
+
+        from django_mongodb_backend import transaction
+
+
+        @transaction.atomic
+        def viewfunc(request):
+            # This code executes inside a transaction.
+            do_stuff()
+
+    and as a :py:term:`context manager`::
+
+        from django_mongodb_backend 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()
+
+    .. admonition:: Avoid catching exceptions inside ``atomic``!
+
+        When exiting an ``atomic`` block, Django looks at whether it's exited
+        normally or with an exception to determine whether to commit or roll
+        back. If you catch and handle exceptions inside an ``atomic`` block,
+        you may hide from Django the fact that a problem has happened. This can
+        result in unexpected behavior.
+
+        This is mostly a concern for :exc:`~django.db.DatabaseError` and its
+        subclasses such as :exc:`~django.db.IntegrityError`. After such an
+        error, the transaction is broken and Django will perform a rollback at
+        the end of the ``atomic`` block.
+
+    .. admonition:: You may need to manually revert app state when rolling back a transaction.
+
+        The values of a model's fields won't be reverted when a transaction
+        rollback happens. This could lead to an inconsistent model state unless
+        you manually restore the original field values.
+
+        For example, given ``MyModel`` with an ``active`` field, this snippet
+        ensures that the ``if obj.active`` check at the end uses the correct
+        value if updating ``active`` to ``True`` fails in the transaction::
+
+            from django_mongodb_backend import transaction
+            from django.db import DatabaseError
+
+            obj = MyModel(active=False)
+            obj.active = True
+            try:
+                with transaction.atomic():
+                    obj.save()
+            except DatabaseError:
+                obj.active = False
+
+            if obj.active:
+                ...
+
+        This also applies to any other mechanism that may hold app state, such
+        as caching or global variables. For example, if the code proactively
+        updates data in the cache after saving an object, it's recommended to
+        use :ref:`transaction.on_commit() <performing-actions-after-commit>`
+        instead, to defer cache alterations until the transaction is actually
+        committed.
+
+    ``atomic`` takes a ``using`` argument which should be the name of a
+    database. If this argument isn't provided, Django uses the ``"default"``
+    database.
+
+.. admonition:: Performance considerations
+
+    Open transactions have a performance cost for your MongoDB server. To
+    minimize this overhead, keep your transactions as short as possible. This
+    is especially important if you're using :func:`atomic` in long-running
+    processes, outside of Django's request / response cycle.
+
+Performing actions after commit
+===============================
+
+The :func:`atomic` function supports Django's
+:func:`~django.db.transaction.on_commit` API to :ref:`perform actions after a
+transaction successfully commits <performing-actions-after-commit>`.
+
+For convenience, :func:`~django.db.transaction.on_commit` is aliased at
+``django_mongodb_backend.transaction.on_commit`` so you can use both::
+
+    from django_mongodb_backend import transaction
+
+
+    transaction.atomic()
+    transaction.on_commit(...)
+
+.. _transactions-limitations:
+
+Limitations
+===========
+
+MongoDB's transaction limitations that are applicable to Django are:
+
+- :meth:`QuerySet.union() <django.db.models.query.QuerySet.union>` is not
+  supported inside a transaction.
+- Savepoints (i.e. nested :func:`~django.db.transaction.atomic` blocks) aren't
+  supported. The outermost :func:`~django.db.transaction.atomic` will start
+  a transaction while any inner :func:`~django.db.transaction.atomic` blocks
+  have no effect.