PYTHON-4008 Support exhaust cursors on mongos 7.1+

doc/changelog.rst

@@ -11,6 +11,14 @@ PyMongo 4.18 brings a number of changes including:
   Session resumption is supported on all Python versions for synchronous clients
   and on Python 3.11+ for async clients.
 - Improved performance for MongoDB 9.0's Intelligent Workload Management (IWM) by only retrying overload errors when doing so is expected to not worsen server conditions.
+- Added support for exhaust cursors (:attr:`~pymongo.cursor.CursorType.EXHAUST`)
+  against mongos 7.1+. An older mongos still raises
+  :class:`~pymongo.errors.InvalidOperation`, now on the first iteration of the
+  cursor rather than from
+  :meth:`~pymongo.synchronous.collection.Collection.find`, since the requirement
+  is checked against the connection in use. Separately, async cursors combining
+  ``limit`` with :attr:`~pymongo.cursor.CursorType.EXHAUST` now raise at ``find``
+  rather than on first iteration, matching the synchronous API.
 - Redacted potentially sensitive authentication mechanism properties, including
   AWS session tokens, from the representations of
   :class:`~pymongo.synchronous.mongo_client.MongoClient` and

doc/contributors.rst

@@ -110,3 +110,4 @@ The following is a list of people who have contributed to
 - Cal Jacobson (cj81499)
 - Sophia Yang (sophiayangDB)
 - Madan Kumar (winklemad)
+- Junyao Dong (carsontung666)

pymongo/asynchronous/collection.py

@@ -1881,8 +1881,14 @@ def find(self, *args: Any, **kwargs: Any) -> AsyncCursor[_DocumentType]:
 
           - The `limit` option can not be used with an exhaust cursor.
 
-          - Exhaust cursors are not supported by mongos and can not be
-            used with a sharded cluster.
+          - When connected to mongos, exhaust cursors require MongoDB 7.1 or
+            newer; an older one raises
+            :class:`~pymongo.errors.InvalidOperation` on the first iteration.
+            On a cluster midway through an upgrade the same call may succeed
+            or raise depending on which mongos it reaches. Behind a load
+            balancer no version check applies; a pre-7.1 mongos there ignores
+            the exhaust request and the cursor falls back to ordinary getMore
+            batches.
 
           - A :class:`~pymongo.cursor.AsyncCursor` instance created with the
             :attr:`~pymongo.cursor.CursorType.EXHAUST` cursor_type requires an
@@ -1892,6 +1898,12 @@ def find(self, *args: Any, **kwargs: Any) -> AsyncCursor[_DocumentType]:
             connection will be closed and discarded without being returned to
             the connection pool.
 
+        .. versionchanged:: 4.18
+           :attr:`~pymongo.cursor.CursorType.EXHAUST` is now permitted when
+           connected to mongos 7.1 or newer. Against an older mongos the
+           resulting :class:`~pymongo.errors.InvalidOperation` is now raised on
+           the first iteration of the cursor rather than by this method.
+
         .. versionchanged:: 4.0
            Removed the ``modifiers`` option.
            Empty projections (eg {} or []) are passed to the server as-is,

pymongo/asynchronous/cursor.py

