Merge pull request #2839 from bagerard/juannyG-transactions-global-session-ver

Juanny g - transactions global session ver

Author: bagerard

.github/workflows/start_mongo.sh

@@ -5,10 +5,12 @@ 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
 if (( $(echo "$MONGODB < 6.0" | bc -l) )); then
-mongo --quiet --eval 'db.runCommand("ping").ok'    # Make sure mongo is awake
+mongo --verbose --eval "rs.initiate()"
+mongo --quiet --eval "rs.status().ok"
 else
-mongosh --quiet  --eval 'db.runCommand("ping").ok'  # Make sure mongo is awake
+mongosh --verbose --eval "rs.initiate()"
+mongosh --quiet --eval "rs.status().ok"
 fi

AUTHORS

@@ -265,3 +265,4 @@ that much better:
  * Ido Shraga (https://github.com/idoshr)
  * Terence Honles (https://github.com/terencehonles)
  * Sean Bermejo (https://github.com/seanbermejo)
+ * Juan Gutierrez (https://github.com/juannyg)

Dockerfile

@@ -1,4 +1,4 @@
-FROM mongo:5
+FROM mongo:4.0
 
 COPY ./entrypoint.sh entrypoint.sh
 RUN chmod u+x entrypoint.sh

docs/changelog.rst

@@ -7,6 +7,11 @@ Changelog
 Development
 ===========
 - (Fill this out as you fix issues and develop your features).
+- Add support for transaction through run_in_transaction (kudos to juannyG for this) #2569
+  Some considerations:
+    - make sure to read https://www.mongodb.com/docs/manual/core/transactions-in-applications/#callback-api-vs-core-api
+    - run_in_transaction context manager relies on Pymongo coreAPI, it will retry automatically in case of `UnknownTransactionCommitResult` but not `TransientTransactionError` exceptions
+    - Using .count() in a transaction will always use Collection.count_document (as estimated_document_count is not supported in transactions)
 
 Changes in 0.29.0
 =================

mongoengine/connection.py

@@ -1,3 +1,5 @@
+import collections
+import threading
 import warnings
 
 from pymongo import MongoClient, ReadPreference, uri_parser
@@ -35,6 +37,7 @@
 _connections = {}
 _dbs = {}
 
+
 READ_PREFERENCE = ReadPreference.PRIMARY
 
 
@@ -470,3 +473,37 @@ 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[-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()

mongoengine/context_managers.py

@@ -1,13 +1,22 @@
 import contextlib
+import logging
 import threading
 from contextlib import contextmanager
 
+from pymongo.errors import ConnectionFailure, OperationFailure
 from pymongo.read_concern import ReadConcern
 from pymongo.write_concern import WriteConcern
 
 from mongoengine.base.fields import _no_dereference_for_fields
 from mongoengine.common import _import_class
-from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db
+from mongoengine.connection import (
+    DEFAULT_CONNECTION_NAME,
+    _clear_session,
+    _get_session,
+    _set_session,
+    get_connection,
+    get_db,
+)
 from mongoengine.pymongo_support import count_documents
 
 __all__ = (
@@ -19,6 +28,7 @@
     "set_write_concern",
     "set_read_write_concern",
     "no_dereferencing_active_for_class",
+    "run_in_transaction",
 )
 
 
@@ -231,11 +241,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})
@@ -311,3 +321,60 @@ def set_read_write_concern(collection, write_concerns, read_concerns):
         write_concern=WriteConcern(**combined_write_concerns),
         read_concern=ReadConcern(**combined_read_concerns),
     )
+
+
+def _commit_with_retry(session):
+    while True:
+        try:
+            # Commit uses write concern set at transaction start.
+            session.commit_transaction()
+            break
+        except (ConnectionFailure, OperationFailure) as exc:
+            # Can retry commit
+            if exc.has_error_label("UnknownTransactionCommitResult"):
+                logging.warning(
+                    "UnknownTransactionCommitResult, retrying commit operation ..."
+                )
+                continue
+            else:
+                # Error during commit
+                raise
+
+
+@contextmanager
+def run_in_transaction(
+    alias=DEFAULT_CONNECTION_NAME, session_kwargs=None, transaction_kwargs=None
+):
+    """run_in_transaction context manager
+    Execute queries within the context in a database transaction.
+
+    Usage:
+
+    .. code-block:: python
+
+        class A(Document):
+            name = StringField()
+
+        with run_in_transaction():
+            a_doc = A.objects.create(name="a")
+            a_doc.update(name="b")
+
+    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. That
+      said, multiple transaction can be nested within the same session for particular connection.
+
+    For more information regarding pymongo transactions: https://pymongo.readthedocs.io/en/stable/api/pymongo/client_session.html#transactions
+    """
+    conn = get_connection(alias)
+    session_kwargs = session_kwargs or {}
+    with conn.start_session(**session_kwargs) as session:
+        transaction_kwargs = transaction_kwargs or {}
+        with session.start_transaction(**transaction_kwargs):
+            try:
+                _set_session(session)
+                yield
+                _commit_with_retry(session)
+            finally:
+                _clear_session()

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,
@@ -271,7 +275,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)
 
     @classmethod
     def _get_timeseries_collection(cls):
@@ -502,17 +506,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
 
@@ -570,7 +578,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(
@@ -873,7 +881,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):
@@ -890,7 +898,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):
@@ -936,7 +946,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
@@ -946,7 +958,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):
@@ -1024,7 +1038,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]