Move re-order logic for client bulkWrite into _client_batched_op_msg_impl

Author: NoahStapp

pymongo/_client_bulk_shared.py

@@ -16,7 +16,8 @@
 """Constants, types, and classes shared across Client Bulk Write API implementations."""
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ChainMap, Mapping, MutableMapping, NoReturn
+from collections import ChainMap
+from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, NoReturn
 
 from pymongo.errors import ClientBulkWriteException, OperationFailure
 from pymongo.helpers_shared import _get_wce_doc

pymongo/asynchronous/client_bulk.py

@@ -21,7 +21,6 @@
 import copy
 import datetime
 import logging
-from collections import ChainMap
 from collections.abc import MutableMapping
 from itertools import islice
 from typing import (
@@ -133,16 +132,8 @@ 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):
-            # Since the data document itself is nested within the insert document
-            # it won't be automatically re-ordered by the BSON conversion.
-            # We use ChainMap here to make the _id field the first field instead.
-            if "_id" in document:
-                document = ChainMap(document, {"_id": document["_id"]})
-            else:
-                id = ObjectId()
-                document["_id"] = id
-                document = ChainMap(document, {"_id": id})
+        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)

pymongo/message.py

@@ -24,6 +24,7 @@
 import datetime
 import random
 import struct
+from collections import ChainMap
 from io import BytesIO as _BytesIO
 from typing import (
     TYPE_CHECKING,
@@ -1111,6 +1112,12 @@ def _check_doc_size_limits(
         # key and the index of its namespace within ns_info as its value.
         op_doc[op_type] = ns_info[namespace]  # type: ignore[index]
 
+        # Since the data document itself is nested within the insert document
+        # it won't be automatically re-ordered by the BSON conversion.
+        # We use ChainMap here to make the _id field the first field instead.
+        if real_op_type == "insert":
+            op_doc["document"] = ChainMap(op_doc["document"], {"_id": op_doc["document"]["_id"]})  # type: ignore[index]
+
         # Encode current operation doc and, if newly added, namespace doc.
         op_doc_encoded = _dict_to_bson(op_doc, False, opts)
         op_length = len(op_doc_encoded)

pymongo/monitoring.py

@@ -189,8 +189,8 @@ def connection_checked_in(self, event):
 from __future__ import annotations
 
 import datetime
-from collections import abc, namedtuple
-from typing import TYPE_CHECKING, Any, ChainMap, Mapping, Optional, Sequence
+from collections import ChainMap, abc, namedtuple
+from typing import TYPE_CHECKING, Any, Mapping, Optional, Sequence
 
 from bson.objectid import ObjectId
 from pymongo.hello import Hello, HelloCompat

pymongo/synchronous/client_bulk.py

@@ -21,7 +21,6 @@
 import copy
 import datetime
 import logging
-from collections import ChainMap
 from collections.abc import MutableMapping
 from itertools import islice
 from typing import (
@@ -133,16 +132,8 @@ 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):
-            # Since the data document itself is nested within the insert document
-            # it won't be automatically re-ordered by the BSON conversion.
-            # We use ChainMap here to make the _id field the first field instead.
-            if "_id" in document:
-                document = ChainMap(document, {"_id": document["_id"]})
-            else:
-                id = ObjectId()
-                document["_id"] = id
-                document = ChainMap(document, {"_id": id})
+        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)