From 083c2474dad0e2ff2130edb5d77e668bb5f15f1e Mon Sep 17 00:00:00 2001
From: Shane Harvey <shnhrv@gmail.com>
Date: Thu, 14 Nov 2019 16:18:12 -0800
Subject: [PATCH] PYTHON-2055 Subtract message header from compressed bulk
 OP_MSG

This change prevents pymongo from generating a bulk OP_COMPRESSED/OP_MSG
with an uncompressed message larger than the server's limit of
maxMessageSizeBytes - 16.
---
 pymongo/message.py |  3 +++
 test/test_bulk.py  | 18 ++++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/pymongo/message.py b/pymongo/message.py
index 6e8c59695..de0a6ea9f 100644
--- a/pymongo/message.py
+++ b/pymongo/message.py
@@ -922,6 +922,9 @@ class _BulkWriteContext(object):
     @property
     def max_message_size(self):
         """A proxy for SockInfo.max_message_size."""
+        if self.compress:
+            # Subtract 16 bytes for the message header.
+            return self.sock_info.max_message_size - 16
         return self.sock_info.max_message_size
 
     @property
diff --git a/test/test_bulk.py b/test/test_bulk.py
index 81bb35c33..ec148a4ac 100644
--- a/test/test_bulk.py
+++ b/test/test_bulk.py
@@ -308,6 +308,24 @@ class TestBulk(BulkTestBase):
         self.assertEqual(n_docs, result.inserted_count)
         self.assertEqual(n_docs, self.coll.count_documents({}))
 
+    @client_context.require_version_min(3, 6)
+    def test_bulk_max_message_size(self):
+        self.coll.delete_many({})
+        self.addCleanup(self.coll.delete_many, {})
+        _16_MB = 16 * 1000 * 1000
+        # Generate a list of documents such that the first batched OP_MSG is
+        # as close as possible to the 48MB limit.
+        docs = [
+            {'_id': 1, 'l': 's' * _16_MB},
+            {'_id': 2, 'l': 's' * _16_MB},
+            {'_id': 3, 'l': 's' * (_16_MB - 10000)},
+        ]
+        # Fill in the remaining ~10000 bytes with small documents.
+        for i in range(4, 10000):
+            docs.append({'_id': i})
+        result = self.coll.insert_many(docs)
+        self.assertEqual(len(docs), len(result.inserted_ids))
+
     def test_generator_insert(self):
         def gen():
             yield {'a': 1, 'b': 1}