@@ -213,19 +213,15 @@ def __init__(
         self._dbname = collection.database.name
         self._collname = collection.name
 
-        # Checking exhaust cursor support requires network IO
-        if _IS_SYNC:
-            self._exhaust_checked = True
-            self._supports_exhaust()  # type: ignore[unused-coroutine]
-        else:
-            self._exhaust = cursor_type == CursorType.EXHAUST
-            self._exhaust_checked = False
+        self._validate_exhaust_handling()
 
-    async def _supports_exhaust(self) -> None:
-        # Exhaust cursor support
+    def _validate_exhaust_handling(self) -> None:
+        """Reject option combinations an exhaust cursor cannot serve.
+
+        Server support is checked against the connection in use, in
+        _check_exhaust_supported.
+        """
         if self._cursor_type == CursorType.EXHAUST:
-            if await self._collection.database.client.is_mongos:
-                raise InvalidOperation("Exhaust cursors are not supported by mongos")
             if self._limit:
                 raise InvalidOperation("Can't use limit and exhaust together.")
             self._exhaust = True
@@ -369,8 +365,6 @@ async def add_option(self, mask: int) -> AsyncCursor[_DocumentType]:
         if mask & _QUERY_OPTIONS["exhaust"]:
             if self._limit:
                 raise InvalidOperation("Can't use limit and exhaust together.")
-            if await self._collection.database.client.is_mongos:
-                raise InvalidOperation("Exhaust cursors are not supported by mongos")
             self._exhaust = True
 
         self._query_flags |= mask
@@ -1125,9 +1119,6 @@ async def rewind(self) -> AsyncCursor[_DocumentType]:
 
     async def next(self) -> _DocumentType:
         """Advance the cursor."""
-        if not self._exhaust_checked:
-            self._exhaust_checked = True
-            await self._supports_exhaust()
         if self._empty:
             raise StopAsyncIteration
         if len(self._data) or await self._refresh():
@@ -1137,9 +1128,6 @@ async def next(self) -> _DocumentType:
 
     async def _next_batch(self, result: list, total: Optional[int] = None) -> bool:  # type: ignore[type-arg]
         """Get all or some documents from the cursor."""
-        if not self._exhaust_checked:
-            self._exhaust_checked = True
-            await self._supports_exhaust()
         if self._empty:
             return False
         if len(self._data) or await self._refresh():

pymongo/common.py

@@ -69,6 +69,9 @@
 # MongoDB 9.0
 MAX_SUPPORTED_WIRE_VERSION = 29
 
+# MongoDB 7.1, the first release whose mongos continues an exhaust getMore stream.
+MONGOS_EXHAUST_WIRE_VERSION = 22
+
 # Frequency to call hello on servers, in seconds.
 HEARTBEAT_FREQUENCY = 10
 

pymongo/message.py

@@ -43,6 +43,7 @@
     RawBSONDocument,
     _inflate_bson,
 )
+from pymongo.common import MONGOS_EXHAUST_WIRE_VERSION
 from pymongo.monitoring import _EventListeners
 
 try:
@@ -52,7 +53,6 @@
 except ImportError:
     _use_c = False
 from pymongo.errors import (
-    ConfigurationError,
     DocumentTooLarge,
     InvalidOperation,
     ProtocolError,
@@ -1220,6 +1220,17 @@ def unpack(cls, msg: bytes | memoryview) -> _OpMsg:
 }
 
 
+def _check_exhaust_supported(conn: _AgnosticConnection) -> None:
+    """Raise if this connection's server will not continue an exhaust stream.
+
+    mongos gained exhaust getMore in 7.1 (SERVER-57297). Load-balanced deployments
+    are left alone: every cursor pins its connection there anyway, so an older mongos
+    costs nothing extra, and refusing would break clients that work today.
+    """
+    if conn.is_mongos and conn.max_wire_version < MONGOS_EXHAUST_WIRE_VERSION:
+        raise InvalidOperation("Exhaust cursors require MongoDB 7.1+ when connected to mongos.")
+
+
 class _Query:
     """A query operation."""
 
@@ -1293,20 +1304,11 @@ def namespace(self) -> str:
         return f"{self.db}.{self.coll}"
 
     def use_command(self, conn: _AgnosticConnection) -> bool:
-        use_find_cmd = False
-        if not self.exhaust:
-            use_find_cmd = True
-        elif conn.max_wire_version >= 8:
-            # OP_MSG supports exhaust on MongoDB 4.2+
-            use_find_cmd = True
-        elif not self.read_concern.ok_for_legacy:
-            raise ConfigurationError(
-                f"read concern level of {self.read_concern.level} is not valid "
-                f"with a max wire version of {conn.max_wire_version}."
-            )
+        if self.exhaust:
+            _check_exhaust_supported(conn)
 
         conn.validate_session(self.client, self.session)  # type: ignore[arg-type]
-        return use_find_cmd
+        return True
 
     def update_command(self, cmd: dict[str, Any]) -> None:
         self._as_command = cmd, self.db
@@ -1426,15 +1428,11 @@ def namespace(self) -> str:
         return f"{self.db}.{self.coll}"
 
     def use_command(self, conn: _AgnosticConnection) -> bool:
-        use_cmd = False
-        if not self.exhaust:
-            use_cmd = True
-        elif conn.max_wire_version >= 8:
-            # OP_MSG supports exhaust on MongoDB 4.2+
-            use_cmd = True
+        if self.exhaust:
+            _check_exhaust_supported(conn)
 
         conn.validate_session(self.client, self.session)  # type: ignore[arg-type]
