PYTHON-5752 - Only retry overload errors if retries are enabled (#2726)

Author: NoahStapp

pymongo/asynchronous/collection.py

@@ -2946,6 +2946,7 @@ async def _aggregate(
             session,
             retryable=not cmd._performs_write,
             operation=_Op.AGGREGATE,
+            is_aggregate_write=cmd._performs_write,
         )
 
     async def aggregate(

pymongo/asynchronous/database.py

@@ -952,7 +952,7 @@ async def inner(
             )
 
         return await self._client._retryable_read(
-            inner, read_preference, session, command_name, None, False
+            inner, read_preference, session, command_name, None, False, is_run_command=True
         )
 
     @_csot.apply

pymongo/asynchronous/mongo_client.py

@@ -2010,6 +2010,8 @@ async def _retry_internal(
         read_pref: Optional[_ServerMode] = None,
         retryable: bool = False,
         operation_id: Optional[int] = None,
+        is_run_command: bool = False,
+        is_aggregate_write: bool = False,
     ) -> T:
         """Internal retryable helper for all client transactions.
 
@@ -2021,6 +2023,8 @@ async def _retry_internal(
         :param address: Server Address, defaults to None
         :param read_pref: Topology of read operation, defaults to None
         :param retryable: If the operation should be retried once, defaults to None
+        :param is_run_command: If this is a runCommand operation, defaults to False
+        :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
 
         :return: Output of the calling func()
         """
@@ -2035,6 +2039,8 @@ async def _retry_internal(
             address=address,
             retryable=retryable,
             operation_id=operation_id,
+            is_run_command=is_run_command,
+            is_aggregate_write=is_aggregate_write,
         ).run()
 
     async def _retryable_read(
@@ -2046,6 +2052,8 @@ async def _retryable_read(
         address: Optional[_Address] = None,
         retryable: bool = True,
         operation_id: Optional[int] = None,
+        is_run_command: bool = False,
+        is_aggregate_write: bool = False,
     ) -> T:
         """Execute an operation with consecutive retries if possible
 
@@ -2061,6 +2069,8 @@ async def _retryable_read(
         :param address: Optional address when sending a message, defaults to None
         :param retryable: if we should attempt retries
             (may not always be supported even if supplied), defaults to False
+        :param is_run_command: If this is a runCommand operation, defaults to False.
+        :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
         """
 
         # Ensure that the client supports retrying on reads and there is no session in
@@ -2079,6 +2089,8 @@ async def _retryable_read(
                 read_pref=read_pref,
                 retryable=retryable,
                 operation_id=operation_id,
+                is_run_command=is_run_command,
+                is_aggregate_write=is_aggregate_write,
             )
 
     async def _retryable_write(
@@ -2748,6 +2760,8 @@ def __init__(
         address: Optional[_Address] = None,
         retryable: bool = False,
         operation_id: Optional[int] = None,
+        is_run_command: bool = False,
+        is_aggregate_write: bool = False,
     ):
         self._last_error: Optional[Exception] = None
         self._retrying = False
@@ -2770,6 +2784,8 @@ def __init__(
         self._operation = operation
         self._operation_id = operation_id
         self._attempt_number = 0
+        self._is_run_command = is_run_command
+        self._is_aggregate_write = is_aggregate_write
 
     async def run(self) -> T:
         """Runs the supplied func() and attempts a retry
@@ -2810,18 +2826,30 @@ async def run(self) -> T:
                 always_retryable = False
                 overloaded = False
                 exc_to_check = exc
+
+                if self._is_run_command and not (
+                    self._client.options.retry_reads and self._client.options.retry_writes
+                ):
+                    raise
+                if self._is_aggregate_write and not self._client.options.retry_writes:
+                    raise
+
                 # Execute specialized catch on read
                 if self._is_read:
                     if isinstance(exc, (ConnectionFailure, OperationFailure)):
                         # ConnectionFailures do not supply a code property
                         exc_code = getattr(exc, "code", None)
                         overloaded = exc.has_error_label("SystemOverloadedError")
                         always_retryable = exc.has_error_label("RetryableError") and overloaded
-                        if not always_retryable and (
-                            self._is_not_eligible_for_retry()
-                            or (
-                                isinstance(exc, OperationFailure)
-                                and exc_code not in helpers_shared._RETRYABLE_ERROR_CODES
+                        if (
+                            not self._client.options.retry_reads
+                            or not always_retryable
+                            and (
+                                self._is_not_eligible_for_retry()
+                                or (
+                                    isinstance(exc, OperationFailure)
+                                    and exc_code not in helpers_shared._RETRYABLE_ERROR_CODES
+                                )
                             )
                         ):
                             raise
@@ -2852,7 +2880,12 @@ async def run(self) -> T:
                     retryable_write_label = exc_to_check.has_error_label("RetryableWriteError")
                     overloaded = exc_to_check.has_error_label("SystemOverloadedError")
                     always_retryable = exc_to_check.has_error_label("RetryableError") and overloaded
-                    if not self._retryable and not always_retryable:
+
+                    # Always retry abortTransaction and commitTransaction up to once
+                    if self._operation not in ["abortTransaction", "commitTransaction"] and (
+                        not self._client.options.retry_writes
+                        or not (self._retryable or always_retryable)
+                    ):
                         raise
                     if retryable_write_label or always_retryable:
                         assert self._session

pymongo/synchronous/collection.py

@@ -2939,6 +2939,7 @@ def _aggregate(
             session,
             retryable=not cmd._performs_write,
             operation=_Op.AGGREGATE,
+            is_aggregate_write=cmd._performs_write,
         )
 
     def aggregate(

pymongo/synchronous/database.py

@@ -952,7 +952,7 @@ def inner(
             )
 
         return self._client._retryable_read(
-            inner, read_preference, session, command_name, None, False
+            inner, read_preference, session, command_name, None, False, is_run_command=True
         )
 
     @_csot.apply

pymongo/synchronous/mongo_client.py

@@ -2006,6 +2006,8 @@ def _retry_internal(
         read_pref: Optional[_ServerMode] = None,
         retryable: bool = False,
         operation_id: Optional[int] = None,
+        is_run_command: bool = False,
+        is_aggregate_write: bool = False,
     ) -> T:
         """Internal retryable helper for all client transactions.
 
@@ -2017,6 +2019,8 @@ def _retry_internal(
         :param address: Server Address, defaults to None
         :param read_pref: Topology of read operation, defaults to None
         :param retryable: If the operation should be retried once, defaults to None
+        :param is_run_command: If this is a runCommand operation, defaults to False
+        :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
 
         :return: Output of the calling func()
         """
@@ -2031,6 +2035,8 @@ def _retry_internal(
             address=address,
             retryable=retryable,
             operation_id=operation_id,
+            is_run_command=is_run_command,
+            is_aggregate_write=is_aggregate_write,
         ).run()
 
     def _retryable_read(
@@ -2042,6 +2048,8 @@ def _retryable_read(
         address: Optional[_Address] = None,
         retryable: bool = True,
         operation_id: Optional[int] = None,
+        is_run_command: bool = False,
+        is_aggregate_write: bool = False,
     ) -> T:
         """Execute an operation with consecutive retries if possible
 
@@ -2057,6 +2065,8 @@ def _retryable_read(
         :param address: Optional address when sending a message, defaults to None
         :param retryable: if we should attempt retries
             (may not always be supported even if supplied), defaults to False
+        :param is_run_command: If this is a runCommand operation, defaults to False.
+        :param is_aggregate_write: If this is a aggregate operation with a write, defaults to False.
         """
 
         # Ensure that the client supports retrying on reads and there is no session in
@@ -2075,6 +2085,8 @@ def _retryable_read(
                 read_pref=read_pref,
                 retryable=retryable,
                 operation_id=operation_id,
+                is_run_command=is_run_command,
+                is_aggregate_write=is_aggregate_write,
             )
 
     def _retryable_write(
@@ -2738,6 +2750,8 @@ def __init__(
         address: Optional[_Address] = None,
         retryable: bool = False,
         operation_id: Optional[int] = None,
+        is_run_command: bool = False,
+        is_aggregate_write: bool = False,
     ):
         self._last_error: Optional[Exception] = None
         self._retrying = False
@@ -2760,6 +2774,8 @@ def __init__(
         self._operation = operation
         self._operation_id = operation_id
         self._attempt_number = 0
+        self._is_run_command = is_run_command
+        self._is_aggregate_write = is_aggregate_write
 
     def run(self) -> T:
         """Runs the supplied func() and attempts a retry
@@ -2800,18 +2816,30 @@ def run(self) -> T:
                 always_retryable = False
                 overloaded = False
                 exc_to_check = exc
+
+                if self._is_run_command and not (
+                    self._client.options.retry_reads and self._client.options.retry_writes
+                ):
+                    raise
+                if self._is_aggregate_write and not self._client.options.retry_writes:
+                    raise
+
                 # Execute specialized catch on read
                 if self._is_read:
                     if isinstance(exc, (ConnectionFailure, OperationFailure)):
                         # ConnectionFailures do not supply a code property
                         exc_code = getattr(exc, "code", None)
                         overloaded = exc.has_error_label("SystemOverloadedError")
                         always_retryable = exc.has_error_label("RetryableError") and overloaded
-                        if not always_retryable and (
-                            self._is_not_eligible_for_retry()
-                            or (
-                                isinstance(exc, OperationFailure)
-                                and exc_code not in helpers_shared._RETRYABLE_ERROR_CODES
+                        if (
+                            not self._client.options.retry_reads
+                            or not always_retryable
+                            and (
+                                self._is_not_eligible_for_retry()
+                                or (
+                                    isinstance(exc, OperationFailure)
+                                    and exc_code not in helpers_shared._RETRYABLE_ERROR_CODES
+                                )
                             )
                         ):
                             raise
@@ -2842,7 +2870,12 @@ def run(self) -> T:
                     retryable_write_label = exc_to_check.has_error_label("RetryableWriteError")
                     overloaded = exc_to_check.has_error_label("SystemOverloadedError")
                     always_retryable = exc_to_check.has_error_label("RetryableError") and overloaded
-                    if not self._retryable and not always_retryable:
+
+                    # Always retry abortTransaction and commitTransaction up to once
+                    if self._operation not in ["abortTransaction", "commitTransaction"] and (
+                        not self._client.options.retry_writes
+                        or not (self._retryable or always_retryable)
+                    ):
                         raise
                     if retryable_write_label or always_retryable:
                         assert self._session

test/asynchronous/test_client_backpressure.py

@@ -227,7 +227,7 @@ async def asyncSetUp(self) -> None:
         self.listener.reset()
         self.app_name = self.__class__.__name__.lower()
         self.client = await self.async_rs_or_single_client(
-            event_listeners=[self.listener], retryWrites=False, appName=self.app_name
+            event_listeners=[self.listener], appName=self.app_name
         )
 
     @patch("random.random")