PYTHON-4666 Fix handling of large documents in client.bulk_write (mongodb#1798)

Author: shruti-sridhar

pymongo/message.py

@@ -1123,6 +1123,8 @@ def _check_doc_size_limits(
         new_message_size = total_ops_length + total_ns_length + op_length + ns_length
         # We have enough data, return this batch.
         if new_message_size > max_doc_sequences_bytes:
+            if idx == 0:
+                _raise_document_too_large(op_type, op_length, max_bson_size + _COMMAND_OVERHEAD)
             break
 
         # Add op and ns documents to this batch.

test/asynchronous/test_client_bulk_write.py

@@ -512,17 +512,15 @@ async def test_returns_error_if_no_writes_can_be_added_to_ops(self):
         # Document too large.
         b_repeated = "b" * self.max_message_size_bytes
         models = [InsertOne(namespace="db.coll", document={"a": b_repeated})]
-        with self.assertRaises(InvalidOperation) as context:
+        with self.assertRaises(DocumentTooLarge):
             await client.bulk_write(models=models)
-        self.assertIn("cannot do an empty bulk write", context.exception._message)
 
         # Namespace too large.
         c_repeated = "c" * self.max_message_size_bytes
         namespace = f"db.{c_repeated}"
         models = [InsertOne(namespace=namespace, document={"a": "b"})]
-        with self.assertRaises(InvalidOperation) as context:
+        with self.assertRaises(DocumentTooLarge):
             await client.bulk_write(models=models)
-        self.assertIn("cannot do an empty bulk write", context.exception._message)
 
     @async_client_context.require_version_min(8, 0, 0, -24)
     @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed")

test/test_client_bulk_write.py

@@ -512,17 +512,15 @@ def test_returns_error_if_no_writes_can_be_added_to_ops(self):
         # Document too large.
         b_repeated = "b" * self.max_message_size_bytes
         models = [InsertOne(namespace="db.coll", document={"a": b_repeated})]
-        with self.assertRaises(InvalidOperation) as context:
+        with self.assertRaises(DocumentTooLarge):
             client.bulk_write(models=models)
-        self.assertIn("cannot do an empty bulk write", context.exception._message)
 
         # Namespace too large.
         c_repeated = "c" * self.max_message_size_bytes
         namespace = f"db.{c_repeated}"
         models = [InsertOne(namespace=namespace, document={"a": "b"})]
-        with self.assertRaises(InvalidOperation) as context:
+        with self.assertRaises(DocumentTooLarge):
             client.bulk_write(models=models)
-        self.assertIn("cannot do an empty bulk write", context.exception._message)
 
     @client_context.require_version_min(8, 0, 0, -24)
     @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed")