PYTHON-1939 Add prose tests for BSON size limits and batch splitting

Author: ShaneHarvey

pymongo/_cmessagemodule.c

@@ -1565,30 +1565,29 @@ _batched_write_command(
         /* We have enough data, return this batch.
          * max_cmd_size accounts for the two trailing null bytes.
          */
+        cur_size = buffer_get_position(buffer) - cur_doc_begin;
         enough_data = (buffer_get_position(buffer) > max_cmd_size);
-        if (enough_data) {
-            cur_size = buffer_get_position(buffer) - cur_doc_begin;
-
-            /* This single document is too large for the command. */
-            if (!idx) {
-                if (op == _INSERT) {
-                    _set_document_too_large(cur_size, max_bson_size);
-                } else {
-                    PyObject* DocumentTooLarge = _error("DocumentTooLarge");
-                    if (DocumentTooLarge) {
-                        /*
-                         * There's nothing intelligent we can say
-                         * about size for update and delete.
-                         */
-                        PyErr_Format(
-                            DocumentTooLarge,
-                            "%s command document too large",
-                            (op == _UPDATE) ? "update": "delete");
-                        Py_DECREF(DocumentTooLarge);
-                    }
+        /* This single document is too large for the command. */
+        if (cur_size > max_cmd_size) {
+            if (op == _INSERT) {
+                _set_document_too_large(cur_size, max_bson_size);
+            } else {
+                PyObject* DocumentTooLarge = _error("DocumentTooLarge");
+                if (DocumentTooLarge) {
+                    /*
+                     * There's nothing intelligent we can say
+                     * about size for update and delete.
+                     */
+                    PyErr_Format(
+                        DocumentTooLarge,
+                        "%s command document too large",
+                        (op == _UPDATE) ? "update": "delete");
+                    Py_DECREF(DocumentTooLarge);
                 }
-                goto fail;
             }
+            goto fail;
+        }
+        if (enough_data) {
             /*
              * Roll the existing buffer back to the beginning
              * of the last document encoded.

pymongo/encryption.py

@@ -49,6 +49,9 @@
                             EncryptionError,
                             InvalidOperation,
                             ServerSelectionTimeoutError)
+from pymongo.message import (_COMMAND_OVERHEAD,
+                             _MAX_ENC_BSON_SIZE,
+                             _raise_document_too_large)
 from pymongo.mongo_client import MongoClient
 from pymongo.pool import _configured_socket, PoolOptions
 from pymongo.read_concern import ReadConcern
@@ -265,11 +268,15 @@ def encrypt(self, database, cmd, check_keys, codec_options):
           The encrypted command to execute.
         """
         self._check_closed()
+        # Workaround for $clusterTime which is incompatible with
+        # check_keys.
+        cluster_time = check_keys and cmd.pop('$clusterTime', None)
+        encoded_cmd = _dict_to_bson(cmd, check_keys, codec_options)
+        max_cmd_size = _MAX_ENC_BSON_SIZE + _COMMAND_OVERHEAD
+        if len(encoded_cmd) > max_cmd_size:
+            raise _raise_document_too_large(
+                next(iter(cmd)), len(encoded_cmd), max_cmd_size)
         with _wrap_encryption_errors():
-            # Workaround for $clusterTime which is incompatible with
-            # check_keys.
-            cluster_time = check_keys and cmd.pop('$clusterTime', None)
-            encoded_cmd = _dict_to_bson(cmd, check_keys, codec_options)
             encrypted_cmd = self._auto_encrypter.encrypt(database, encoded_cmd)
             # TODO: PYTHON-1922 avoid decoding the encrypted_cmd.
             encrypt_cmd = _inflate_bson(

pymongo/message.py

@@ -1421,13 +1421,14 @@ def _batched_write_command_impl(
         value = bson.BSON.encode(doc, check_keys, opts)
         # Is there enough room to add this document? max_cmd_size accounts for
         # the two trailing null bytes.
+        doc_too_large = len(value) > max_cmd_size
         enough_data = (buf.tell() + len(key) + len(value)) >= max_cmd_size
         enough_documents = (idx >= max_write_batch_size)
+        if doc_too_large:
+            write_op = list(_FIELD_MAP.keys())[operation]
+            _raise_document_too_large(
+                write_op, len(value), max_bson_size)
         if enough_data or enough_documents:
-            if not idx:
-                write_op = list(_FIELD_MAP.keys())[operation]
-                _raise_document_too_large(
-                    write_op, len(value), max_bson_size)
             break
         buf.write(_BSONOBJ)
         buf.write(key)

test/client-side-encryption/limits/limits-doc.json

@@ -0,0 +1,102 @@
+{
+    "00": "a", 
+    "01": "a", 
+    "02": "a", 
+    "03": "a", 
+    "04": "a", 
+    "05": "a", 
+    "06": "a", 
+    "07": "a", 
+    "08": "a", 
+    "09": "a", 
+    "10": "a", 
+    "11": "a", 
+    "12": "a", 
+    "13": "a", 
+    "14": "a", 
+    "15": "a", 
+    "16": "a", 
+    "17": "a", 
+    "18": "a", 
+    "19": "a", 
+    "20": "a", 
+    "21": "a", 
+    "22": "a", 
+    "23": "a", 
+    "24": "a", 
+    "25": "a", 
+    "26": "a", 
+    "27": "a", 
+    "28": "a", 
+    "29": "a", 
+    "30": "a", 
+    "31": "a", 
+    "32": "a", 
+    "33": "a", 
+    "34": "a", 
+    "35": "a", 
+    "36": "a", 
+    "37": "a", 
+    "38": "a", 
+    "39": "a", 
+    "40": "a", 
+    "41": "a", 
+    "42": "a", 
+    "43": "a", 
+    "44": "a", 
+    "45": "a", 
+    "46": "a", 
+    "47": "a", 
+    "48": "a", 
+    "49": "a", 
+    "50": "a", 
+    "51": "a", 
+    "52": "a", 
+    "53": "a", 
+    "54": "a", 
+    "55": "a", 
+    "56": "a", 
+    "57": "a", 
+    "58": "a", 
+    "59": "a", 
+    "60": "a", 
+    "61": "a", 
+    "62": "a", 
+    "63": "a", 
+    "64": "a", 
+    "65": "a", 
+    "66": "a", 
+    "67": "a", 
+    "68": "a", 
+    "69": "a", 
+    "70": "a", 
+    "71": "a", 
+    "72": "a", 
+    "73": "a", 
+    "74": "a", 
+    "75": "a", 
+    "76": "a", 
+    "77": "a", 
+    "78": "a", 
+    "79": "a", 
+    "80": "a", 
+    "81": "a", 
+    "82": "a", 
+    "83": "a", 
+    "84": "a", 
+    "85": "a", 
+    "86": "a", 
+    "87": "a", 
+    "88": "a", 
+    "89": "a", 
+    "90": "a", 
+    "91": "a", 
+    "92": "a", 
+    "93": "a", 
+    "94": "a", 
+    "95": "a", 
+    "96": "a", 
+    "97": "a", 
+    "98": "a", 
+    "99": "a"
+}

test/client-side-encryption/limits/limits-key.json

@@ -0,0 +1,31 @@
+{
+    "status": {
+        "$numberInt": "1"
+    }, 
+    "_id": {
+        "$binary": {
+            "base64": "LOCALAAAAAAAAAAAAAAAAA==", 
+            "subType": "04"
+        }
+    }, 
+    "masterKey": {
+        "provider": "local"
+    }, 
+    "updateDate": {
+        "$date": {
+            "$numberLong": "1557827033449"
+        }
+    }, 
+    "keyMaterial": {
+        "$binary": {
+            "base64": "Ce9HSz/HKKGkIt4uyy+jDuKGA+rLC2cycykMo6vc8jXxqa1UVDYHWq1r+vZKbnnSRBfB981akzRKZCFpC05CTyFqDhXv6OnMjpG97OZEREGIsHEYiJkBW0jJJvfLLgeLsEpBzsro9FztGGXASxyxFRZFhXvHxyiLOKrdWfs7X1O/iK3pEoHMx6uSNSfUOgbebLfIqW7TO++iQS5g1xovXA==", 
+            "subType": "00"
+        }
+    }, 
+    "creationDate": {
+        "$date": {
+            "$numberLong": "1557827033449"
+        }
+    },
+    "keyAltNames": [ "local" ]
+}