-        return use_cmd
+        return True
 
     def update_command(self, cmd: dict[str, Any]) -> None:
         self._as_command = cmd, self.db

pymongo/synchronous/collection.py

@@ -1881,8 +1881,14 @@ def find(self, *args: Any, **kwargs: Any) -> Cursor[_DocumentType]:
 
           - The `limit` option can not be used with an exhaust cursor.
 
-          - Exhaust cursors are not supported by mongos and can not be
-            used with a sharded cluster.
+          - When connected to mongos, exhaust cursors require MongoDB 7.1 or
+            newer; an older one raises
+            :class:`~pymongo.errors.InvalidOperation` on the first iteration.
+            On a cluster midway through an upgrade the same call may succeed
+            or raise depending on which mongos it reaches. Behind a load
+            balancer no version check applies; a pre-7.1 mongos there ignores
+            the exhaust request and the cursor falls back to ordinary getMore
+            batches.
 
           - A :class:`~pymongo.cursor.Cursor` instance created with the
             :attr:`~pymongo.cursor.CursorType.EXHAUST` cursor_type requires an
@@ -1892,6 +1898,12 @@ def find(self, *args: Any, **kwargs: Any) -> Cursor[_DocumentType]:
             connection will be closed and discarded without being returned to
             the connection pool.
 
+        .. versionchanged:: 4.18
+           :attr:`~pymongo.cursor.CursorType.EXHAUST` is now permitted when
+           connected to mongos 7.1 or newer. Against an older mongos the
+           resulting :class:`~pymongo.errors.InvalidOperation` is now raised on
+           the first iteration of the cursor rather than by this method.
+
         .. versionchanged:: 4.0
            Removed the ``modifiers`` option.
            Empty projections (eg {} or []) are passed to the server as-is,

pymongo/synchronous/cursor.py

@@ -213,19 +213,15 @@ def __init__(
         self._dbname = collection.database.name
         self._collname = collection.name
 
-        # Checking exhaust cursor support requires network IO
-        if _IS_SYNC:
-            self._exhaust_checked = True
-            self._supports_exhaust()  # type: ignore[unused-coroutine]
-        else:
-            self._exhaust = cursor_type == CursorType.EXHAUST
-            self._exhaust_checked = False
+        self._validate_exhaust_handling()
 
-    def _supports_exhaust(self) -> None:
-        # Exhaust cursor support
+    def _validate_exhaust_handling(self) -> None:
+        """Reject option combinations an exhaust cursor cannot serve.
+
+        Server support is checked against the connection in use, in
+        _check_exhaust_supported.
+        """
         if self._cursor_type == CursorType.EXHAUST:
-            if self._collection.database.client.is_mongos:
-                raise InvalidOperation("Exhaust cursors are not supported by mongos")
             if self._limit:
                 raise InvalidOperation("Can't use limit and exhaust together.")
             self._exhaust = True
@@ -369,8 +365,6 @@ def add_option(self, mask: int) -> Cursor[_DocumentType]:
         if mask & _QUERY_OPTIONS["exhaust"]:
             if self._limit:
                 raise InvalidOperation("Can't use limit and exhaust together.")
-            if self._collection.database.client.is_mongos:
-                raise InvalidOperation("Exhaust cursors are not supported by mongos")
             self._exhaust = True
 
         self._query_flags |= mask
@@ -1121,9 +1115,6 @@ def rewind(self) -> Cursor[_DocumentType]:
 
     def next(self) -> _DocumentType:
         """Advance the cursor."""
-        if not self._exhaust_checked:
-            self._exhaust_checked = True
-            self._supports_exhaust()
         if self._empty:
             raise StopIteration
         if len(self._data) or self._refresh():
@@ -1133,9 +1124,6 @@ def next(self) -> _DocumentType:
 
     def _next_batch(self, result: list, total: Optional[int] = None) -> bool:  # type: ignore[type-arg]
         """Get all or some documents from the cursor."""
-        if not self._exhaust_checked:
-            self._exhaust_checked = True
-            self._supports_exhaust()
         if self._empty:
             return False
         if len(self._data) or self._refresh():

test/__init__.py

@@ -763,6 +763,20 @@ def require_retryable_writes(self, func):
             func=func,
         )
 
