Fix tests and Improve commit handler so that it supports UnknownTransactionCommitResult

Author: bagerard

mongoengine/context_managers.py

@@ -2,6 +2,7 @@
 import threading
 from contextlib import contextmanager
 
+from pymongo.errors import ConnectionFailure, OperationFailure
 from pymongo.read_concern import ReadConcern
 from pymongo.write_concern import WriteConcern
 
@@ -321,6 +322,23 @@ def set_read_write_concern(collection, write_concerns, read_concerns):
     )
 
 
+def _commit_with_retry(session):
+    while True:
+        try:
+            # Commit uses write concern set at transaction start.
+            session.commit_transaction()
+            print("Transaction committed.")
+            break
+        except (ConnectionFailure, OperationFailure) as exc:
+            # Can retry commit
+            if exc.has_error_label("UnknownTransactionCommitResult"):
+                print("UnknownTransactionCommitResult, retrying commit operation ...")
+                continue
+            else:
+                print("Error during commit ...")
+                raise
+
+
 @contextmanager
 def run_in_transaction(
     alias=DEFAULT_CONNECTION_NAME, session_kwargs=None, transaction_kwargs=None
@@ -355,5 +373,6 @@ class A(Document):
             try:
                 _set_session(session)
                 yield
+                _commit_with_retry(session)
             finally:
                 _clear_session()

tests/test_connection.py

@@ -658,9 +658,10 @@ def test_multiple_connection_settings(self):
         mongo_connections["t2"].server_info()
 
         assert mongo_connections["t1"].address[0] == "localhost"
-        assert (
-            mongo_connections["t2"].address[0] == "localhost"
-        )  # weird but we have this with replicaset
+        assert mongo_connections["t2"].address[0] in (
+            "localhost",
+            "127.0.0.1",
+        )  # weird but there is a discrepancy in the address in replicaset setup
         assert mongo_connections["t1"].read_preference == ReadPreference.PRIMARY
         assert (
             mongo_connections["t2"].read_preference == ReadPreference.PRIMARY_PREFERRED

tests/test_context_managers.py

@@ -1,3 +1,4 @@
+import logging
 import random
 import time
 import unittest
@@ -28,7 +29,7 @@
 
 
 class TestRollbackError(Exception):
-    pass
+    __test__ = False  # Silence pytest warning
 
 
 class TestableThread(Thread):
@@ -38,6 +39,8 @@ class TestableThread(Thread):
     REF: https://gist.github.com/sbrugman/59b3535ebcd5aa0e2598293cfa58b6ab
     """
 
+    __test__ = False  # Silence pytest warning
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.exc = None
@@ -753,14 +756,15 @@ def run_tx():
 
             except TestRollbackError:
                 pass
-            # except OperationError as op_failure:
-            #     """
-            #     See thread safety test below for more details about TransientTransactionError handling
-            #     """
-            #     if "TransientTransactionError" in str(op_failure):
-            #         run_tx()
-            #     else:
-            #         raise op_failure
+            except OperationError as op_failure:
+                """
+                See thread safety test below for more details about TransientTransactionError handling
+                """
+                if "TransientTransactionError" in str(op_failure):
+                    logging.warning("TransientTransactionError - retrying...")
+                    run_tx()
+                else:
+                    raise op_failure
 
         run_tx()
         assert "a" == A.objects.get(id=a_doc.id).name
@@ -789,10 +793,10 @@ def test_thread_safety_of_transactions(self):
         case, then no amount of runtime variability should have
         an effect on the output.
 
-        This test sets up 10 records, each with an integer field
+        This test sets up e.g 10 records, each with an integer field
         of value 0 - 9.
 
-        We then spin up 10 threads and attempt to update a target
+        We then spin up e.g 10 threads and attempt to update a target
         record by multiplying its integer value by 10. Then, if
         the target record is even, throw an exception, which
         should then roll the transaction back. The odd rows always
@@ -807,24 +811,26 @@ def test_thread_safety_of_transactions(self):
         connect("mongoenginetest")
 
         class A(Document):
-            i = IntField()
+            i = IntField(unique=True)
 
         A.drop_collection()
         # Ensure the collection is created
-        A.objects.create(i=0)
+        _ = A.objects.first()
+
+        thread_count = 20
 
         def thread_fn(idx):
             # Open the transaction at some unknown interval
             time.sleep(random.uniform(0.1, 0.5))
             try:
                 with run_in_transaction():
                     a = A.objects.get(i=idx)
-                    a.i = idx * 10
+                    a.i = idx * thread_count
                     # Save at some unknown interval
                     time.sleep(random.uniform(0.1, 0.5))
                     a.save()
 
-                    # Force roll backs for the even runs...
+                    # Force rollbacks for the even runs...
                     if idx % 2 == 0:
                         raise TestRollbackError()
             except TestRollbackError:
@@ -841,6 +847,7 @@ def thread_fn(idx):
                 """
                 error_labels = op_failure.details.get("errorLabels", [])
                 if "TransientTransactionError" in error_labels:
+                    logging.warning("TransientTransactionError - retrying...")
                     thread_fn(idx)
                 else:
                     raise op_failure
@@ -854,7 +861,6 @@ def thread_fn(idx):
             A.objects.all().delete()
 
             # Prepopulate the data for reads
-            thread_count = 20
             for i in range(thread_count):
                 A.objects.create(i=i)
 
@@ -868,13 +874,10 @@ def thread_fn(idx):
                 t.join()
 
             # Check the sum
-            expected_sum = 0
-            for i in range(thread_count):
-                if i % 2 == 0:
-                    expected_sum += i
-                else:
-                    expected_sum += i * 10
-            assert expected_sum == 1090
+            expected_sum = sum(
+                i if i % 2 == 0 else i * thread_count for i in range(thread_count)
+            )
+            assert expected_sum == 2090
             assert expected_sum == sum(a.i for a in A.objects.all())