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.