Unwrap ChainMaps during logging

Author: NoahStapp

pymongo/_client_bulk_shared.py

@@ -16,7 +16,7 @@
 """Constants, types, and classes shared across Client Bulk Write API implementations."""
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, NoReturn
+from typing import TYPE_CHECKING, Any, ChainMap, Mapping, MutableMapping, NoReturn
 
 from pymongo.errors import ClientBulkWriteException, OperationFailure
 from pymongo.helpers_shared import _get_wce_doc
@@ -63,6 +63,10 @@ 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/asynchronous/client_bulk.py

@@ -133,10 +133,16 @@ 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 = ChainMap(document, {"_id": ObjectId()})
-        elif not isinstance(document, RawBSONDocument) and "_id" in document:
-            document = ChainMap(document, {"_id": document["_id"]})
+        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})
         cmd = {"insert": -1, "document": document}
         self.ops.append(("insert", cmd))
         self.namespaces.append(namespace)

pymongo/monitoring.py

@@ -190,7 +190,7 @@ def connection_checked_in(self, event):
 
 import datetime
 from collections import abc, namedtuple
-from typing import TYPE_CHECKING, Any, Mapping, Optional, Sequence
+from typing import TYPE_CHECKING, Any, ChainMap, Mapping, Optional, Sequence
 
 from bson.objectid import ObjectId
 from pymongo.hello import Hello, HelloCompat
@@ -625,6 +625,11 @@ 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,

pymongo/synchronous/client_bulk.py

@@ -133,10 +133,16 @@ 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 = ChainMap(document, {"_id": ObjectId()})
-        elif not isinstance(document, RawBSONDocument) and "_id" in document:
-            document = ChainMap(document, {"_id": document["_id"]})
+        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})
         cmd = {"insert": -1, "document": document}
         self.ops.append(("insert", cmd))
         self.namespaces.append(namespace)

test/mockupdb/test_id_ordering.py

@@ -48,22 +48,31 @@ def test_id_ordering(self):
         server.run()
         self.addCleanup(server.stop)
 
+        # We also verify that the original document contains an _id field after each insert
+        document = {"x": 1}
+
         client = self.simple_client(server.uri, loadBalanced=True)
         collection = client.db.coll
-        with going(collection.insert_one, {"x": 1}):
+        with going(collection.insert_one, document):
             request = server.receives()
             self.assertEqual("_id", next(iter(request["documents"][0])))
             request.reply({"ok": 1})
+        self.assertIn("_id", document)
+
+        document = {"x1": 1}
 
-        with going(collection.bulk_write, [InsertOne({"x1": 1})]):
+        with going(collection.bulk_write, [InsertOne(document)]):
             request = server.receives()
             self.assertEqual("_id", next(iter(request["documents"][0])))
             request.reply({"ok": 1})
+        self.assertIn("_id", document)
 
-        with going(client.bulk_write, [InsertOne(namespace="db.coll", document={"x2": 1})]):
+        document = {"x2": 1}
+        with going(client.bulk_write, [InsertOne(namespace="db.coll", document=document)]):
             request = server.receives()
             self.assertEqual("_id", next(iter(request["ops"][0]["document"])))
             request.reply({"ok": 1})
+        self.assertIn("_id", document)
 
         # Re-ordering user-supplied _id fields is not required by the spec, but PyMongo does it for performance reasons
         with going(collection.insert_one, {"x": 1, "_id": 111}):