Deprioritize previous servers for overload retry attempts

Author: blink1073

pymongo/asynchronous/collection.py

@@ -20,7 +20,6 @@
 from typing import (
     TYPE_CHECKING,
     Any,
-    AsyncContextManager,
     Callable,
     Coroutine,
     Generic,
@@ -58,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 (
@@ -573,11 +571,6 @@ async def watch(
         await change_stream._initialize_cursor()
         return change_stream
 
-    async def _conn_for_writes(
-        self, session: Optional[AsyncClientSession], operation: str
-    ) -> AsyncContextManager[AsyncConnection]:
-        return await self._database.client._conn_for_writes(session, operation)
-
     async def _command(
         self,
         conn: AsyncConnection,
@@ -654,7 +647,10 @@ async def _create_helper(
             if "size" in options:
                 options["size"] = float(options["size"])
             cmd.update(options)
-        async with await self._conn_for_writes(session, operation=_Op.CREATE) as conn:
+
+        async def inner(
+            session: Optional[AsyncClientSession], conn: AsyncConnection, _retryable_write: bool
+        ) -> None:
             if qev2_required and conn.max_wire_version < 21:
                 raise ConfigurationError(
                     "Driver support of Queryable Encryption is incompatible with server. "
@@ -671,6 +667,8 @@ async def _create_helper(
                 session=session,
             )
 
+        await self.database.client._retryable_write(False, inner, session, _Op.CREATE)
+
     async def _create(
         self,
         options: MutableMapping[str, Any],
@@ -2229,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]:
@@ -2243,7 +2240,10 @@ async def _create_indexes(
             command (like maxTimeMS) can be passed as keyword arguments.
         """
         names = []
-        async with await self._conn_for_writes(session, operation=_Op.CREATE_INDEXES) as conn:
+
+        async def inner(
+            session: Optional[AsyncClientSession], conn: AsyncConnection, _retryable_write: bool
+        ) -> list[str]:
             supports_quorum = conn.max_wire_version >= 9
 
             def gen_indexes() -> Iterator[Mapping[str, Any]]:
@@ -2272,7 +2272,11 @@ def gen_indexes() -> Iterator[Mapping[str, Any]]:
                 write_concern=self._write_concern_for(session),
                 session=session,
             )
-        return names
+            return names
+
+        return await self.database.client._retryable_write(
+            False, inner, session, _Op.CREATE_INDEXES
+        )
 
     async def create_index(
         self,
@@ -2474,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,
@@ -2493,7 +2496,10 @@ async def _drop_index(
         cmd.update(kwargs)
         if comment is not None:
             cmd["comment"] = comment
-        async with await self._conn_for_writes(session, operation=_Op.DROP_INDEXES) as conn:
+
+        async def inner(
+            session: Optional[AsyncClientSession], conn: AsyncConnection, _retryable_write: bool
+        ) -> None:
             await self._command(
                 conn,
                 cmd,
@@ -2503,6 +2509,8 @@ async def _drop_index(
                 session=session,
             )
 
+        await self.database.client._retryable_write(False, inner, session, _Op.DROP_INDEXES)
+
     async def list_indexes(
         self,
         session: Optional[AsyncClientSession] = None,
@@ -2766,17 +2774,22 @@ def gen_indexes() -> Iterator[Mapping[str, Any]]:
         cmd = {"createSearchIndexes": self.name, "indexes": list(gen_indexes())}
         cmd.update(kwargs)
 
-        async with await self._conn_for_writes(
-            session, operation=_Op.CREATE_SEARCH_INDEXES
-        ) as conn:
+        async def inner(
+            session: Optional[AsyncClientSession], conn: AsyncConnection, _retryable_write: bool
+        ) -> list[str]:
             resp = await self._command(
                 conn,
                 cmd,
                 read_preference=ReadPreference.PRIMARY,
                 codec_options=_UNICODE_REPLACE_CODEC_OPTIONS,
+                session=session,
             )
             return [index["name"] for index in resp["indexesCreated"]]
 
+        return self.database.client._retryable_write(
+            False, inner, session, _Op.CREATE_SEARCH_INDEXES
+        )
+
     async def drop_search_index(
         self,
         name: str,
@@ -2802,15 +2815,21 @@ async def drop_search_index(
         cmd.update(kwargs)
         if comment is not None:
             cmd["comment"] = comment
-        async with await self._conn_for_writes(session, operation=_Op.DROP_SEARCH_INDEXES) as conn:
+
+        async def inner(
+            session: Optional[AsyncClientSession], conn: AsyncConnection, _retryable_write: bool
+        ) -> None:
             await self._command(
                 conn,
                 cmd,
                 read_preference=ReadPreference.PRIMARY,
                 allowable_errors=["ns not found", 26],
                 codec_options=_UNICODE_REPLACE_CODEC_OPTIONS,
+                session=session,
             )
 
+        return self.database.client._retryable_write(False, inner, session, _Op.DROP_SEARCH_INDEXES)
+
     async def update_search_index(
         self,
         name: str,
@@ -2838,15 +2857,21 @@ async def update_search_index(
         cmd.update(kwargs)
         if comment is not None:
             cmd["comment"] = comment
-        async with await self._conn_for_writes(session, operation=_Op.UPDATE_SEARCH_INDEX) as conn:
+
+        async def inner(
+            session: Optional[AsyncClientSession], conn: AsyncConnection, _retryable_write: bool
+        ) -> None:
             await self._command(
                 conn,
                 cmd,
                 read_preference=ReadPreference.PRIMARY,
                 allowable_errors=["ns not found", 26],
                 codec_options=_UNICODE_REPLACE_CODEC_OPTIONS,
+                session=session,
             )
 
+        return self.database.client._retryable_write(False, inner, session, _Op.UPDATE_SEARCH_INDEX)
+
     async def options(
         self,
         session: Optional[AsyncClientSession] = None,
@@ -3075,7 +3100,6 @@ async def aggregate_raw_batches(
             )
 
     @_csot.apply
-    @_retry_overload
     async def rename(
         self,
         new_name: str,
@@ -3127,17 +3151,21 @@ async def rename(
         if comment is not None:
             cmd["comment"] = comment
         write_concern = self._write_concern_for_cmd(cmd, session)
+        client = self._database.client
 
-        async with await self._conn_for_writes(session, operation=_Op.RENAME) as conn:
-            async with self._database.client._tmp_session(session) as s:
-                return await conn.command(
-                    "admin",
-                    cmd,
-                    write_concern=write_concern,
-                    parse_write_concern_error=True,
-                    session=s,
-                    client=self._database.client,
-                )
+        async def inner(
+            session: Optional[AsyncClientSession], conn: AsyncConnection, _retryable_write: bool
+        ) -> MutableMapping[str, Any]:
+            return await conn.command(
+                "admin",
+                cmd,
+                write_concern=write_concern,
+                parse_write_concern_error=True,
+                session=session,
+                client=client,
+            )
+
+        return client._retryable_write(False, inner, session, _Op.RENAME)
 
     async def distinct(
         self,

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]],
@@ -935,14 +932,15 @@ async def command(
 
         if read_preference is None:
             read_preference = (session and session._txn_read_preference()) or ReadPreference.PRIMARY
-        async with await self._client._conn_for_reads(
-            read_preference, session, operation=command_name
-        ) as (
-            connection,
-            read_preference,
-        ):
+
+        async def inner(
+            session: Optional[AsyncClientSession],
+            _server: Server,
+            conn: AsyncConnection,
+            read_preference: _ServerMode,
+        ) -> Union[dict[str, Any], _CodecDocumentType]:
             return await self._command(
-                connection,
+                conn,
                 command,
                 value,
                 check,
@@ -953,8 +951,11 @@ async def command(
                 **kwargs,
             )
 
+        return await self._client._retryable_read(
+            inner, read_preference, session, command_name, None, False
+        )
+
     @_csot.apply
-    @_retry_overload
     async def cursor_command(
         self,
         command: Union[str, MutableMapping[str, Any]],
@@ -1021,17 +1022,17 @@ async def cursor_command(
 
         async with self._client._tmp_session(session) as tmp_session:
             opts = codec_options or DEFAULT_CODEC_OPTIONS
-
             if read_preference is None:
                 read_preference = (
                     tmp_session and tmp_session._txn_read_preference()
                 ) or ReadPreference.PRIMARY
-            async with await self._client._conn_for_reads(
-                read_preference, tmp_session, command_name
-            ) as (
-                conn,
-                read_preference,
-            ):
+
+            async def inner(
+                session: Optional[AsyncClientSession],
+                _server: Server,
+                conn: AsyncConnection,
+                read_preference: _ServerMode,
+            ) -> AsyncCommandCursor[_DocumentType]:
                 response = await self._command(
                     conn,
                     command,
@@ -1040,7 +1041,7 @@ async def cursor_command(
                     None,
                     read_preference,
                     opts,
-                    session=tmp_session,
+                    session=session,
                     **kwargs,
                 )
                 coll = self.get_collection("$cmd", read_preference=read_preference)
@@ -1050,14 +1051,18 @@ async def cursor_command(
                         response["cursor"],
                         conn.address,
                         max_await_time_ms=max_await_time_ms,
-                        session=tmp_session,
+                        session=session,
                         comment=comment,
                     )
                     await cmd_cursor._maybe_pin_connection(conn)
                     return cmd_cursor
                 else:
                     raise InvalidOperation("Command does not return a cursor.")
 
+            return await self.client._retryable_read(
+                inner, read_preference, tmp_session, command_name, False
+            )
+
     async def _retryable_read_command(
         self,
         command: Union[str, MutableMapping[str, Any]],
@@ -1259,18 +1264,21 @@ async def _drop_helper(
         if comment is not None:
             command["comment"] = comment
 
-        async with await self._client._conn_for_writes(session, operation=_Op.DROP) as connection:
+        async def inner(
+            session: Optional[AsyncClientSession], conn: AsyncConnection, _retryable_write: bool
+        ) -> dict[str, Any]:
             return await self._command(
-                connection,
+                conn,
                 command,
                 allowable_errors=["ns not found", 26],
                 write_concern=self._write_concern_for(session),
                 parse_write_concern_error=True,
                 session=session,
             )
 
+        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]],