+    def supports_exhaust_cursors(self):
+        """Whether this deployment supports exhaust cursors."""
+        if self.is_mongos:
+            return self.version.at_least(7, 1)
+        return True
+
+    def require_exhaust_cursors(self, func):
+        """Run a test only if the deployment supports exhaust cursors."""
+        return self._require(
+            self.supports_exhaust_cursors,
+            "This server does not support exhaust cursors",
+            func=func,
+        )
+
     def supports_transactions(self):
         if self.version.at_least(4, 1, 8):
             return self.is_mongos or self.is_rs

test/asynchronous/__init__.py

@@ -763,6 +763,20 @@ def require_retryable_writes(self, func):
             func=func,
         )
 
+    def supports_exhaust_cursors(self):
+        """Whether this deployment supports exhaust cursors."""
+        if self.is_mongos:
+            return self.version.at_least(7, 1)
+        return True
+
+    def require_exhaust_cursors(self, func):
+        """Run a test only if the deployment supports exhaust cursors."""
+        return self._require(
+            self.supports_exhaust_cursors,
+            "This server does not support exhaust cursors",
+            func=func,
+        )
+
     def supports_transactions(self):
         if self.version.at_least(4, 1, 8):
             return self.is_mongos or self.is_rs

test/asynchronous/test_client.py

@@ -2417,8 +2417,8 @@ class TestExhaustCursor(AsyncIntegrationTest):
 
     def setUp(self):
         super().setUp()
-        if async_client_context.is_mongos:
-            raise SkipTest("mongos doesn't support exhaust, SERVER-2627")
+        if not async_client_context.supports_exhaust_cursors():
+            raise SkipTest("mongos serves exhaust cursors only from 7.1, SERVER-57297")
 
     async def test_exhaust_query_server_error(self):
         # When doing an exhaust query, the socket stays checked out on success

test/asynchronous/test_collection.py

@@ -27,7 +27,7 @@
 
 from pymongo.asynchronous.database import AsyncDatabase
 from pymongo.asynchronous.helpers import anext
-from test.asynchronous.utils import async_get_pool, async_is_mongos
+from test.asynchronous.utils import async_get_pool
 
 sys.path[0:0] = [""]
 
@@ -1796,8 +1796,15 @@ async def test_cursor_timeout(self):
         await self.db.test.find(no_cursor_timeout=True).to_list()
         await self.db.test.find(no_cursor_timeout=False).to_list()
 
+    async def test_exhaust_limit_raises_without_iterating(self):
+        # The limit conflict is settled at find(); the mongos wire version is not.
+        with self.assertRaises(InvalidOperation):
+            self.db.test.find(cursor_type=CursorType.EXHAUST, limit=5)
+        self.db.test.find(cursor_type=CursorType.EXHAUST)
+
     async def test_exhaust(self):
-        if await async_is_mongos(self.db.client):
+        # mongos only serves exhaust cursors from 7.1 onwards (SERVER-57297).
+        if not async_client_context.supports_exhaust_cursors():
             with self.assertRaises(InvalidOperation):
                 await anext(self.db.test.find(cursor_type=CursorType.EXHAUST))
             return

test/asynchronous/test_cursor.py

@@ -130,8 +130,8 @@ async def test_add_remove_option(self):
         self.assertEqual(0, cursor._query_flags)
 
     async def test_add_remove_option_exhaust(self):
-        # Exhaust - which mongos doesn't support
-        if async_client_context.is_mongos:
+        # mongos only serves exhaust cursors from 7.1 onwards (SERVER-57297).
+        if not async_client_context.supports_exhaust_cursors():
             with self.assertRaises(InvalidOperation):
                 await anext(self.db.test.find(cursor_type=CursorType.EXHAUST))
         else:
@@ -1615,7 +1615,7 @@ async def test_clone(self):
         self.assertIsInstance(await anext(cursor.clone()), bytes)
         self.assertIsInstance(await anext(copy.copy(cursor)), bytes)
 
-    @async_client_context.require_no_mongos
+    @async_client_context.require_exhaust_cursors
     async def test_exhaust(self):
         c = self.db.test
         await c.insert_many({"_id": i} for i in range(200))
@@ -1849,7 +1849,7 @@ async def test_monitoring(self):
             listener.reset()
 
     @async_client_context.require_version_min(5, 0, -1)
-    @async_client_context.require_no_mongos
+    @async_client_context.require_exhaust_cursors
     @async_client_context.require_sync
     async def test_exhaust_cursor_db_set(self):
         listener = OvertCommandListener()