Isolate ChainMap + add RawBSONDocument not inflated test

Author: NoahStapp

pymongo/_client_bulk_shared.py

@@ -16,7 +16,6 @@
 """Constants, types, and classes shared across Client Bulk Write API implementations."""
 from __future__ import annotations
 
-from collections import ChainMap
 from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, NoReturn
 
 from pymongo.errors import ClientBulkWriteException, OperationFailure
@@ -64,10 +63,6 @@ def _throw_client_bulk_write_exception(
     """Raise a ClientBulkWriteException from the full result."""
     # retryWrites on MMAPv1 should raise an actionable error.
     if full_result["writeErrors"]:
-        # Unpack ChainMaps into the original document only
-        for doc in full_result["writeErrors"]:
-            if "document" in doc["op"] and isinstance(doc["op"]["document"], ChainMap):
-                doc["op"]["document"] = doc["op"]["document"].maps[0]
         full_result["writeErrors"].sort(key=lambda error: error["idx"])
         err = full_result["writeErrors"][0]
         code = err["code"]

pymongo/message.py

@@ -1115,11 +1115,15 @@ def _check_doc_size_limits(
         # 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.
+        doc_to_encode = op_doc
         if real_op_type == "insert":
-            op_doc["document"] = ChainMap(op_doc["document"], {"_id": op_doc["document"]["_id"]})  # type: ignore[index]
+            doc = op_doc["document"]
+            if not isinstance(doc, RawBSONDocument):
+                doc_to_encode = op_doc.copy()  # Shallow copy
+                doc_to_encode["document"] = ChainMap(doc, {"_id": doc["_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_doc_encoded = _dict_to_bson(doc_to_encode, False, opts)
         op_length = len(op_doc_encoded)
         if ns_doc:
             ns_doc_encoded = _dict_to_bson(ns_doc, False, opts)

pymongo/monitoring.py

@@ -189,7 +189,7 @@ def connection_checked_in(self, event):
 from __future__ import annotations
 
 import datetime
-from collections import ChainMap, abc, namedtuple
+from collections import abc, namedtuple
 from typing import TYPE_CHECKING, Any, Mapping, Optional, Sequence
 
 from bson.objectid import ObjectId
@@ -625,11 +625,6 @@ def __init__(
             raise ValueError(f"{command!r} is not a valid command")
         # Command name must be first key.
         command_name = next(iter(command))
-        # Unpack ChainMaps into the original document only
-        if command_name == "bulkWrite" and "ops" in command:
-            for doc in command["ops"]:
-                if "document" in doc and isinstance(doc["document"], ChainMap):
-                    doc["document"] = doc["document"].maps[0]
         super().__init__(
             command_name,
             request_id,

test/asynchronous/test_client_bulk_write.py

@@ -18,6 +18,8 @@
 import os
 import sys
 
+from bson import encode
+
 sys.path[0:0] = [""]
 
 from test.asynchronous import (
@@ -84,6 +86,17 @@ async def test_formats_write_error_correctly(self):
         self.assertEqual(write_error["idx"], 1)
         self.assertEqual(write_error["op"], {"insert": 0, "document": {"_id": 1}})
 
+    @async_client_context.require_version_min(8, 0, 0, -24)
+    @async_client_context.require_no_serverless
+    async def test_raw_bson_not_inflated(self):
+        doc = RawBSONDocument(encode({"a": "b" * 100}))
+        models = [
+            InsertOne(namespace="db.coll", document=doc),
+        ]
+        await self.client.bulk_write(models=models)
+
+        self.assertIsNone(doc._RawBSONDocument__inflated_doc)
+
 
 # https://github.com/mongodb/specifications/tree/master/source/crud/tests
 class TestClientBulkWriteCRUD(AsyncIntegrationTest):

test/test_client_bulk_write.py

@@ -18,6 +18,8 @@
 import os
 import sys
 
+from bson import encode
+
 sys.path[0:0] = [""]
 
 from test import (
@@ -84,6 +86,17 @@ def test_formats_write_error_correctly(self):
         self.assertEqual(write_error["idx"], 1)
         self.assertEqual(write_error["op"], {"insert": 0, "document": {"_id": 1}})
 
+    @client_context.require_version_min(8, 0, 0, -24)
+    @client_context.require_no_serverless
+    def test_raw_bson_not_inflated(self):
+        doc = RawBSONDocument(encode({"a": "b" * 100}))
+        models = [
+            InsertOne(namespace="db.coll", document=doc),
+        ]
+        self.client.bulk_write(models=models)
+
+        self.assertIsNone(doc._RawBSONDocument__inflated_doc)
+
 
 # https://github.com/mongodb/specifications/tree/master/source/crud/tests
 class TestClientBulkWriteCRUD(IntegrationTest):