diff --git a/pymongo/_client_bulk_shared.py b/pymongo/_client_bulk_shared.py index f4f7f8fe7f..0019d0c0bb 100644 --- a/pymongo/_client_bulk_shared.py +++ b/pymongo/_client_bulk_shared.py @@ -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( @@ -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 diff --git a/pymongo/asynchronous/bulk.py b/pymongo/asynchronous/bulk.py index 86b9528b6e..81186499f6 100644 --- a/pymongo/asynchronous/bulk.py +++ b/pymongo/asynchronous/bulk.py @@ -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 ( @@ -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, @@ -59,7 +51,6 @@ from pymongo.helpers_shared import _RETRYABLE_ERROR_CODES from pymongo.message import ( _DELETE, - _INSERT, _UPDATE, _BulkWriteContext, _EncryptedBulkWriteContext, @@ -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( diff --git a/pymongo/asynchronous/client_bulk.py b/pymongo/asynchronous/client_bulk.py index 2847dcb2c0..bfbe9e2c6b 100644 --- a/pymongo/asynchronous/client_bulk.py +++ b/pymongo/asynchronous/client_bulk.py @@ -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.collection import AsyncCollection @@ -46,15 +44,11 @@ from pymongo.asynchronous.mongo_client import AsyncMongoClient from pymongo.asynchronous.pool import AsyncConnection from pymongo._client_bulk_shared import ( + _AgnosticClientBulkBase, _merge_command, _throw_client_bulk_write_exception, ) 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, ConnectionFailure, @@ -76,152 +70,15 @@ InsertOneResult, UpdateResult, ) -from pymongo.typings import _DocumentOut, _Pipeline from pymongo.write_concern import WriteConcern _IS_SYNC = False -class _AsyncClientBulk: +class _AsyncClientBulk(_AgnosticClientBulkBase): """The private guts of the client-level bulk write API.""" - def __init__( - self, - client: AsyncMongoClient[Any], - 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 _AsyncClientBulk instance.""" - self.client = client - self.write_concern = write_concern - self.let = let - if self.let is not None: - common.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 an update_many is not retryable. - self.is_retryable = False - self.ops.append(("delete", cmd)) - self.namespaces.append(namespace) - self.total_ops += 1 + client: AsyncMongoClient[Any] @_handle_reauth async def write_command( diff --git a/pymongo/bulk_shared.py b/pymongo/bulk_shared.py index a086611be5..54f0d4cd22 100644 --- a/pymongo/bulk_shared.py +++ b/pymongo/bulk_shared.py @@ -17,19 +17,28 @@ from __future__ import annotations -from collections.abc import Mapping, MutableMapping -from typing import TYPE_CHECKING, Any, NoReturn +from collections.abc import Iterator, Mapping, MutableMapping +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 BulkWriteError, OperationFailure from pymongo.helpers_shared import _get_wce_doc from pymongo.message import ( _DELETE, _INSERT, _UPDATE, + _BulkWriteContext, + _EncryptedBulkWriteContext, ) if TYPE_CHECKING: - from pymongo.typings import _DocumentOut + from pymongo.typings import _AgnosticCollection, _DocumentOut, _DocumentType, _Pipeline _DELETE_ALL: int = 0 @@ -131,3 +140,159 @@ def _raise_bulk_write_error(full_result: _DocumentOut) -> NoReturn: ) raise OperationFailure(errmsg, code, full_result) raise BulkWriteError(full_result) + + +class _AgnosticBulkBase: + """The private guts of the bulk write API.""" + + def __init__( + self, + collection: _AgnosticCollection[_DocumentType], + ordered: bool, + bypass_document_validation: Optional[bool], + comment: Optional[str] = None, + let: Optional[Any] = None, + ) -> None: + """Initialize an _AgnosticBulkBase 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: + 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 diff --git a/pymongo/synchronous/bulk.py b/pymongo/synchronous/bulk.py index be12a72127..291965fca1 100644 --- a/pymongo/synchronous/bulk.py +++ b/pymongo/synchronous/bulk.py @@ -29,23 +29,15 @@ 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.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, @@ -54,7 +46,6 @@ from pymongo.helpers_shared import _RETRYABLE_ERROR_CODES from pymongo.message import ( _DELETE, - _INSERT, _UPDATE, _BulkWriteContext, _EncryptedBulkWriteContext, @@ -71,165 +62,14 @@ from pymongo.synchronous.collection import Collection from pymongo.synchronous.mongo_client import MongoClient from pymongo.synchronous.pool import Connection - from pymongo.typings import _DocumentOut, _DocumentType, _Pipeline _IS_SYNC = True -class _Bulk: +class _Bulk(_AgnosticBulkBase): """The private guts of the bulk write API.""" - def __init__( - self, - collection: Collection[_DocumentType], - ordered: bool, - bypass_document_validation: Optional[bool], - comment: Optional[str] = None, - let: Optional[Any] = None, - ) -> None: - """Initialize a _Bulk 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: Collection[Any] @_handle_reauth def write_command( diff --git a/pymongo/synchronous/client_bulk.py b/pymongo/synchronous/client_bulk.py index 5e5d1a76e8..549c9d59ca 100644 --- a/pymongo/synchronous/client_bulk.py +++ b/pymongo/synchronous/client_bulk.py @@ -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.synchronous.client_session import ClientSession from pymongo.synchronous.collection import Collection @@ -46,15 +44,11 @@ from pymongo.synchronous.mongo_client import MongoClient from pymongo.synchronous.pool import Connection from pymongo._client_bulk_shared import ( + _AgnosticClientBulkBase, _merge_command, _throw_client_bulk_write_exception, ) 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, ConnectionFailure, @@ -76,152 +70,15 @@ InsertOneResult, UpdateResult, ) -from pymongo.typings import _DocumentOut, _Pipeline from pymongo.write_concern import WriteConcern _IS_SYNC = True -class _ClientBulk: +class _ClientBulk(_AgnosticClientBulkBase): """The private guts of the client-level bulk write API.""" - def __init__( - self, - client: MongoClient[Any], - 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 _ClientBulk instance.""" - self.client = client - self.write_concern = write_concern - self.let = let - if self.let is not None: - common.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 an update_many is not retryable. - self.is_retryable = False - self.ops.append(("delete", cmd)) - self.namespaces.append(namespace) - self.total_ops += 1 + client: MongoClient[Any] @_handle_reauth def write_command( diff --git a/pymongo/typings.py b/pymongo/typings.py index 0281c57c59..e0777e07e1 100644 --- a/pymongo/typings.py +++ b/pymongo/typings.py @@ -70,7 +70,6 @@ def strip_optional(elem: Optional[_T]) -> _T: __all__ = [ "_Address", - "_AgnosticMongoClient", "_CollationIn", "_DocumentOut", "_DocumentType",