diff --git a/pymongo/asynchronous/command_cursor.py b/pymongo/asynchronous/command_cursor.py index 71404281a4..c05c4bd49d 100644 --- a/pymongo/asynchronous/command_cursor.py +++ b/pymongo/asynchronous/command_cursor.py @@ -17,7 +17,7 @@ from __future__ import annotations from collections import deque -from collections.abc import AsyncIterator, Mapping, Sequence +from collections.abc import AsyncIterator, Mapping from typing import ( TYPE_CHECKING, Any, @@ -27,11 +27,11 @@ from bson import CodecOptions, _convert_raw_document_lists_to_streams from pymongo.asynchronous.cursor_base import _AsyncCursorBase, _ConnectionManager -from pymongo.cursor_shared import _CURSOR_CLOSED_ERRORS +from pymongo.cursor_shared import _CURSOR_CLOSED_ERRORS, _AgnosticCommandCursorBase from pymongo.errors import ConnectionFailure, InvalidOperation, OperationFailure from pymongo.message import _GetMore, _OpMsg, _RawBatchGetMore from pymongo.response import PinnedResponse -from pymongo.typings import _Address, _DocumentOut, _DocumentType +from pymongo.typings import _Address, _DocumentType if TYPE_CHECKING: from pymongo.asynchronous.client_session import AsyncClientSession @@ -41,7 +41,9 @@ _IS_SYNC = False -class AsyncCommandCursor(_AsyncCursorBase[_DocumentType]): +class AsyncCommandCursor( + _AgnosticCommandCursorBase[_DocumentType], _AsyncCursorBase[_DocumentType] +): """An asynchronous cursor / iterator over command cursors. Used by :meth:`~pymongo.asynchronous.collection.AsyncCollection.aggregate`, :meth:`~pymongo.asynchronous.database.AsyncDatabase.aggregate`, @@ -55,89 +57,6 @@ class AsyncCommandCursor(_AsyncCursorBase[_DocumentType]): Should not be called directly by application developers. """ - _getmore_class = _GetMore - - def __init__( - self, - collection: AsyncCollection[_DocumentType], - cursor_info: Mapping[str, Any], - address: Optional[_Address], - batch_size: int = 0, - max_await_time_ms: Optional[int] = None, - session: Optional[AsyncClientSession] = None, - comment: Any = None, - ) -> None: - """Create a new command cursor.""" - self._sock_mgr: Any = None - self._collection: AsyncCollection[_DocumentType] = collection - self._id = cursor_info["id"] - self._data = deque(cursor_info["firstBatch"]) - self._postbatchresumetoken: Optional[Mapping[str, Any]] = cursor_info.get( - "postBatchResumeToken" - ) - self._address = address - self._batch_size = batch_size - self._max_await_time_ms = max_await_time_ms - self._timeout = self._collection.database.client.options.timeout - self._session = session - if self._session is not None: - self._session._attached_to_cursor = True - self._killed = self._id == 0 - self._comment = comment - if self._killed: - self._end_session() - - if "ns" in cursor_info: - self._ns = cursor_info["ns"] - else: - self._ns = collection.full_name - - self.batch_size(batch_size) - - if not isinstance(max_await_time_ms, int) and max_await_time_ms is not None: - raise TypeError( - f"max_await_time_ms must be an integer or None, not {type(max_await_time_ms)}" - ) - - def _get_namespace(self) -> str: - return self._ns - - def batch_size(self, batch_size: int) -> AsyncCommandCursor[_DocumentType]: - """Limits the number of documents returned in one batch. Each batch - requires a round trip to the server. It can be adjusted to optimize - performance and limit data transfer. - - .. note:: batch_size can not override MongoDB's internal limits on the - amount of data it will return to the client in a single batch (i.e - if you set batch size to 1,000,000,000, MongoDB will currently only - return 4-16MB of results per batch). - - Raises :exc:`TypeError` if `batch_size` is not an integer. - Raises :exc:`ValueError` if `batch_size` is less than ``0``. - - :param batch_size: The size of each batch of results requested. - """ - if not isinstance(batch_size, int): - raise TypeError(f"batch_size must be an integer, not {type(batch_size)}") - if batch_size < 0: - raise ValueError("batch_size must be >= 0") - - self._batch_size = (batch_size == 1 and 2) or batch_size - return self - - def _has_next(self) -> bool: - """Returns `True` if the cursor has documents remaining from the - previous batch. - """ - return len(self._data) > 0 - - @property - def _post_batch_resume_token(self) -> Optional[Mapping[str, Any]]: - """Retrieve the postBatchResumeToken from the response to a - changeStream aggregate or getMore. - """ - return self._postbatchresumetoken - async def _maybe_pin_connection(self, conn: AsyncConnection) -> None: client = self._collection.database.client if not client._should_pin_cursor(self._session): @@ -152,22 +71,6 @@ async def _maybe_pin_connection(self, conn: AsyncConnection) -> None: else: self._sock_mgr = conn_mgr - def _unpack_response( - self, - response: _OpMsg, - cursor_id: Optional[int], - codec_options: CodecOptions[Mapping[str, Any]], - user_fields: Optional[Mapping[str, Any]] = None, - legacy_response: bool = False, - ) -> Sequence[_DocumentOut]: - return response.unpack_response(cursor_id, codec_options, user_fields, legacy_response) - - def _end_session(self) -> None: - if self._session and self._session._implicit: - self._session._attached_to_cursor = False - self._session._end_implicit_session() - self._session = None - async def _send_message(self, operation: _GetMore) -> None: """Send a getmore message and handle the response.""" client = self._collection.database.client diff --git a/pymongo/cursor_shared.py b/pymongo/cursor_shared.py index df0e1e2f58..115c7a51b7 100644 --- a/pymongo/cursor_shared.py +++ b/pymongo/cursor_shared.py @@ -18,11 +18,22 @@ from __future__ import annotations from abc import ABC, abstractmethod +from collections import deque from collections.abc import Mapping, Sequence -from typing import Any, Generic, Optional, Union +from typing import TYPE_CHECKING, Any, Generic, Optional, Union + +from bson import CodecOptions +from pymongo.message import _CursorAddress, _GetMore, _OpMsg +from pymongo.typings import ( + _Address, + _AgnosticClientSession, + _AgnosticCollection, + _DocumentOut, + _DocumentType, +) -from pymongo.message import _CursorAddress -from pymongo.typings import _Address, _DocumentType +if TYPE_CHECKING: + from typing_extensions import Self _CURSOR_DOC_FIELDS = {"cursor": {"firstBatch": 1, "nextBatch": 1}} @@ -133,6 +144,115 @@ def _die_no_lock(self) -> None: self._sock_mgr = None +class _AgnosticCommandCursorBase(_AgnosticCursorBase[_DocumentType]): + """An agnostic cursor / iterator over command cursors. + Used by aggregate, list_indexes, list_search_indexes, list_collections, cursor_command, + and list_databases helpers on both synchronous and asynchronous APIs to iterate MongoDB + command results. + + Should not be called directly by application developers. + """ + + _getmore_class = _GetMore + + def __init__( + self, + collection: _AgnosticCollection[_DocumentType], + cursor_info: Mapping[str, Any], + address: Optional[_Address], + batch_size: int = 0, + max_await_time_ms: Optional[int] = None, + session: Optional[_AgnosticClientSession] = None, + comment: Any = None, + ) -> None: + """Create a new command cursor.""" + self._sock_mgr: Any = None + self._collection = collection + self._id = cursor_info["id"] + self._data = deque(cursor_info["firstBatch"]) + self._postbatchresumetoken: Optional[Mapping[str, Any]] = cursor_info.get( + "postBatchResumeToken" + ) + self._address = address + self._batch_size = batch_size + self._max_await_time_ms = max_await_time_ms + self._timeout = self._collection.database.client.options.timeout + self._session = session + if self._session is not None: + self._session._attached_to_cursor = True + self._killed = self._id == 0 + self._comment = comment + if self._killed: + self._end_session() + + if "ns" in cursor_info: + self._ns = cursor_info["ns"] + else: + self._ns = collection.full_name + + self.batch_size(batch_size) + + if not isinstance(max_await_time_ms, int) and max_await_time_ms is not None: + raise TypeError( + f"max_await_time_ms must be an integer or None, not {type(max_await_time_ms)}" + ) + + def _get_namespace(self) -> str: + return self._ns + + def batch_size(self, batch_size: int) -> Self: + """Limits the number of documents returned in one batch. Each batch + requires a round trip to the server. It can be adjusted to optimize + performance and limit data transfer. + + .. note:: batch_size can not override MongoDB's internal limits on the + amount of data it will return to the client in a single batch (i.e + if you set batch size to 1,000,000,000, MongoDB will currently only + return 4-16MB of results per batch). + + Raises :exc:`TypeError` if `batch_size` is not an integer. + Raises :exc:`ValueError` if `batch_size` is less than ``0``. + + :param batch_size: The size of each batch of results requested. + """ + if not isinstance(batch_size, int): + raise TypeError(f"batch_size must be an integer, not {type(batch_size)}") + if batch_size < 0: + raise ValueError("batch_size must be >= 0") + + self._batch_size = (batch_size == 1 and 2) or batch_size + return self + + def _has_next(self) -> bool: + """Returns `True` if the cursor has documents remaining from the + previous batch. + """ + return len(self._data) > 0 + + @property + def _post_batch_resume_token(self) -> Optional[Mapping[str, Any]]: + """Retrieve the postBatchResumeToken from the response to a + changeStream aggregate or getMore. + """ + return self._postbatchresumetoken + + def _unpack_response( + self, + response: _OpMsg, + cursor_id: Optional[int], + codec_options: CodecOptions[Mapping[str, Any]], + user_fields: Optional[Mapping[str, Any]] = None, + legacy_response: bool = False, + ) -> Sequence[_DocumentOut]: + return response.unpack_response(cursor_id, codec_options, user_fields, legacy_response) + + def _end_session(self) -> None: + if self._session and self._session._implicit: + self._session._attached_to_cursor = False + self._session._end_implicit_session() + self._session = None + + # These errors mean that the server has already killed the cursor so there is # no need to send killCursors. _CURSOR_CLOSED_ERRORS = frozenset( diff --git a/pymongo/synchronous/command_cursor.py b/pymongo/synchronous/command_cursor.py index 8868d87939..b427fed844 100644 --- a/pymongo/synchronous/command_cursor.py +++ b/pymongo/synchronous/command_cursor.py @@ -17,7 +17,7 @@ from __future__ import annotations from collections import deque -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Iterator, Mapping from typing import ( TYPE_CHECKING, Any, @@ -26,12 +26,12 @@ ) from bson import CodecOptions, _convert_raw_document_lists_to_streams -from pymongo.cursor_shared import _CURSOR_CLOSED_ERRORS +from pymongo.cursor_shared import _CURSOR_CLOSED_ERRORS, _AgnosticCommandCursorBase from pymongo.errors import ConnectionFailure, InvalidOperation, OperationFailure from pymongo.message import _GetMore, _OpMsg, _RawBatchGetMore from pymongo.response import PinnedResponse from pymongo.synchronous.cursor_base import _ConnectionManager, _CursorBase -from pymongo.typings import _Address, _DocumentOut, _DocumentType +from pymongo.typings import _Address, _DocumentType if TYPE_CHECKING: from pymongo.synchronous.client_session import ClientSession @@ -41,7 +41,7 @@ _IS_SYNC = True -class CommandCursor(_CursorBase[_DocumentType]): +class CommandCursor(_AgnosticCommandCursorBase[_DocumentType], _CursorBase[_DocumentType]): """A cursor / iterator over command cursors. Used by :meth:`~pymongo.collection.Collection.aggregate`, :meth:`~pymongo.database.Database.aggregate`, @@ -55,89 +55,6 @@ class CommandCursor(_CursorBase[_DocumentType]): Should not be called directly by application developers. """ - _getmore_class = _GetMore - - def __init__( - self, - collection: Collection[_DocumentType], - cursor_info: Mapping[str, Any], - address: Optional[_Address], - batch_size: int = 0, - max_await_time_ms: Optional[int] = None, - session: Optional[ClientSession] = None, - comment: Any = None, - ) -> None: - """Create a new command cursor.""" - self._sock_mgr: Any = None - self._collection: Collection[_DocumentType] = collection - self._id = cursor_info["id"] - self._data = deque(cursor_info["firstBatch"]) - self._postbatchresumetoken: Optional[Mapping[str, Any]] = cursor_info.get( - "postBatchResumeToken" - ) - self._address = address - self._batch_size = batch_size - self._max_await_time_ms = max_await_time_ms - self._timeout = self._collection.database.client.options.timeout - self._session = session - if self._session is not None: - self._session._attached_to_cursor = True - self._killed = self._id == 0 - self._comment = comment - if self._killed: - self._end_session() - - if "ns" in cursor_info: - self._ns = cursor_info["ns"] - else: - self._ns = collection.full_name - - self.batch_size(batch_size) - - if not isinstance(max_await_time_ms, int) and max_await_time_ms is not None: - raise TypeError( - f"max_await_time_ms must be an integer or None, not {type(max_await_time_ms)}" - ) - - def _get_namespace(self) -> str: - return self._ns - - def batch_size(self, batch_size: int) -> CommandCursor[_DocumentType]: - """Limits the number of documents returned in one batch. Each batch - requires a round trip to the server. It can be adjusted to optimize - performance and limit data transfer. - - .. note:: batch_size can not override MongoDB's internal limits on the - amount of data it will return to the client in a single batch (i.e - if you set batch size to 1,000,000,000, MongoDB will currently only - return 4-16MB of results per batch). - - Raises :exc:`TypeError` if `batch_size` is not an integer. - Raises :exc:`ValueError` if `batch_size` is less than ``0``. - - :param batch_size: The size of each batch of results requested. - """ - if not isinstance(batch_size, int): - raise TypeError(f"batch_size must be an integer, not {type(batch_size)}") - if batch_size < 0: - raise ValueError("batch_size must be >= 0") - - self._batch_size = (batch_size == 1 and 2) or batch_size - return self - - def _has_next(self) -> bool: - """Returns `True` if the cursor has documents remaining from the - previous batch. - """ - return len(self._data) > 0 - - @property - def _post_batch_resume_token(self) -> Optional[Mapping[str, Any]]: - """Retrieve the postBatchResumeToken from the response to a - changeStream aggregate or getMore. - """ - return self._postbatchresumetoken - def _maybe_pin_connection(self, conn: Connection) -> None: client = self._collection.database.client if not client._should_pin_cursor(self._session): @@ -152,22 +69,6 @@ def _maybe_pin_connection(self, conn: Connection) -> None: else: self._sock_mgr = conn_mgr - def _unpack_response( - self, - response: _OpMsg, - cursor_id: Optional[int], - codec_options: CodecOptions[Mapping[str, Any]], - user_fields: Optional[Mapping[str, Any]] = None, - legacy_response: bool = False, - ) -> Sequence[_DocumentOut]: - return response.unpack_response(cursor_id, codec_options, user_fields, legacy_response) - - def _end_session(self) -> None: - if self._session and self._session._implicit: - self._session._attached_to_cursor = False - self._session._end_implicit_session() - self._session = None - def _send_message(self, operation: _GetMore) -> None: """Send a getmore message and handle the response.""" client = self._collection.database.client diff --git a/pymongo/typings.py b/pymongo/typings.py index 6ed660e57a..0281c57c59 100644 --- a/pymongo/typings.py +++ b/pymongo/typings.py @@ -31,12 +31,14 @@ from pymongo.asynchronous.bulk import _AsyncBulk from pymongo.asynchronous.client_bulk import _AsyncClientBulk from pymongo.asynchronous.client_session import AsyncClientSession + from pymongo.asynchronous.collection import AsyncCollection from pymongo.asynchronous.mongo_client import AsyncMongoClient from pymongo.asynchronous.pool import AsyncConnection from pymongo.collation import Collation from pymongo.synchronous.bulk import _Bulk from pymongo.synchronous.client_bulk import _ClientBulk from pymongo.synchronous.client_session import ClientSession + from pymongo.synchronous.collection import Collection from pymongo.synchronous.mongo_client import MongoClient from pymongo.synchronous.pool import Connection @@ -51,6 +53,7 @@ # Type hinting types for compatibility between async and sync classes _AgnosticMongoClient = Union["AsyncMongoClient", "MongoClient"] # type: ignore[type-arg] +_AgnosticCollection = Union["AsyncCollection[_DocumentType]", "Collection[_DocumentType]"] _AgnosticConnection = Union["AsyncConnection", "Connection"] _AgnosticClientSession = Union["AsyncClientSession", "ClientSession"] _AgnosticBulk = Union["_AsyncBulk", "_Bulk"]