Skip to content

PYTHON-6023 Consolidate all non-IO code in AsyncCursor/Cursor into _A… #3045

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sleepyStick wants to merge 5 commits into
base: main
Choose a base branch
from
Open

PYTHON-6023 Consolidate all non-IO code in AsyncCursor/Cursor into _A… #3045

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
92f4edd
PYTHON-6023 Consolidate all non-IO code in AsyncCursor/Cursor into _A…
sleepyStick Sep 3, 2026
4e58919
Merge branch 'main' into PYTHON-6023
sleepyStick Sep 8, 2026
ec0fc6a
fix typing
sleepyStick Sep 8, 2026
a5e4e00
Merge branch 'PYTHON-6023' of github.com:sleepyStick/mongo-python-dri…
sleepyStick Sep 8, 2026
1613dc2
Potential fix for pull request finding
sleepyStick Sep 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 6 additions & 103 deletions pymongo/asynchronous/command_cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -41,7 +41,9 @@
_IS_SYNC = False


class AsyncCommandCursor(_AsyncCursorBase[_DocumentType]):
class AsyncCommandCursor(
_AgnosticCommandCursorBase[_DocumentType], _AsyncCursorBase[_DocumentType]
):
Comment on lines +44 to +46

@sleepyStick sleepyStick Sep 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe i'm misunderstanding the comment, but I if the non-io code is moved into _AgnosticCommandCursorBase, there is no sync vs async so i don't think this matters?

@aclark4life aclark4life Sep 9, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the type-checking value is lost for callers of AsyncCommandCursor / CommandCursor who get  Any / union-typed self._collection instead of the collection doc type. So maybe we do something like:

# pymongo/asynchronous/command_cursor.py
class AsyncCommandCursor(
    _AgnosticCommandCursorBase[_DocumentType], _AsyncCursorBase[_DocumentType]
):
    _collection: AsyncCollection[_DocumentType]
    ...

# pymongo/synchronous/command_cursor.py
class CommandCursor(_AgnosticCommandCursorBase[_DocumentType], _CursorBase[_DocumentType]):
    _collection: Collection[_DocumentType]
    ...

"""An asynchronous cursor / iterator over command cursors.
Used by :meth:`~pymongo.asynchronous.collection.AsyncCollection.aggregate`,
:meth:`~pymongo.asynchronous.database.AsyncDatabase.aggregate`,
Expand All @@ -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):
Expand All @@ -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
Expand Down
126 changes: 123 additions & 3 deletions pymongo/cursor_shared.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}}

Expand Down Expand Up @@ -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(
Expand Down
Loading
Loading