Roll back with_transaction abstraction back to context mgr

Author: juannyG

mongoengine/connection.py

@@ -25,7 +25,6 @@
     "get_connection",
     "get_db",
     "register_connection",
-    "run_in_transaction",
 ]
 
 
@@ -508,55 +507,3 @@ def _get_session():
 
 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:
-            return session.with_transaction(**transaction_kwargs)
-        finally:
-            _clear_session()

mongoengine/context_managers.py

@@ -7,7 +7,10 @@
 from mongoengine.common import _import_class
 from mongoengine.connection import (
     DEFAULT_CONNECTION_NAME,
+    _clear_session,
     _get_session,
+    _set_session,
+    get_connection,
     get_db,
 )
 from mongoengine.pymongo_support import count_documents
@@ -21,6 +24,7 @@
     "set_write_concern",
     "set_read_write_concern",
     "no_dereferencing_active_for_class",
+    "run_in_transaction",
 )
 
 
@@ -321,3 +325,41 @@ def set_read_write_concern(collection, write_concerns, read_concerns):
         write_concern=WriteConcern(**combined_write_concerns),
         read_concern=ReadConcern(**combined_read_concerns),
     )
+
+
+@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
+            finally:
+                _clear_session()