PYTHON-5767 Implement DRIVERS-3427 phase 1 rollout changes

- Change MAX_RETRIES default from 5 to 2
- Add configurable maxAdaptiveRetries client option (default: 2)
- Add enableOverloadRetargeting client option (default: false)
- Add retry metadata ("retry": N) to outgoing command bodies on retry
- Remove _retry_overload decorator from collection/database methods

Co-Authored-By: Claude Code <noreply@anthropic.com>

Author: Jibola

pymongo/asynchronous/collection.py

@@ -57,7 +57,6 @@
     AsyncCursor,
     AsyncRawBatchCursor,
 )
-from pymongo.asynchronous.helpers import _retry_overload
 from pymongo.collation import validate_collation_or_none
 from pymongo.common import _ecoc_coll_name, _esc_coll_name
 from pymongo.errors import (
@@ -2228,7 +2227,6 @@ async def create_indexes(
         return await self._create_indexes(indexes, session, **kwargs)
 
     @_csot.apply
-    @_retry_overload
     async def _create_indexes(
         self, indexes: Sequence[IndexModel], session: Optional[AsyncClientSession], **kwargs: Any
     ) -> list[str]:
@@ -2480,7 +2478,6 @@ async def drop_index(
         await self._drop_index(index_or_name, session, comment, **kwargs)
 
     @_csot.apply
-    @_retry_overload
     async def _drop_index(
         self,
         index_or_name: _IndexKeyHint,
@@ -3104,7 +3101,6 @@ async def aggregate_raw_batches(
             )
 
     @_csot.apply
-    @_retry_overload
     async def rename(
         self,
         new_name: str,

pymongo/asynchronous/database.py

@@ -38,7 +38,6 @@
 from pymongo.asynchronous.change_stream import AsyncDatabaseChangeStream
 from pymongo.asynchronous.collection import AsyncCollection
 from pymongo.asynchronous.command_cursor import AsyncCommandCursor
-from pymongo.asynchronous.helpers import _retry_overload
 from pymongo.common import _ecoc_coll_name, _esc_coll_name
 from pymongo.database_shared import _check_name, _CodecDocumentType
 from pymongo.errors import CollectionInvalid, InvalidOperation
@@ -479,7 +478,6 @@ async def watch(
         return change_stream
 
     @_csot.apply
-    @_retry_overload
     async def create_collection(
         self,
         name: str,
@@ -822,7 +820,6 @@ async def command(
         ...
 
     @_csot.apply
-    @_retry_overload
     async def command(
         self,
         command: Union[str, MutableMapping[str, Any]],
@@ -959,7 +956,6 @@ async def inner(
         )
 
     @_csot.apply
-    @_retry_overload
     async def cursor_command(
         self,
         command: Union[str, MutableMapping[str, Any]],
@@ -1283,7 +1279,6 @@ async def inner(
         return await self.client._retryable_write(False, inner, session, _Op.DROP)
 
     @_csot.apply
-    @_retry_overload
     async def drop_collection(
         self,
         name_or_collection: Union[str, AsyncCollection[_DocumentTypeArg]],

pymongo/asynchronous/helpers.py

@@ -22,6 +22,7 @@
 import socket
 import sys
 import time as time  # noqa: PLC0414 # needed in sync version
+from contextvars import ContextVar
 from typing import (
     Any,
     Callable,
@@ -76,9 +77,13 @@ async def inner(*args: Any, **kwargs: Any) -> Any:
     return cast(F, inner)
 
 
-_MAX_RETRIES = 5
+_MAX_RETRIES = 2
 _BACKOFF_INITIAL = 0.1
 _BACKOFF_MAX = 10
+
+# Context variable used to pass the current retry attempt number to conn.command()
+# so that retry metadata can be injected into outgoing command bodies.
+_RETRY_ATTEMPT: ContextVar[int] = ContextVar("_retry_attempt", default=0)
 DEFAULT_RETRY_TOKEN_CAPACITY = 1000.0
 DEFAULT_RETRY_TOKEN_RETURN = 0.1
 

pymongo/asynchronous/mongo_client.py

@@ -69,6 +69,7 @@
 from pymongo.asynchronous.client_session import _SESSION, _EmptyServerSession
 from pymongo.asynchronous.command_cursor import AsyncCommandCursor
 from pymongo.asynchronous.helpers import (
+    _RETRY_ATTEMPT,
     _RetryPolicy,
     _TokenBucket,
 )
@@ -895,7 +896,9 @@ def __init__(
         )
 
         self._retry_policy = _RetryPolicy(
-            _TokenBucket(), adaptive_retry=self._options.adaptive_retries
+            _TokenBucket(),
+            attempts=self._options.max_adaptive_retries,
+            adaptive_retry=self._options.adaptive_retries,
         )
 
         self._init_based_on_options(self._seeds, srv_max_hosts, srv_service_name)
@@ -2820,6 +2823,7 @@ async def run(self) -> T:
 
         while True:
             self._check_last_error(check_csot=True)
+            retry_token = _RETRY_ATTEMPT.set(self._attempt_number)
             try:
                 res = await self._read() if self._is_read else await self._write()
                 await self._retry_policy.record_success(self._attempt_number > 0)
@@ -2930,7 +2934,7 @@ async def run(self) -> T:
                             transaction.set_starting()
                         transaction.attempt = 0
 
-                if (
+                if self._client.options.enable_overload_retargeting and (
                     self._server is not None
                     and self._client.topology_description.topology_type_name == "Sharded"
                     or exc.has_error_label("SystemOverloadedError")
@@ -2946,6 +2950,8 @@ async def run(self) -> T:
                         else:
                             raise
                     await asyncio.sleep(delay)
+            finally:
+                _RETRY_ATTEMPT.reset(retry_token)
 
     def _is_not_eligible_for_retry(self) -> bool:
         """Checks if the exchange is not eligible for retry"""

pymongo/asynchronous/pool.py

@@ -39,7 +39,7 @@
 from bson import DEFAULT_CODEC_OPTIONS
 from pymongo import _csot, helpers_shared
 from pymongo.asynchronous.client_session import _validate_session_write_concern
-from pymongo.asynchronous.helpers import _handle_reauth
+from pymongo.asynchronous.helpers import _RETRY_ATTEMPT, _handle_reauth
 from pymongo.asynchronous.network import command
 from pymongo.common import (
     MAX_BSON_SIZE,
@@ -395,6 +395,9 @@ async def command(
         if session:
             session._apply_to(spec, retryable_write, read_preference, self)
         self.send_cluster_time(spec, session, client)
+        retry_attempt = _RETRY_ATTEMPT.get()
+        if retry_attempt > 0:
+            spec["retry"] = retry_attempt
         listeners = self.listeners if publish_events else None
         unacknowledged = bool(write_concern and not write_concern.acknowledged)
         if self.op_msg_enabled:

pymongo/client_options.py

@@ -240,6 +240,16 @@ def __init__(
             if "adaptive_retries" in options
             else options.get("adaptiveretries", common.ADAPTIVE_RETRIES)
         )
+        self.__max_adaptive_retries = (
+            options.get("max_adaptive_retries", common.MAX_ADAPTIVE_RETRIES)
+            if "max_adaptive_retries" in options
+            else options.get("maxadaptiveretries", common.MAX_ADAPTIVE_RETRIES)
+        )
+        self.__enable_overload_retargeting = (
+            options.get("enable_overload_retargeting", common.ENABLE_OVERLOAD_RETARGETING)
+            if "enable_overload_retargeting" in options
+            else options.get("enableoverloadretargeting", common.ENABLE_OVERLOAD_RETARGETING)
+        )
 
     @property
     def _options(self) -> Mapping[str, Any]:
@@ -359,3 +369,19 @@ def adaptive_retries(self) -> bool:
         .. versionadded:: 4.XX
         """
         return self.__adaptive_retries
+
+    @property
+    def max_adaptive_retries(self) -> int:
+        """The configured maxAdaptiveRetries option.
+
+        .. versionadded:: 4.XX
+        """
+        return self.__max_adaptive_retries
+
+    @property
+    def enable_overload_retargeting(self) -> bool:
+        """The configured enableOverloadRetargeting option.
+
+        .. versionadded:: 4.XX
+        """
+        return self.__enable_overload_retargeting

pymongo/common.py

@@ -143,6 +143,12 @@
 # Default value for adaptiveRetries
 ADAPTIVE_RETRIES = False
 
+# Default value for maxAdaptiveRetries
+MAX_ADAPTIVE_RETRIES = 2
+
+# Default value for enableOverloadRetargeting
+ENABLE_OVERLOAD_RETARGETING = False
+
 # Auth mechanism properties that must raise an error instead of warning if they invalidate.
 _MECH_PROP_MUST_RAISE = ["CANONICALIZE_HOST_NAME"]
 
@@ -776,6 +782,8 @@ def validate_server_monitoring_mode(option: str, value: str) -> str:
     "auto_encryption_opts": validate_auto_encryption_opts_or_none,
     "authoidcallowedhosts": validate_list,
     "adaptive_retries": validate_boolean_or_string,
+    "max_adaptive_retries": validate_non_negative_integer,
+    "enable_overload_retargeting": validate_boolean_or_string,
 }
 
 # Dictionary where keys are any URI option name, and values are the

pymongo/synchronous/collection.py

@@ -88,7 +88,6 @@
     Cursor,
     RawBatchCursor,
 )
-from pymongo.synchronous.helpers import _retry_overload
 from pymongo.typings import _CollationIn, _DocumentType, _DocumentTypeArg, _Pipeline
 from pymongo.write_concern import DEFAULT_WRITE_CONCERN, WriteConcern, validate_boolean
 
@@ -2225,7 +2224,6 @@ def create_indexes(
         return self._create_indexes(indexes, session, **kwargs)
 
     @_csot.apply
-    @_retry_overload
     def _create_indexes(
         self, indexes: Sequence[IndexModel], session: Optional[ClientSession], **kwargs: Any
     ) -> list[str]:
@@ -2475,7 +2473,6 @@ def drop_index(
         self._drop_index(index_or_name, session, comment, **kwargs)
 
     @_csot.apply
-    @_retry_overload
     def _drop_index(
         self,
         index_or_name: _IndexKeyHint,
@@ -3097,7 +3094,6 @@ def aggregate_raw_batches(
             )
 
     @_csot.apply
-    @_retry_overload
     def rename(
         self,
         new_name: str,

pymongo/synchronous/database.py

@@ -43,7 +43,6 @@
 from pymongo.synchronous.change_stream import DatabaseChangeStream
 from pymongo.synchronous.collection import Collection
 from pymongo.synchronous.command_cursor import CommandCursor
-from pymongo.synchronous.helpers import _retry_overload
 from pymongo.typings import _CollationIn, _DocumentType, _DocumentTypeArg, _Pipeline
 
 if TYPE_CHECKING:
@@ -479,7 +478,6 @@ def watch(
         return change_stream
 
     @_csot.apply
-    @_retry_overload
     def create_collection(
         self,
         name: str,
@@ -822,7 +820,6 @@ def command(
         ...
 
     @_csot.apply
-    @_retry_overload
     def command(
         self,
         command: Union[str, MutableMapping[str, Any]],
@@ -959,7 +956,6 @@ def inner(
         )
 
     @_csot.apply
-    @_retry_overload
     def cursor_command(
         self,
         command: Union[str, MutableMapping[str, Any]],
@@ -1280,7 +1276,6 @@ def inner(
         return self.client._retryable_write(False, inner, session, _Op.DROP)
 
     @_csot.apply
-    @_retry_overload
     def drop_collection(
         self,
         name_or_collection: Union[str, Collection[_DocumentTypeArg]],

pymongo/synchronous/helpers.py

@@ -22,6 +22,7 @@
 import socket
 import sys
 import time as time  # noqa: PLC0414 # needed in sync version
+from contextvars import ContextVar
 from typing import (
     Any,
     Callable,
@@ -76,9 +77,13 @@ def inner(*args: Any, **kwargs: Any) -> Any:
     return cast(F, inner)
 
 
-_MAX_RETRIES = 5
+_MAX_RETRIES = 2
 _BACKOFF_INITIAL = 0.1
 _BACKOFF_MAX = 10
+
+# Context variable used to pass the current retry attempt number to conn.command()
+# so that retry metadata can be injected into outgoing command bodies.
+_RETRY_ATTEMPT: ContextVar[int] = ContextVar("_retry_attempt", default=0)
 DEFAULT_RETRY_TOKEN_CAPACITY = 1000.0
 DEFAULT_RETRY_TOKEN_RETURN = 0.1