Merge branch 'master' of github.com:MongoEngine/mongoengine into remove_deprecated_warnings

Author: bagerard

.github/workflows/github-actions.yml

@@ -20,17 +20,16 @@ env:
   MONGODB_6_0: "6.0"
   MONGODB_7_0: "7.0"
 
-  PYMONGO_3_4: 3.4
-  PYMONGO_3_6: 3.6
-  PYMONGO_3_9: 3.9
   PYMONGO_3_11: 3.11
   PYMONGO_3_12: 3.12
+  PYMONGO_3_13: 3.13
   PYMONGO_4_0: 4.0
   PYMONGO_4_3: 4.3.3
   PYMONGO_4_4: 4.4.1
   PYMONGO_4_6: 4.6.2
   PYMONGO_4_7: 4.7.3
   PYMONGO_4_8: 4.8.0
+  PYMONGO_4_9: 4.9
 
   MAIN_PYTHON_VERSION: 3.9
 
@@ -61,15 +60,12 @@ jobs:
         MONGODB: [$MONGODB_4_0]
         PYMONGO: [$PYMONGO_3_11]
         include:
-          - python-version: 3.7
-            MONGODB: $MONGODB_3_6
-            PYMONGO: $PYMONGO_3_9
           - python-version: 3.8
-            MONGODB: $MONGODB_4_4
-            PYMONGO: $PYMONGO_3_11
+            MONGODB: $MONGODB_3_6
+            PYMONGO: $PYMONGO_3_12
           - python-version: 3.9
             MONGODB: $MONGODB_4_4
-            PYMONGO: $PYMONGO_3_12
+            PYMONGO: $PYMONGO_3_13
           - python-version: "3.10"
             MONGODB: $MONGODB_4_4
             PYMONGO: $PYMONGO_4_0
@@ -88,6 +84,9 @@ jobs:
           - python-version: "3.11"
             MONGODB: $MONGODB_7_0
             PYMONGO: $PYMONGO_4_8
+          - python-version: "3.11"
+            MONGODB: $MONGODB_7_0
+            PYMONGO: $PYMONGO_4_9
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}

.github/workflows/start_mongo.sh

@@ -5,10 +5,18 @@ 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
+
+args=(--dbpath $mongodb_dir/data --logpath $mongodb_dir/mongodb.log --fork --replSet mongoengine)
+if (( $(echo "$MONGODB > 3.8" | bc -l) )); then
+    args+=(--setParameter maxTransactionLockRequestTimeoutMillis=1000)
+fi
+
+$mongodb_dir/bin/mongod "${args[@]}"
 
 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)
 - Further to the deprecation warning, remove ability to use an unpacked list to `Queryset.aggregate(*pipeline)`, a plain list must be provided instead `Queryset.aggregate(pipeline)`, as it's closer to pymongo interface
 - Further to the deprecation warning, remove `full_response` from `QuerySet.modify` as it wasn't supported with Pymongo 3+
 - Fixed stacklevel of many warnings (to point places emitting the warning more accurately)

entrypoint.sh

@@ -3,4 +3,4 @@
 mongod --replSet mongoengine --fork --logpath=/var/log/mongodb.log
 mongo db --eval "rs.initiate()"
 mongod --shutdown
-mongod --replSet mongoengine --bind_ip 0.0.0.0
+mongod --replSet mongoengine --bind_ip 0.0.0.0 --setParameter maxTransactionLockRequestTimeoutMillis=1000

mongoengine/connection.py

@@ -1,8 +1,14 @@
+import collections
+import threading
 import warnings
 
 from pymongo import MongoClient, ReadPreference, uri_parser
 from pymongo.common import _UUID_REPRESENTATIONS
-from pymongo.database import _check_name
+
+try:
+    from pymongo.database_shared import _check_name
+except ImportError:
+    from pymongo.database import _check_name
 
 # DriverInfo was added in PyMongo 3.7.
 try:
@@ -35,6 +41,7 @@
 _connections = {}
 _dbs = {}
 
+
 READ_PREFERENCE = ReadPreference.PRIMARY
 
 
@@ -471,3 +478,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)