Skip to content

Commit 55c8bdd

Browse files
ShaneHarveyShaneHarvey
committed
PYTHON-1939 Add prose tests for BSON size limits and batch splitting
1 parent 94e2b10 commit 55c8bdd

File tree

8 files changed

+1672
-29
lines changed

8 files changed

+1672
-29
lines changed

‎pymongo/_cmessagemodule.c

Lines changed: 20 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1565,30 +1565,29 @@ _batched_write_command(
15651565
/* We have enough data, return this batch.
15661566
* max_cmd_size accounts for the two trailing null bytes.
15671567
*/
1568+
cur_size = buffer_get_position(buffer) - cur_doc_begin;
15681569
enough_data = (buffer_get_position(buffer) > max_cmd_size);
1569-
if (enough_data) {
1570-
cur_size = buffer_get_position(buffer) - cur_doc_begin;
1571-
1572-
/* This single document is too large for the command. */
1573-
if (!idx) {
1574-
if (op == _INSERT) {
1575-
_set_document_too_large(cur_size, max_bson_size);
1576-
} else {
1577-
PyObject* DocumentTooLarge = _error("DocumentTooLarge");
1578-
if (DocumentTooLarge) {
1579-
/*
1580-
* There's nothing intelligent we can say
1581-
* about size for update and delete.
1582-
*/
1583-
PyErr_Format(
1584-
DocumentTooLarge,
1585-
"%s command document too large",
1586-
(op == _UPDATE) ? "update": "delete");
1587-
Py_DECREF(DocumentTooLarge);
1588-
}
1570+
/* This single document is too large for the command. */
1571+
if (cur_size > max_cmd_size) {
1572+
if (op == _INSERT) {
1573+
_set_document_too_large(cur_size, max_bson_size);
1574+
} else {
1575+
PyObject* DocumentTooLarge = _error("DocumentTooLarge");
1576+
if (DocumentTooLarge) {
1577+
/*
1578+
* There's nothing intelligent we can say
1579+
* about size for update and delete.
1580+
*/
1581+
PyErr_Format(
1582+
DocumentTooLarge,
1583+
"%s command document too large",
1584+
(op == _UPDATE) ? "update": "delete");
1585+
Py_DECREF(DocumentTooLarge);
15891586
}
1590-
goto fail;
15911587
}
1588+
goto fail;
1589+
}
1590+
if (enough_data) {
15921591
/*
15931592
* Roll the existing buffer back to the beginning
15941593
* of the last document encoded.

‎pymongo/encryption.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,9 @@
4949
EncryptionError,
5050
InvalidOperation,
5151
ServerSelectionTimeoutError)
52+
from pymongo.message import (_COMMAND_OVERHEAD,
53+
_MAX_ENC_BSON_SIZE,
54+
_raise_document_too_large)
5255
from pymongo.mongo_client import MongoClient
5356
from pymongo.pool import _configured_socket, PoolOptions
5457
from pymongo.read_concern import ReadConcern
@@ -265,11 +268,15 @@ def encrypt(self, database, cmd, check_keys, codec_options):
265268
The encrypted command to execute.
266269
"""
267270
self._check_closed()
271+
# Workaround for $clusterTime which is incompatible with
272+
# check_keys.
273+
cluster_time = check_keys and cmd.pop('$clusterTime', None)
274+
encoded_cmd = _dict_to_bson(cmd, check_keys, codec_options)
275+
max_cmd_size = _MAX_ENC_BSON_SIZE + _COMMAND_OVERHEAD
276+
if len(encoded_cmd) > max_cmd_size:
277+
raise _raise_document_too_large(
278+
next(iter(cmd)), len(encoded_cmd), max_cmd_size)
268279
with _wrap_encryption_errors():
269-
# Workaround for $clusterTime which is incompatible with
270-
# check_keys.
271-
cluster_time = check_keys and cmd.pop('$clusterTime', None)
272-
encoded_cmd = _dict_to_bson(cmd, check_keys, codec_options)
273280
encrypted_cmd = self._auto_encrypter.encrypt(database, encoded_cmd)
274281
# TODO: PYTHON-1922 avoid decoding the encrypted_cmd.
275282
encrypt_cmd = _inflate_bson(

‎pymongo/message.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1421,13 +1421,14 @@ def _batched_write_command_impl(
14211421
value = bson.BSON.encode(doc, check_keys, opts)
14221422
# Is there enough room to add this document? max_cmd_size accounts for
14231423
# the two trailing null bytes.
1424+
doc_too_large = len(value) > max_cmd_size
14241425
enough_data = (buf.tell() + len(key) + len(value)) >= max_cmd_size
14251426
enough_documents = (idx >= max_write_batch_size)
1427+
if doc_too_large:
1428+
write_op = list(_FIELD_MAP.keys())[operation]
1429+
_raise_document_too_large(
1430+
write_op, len(value), max_bson_size)
14261431
if enough_data or enough_documents:
1427-
if not idx:
1428-
write_op = list(_FIELD_MAP.keys())[operation]
1429-
_raise_document_too_large(
1430-
write_op, len(value), max_bson_size)
14311432
break
14321433
buf.write(_BSONOBJ)
14331434
buf.write(key)

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

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

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
{
2+
"status": {
3+
"$numberInt": "1"
4+
},
5+
"_id": {
6+
"$binary": {
7+
"base64": "LOCALAAAAAAAAAAAAAAAAA==",
8+
"subType": "04"
9+
}
10+
},
11+
"masterKey": {
12+
"provider": "local"
13+
},
14+
"updateDate": {
15+
"$date": {
16+
"$numberLong": "1557827033449"
17+
}
18+
},
19+
"keyMaterial": {
20+
"$binary": {
21+
"base64": "Ce9HSz/HKKGkIt4uyy+jDuKGA+rLC2cycykMo6vc8jXxqa1UVDYHWq1r+vZKbnnSRBfB981akzRKZCFpC05CTyFqDhXv6OnMjpG97OZEREGIsHEYiJkBW0jJJvfLLgeLsEpBzsro9FztGGXASxyxFRZFhXvHxyiLOKrdWfs7X1O/iK3pEoHMx6uSNSfUOgbebLfIqW7TO++iQS5g1xovXA==",
22+
"subType": "00"
23+
}
24+
},
25+
"creationDate": {
26+
"$date": {
27+
"$numberLong": "1557827033449"
28+
}
29+
},
30+
"keyAltNames": [ "local" ]
31+
}

0 commit comments

Comments
 (0)