Skip to content

PYTHON-6024 Consolidate all non-IO code in _AsyncClientBulk/_AsyncBul… #3087

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

Merged
sleepyStick merged 1 commit into from
Oct 1, 2026
Merged

PYTHON-6024 Consolidate all non-IO code in _AsyncClientBulk/_AsyncBul… #3087

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
155 changes: 153 additions & 2 deletions pymongo/_client_bulk_shared.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,22 @@
from __future__ import annotations

from collections.abc import Mapping, MutableMapping
from typing import TYPE_CHECKING, Any, NoReturn
from typing import TYPE_CHECKING, Any, NoReturn, Optional, Union

from bson.objectid import ObjectId
from bson.raw_bson import RawBSONDocument
from pymongo.common import (
validate_is_document_type,
validate_ok_for_replace,
validate_ok_for_update,
)
from pymongo.errors import ClientBulkWriteException, OperationFailure
from pymongo.helpers_shared import _get_wce_doc
from pymongo.message import _ClientBulkWriteContext
from pymongo.write_concern import WriteConcern

if TYPE_CHECKING:
from pymongo.typings import _DocumentOut
from pymongo.typings import _AgnosticMongoClient, _DocumentOut, _Pipeline


def _merge_command(
Expand Down Expand Up @@ -79,3 +88,145 @@ def _throw_client_bulk_write_exception(
if isinstance(full_result["error"], BaseException):
raise ClientBulkWriteException(full_result, verbose_results) from full_result["error"]
raise ClientBulkWriteException(full_result, verbose_results)


class _AgnosticClientBulkBase:
"""The private guts of the client-level bulk write API."""

def __init__(
self,
client: _AgnosticMongoClient,
write_concern: WriteConcern,
ordered: bool = True,
bypass_document_validation: Optional[bool] = None,
comment: Optional[str] = None,
let: Optional[Any] = None,
verbose_results: bool = False,
) -> None:
"""Initialize a _AgnosticClientBulkBase instance."""
self.client = client
self.write_concern = write_concern
self.let = let
if self.let is not None:
validate_is_document_type("let", self.let)
self.ordered = ordered
self.bypass_doc_val = bypass_document_validation
self.comment = comment
self.verbose_results = verbose_results
self.ops: list[tuple[str, Mapping[str, Any]]] = []
self.namespaces: list[str] = []
self.idx_offset: int = 0
self.total_ops: int = 0
self.executed = False
self.uses_collation = False
self.uses_array_filters = False
self.is_retryable = self.client.options.retry_writes
self.retrying = False
self.started_retryable_write = False

@property
def bulk_ctx_class(self) -> type[_ClientBulkWriteContext]:
return _ClientBulkWriteContext

def add_insert(self, namespace: str, document: _DocumentOut) -> None:
"""Add an insert document to the list of ops."""
validate_is_document_type("document", document)
# Generate ObjectId client side.
if not (isinstance(document, RawBSONDocument) or "_id" in document):
document["_id"] = ObjectId()
cmd = {"insert": -1, "document": document}
self.ops.append(("insert", cmd))
self.namespaces.append(namespace)
self.total_ops += 1

def add_update(
self,
namespace: str,
selector: Mapping[str, Any],
update: Union[Mapping[str, Any], _Pipeline],
multi: bool,
upsert: Optional[bool] = None,
collation: Optional[Mapping[str, Any]] = None,
array_filters: Optional[list[Mapping[str, Any]]] = None,
hint: Union[str, dict[str, Any], None] = None,
sort: Optional[Mapping[str, Any]] = None,
) -> None:
"""Create an update document and add it to the list of ops."""
validate_ok_for_update(update)
cmd = {
"update": -1,
"filter": selector,
"updateMods": update,
"multi": multi,
}
if upsert is not None:
cmd["upsert"] = upsert
if array_filters is not None:
self.uses_array_filters = True
cmd["arrayFilters"] = array_filters
if hint is not None:
cmd["hint"] = hint
if collation is not None:
self.uses_collation = True
cmd["collation"] = collation
if sort is not None:
cmd["sort"] = sort
if multi:
# A bulk_write containing an update_many is not retryable.
self.is_retryable = False
self.ops.append(("update", cmd))
self.namespaces.append(namespace)
self.total_ops += 1

def add_replace(
self,
namespace: str,
selector: Mapping[str, Any],
replacement: Mapping[str, Any],
upsert: Optional[bool] = None,
collation: Optional[Mapping[str, Any]] = None,
hint: Union[str, dict[str, Any], None] = None,
sort: Optional[Mapping[str, Any]] = None,
) -> None:
"""Create a replace document and add it to the list of ops."""
validate_ok_for_replace(replacement)
cmd = {
"update": -1,
"filter": selector,
"updateMods": replacement,
"multi": False,
}
if upsert is not None:
cmd["upsert"] = upsert
if hint is not None:
cmd["hint"] = hint
if collation is not None:
self.uses_collation = True
cmd["collation"] = collation
if sort is not None:
cmd["sort"] = sort
self.ops.append(("replace", cmd))
self.namespaces.append(namespace)
self.total_ops += 1

def add_delete(
self,
namespace: str,
selector: Mapping[str, Any],
multi: bool,
collation: Optional[Mapping[str, Any]] = None,
hint: Union[str, dict[str, Any], None] = None,
) -> None:
"""Create a delete document and add it to the list of ops."""
cmd = {"delete": -1, "filter": selector, "multi": multi}
if hint is not None:
cmd["hint"] = hint
if collation is not None:
self.uses_collation = True
cmd["collation"] = collation
if multi:
# A bulk_write containing a delete_many is not retryable.
self.is_retryable = False
self.ops.append(("delete", cmd))
self.namespaces.append(namespace)
self.total_ops += 1
168 changes: 4 additions & 164 deletions pymongo/asynchronous/bulk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,7 @@
Union,
)

from bson.objectid import ObjectId
from bson.raw_bson import RawBSONDocument
from pymongo import _csot, common
from pymongo import _csot
from pymongo._telemetry import _generate_op_id_or_none
from pymongo.asynchronous.client_session import AsyncClientSession
from pymongo.asynchronous.command_runner import (
Expand All @@ -40,17 +38,11 @@
from pymongo.asynchronous.helpers import _handle_reauth
from pymongo.bulk_shared import (
_COMMANDS,
_DELETE_ALL,
_AgnosticBulkBase,
_merge_command,
_raise_bulk_write_error,
_Run,
)
from pymongo.client_session_shared import _validate_session_write_concern
from pymongo.common import (
validate_is_document_type,
validate_ok_for_replace,
validate_ok_for_update,
)
from pymongo.errors import (
ConfigurationError,
InvalidOperation,
Expand All @@ -59,7 +51,6 @@
from pymongo.helpers_shared import _RETRYABLE_ERROR_CODES
from pymongo.message import (
_DELETE,
_INSERT,
_UPDATE,
_BulkWriteContext,
_EncryptedBulkWriteContext,
Expand All @@ -71,165 +62,14 @@
from pymongo.asynchronous.collection import AsyncCollection
from pymongo.asynchronous.mongo_client import AsyncMongoClient
from pymongo.asynchronous.pool import AsyncConnection
from pymongo.typings import _DocumentOut, _DocumentType, _Pipeline

_IS_SYNC = False


class _AsyncBulk:
class _AsyncBulk(_AgnosticBulkBase):
"""The private guts of the bulk write API."""

def __init__(
self,
collection: AsyncCollection[_DocumentType],
ordered: bool,
bypass_document_validation: Optional[bool],
comment: Optional[str] = None,
let: Optional[Any] = None,
) -> None:
"""Initialize a _AsyncBulk instance."""
self.collection = collection.with_options(
codec_options=collection.codec_options._replace(
unicode_decode_error_handler="replace", document_class=dict
)
)
self.let = let
if self.let is not None:
common.validate_is_document_type("let", self.let)
self.comment: Optional[str] = comment
self.ordered = ordered
self.ops: list[tuple[int, Mapping[str, Any]]] = []
self.executed = False
self.bypass_doc_val = bypass_document_validation
self.uses_collation = False
self.uses_array_filters = False
self.uses_sort = False
self.is_retryable = True
self.retrying = False
self.started_retryable_write = False
# Extra state so that we know where to pick up on a retry attempt.
self.current_run = None
self.next_run = None
self.is_encrypted = False

@property
def bulk_ctx_class(self) -> type[_BulkWriteContext]:
encrypter = self.collection.database.client._encrypter
if encrypter and not encrypter._bypass_auto_encryption:
self.is_encrypted = True
return _EncryptedBulkWriteContext
else:
self.is_encrypted = False
return _BulkWriteContext

def add_insert(self, document: _DocumentOut) -> None:
"""Add an insert document to the list of ops."""
validate_is_document_type("document", document)
# Generate ObjectId client side.
if not (isinstance(document, RawBSONDocument) or "_id" in document):
document["_id"] = ObjectId()
self.ops.append((_INSERT, document))

def add_update(
self,
selector: Mapping[str, Any],
update: Union[Mapping[str, Any], _Pipeline],
multi: bool,
upsert: Optional[bool],
collation: Optional[Mapping[str, Any]] = None,
array_filters: Optional[list[Mapping[str, Any]]] = None,
hint: Union[str, dict[str, Any], None] = None,
sort: Optional[Mapping[str, Any]] = None,
) -> None:
"""Create an update document and add it to the list of ops."""
validate_ok_for_update(update)
cmd: dict[str, Any] = {"q": selector, "u": update, "multi": multi}
if upsert is not None:
cmd["upsert"] = upsert
if collation is not None:
self.uses_collation = True
cmd["collation"] = collation
if array_filters is not None:
self.uses_array_filters = True
cmd["arrayFilters"] = array_filters
if hint is not None:
cmd["hint"] = hint
if sort is not None:
self.uses_sort = True
cmd["sort"] = sort
if multi:
# A bulk_write containing an update_many is not retryable.
self.is_retryable = False
self.ops.append((_UPDATE, cmd))

def add_replace(
self,
selector: Mapping[str, Any],
replacement: Mapping[str, Any],
upsert: Optional[bool],
collation: Optional[Mapping[str, Any]] = None,
hint: Union[str, dict[str, Any], None] = None,
sort: Optional[Mapping[str, Any]] = None,
) -> None:
"""Create a replace document and add it to the list of ops."""
validate_ok_for_replace(replacement)
cmd: dict[str, Any] = {"q": selector, "u": replacement}
if upsert is not None:
cmd["upsert"] = upsert
if collation is not None:
self.uses_collation = True
cmd["collation"] = collation
if hint is not None:
cmd["hint"] = hint
if sort is not None:
self.uses_sort = True
cmd["sort"] = sort
self.ops.append((_UPDATE, cmd))

def add_delete(
self,
selector: Mapping[str, Any],
limit: int,
collation: Optional[Mapping[str, Any]] = None,
hint: Union[str, dict[str, Any], None] = None,
) -> None:
"""Create a delete document and add it to the list of ops."""
cmd: dict[str, Any] = {"q": selector, "limit": limit}
if collation is not None:
self.uses_collation = True
cmd["collation"] = collation
if hint is not None:
cmd["hint"] = hint
if limit == _DELETE_ALL:
# A bulk_write containing a delete_many is not retryable.
self.is_retryable = False
self.ops.append((_DELETE, cmd))

def gen_ordered(self) -> Iterator[Optional[_Run]]:
"""Generate batches of operations, batched by type of
operation, in the order **provided**.
"""
run = None
for idx, (op_type, operation) in enumerate(self.ops):
if run is None:
run = _Run(op_type)
elif run.op_type != op_type:
yield run
run = _Run(op_type)
run.add(idx, operation)
yield run

def gen_unordered(self) -> Iterator[_Run]:
"""Generate batches of operations, batched by type of
operation, in arbitrary order.
"""
operations = [_Run(_INSERT), _Run(_UPDATE), _Run(_DELETE)]
for idx, (op_type, operation) in enumerate(self.ops):
operations[op_type].add(idx, operation)

for run in operations:
if run.ops:
yield run
collection: AsyncCollection[Any]

@_handle_reauth
async def write_command(
Expand Down
Loading
Loading