PYTHON-5751 - Only retry overload errors if retries are enabled

Author: NoahStapp

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

@@ -2009,6 +2009,7 @@ async def _retry_internal(
         read_pref: Optional[_ServerMode] = None,
         retryable: bool = False,
         operation_id: Optional[int] = None,
+        is_run_command: bool = False,
     ) -> T:
         """Internal retryable helper for all client transactions.
 
@@ -2020,6 +2021,7 @@ 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
 
         :return: Output of the calling func()
         """
@@ -2034,6 +2036,7 @@ async def _retry_internal(
             address=address,
             retryable=retryable,
             operation_id=operation_id,
+            is_run_command=is_run_command,
         ).run()
 
     async def _retryable_read(
@@ -2045,6 +2048,7 @@ async def _retryable_read(
         address: Optional[_Address] = None,
         retryable: bool = True,
         operation_id: Optional[int] = None,
+        is_run_command: bool = False,
     ) -> T:
         """Execute an operation with consecutive retries if possible
 
@@ -2060,6 +2064,7 @@ 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.
         """
 
         # Ensure that the client supports retrying on reads and there is no session in
@@ -2078,6 +2083,7 @@ async def _retryable_read(
                 read_pref=read_pref,
                 retryable=retryable,
                 operation_id=operation_id,
+                is_run_command=is_run_command,
             )
 
     async def _retryable_write(
@@ -2747,6 +2753,7 @@ def __init__(
         address: Optional[_Address] = None,
         retryable: bool = False,
         operation_id: Optional[int] = None,
+        is_run_command: bool = False,
     ):
         self._last_error: Optional[Exception] = None
         self._retrying = False
@@ -2769,6 +2776,7 @@ def __init__(
         self._operation = operation
         self._operation_id = operation_id
         self._attempt_number = 0
+        self._is_run_command = is_run_command
 
     async def run(self) -> T:
         """Runs the supplied func() and attempts a retry
@@ -2809,18 +2817,27 @@ 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
                 # 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
@@ -2851,7 +2868,11 @@ 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:
+                    if (
+                        not self._client.options.retry_writes
+                        or not self._retryable
+                        and not always_retryable
+                    ):
                         raise
                     if retryable_write_label or always_retryable:
                         assert self._session

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

@@ -2005,6 +2005,7 @@ def _retry_internal(
         read_pref: Optional[_ServerMode] = None,
         retryable: bool = False,
         operation_id: Optional[int] = None,
+        is_run_command: bool = False,
     ) -> T:
         """Internal retryable helper for all client transactions.
 
@@ -2016,6 +2017,7 @@ 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
 
         :return: Output of the calling func()
         """
@@ -2030,6 +2032,7 @@ def _retry_internal(
             address=address,
             retryable=retryable,
             operation_id=operation_id,
+            is_run_command=is_run_command,
         ).run()
 
     def _retryable_read(
@@ -2041,6 +2044,7 @@ def _retryable_read(
         address: Optional[_Address] = None,
         retryable: bool = True,
         operation_id: Optional[int] = None,
+        is_run_command: bool = False,
     ) -> T:
         """Execute an operation with consecutive retries if possible
 
@@ -2056,6 +2060,7 @@ 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.
         """
 
         # Ensure that the client supports retrying on reads and there is no session in
@@ -2074,6 +2079,7 @@ def _retryable_read(
                 read_pref=read_pref,
                 retryable=retryable,
                 operation_id=operation_id,
+                is_run_command=is_run_command,
             )
 
     def _retryable_write(
@@ -2737,6 +2743,7 @@ def __init__(
         address: Optional[_Address] = None,
         retryable: bool = False,
         operation_id: Optional[int] = None,
+        is_run_command: bool = False,
     ):
         self._last_error: Optional[Exception] = None
         self._retrying = False
@@ -2759,6 +2766,7 @@ def __init__(
         self._operation = operation
         self._operation_id = operation_id
         self._attempt_number = 0
+        self._is_run_command = is_run_command
 
     def run(self) -> T:
         """Runs the supplied func() and attempts a retry
@@ -2799,18 +2807,27 @@ 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
                 # 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
@@ -2841,7 +2858,11 @@ 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:
+                    if (
+                        not self._client.options.retry_writes
+                        or not self._retryable
+                        and not 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")