Transactions via an internal global session var

.github/workflows/start_mongo.sh

@@ -5,5 +5,6 @@ MONGODB=$1
 mongodb_dir=$(find ${PWD}/ -type d -name "mongodb-linux-x86_64*")
 
 mkdir $mongodb_dir/data
-$mongodb_dir/bin/mongod --dbpath $mongodb_dir/data --logpath $mongodb_dir/mongodb.log --fork
+$mongodb_dir/bin/mongod --dbpath $mongodb_dir/data --logpath $mongodb_dir/mongodb.log --fork --replSet mongoengine
 mongo --eval 'db.version();'    # Make sure mongo is awake
+mongo --eval "rs.initiate()"

AUTHORS

@@ -264,3 +264,4 @@ that much better:
  * oleksandr-l5 (https://github.com/oleksandr-l5)
  * Ido Shraga (https://github.com/idoshr)
  * Terence Honles (https://github.com/terencehonles)
+ * Juan Gutierrez (https://github.com/juannyg)

mongoengine/connection.py

@@ -1,3 +1,5 @@
+import collections
+import threading
 import warnings
 
 from pymongo import MongoClient, ReadPreference, uri_parser
@@ -26,6 +28,7 @@
     "get_connection",
     "get_db",
     "register_connection",
+    "run_in_transaction",
 ]
 
 
@@ -38,6 +41,7 @@
 _connections = {}
 _dbs = {}
 
+
 READ_PREFERENCE = ReadPreference.PRIMARY
 
 
@@ -463,3 +467,89 @@ def connect(db=None, alias=DEFAULT_CONNECTION_NAME, **kwargs):
 # Support old naming convention
 _get_connection = get_connection
 _get_db = get_db
+
+
+class _LocalSessions(threading.local):
+    def __init__(self):
+        self.sessions = collections.deque()
+
+    def append(self, session):
+        self.sessions.append(session)
+
+    def get_current(self):
+        if len(self.sessions):
+            return self.sessions[len(self.sessions) - 1]
+
+    def clear_current(self):
+        if len(self.sessions):
+            self.sessions.pop()
+
+    def clear_all(self):
+        self.sessions.clear()
+
+
+_local_sessions = _LocalSessions()
+
+
+def _set_session(session):
+    _local_sessions.append(session)
+
+
+def _get_session():
+    return _local_sessions.get_current()
+
+
+def _clear_session():
+    return _local_sessions.clear_current()
+
+
+def run_in_transaction(
+    callback,
+    alias=DEFAULT_CONNECTION_NAME,
+    session_kwargs=None,
+    transaction_kwargs=None,
+):
+    """Execute queries within a MongoDB transaction.
+
+    Usage:
+
+    .. code-block:: python
+
+        class A(Document):
+            name = StringField()
+
+        def callback(session):
+            a_doc = A.objects.create(name="a")
+            a_doc.update(name="b")
+        run_in_transaction(callback)
+
+        # With custom args/kwargs
+        def callback(session, custom_arg, customer_kwarg=None):
+            a_doc.update(name=f'{custom_arg}-{custom_kwarg}')
+        run_in_transaction(
+            lambda s: callback(s, 'arg', custom_kwarg='kwarg')
+        )
+
+    Be aware that:
+    - Mongo transactions run inside a session which is bound to a connection. If you attempt to
+      execute a transaction across a different connection alias, pymongo will raise an exception. In
+      other words: you cannot create a transaction that crosses different database connections.
+
+    For more information regarding pymongo transactions: https://pymongo.readthedocs.io/en/stable/api/pymongo/client_session.html#transactions
+    """
+
+    if PYMONGO_VERSION < (3, 9):
+        raise mongoengine.errors.OperationError(
+            "pymongo>=3.9 is required to use transactions"
+        )
+
+    conn = get_connection(alias)
+    session_kwargs = session_kwargs or {}
+    with conn.start_session(**session_kwargs) as session:
+        transaction_kwargs = transaction_kwargs or {}
+        transaction_kwargs["callback"] = callback
+        _set_session(session)
+        try:
+            session.with_transaction(**transaction_kwargs)
+        finally:
+            _clear_session()

mongoengine/context_managers.py

@@ -4,7 +4,11 @@
 from pymongo.write_concern import WriteConcern
 
 from mongoengine.common import _import_class
-from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db
+from mongoengine.connection import (
+    DEFAULT_CONNECTION_NAME,
+    _get_session,
+    get_db,
+)
 from mongoengine.pymongo_support import count_documents
 
 __all__ = (
@@ -210,11 +214,11 @@ def __init__(self, alias=DEFAULT_CONNECTION_NAME):
         }
 
     def _turn_on_profiling(self):
-        profile_update_res = self.db.command({"profile": 0})
+        profile_update_res = self.db.command({"profile": 0}, session=_get_session())
         self.initial_profiling_level = profile_update_res["was"]
 
         self.db.system.profile.drop()
-        self.db.command({"profile": 2})
+        self.db.command({"profile": 2}, session=_get_session())
 
     def _resets_profiling(self):
         self.db.command({"profile": self.initial_profiling_level})

mongoengine/dereference.py

@@ -8,7 +8,7 @@
     get_document,
 )
 from mongoengine.base.datastructures import LazyReference
-from mongoengine.connection import get_db
+from mongoengine.connection import _get_session, get_db
 from mongoengine.document import Document, EmbeddedDocument
 from mongoengine.fields import (
     DictField,
@@ -187,13 +187,15 @@ def _fetch_objects(self, doc_type=None):
 
                 if doc_type:
                     references = doc_type._get_db()[collection].find(
-                        {"_id": {"$in": refs}}
+                        {"_id": {"$in": refs}}, session=_get_session()
                     )
                     for ref in references:
                         doc = doc_type._from_son(ref)
                         object_map[(collection, doc.id)] = doc
                 else:
-                    references = get_db()[collection].find({"_id": {"$in": refs}})
+                    references = get_db()[collection].find(
+                        {"_id": {"$in": refs}}, session=_get_session()
+                    )
                     for ref in references:
                         if "_cls" in ref:
                             doc = get_document(ref["_cls"])._from_son(ref)

mongoengine/document.py

@@ -15,7 +15,11 @@
     get_document,
 )
 from mongoengine.common import _import_class
-from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db
+from mongoengine.connection import (
+    DEFAULT_CONNECTION_NAME,
+    _get_session,
+    get_db,
+)
 from mongoengine.context_managers import (
     set_write_concern,
     switch_collection,
@@ -269,7 +273,7 @@ def _get_capped_collection(cls):
         if max_documents:
             opts["max"] = max_documents
 
-        return db.create_collection(collection_name, **opts)
+        return db.create_collection(collection_name, session=_get_session(), **opts)
 
     def to_mongo(self, *args, **kwargs):
         data = super().to_mongo(*args, **kwargs)
@@ -479,17 +483,21 @@ def _save_create(self, doc, force_insert, write_concern):
         collection = self._get_collection()
         with set_write_concern(collection, write_concern) as wc_collection:
             if force_insert:
-                return wc_collection.insert_one(doc).inserted_id
+                return wc_collection.insert_one(doc, session=_get_session()).inserted_id
             # insert_one will provoke UniqueError alongside save does not
             # therefore, it need to catch and call replace_one.
             if "_id" in doc:
                 select_dict = {"_id": doc["_id"]}
                 select_dict = self._integrate_shard_key(doc, select_dict)
-                raw_object = wc_collection.find_one_and_replace(select_dict, doc)
+                raw_object = wc_collection.find_one_and_replace(
+                    select_dict, doc, session=_get_session()
+                )
                 if raw_object:
                     return doc["_id"]
 
-            object_id = wc_collection.insert_one(doc).inserted_id
+            object_id = wc_collection.insert_one(
+                doc, session=_get_session()
+            ).inserted_id
 
         return object_id
 
@@ -547,7 +555,7 @@ def _save_update(self, doc, save_condition, write_concern):
             upsert = save_condition is None
             with set_write_concern(collection, write_concern) as wc_collection:
                 last_error = wc_collection.update_one(
-                    select_dict, update_doc, upsert=upsert
+                    select_dict, update_doc, upsert=upsert, session=_get_session()
                 ).raw_result
             if not upsert and last_error["n"] == 0:
                 raise SaveConditionError(
@@ -850,7 +858,7 @@ def drop_collection(cls):
             )
         cls._collection = None
         db = cls._get_db()
-        db.drop_collection(coll_name)
+        db.drop_collection(coll_name, session=_get_session())
 
     @classmethod
     def create_index(cls, keys, background=False, **kwargs):
@@ -867,7 +875,9 @@ def create_index(cls, keys, background=False, **kwargs):
         index_spec["background"] = background
         index_spec.update(kwargs)
 
-        return cls._get_collection().create_index(fields, **index_spec)
+        return cls._get_collection().create_index(
+            fields, session=_get_session(), **index_spec
+        )
 
     @classmethod
     def ensure_indexes(cls):
@@ -913,7 +923,9 @@ def ensure_indexes(cls):
                 if "cls" in opts:
                     del opts["cls"]
 
-                collection.create_index(fields, background=background, **opts)
+                collection.create_index(
+                    fields, background=background, session=_get_session(), **opts
+                )
 
         # If _cls is being used (for polymorphism), it needs an index,
         # only if another index doesn't begin with _cls
@@ -923,7 +935,9 @@ def ensure_indexes(cls):
             if "cls" in index_opts:
                 del index_opts["cls"]
 
-            collection.create_index("_cls", background=background, **index_opts)
+            collection.create_index(
+                "_cls", background=background, session=_get_session(), **index_opts
+            )
 
     @classmethod
     def list_indexes(cls):
@@ -1001,7 +1015,7 @@ def compare_indexes(cls):
 
         existing = []
         collection = cls._get_collection()
-        for info in collection.index_information().values():
+        for info in collection.index_information(session=_get_session()).values():
             if "_fts" in info["key"][0]:
                 # Useful for text indexes (but not only)
                 index_type = info["key"][0][1]