retrying vars back in bulk

Author: sleepyStick

pymongo/asynchronous/bulk.py

@@ -112,6 +112,9 @@ def __init__(
         self.uses_hint_update = False
         self.uses_hint_delete = False
         self.uses_sort = False
+        self.is_retryable = True
+        self.retrying = False
+        self.started_retryable_write = False
         # Extra state so that we know where to pick up on a retry attempt.
         self.current_run = None
         self.next_run = None
@@ -127,23 +130,23 @@ def bulk_ctx_class(self) -> Type[_BulkWriteContext]:
             self.is_encrypted = False
             return _BulkWriteContext
 
-    @property
-    def is_retryable(self) -> bool:
-        if self.current_run:
-            return self.current_run.is_retryable
-        return True
-
-    @property
-    def retrying(self) -> bool:
-        if self.current_run:
-            return self.current_run.retrying
-        return False
-
-    @property
-    def started_retryable_write(self) -> bool:
-        if self.current_run:
-            return self.current_run.started_retryable_write
-        return False
+    # @property
+    # def is_retryable(self) -> bool:
+    #     if self.current_run:
+    #         return self.current_run.is_retryable
+    #     return True
+    #
+    # @property
+    # def retrying(self) -> bool:
+    #     if self.current_run:
+    #         return self.current_run.retrying
+    #     return False
+    #
+    # @property
+    # def started_retryable_write(self) -> bool:
+    #     if self.current_run:
+    #         return self.current_run.started_retryable_write
+    #     return False
 
     def add_insert(self, document: _DocumentOut) -> bool:
         """Add an insert document to the list of ops."""
@@ -255,7 +258,7 @@ def gen_ordered(
                 yield run
                 run = _Run(op_type)
             run.add(idx, operation)
-            run.is_retryable = run.is_retryable and retryable
+            self.is_retryable = self.is_retryable and retryable
         if run is None:
             raise InvalidOperation("No operations to execute")
         yield run
@@ -273,7 +276,7 @@ def gen_unordered(
             retryable = process(request)
             (op_type, operation) = self.ops[idx]
             operations[op_type].add(idx, operation)
-            operations[op_type].is_retryable = operations[op_type].is_retryable and retryable
+            self.is_retryable = self.is_retryable and retryable
         if (
             len(operations[_INSERT].ops) == 0
             and len(operations[_UPDATE].ops) == 0
@@ -533,7 +536,7 @@ async def _execute_command(
         last_run = False
 
         while run:
-            if not run.retrying:
+            if not self.retrying:
                 self.next_run = next(generator, None)
                 if self.next_run is None:
                     last_run = True
@@ -567,10 +570,10 @@ async def _execute_command(
                 if session:
                     # Start a new retryable write unless one was already
                     # started for this command.
-                    if run.is_retryable and not run.started_retryable_write:
+                    if self.is_retryable and not self.started_retryable_write:
                         session._start_retryable_write()
                         self.started_retryable_write = True
-                    session._apply_to(cmd, run.is_retryable, ReadPreference.PRIMARY, conn)
+                    session._apply_to(cmd, self.is_retryable, ReadPreference.PRIMARY, conn)
                 conn.send_cluster_time(cmd, session, client)
                 conn.add_server_api(cmd)
                 # CSOT: apply timeout before encoding the command.

pymongo/asynchronous/mongo_client.py

@@ -2674,8 +2674,8 @@ def __init__(
         self._operation_id = operation_id
 
     def _bulk_retryable(self) -> bool:
-        if self._bulk is not None and self._bulk.current_run is not None:
-            return self._bulk.current_run.is_retryable
+        if self._bulk is not None:
+            return self._bulk.is_retryable
         return True
 
     async def run(self) -> T:
@@ -2695,8 +2695,8 @@ async def run(self) -> T:
             and not self._is_read
         ):
             self._session._start_retryable_write()  # type: ignore
-            if self._bulk and self._bulk.current_run:
-                self._bulk.current_run.started_retryable_write = True
+            if self._bulk:
+                self._bulk.started_retryable_write = True
 
         while True:
             self._check_last_error(check_csot=True)
@@ -2746,10 +2746,7 @@ async def run(self) -> T:
                         else:
                             raise
                     if self._bulk:
-                        if self._bulk.current_run:
-                            self._bulk.current_run.retrying = True
-                        else:
-                            self._bulk.retrying = True
+                        self._bulk.retrying = True
                     else:
                         self._retrying = True
                     if not exc.has_error_label("NoWritesPerformed"):
@@ -2770,11 +2767,7 @@ def _is_not_eligible_for_retry(self) -> bool:
 
     def _is_retrying(self) -> bool:
         """Checks if the exchange is currently undergoing a retry"""
-        return (
-            self._bulk.current_run.retrying
-            if self._bulk is not None and self._bulk.current_run is not None
-            else self._retrying
-        )
+        return self._bulk.retrying if self._bulk is not None else self._retrying
 
     def _is_session_state_retryable(self) -> bool:
         """Checks if provided session is eligible for retry
@@ -2834,8 +2827,8 @@ async def _write(self) -> T:
                     # not support sessions raise the last error.
                     self._check_last_error()
                     self._retryable = False
-                    if self._bulk and self._bulk.current_run:
-                        self._bulk.current_run.is_retryable = False
+                    if self._bulk:
+                        self._bulk.is_retryable = False
                 return await self._func(self._session, conn)  # type: ignore
         except PyMongoError as exc:
             if not self._retryable or not self._bulk_retryable():

pymongo/bulk_shared.py

@@ -50,9 +50,6 @@ def __init__(self, op_type: int) -> None:
         self.index_map: list[int] = []
         self.ops: list[Any] = []
         self.idx_offset: int = 0
-        self.is_retryable = True
-        self.retrying = False
-        self.started_retryable_write = False
 
     def index(self, idx: int) -> int:
         """Get the original index of an operation in this run.

pymongo/synchronous/bulk.py

@@ -112,6 +112,9 @@ def __init__(
         self.uses_hint_update = False
         self.uses_hint_delete = False
         self.uses_sort = False
+        self.is_retryable = True
+        self.retrying = False
+        self.started_retryable_write = False
         # Extra state so that we know where to pick up on a retry attempt.
         self.current_run = None
         self.next_run = None
@@ -127,23 +130,23 @@ def bulk_ctx_class(self) -> Type[_BulkWriteContext]:
             self.is_encrypted = False
             return _BulkWriteContext
 
-    @property
-    def is_retryable(self) -> bool:
-        if self.current_run:
-            return self.current_run.is_retryable
-        return True
-
-    @property
-    def retrying(self) -> bool:
-        if self.current_run:
-            return self.current_run.retrying
-        return False
-
-    @property
-    def started_retryable_write(self) -> bool:
-        if self.current_run:
-            return self.current_run.started_retryable_write
-        return False
+    # @property
+    # def is_retryable(self) -> bool:
+    #     if self.current_run:
+    #         return self.current_run.is_retryable
+    #     return True
+    #
+    # @property
+    # def retrying(self) -> bool:
+    #     if self.current_run:
+    #         return self.current_run.retrying
+    #     return False
+    #
+    # @property
+    # def started_retryable_write(self) -> bool:
+    #     if self.current_run:
+    #         return self.current_run.started_retryable_write
+    #     return False
 
     def add_insert(self, document: _DocumentOut) -> bool:
         """Add an insert document to the list of ops."""
@@ -255,7 +258,7 @@ def gen_ordered(
                 yield run
                 run = _Run(op_type)
             run.add(idx, operation)
-            run.is_retryable = run.is_retryable and retryable
+            self.is_retryable = self.is_retryable and retryable
         if run is None:
             raise InvalidOperation("No operations to execute")
         yield run
@@ -273,7 +276,7 @@ def gen_unordered(
             retryable = process(request)
             (op_type, operation) = self.ops[idx]
             operations[op_type].add(idx, operation)
-            operations[op_type].is_retryable = operations[op_type].is_retryable and retryable
+            self.is_retryable = self.is_retryable and retryable
         if (
             len(operations[_INSERT].ops) == 0
             and len(operations[_UPDATE].ops) == 0
@@ -533,7 +536,7 @@ def _execute_command(
         last_run = False
 
         while run:
-            if not run.retrying:
+            if not self.retrying:
                 self.next_run = next(generator, None)
                 if self.next_run is None:
                     last_run = True
@@ -567,10 +570,10 @@ def _execute_command(
                 if session:
                     # Start a new retryable write unless one was already
                     # started for this command.
-                    if run.is_retryable and not run.started_retryable_write:
+                    if self.is_retryable and not self.started_retryable_write:
                         session._start_retryable_write()
                         self.started_retryable_write = True
-                    session._apply_to(cmd, run.is_retryable, ReadPreference.PRIMARY, conn)
+                    session._apply_to(cmd, self.is_retryable, ReadPreference.PRIMARY, conn)
                 conn.send_cluster_time(cmd, session, client)
                 conn.add_server_api(cmd)
                 # CSOT: apply timeout before encoding the command.

pymongo/synchronous/mongo_client.py

@@ -2662,8 +2662,8 @@ def __init__(
         self._operation_id = operation_id
 
     def _bulk_retryable(self) -> bool:
-        if self._bulk is not None and self._bulk.current_run is not None:
-            return self._bulk.current_run.is_retryable
+        if self._bulk is not None:
+            return self._bulk.is_retryable
         return True
 
     def run(self) -> T:
@@ -2683,8 +2683,8 @@ def run(self) -> T:
             and not self._is_read
         ):
             self._session._start_retryable_write()  # type: ignore
-            if self._bulk and self._bulk.current_run:
-                self._bulk.current_run.started_retryable_write = True
+            if self._bulk:
+                self._bulk.started_retryable_write = True
 
         while True:
             self._check_last_error(check_csot=True)
@@ -2734,10 +2734,7 @@ def run(self) -> T:
                         else:
                             raise
                     if self._bulk:
-                        if self._bulk.current_run:
-                            self._bulk.current_run.retrying = True
-                        else:
-                            self._bulk.retrying = True
+                        self._bulk.retrying = True
                     else:
                         self._retrying = True
                     if not exc.has_error_label("NoWritesPerformed"):
@@ -2758,11 +2755,7 @@ def _is_not_eligible_for_retry(self) -> bool:
 
     def _is_retrying(self) -> bool:
         """Checks if the exchange is currently undergoing a retry"""
-        return (
-            self._bulk.current_run.retrying
-            if self._bulk is not None and self._bulk.current_run is not None
-            else self._retrying
-        )
+        return self._bulk.retrying if self._bulk is not None else self._retrying
 
     def _is_session_state_retryable(self) -> bool:
         """Checks if provided session is eligible for retry
@@ -2822,8 +2815,8 @@ def _write(self) -> T:
                     # not support sessions raise the last error.
                     self._check_last_error()
                     self._retryable = False
-                    if self._bulk and self._bulk.current_run:
-                        self._bulk.current_run.is_retryable = False
+                    if self._bulk:
+                        self._bulk.is_retryable = False
                 return self._func(self._session, conn)  # type: ignore
         except PyMongoError as exc:
             if not self._retryable or not self._bulk_retryable():