diff --git a/bson/__init__.py b/bson/__init__.py index 793c2bbd8f..200b34e119 100644 --- a/bson/__init__.py +++ b/bson/__init__.py @@ -109,6 +109,7 @@ ) from bson.code import Code from bson.codec_options import ( + _RAW_BSON_VIEW_THRESHOLD, DEFAULT_CODEC_OPTIONS, CodecOptions, DatetimeConversion, @@ -242,8 +243,27 @@ def get_data_and_view(data: Any) -> tuple[Any, memoryview]: if isinstance(data, (bytes, bytearray)): return data, memoryview(data) - view = memoryview(data) - return view.tobytes(), view + # Copy other inputs so decoding is immutable + data = memoryview(data).tobytes() + return data, memoryview(data) + + +def _raw_as_bytes(raw: Union[bytes, bytearray, memoryview]) -> bytes: + """Convert a raw BSON buffer (bytes, bytearray, or memoryview) to bytes.""" + return raw if isinstance(raw, bytes) else bytes(raw) + + +def _raw_slice(data: Any, view: memoryview, position: int, end: int, obj_size: int) -> Any: + """Return the raw BSON document spanning ``position`` to ``end`` for use as a buffer.""" + whole_span = position == 0 and obj_size == len(data) + if view.readonly: # data is immutable (bytes). + if whole_span: + return data + if obj_size >= _RAW_BSON_VIEW_THRESHOLD: + return view[position : end + 1] + return data[position : end + 1] + # Mutable buffer, must copy. + return bytes(data) if whole_span else _raw_as_bytes(data[position : end + 1]) def _raise_unknown_type(element_type: int, element_name: str) -> NoReturn: @@ -311,7 +331,8 @@ def _get_object( """Decode a BSON subdocument to opts.document_class or bson.dbref.DBRef.""" obj_size, end = _get_object_size(data, position, obj_end) if _raw_document_class(opts.document_class): - return (opts.document_class(data[position : end + 1], opts), position + obj_size) + buf = _raw_slice(data, view, position, end, obj_size) + return (opts.document_class(buf, opts), position + obj_size) obj = _elements_to_dict(data, view, position + 4, end, opts) @@ -618,7 +639,8 @@ def _bson_to_dict(data: Any, opts: CodecOptions[_DocumentType]) -> _DocumentType data, view = get_data_and_view(data) try: if _raw_document_class(opts.document_class): - return opts.document_class(data, opts) # type:ignore[call-arg] + buf = _raw_slice(data, view, 0, len(data) - 1, len(data)) + return opts.document_class(buf, opts) # type:ignore[call-arg] _, end = _get_object_size(data, 0, len(data)) return cast("_DocumentType", _elements_to_dict(data, view, 4, end, opts)) except InvalidBSON: @@ -708,7 +730,8 @@ def _encode_bytes(name: bytes, value: bytes, dummy0: Any, dummy1: Any) -> bytes: def _encode_mapping(name: bytes, value: Any, check_keys: bool, opts: CodecOptions[Any]) -> bytes: """Encode a mapping type.""" if _raw_document_class(value): - return b"\x03" + name + cast(bytes, value.raw) + # join avoids a copy by consuming a memoryview raw directly. + return b"".join((b"\x03", name, value.raw)) data = b"".join([_element_to_bson(key, val, check_keys, opts) for key, val in value.items()]) return b"\x03" + name + _PACK_INT(len(data) + 5) + data + b"\x00" @@ -994,7 +1017,7 @@ def _dict_to_bson( ) -> bytes: """Encode a document to BSON.""" if _raw_document_class(doc): - return cast(bytes, doc.raw) + return _raw_as_bytes(doc.raw) try: elements = [] if top_level and "_id" in doc: @@ -1109,7 +1132,8 @@ def _decode_all(data: _ReadableBuffer, opts: CodecOptions[_DocumentType]) -> lis if data[obj_end] != 0: raise InvalidBSON("bad eoo") if use_raw: - docs.append(opts.document_class(data[position : obj_end + 1], opts)) # type: ignore + raw_buf = _raw_slice(data, view, position, obj_end, obj_size) + docs.append(opts.document_class(raw_buf, opts)) # type: ignore else: docs.append(_elements_to_dict(data, view, position + 4, obj_end, opts)) position += obj_size diff --git a/bson/_cbsonmodule.c b/bson/_cbsonmodule.c index a9ef25e01f..7efd0cf02a 100644 --- a/bson/_cbsonmodule.c +++ b/bson/_cbsonmodule.c @@ -84,6 +84,7 @@ struct module_state { PyObject* _from_bid_str; int64_t min_millis; int64_t max_millis; + Py_ssize_t raw_bson_view_threshold; }; #define GETSTATE(m) ((struct module_state*)PyModule_GetState(m)) @@ -250,6 +251,11 @@ static int _write_element_to_buffer(PyObject* self, buffer_t buffer, */ static int write_raw_doc(buffer_t buffer, PyObject* raw, PyObject* _raw); +/* Get a read-only buffer view of a bytes-like object. + * Returns 1 on success or 0 on failure with an exception set. + */ +static int _get_buffer(PyObject *exporter, Py_buffer *view); + #if PY_VERSION_HEX >= PYTHON_3_12 /* Transfer traceback from old_exc to new_exc. * Steals reference to old_exc. */ @@ -691,6 +697,7 @@ static int _load_python_objects(PyObject* module) { PyObject* compiled = NULL; PyObject* min_datetime_ms = NULL; PyObject* max_datetime_ms = NULL; + PyObject* raw_bson_view_threshold = NULL; struct module_state *state = GETSTATE(module); if (!state) { return 1; @@ -742,15 +749,19 @@ static int _load_python_objects(PyObject* module) { _load_object(&min_datetime_ms, "bson.datetime_ms", "_MIN_UTC_MS") || _load_object(&max_datetime_ms, "bson.datetime_ms", "_MAX_UTC_MS") || _load_object(&state->min_datetime, "bson.datetime_ms", "_MIN_UTC") || - _load_object(&state->max_datetime, "bson.datetime_ms", "_MAX_UTC")) { + _load_object(&state->max_datetime, "bson.datetime_ms", "_MAX_UTC") || + _load_object(&raw_bson_view_threshold, "bson.codec_options", "_RAW_BSON_VIEW_THRESHOLD")) { return 1; } state->min_millis = PyLong_AsLongLong(min_datetime_ms); state->max_millis = PyLong_AsLongLong(max_datetime_ms); + state->raw_bson_view_threshold = PyLong_AsSsize_t(raw_bson_view_threshold); Py_DECREF(min_datetime_ms); Py_DECREF(max_datetime_ms); - if ((state->min_millis == -1 || state->max_millis == -1) && PyErr_Occurred()) { + Py_DECREF(raw_bson_view_threshold); + if ((state->min_millis == -1 || state->max_millis == -1 || + state->raw_bson_view_threshold == -1) && PyErr_Occurred()) { return 1; } @@ -914,6 +925,7 @@ int convert_codec_options(PyObject* self, PyObject* options_obj, codec_options_t options->is_raw_bson = (101 == type_marker); options->is_dict_class = (options->document_class == (PyObject*)&PyDict_Type); + options->buffer_owner = NULL; options->options_obj = options_obj; Py_INCREF(options->options_obj); @@ -1735,29 +1747,45 @@ int decode_and_write_pair(PyObject* self, buffer_t buffer, * Returns the number of bytes written or 0 on failure. */ static int write_raw_doc(buffer_t buffer, PyObject* raw, PyObject* _raw_str) { - char* bytes; + char* data; Py_ssize_t len; int len_int; int bytes_written = 0; PyObject* bytes_obj = NULL; + Py_buffer view = {0}; bytes_obj = PyObject_GetAttr(raw, _raw_str); if (!bytes_obj) { goto fail; } - if (-1 == PyBytes_AsStringAndSize(bytes_obj, &bytes, &len)) { + if (PyBytes_Check(bytes_obj)) { + /* The common case: raw is bytes. */ + data = PyBytes_AS_STRING(bytes_obj); + len = PyBytes_GET_SIZE(bytes_obj); + } else if (PyMemoryView_Check(bytes_obj)) { + /* raw may also be a memoryview of the decode buffer. */ + if (!_get_buffer(bytes_obj, &view)) { + goto fail; + } + data = (char*)view.buf; + len = view.len; + } else { + PyErr_Format(PyExc_TypeError, + "RawBSONDocument.raw must be bytes or memoryview, not %.200s", + Py_TYPE(bytes_obj)->tp_name); goto fail; } len_int = _downcast_and_check(len, 0); if (-1 == len_int) { goto fail; } - if (!buffer_write_bytes(buffer, bytes, len_int)) { + if (!buffer_write_bytes(buffer, data, len_int)) { goto fail; } bytes_written = len_int; fail: + PyBuffer_Release(&view); Py_XDECREF(bytes_obj); return bytes_written; } @@ -2031,6 +2059,13 @@ static PyObject* _cbson_dict_to_bson(PyObject* self, PyObject* args) { if (NULL == raw_bson_document_bytes_obj) { return NULL; } + /* raw may be a memoryview but + * encoding must always produce bytes. */ + if (!PyBytes_Check(raw_bson_document_bytes_obj)) { + PyObject* as_bytes = PyBytes_FromObject(raw_bson_document_bytes_obj); + Py_DECREF(raw_bson_document_bytes_obj); + return as_bytes; + } return raw_bson_document_bytes_obj; } @@ -2832,7 +2867,9 @@ static int _element_to_dict(PyObject* self, const char* string, } static PyObject* _cbson_element_to_dict(PyObject* self, PyObject* args) { - /* TODO: Support buffer protocol */ + /* TODO(PYTHON-6038): buffer-protocol inputs are copied + * upstream in get_data_and_view. Native RawBSONDocument inflation in C + * should accept the buffer directly. */ char* string; PyObject* bson; PyObject* options_obj = NULL; @@ -2843,7 +2880,7 @@ static PyObject* _cbson_element_to_dict(PyObject* self, PyObject* args) { int raw_array = 0; PyObject* name; PyObject* value; - PyObject* result_tuple; + PyObject* result_tuple = NULL; if (!(PyArg_ParseTuple(args, "OIIOp", &bson, &position, &max, &options_obj, &raw_array) && @@ -2853,22 +2890,23 @@ static PyObject* _cbson_element_to_dict(PyObject* self, PyObject* args) { if (!PyBytes_Check(bson)) { PyErr_SetString(PyExc_TypeError, "argument to _element_to_dict must be a bytes object"); - return NULL; + goto done; } string = PyBytes_AS_STRING(bson); + options.buffer_owner = bson; new_position = _element_to_dict(self, string, position, max, &options, raw_array, &name, &value); if (new_position < 0) { - return NULL; + goto done; } result_tuple = Py_BuildValue("NNi", name, value, new_position); if (!result_tuple) { Py_DECREF(name); Py_DECREF(value); - return NULL; } +done: destroy_codec_options(&options); return result_tuple; } @@ -2934,7 +2972,44 @@ static PyObject* elements_to_dict(PyObject* self, const char* string, const codec_options_t* options) { PyObject* result; if (options->is_raw_bson) { - PyObject* bson_bytes = PyBytes_FromStringAndSize(string, max); + PyObject* bson_bytes; + PyObject* buffer_owner = options->buffer_owner; + int owner_is_bytes = buffer_owner && PyBytes_Check(buffer_owner); + Py_ssize_t offset = 0; + if (owner_is_bytes) { + offset = string - PyBytes_AS_STRING(buffer_owner); + if (offset < 0 || + offset + (Py_ssize_t)max > PyBytes_GET_SIZE(buffer_owner)) { + PyObject* InvalidBSON = _error("InvalidBSON"); + if (InvalidBSON) { + PyErr_SetString(InvalidBSON, "invalid buffer offset"); + Py_DECREF(InvalidBSON); + } + return NULL; + } + } + if (owner_is_bytes && offset == 0 && + (Py_ssize_t)max == PyBytes_GET_SIZE(buffer_owner)) { + /* The document spans the entire buffer, pass the buffer + * itself through. */ + bson_bytes = buffer_owner; + Py_INCREF(bson_bytes); + } else if (owner_is_bytes && + (Py_ssize_t)max >= GETSTATE(self)->raw_bson_view_threshold) { + /* Zero-copy: pass a read-only slice of the buffer instead of a + * bytes copy. Only immutable (bytes) buffers may be sliced this + * way, mutable buffers must be copied so + * the caller can't mutate the document after decoding. */ + PyObject* top_view = PyMemoryView_FromObject(buffer_owner); + if (!top_view) { + return NULL; + } + bson_bytes = PySequence_GetSlice(top_view, offset, + offset + (Py_ssize_t)max); + Py_DECREF(top_view); + } else { + bson_bytes = PyBytes_FromStringAndSize(string, max); + } if (!bson_bytes) { return NULL; } @@ -2974,6 +3049,51 @@ static int _get_buffer(PyObject *exporter, Py_buffer *view) { return 0; } +/* Return 1 if any document in a stream of BSON documents is at least + * `threshold` bytes. Malformed lengths return 0. */ +static int _contains_view_eligible_doc(const char* data, Py_ssize_t len, + Py_ssize_t threshold) { + Py_ssize_t position = 0; + while (len - position >= 4) { + int32_t size; + memcpy(&size, data + position, 4); + size = (int32_t)BSON_UINT32_FROM_LE(size); + if (size < BSON_MIN_SIZE || (Py_ssize_t)size > len - position) { + return 0; + } + if ((Py_ssize_t)size >= threshold) { + return 1; + } + position += size; + } + return 0; +} + +/* Prepare an input buffer for raw-document decoding: bytes and + * bytearray are returned as-is, with other inputs copied to + * bytes only if the stream contains a document large enough for a zero-copy + * view. Returns a new reference or NULL on failure with an exception set. */ +static PyObject* _prepare_input_buffer(PyObject* self, PyObject* bson, + const codec_options_t* options) { + Py_buffer tmp = {0}; + int needs_copy; + if (!options->is_raw_bson || PyBytes_Check(bson) || PyByteArray_Check(bson)) { + Py_INCREF(bson); + return bson; + } + if (!_get_buffer(bson, &tmp)) { + return NULL; + } + needs_copy = _contains_view_eligible_doc( + (const char*)tmp.buf, tmp.len, GETSTATE(self)->raw_bson_view_threshold); + PyBuffer_Release(&tmp); + if (!needs_copy) { + Py_INCREF(bson); + return bson; + } + return PyBytes_FromObject(bson); +} + static PyObject* _cbson_bson_to_dict(PyObject* self, PyObject* args) { int32_t size; Py_ssize_t total_size; @@ -2989,7 +3109,14 @@ static PyObject* _cbson_bson_to_dict(PyObject* self, PyObject* args) { return result; } + bson = _prepare_input_buffer(self, bson, &options); + if (!bson) { + destroy_codec_options(&options); + return result; + } + if (!_get_buffer(bson, &view)) { + Py_DECREF(bson); destroy_codec_options(&options); return result; } @@ -3007,6 +3134,7 @@ static PyObject* _cbson_bson_to_dict(PyObject* self, PyObject* args) { } string = (char*)view.buf; + options.buffer_owner = bson; memcpy(&size, string, 4); size = (int32_t)BSON_UINT32_FROM_LE(size); if (size < BSON_MIN_SIZE) { @@ -3039,6 +3167,7 @@ static PyObject* _cbson_bson_to_dict(PyObject* self, PyObject* args) { result = elements_to_dict(self, string, (unsigned)size, &options); done: PyBuffer_Release(&view); + Py_DECREF(bson); destroy_codec_options(&options); return result; } @@ -3059,12 +3188,20 @@ static PyObject* _cbson_decode_all(PyObject* self, PyObject* args) { return NULL; } + bson = _prepare_input_buffer(self, bson, &options); + if (!bson) { + destroy_codec_options(&options); + return NULL; + } + if (!_get_buffer(bson, &view)) { + Py_DECREF(bson); destroy_codec_options(&options); return NULL; } total_size = view.len; string = (char*)view.buf; + options.buffer_owner = bson; if (!(result = PyList_New(0))) { goto fail; @@ -3133,6 +3270,7 @@ static PyObject* _cbson_decode_all(PyObject* self, PyObject* args) { result = NULL; done: PyBuffer_Release(&view); + Py_DECREF(bson); destroy_codec_options(&options); return result; } diff --git a/bson/_cbsonmodule.h b/bson/_cbsonmodule.h index a9bee24b8d..2cf357bd5d 100644 --- a/bson/_cbsonmodule.h +++ b/bson/_cbsonmodule.h @@ -73,6 +73,7 @@ typedef struct codec_options_t { PyObject* options_obj; unsigned char is_raw_bson; unsigned char is_dict_class; + PyObject* buffer_owner; /* The owning decode input buffer for RawBSONDocument */ } codec_options_t; /* C API functions */ diff --git a/bson/codec_options.py b/bson/codec_options.py index 848f672f54..89d87a74b6 100644 --- a/bson/codec_options.py +++ b/bson/codec_options.py @@ -41,6 +41,10 @@ _RAW_BSON_DOCUMENT_MARKER = 101 +# Raw BSON documents at least this many bytes are exposed as read-only +# memoryview slices of the decode buffer instead of bytes copies. +_RAW_BSON_VIEW_THRESHOLD = 4096 + def _raw_document_class(document_class: Any) -> bool: """Determine if a document_class is a RawBSONDocument class.""" diff --git a/bson/json_util.py b/bson/json_util.py index e12e04ccb8..8d84184487 100644 --- a/bson/json_util.py +++ b/bson/json_util.py @@ -97,7 +97,10 @@ performance improvement. `python-bsonjs` is a fast BSON to MongoDB Extended JSON converter for Python built on top of `libbson `_. `python-bsonjs` works best - with PyMongo when using :class:`~bson.raw_bson.RawBSONDocument`. + with PyMongo when using :class:`~bson.raw_bson.RawBSONDocument`. Note that + `python-bsonjs` requires an exact :class:`bytes` instance, while + :attr:`~bson.raw_bson.RawBSONDocument.raw` may be a :class:`memoryview`, + so pass ``bytes(doc.raw)``. """ from __future__ import annotations diff --git a/bson/raw_bson.py b/bson/raw_bson.py index 42bd19cab4..411916fc7a 100644 --- a/bson/raw_bson.py +++ b/bson/raw_bson.py @@ -53,10 +53,11 @@ from __future__ import annotations +import copyreg from collections.abc import ItemsView, Iterator, Mapping from typing import Any, Optional -from bson import _get_object_size, _raw_to_dict +from bson import _get_object_size, _raw_as_bytes, _raw_to_dict from bson.codec_options import _RAW_BSON_DOCUMENT_MARKER, CodecOptions from bson.codec_options import DEFAULT_CODEC_OPTIONS as DEFAULT @@ -142,7 +143,16 @@ class from the standard library so it can be used like a read-only @property def raw(self) -> bytes | memoryview: - """The raw BSON bytes composing this document.""" + """The raw BSON bytes composing this document. + + .. versionchanged:: 4.18 + Documents and subdocuments 4KB and larger decoded from an + immutable buffer are returned as read-only :class:`memoryview` + slices of that buffer instead of :class:`bytes` copies. Documents + decoded from mutable buffers such as :class:`bytearray` are + always :class:`bytes` copies. Call ``bytes(doc.raw)`` + to get an independent copy. + """ return self.__raw def items(self) -> ItemsView[str, Any]: @@ -179,8 +189,21 @@ def __eq__(self, other: Any) -> bool: __hash__ = None # type: ignore[assignment] + def __getstate__(self) -> tuple[Optional[dict[str, Any]], dict[str, Any]]: + slots_state: dict[str, Any] = { + name: getattr(self, name) + for name in copyreg._slotnames(type(self)) # type: ignore[attr-defined] + if hasattr(self, name) + } + slots_state["_RawBSONDocument__raw"] = _raw_as_bytes(self.__raw) + slots_state["_RawBSONDocument__inflated_doc"] = None + return getattr(self, "__dict__", None), slots_state + def __repr__(self) -> str: - return f"{self.__class__.__name__}({self.raw!r}, codec_options={self.__codec_options!r})" + return ( + f"{self.__class__.__name__}({_raw_as_bytes(self.__raw)!r}, " + f"codec_options={self.__codec_options!r})" + ) class _RawArrayBSONDocument(RawBSONDocument): diff --git a/doc/changelog.rst b/doc/changelog.rst index e752f9bcb2..fb9d635990 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -29,6 +29,13 @@ PyMongo 4.18 brings a number of changes including: attempts, so consumers can correlate a retried operation's events. As a result, ``operation_id`` is no longer equal to the per-attempt ``request_id`` for these operations. +- Improved the performance and memory usage of decoding large documents to + :class:`~bson.raw_bson.RawBSONDocument`. Documents and subdocuments that are 4KB or greater + and decoded from an immutable buffer are now exposed as read-only :class:`memoryview` + slices instead of :class:`bytes` copies. Documents decoded from mutable buffers such as a + :class:`bytearray` are always :class:`bytes` copies. +- :func:`bson.get_data_and_view` now returns a view of a private :class:`bytes` copy + for buffer-protocol inputs other than :class:`bytes` or :class:`bytearray`. - Fixed a potential out-of-bounds read in the C extension when decoding an array of BSON documents. An embedded document whose declared length exceeds the bytes remaining in the array now raises diff --git a/test/asynchronous/test_raw_bson.py b/test/asynchronous/test_raw_bson.py index 688da7a670..6675c30928 100644 --- a/test/asynchronous/test_raw_bson.py +++ b/test/asynchronous/test_raw_bson.py @@ -19,13 +19,12 @@ sys.path[0:0] = [""] -from bson import Code, DBRef, decode, encode +from bson import decode, encode from bson.binary import JAVA_LEGACY, Binary, UuidRepresentation from bson.codec_options import CodecOptions -from bson.errors import InvalidBSON from bson.raw_bson import DEFAULT_RAW_BSON_OPTIONS, RawBSONDocument -from bson.son import SON from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest +from test.test_raw_bson_shared import TEST_RAW_BSON _IS_SYNC = False @@ -34,40 +33,22 @@ class TestRawBSONDocument(AsyncIntegrationTest): # {'_id': ObjectId('556df68b6e32ab21a95e0785'), # 'name': 'Sherlock', # 'addresses': [{'street': 'Baker Street'}]} - bson_string = ( - b"Z\x00\x00\x00\x07_id\x00Um\xf6\x8bn2\xab!\xa9^\x07\x85\x02name\x00\t" - b"\x00\x00\x00Sherlock\x00\x04addresses\x00&\x00\x00\x00\x030\x00\x1e" - b"\x00\x00\x00\x02street\x00\r\x00\x00\x00Baker Street\x00\x00\x00\x00" - ) + bson_string = TEST_RAW_BSON document = RawBSONDocument(bson_string) async def asyncTearDown(self): if async_client_context.connected: await self.client.pymongo_test.test_raw.drop() - def test_decode(self): - self.assertEqual("Sherlock", self.document["name"]) - first_address = self.document["addresses"][0] - self.assertIsInstance(first_address, RawBSONDocument) - self.assertEqual("Baker Street", first_address["street"]) - - def test_raw(self): - self.assertEqual(self.bson_string, self.document.raw) - - def test_empty_doc(self): - doc = RawBSONDocument(encode({})) - with self.assertRaises(KeyError): - doc["does-not-exist"] - - def test_invalid_bson_sequence(self): - bson_byte_sequence = encode({"a": 1}) + encode({}) - with self.assertRaisesRegex(InvalidBSON, "invalid object length"): - RawBSONDocument(bson_byte_sequence) - - def test_invalid_bson_eoo(self): - invalid_bson_eoo = encode({"a": 1})[:-1] + b"\x01" - with self.assertRaisesRegex(InvalidBSON, "bad eoo"): - RawBSONDocument(invalid_bson_eoo) + @async_client_context.require_connection + async def test_round_trip_view_backed_document(self): + inner = {"payload": "x" * 8000, "marker": 1} + subdoc = RawBSONDocument(encode({"big": inner}))["big"] + self.assertIsInstance(subdoc.raw, memoryview) + coll = self.client.pymongo_test.test_raw + await coll.insert_one(subdoc) + result = await coll.find_one({"marker": 1}, {"_id": False}) + self.assertEqual(inner, result) @async_client_context.require_connection async def test_round_trip(self): @@ -101,24 +82,6 @@ async def test_round_trip_raw_uuid(self): raw_coll = coll.with_options(codec_options=DEFAULT_RAW_BSON_OPTIONS) self.assertEqual(await raw_coll.find_one(), raw) - def test_with_codec_options(self): - # {'date': datetime.datetime(2015, 6, 3, 18, 40, 50, 826000), - # '_id': UUID('026fab8f-975f-4965-9fbf-85ad874c60ff')} - # encoded with JAVA_LEGACY uuid representation. - bson_string = ( - b"-\x00\x00\x00\x05_id\x00\x10\x00\x00\x00\x03eI_\x97\x8f\xabo\x02" - b"\xff`L\x87\xad\x85\xbf\x9f\tdate\x00\x8a\xd6\xb9\xbaM" - b"\x01\x00\x00\x00" - ) - document = RawBSONDocument( - bson_string, - codec_options=CodecOptions( - uuid_representation=JAVA_LEGACY, document_class=RawBSONDocument - ), - ) - - self.assertEqual(uuid.UUID("026fab8f-975f-4965-9fbf-85ad874c60ff"), document["_id"]) - @async_client_context.require_connection async def test_round_trip_codec_options(self): doc = { @@ -188,31 +151,6 @@ async def test_write_response_raw_bson(self): await coll.update_one(self.document, {"$set": {"a": "b"}}, upsert=True) await coll.update_many(self.document, {"$set": {"b": "c"}}) - def test_preserve_key_ordering(self): - keyvaluepairs = [ - ("a", 1), - ("b", 2), - ("c", 3), - ] - rawdoc = RawBSONDocument(encode(SON(keyvaluepairs))) - - for rkey, elt in zip(rawdoc, keyvaluepairs): - self.assertEqual(rkey, elt[0]) - - def test_contains_code_with_scope(self): - doc = RawBSONDocument(encode({"value": Code("x=1", scope={})})) - - self.assertEqual(decode(encode(doc)), {"value": Code("x=1", {})}) - self.assertEqual(doc["value"].scope, RawBSONDocument(encode({}))) - - def test_contains_dbref(self): - doc = RawBSONDocument(encode({"value": DBRef("test", "id")})) - raw = {"$ref": "test", "$id": "id"} - raw_encoded = encode(decode(encode(raw))) - - self.assertEqual(decode(encode(doc)), {"value": DBRef("test", "id")}) - self.assertEqual(doc["value"].raw, raw_encoded) - if __name__ == "__main__": unittest.main() diff --git a/test/test_raw_bson.py b/test/test_raw_bson.py index 0c068d7dbf..b9f0d6239c 100644 --- a/test/test_raw_bson.py +++ b/test/test_raw_bson.py @@ -19,13 +19,12 @@ sys.path[0:0] = [""] -from bson import Code, DBRef, decode, encode +from bson import decode, encode from bson.binary import JAVA_LEGACY, Binary, UuidRepresentation from bson.codec_options import CodecOptions -from bson.errors import InvalidBSON from bson.raw_bson import DEFAULT_RAW_BSON_OPTIONS, RawBSONDocument -from bson.son import SON from test import IntegrationTest, client_context, unittest +from test.test_raw_bson_shared import TEST_RAW_BSON _IS_SYNC = True @@ -34,40 +33,22 @@ class TestRawBSONDocument(IntegrationTest): # {'_id': ObjectId('556df68b6e32ab21a95e0785'), # 'name': 'Sherlock', # 'addresses': [{'street': 'Baker Street'}]} - bson_string = ( - b"Z\x00\x00\x00\x07_id\x00Um\xf6\x8bn2\xab!\xa9^\x07\x85\x02name\x00\t" - b"\x00\x00\x00Sherlock\x00\x04addresses\x00&\x00\x00\x00\x030\x00\x1e" - b"\x00\x00\x00\x02street\x00\r\x00\x00\x00Baker Street\x00\x00\x00\x00" - ) + bson_string = TEST_RAW_BSON document = RawBSONDocument(bson_string) def tearDown(self): if client_context.connected: self.client.pymongo_test.test_raw.drop() - def test_decode(self): - self.assertEqual("Sherlock", self.document["name"]) - first_address = self.document["addresses"][0] - self.assertIsInstance(first_address, RawBSONDocument) - self.assertEqual("Baker Street", first_address["street"]) - - def test_raw(self): - self.assertEqual(self.bson_string, self.document.raw) - - def test_empty_doc(self): - doc = RawBSONDocument(encode({})) - with self.assertRaises(KeyError): - doc["does-not-exist"] - - def test_invalid_bson_sequence(self): - bson_byte_sequence = encode({"a": 1}) + encode({}) - with self.assertRaisesRegex(InvalidBSON, "invalid object length"): - RawBSONDocument(bson_byte_sequence) - - def test_invalid_bson_eoo(self): - invalid_bson_eoo = encode({"a": 1})[:-1] + b"\x01" - with self.assertRaisesRegex(InvalidBSON, "bad eoo"): - RawBSONDocument(invalid_bson_eoo) + @client_context.require_connection + def test_round_trip_view_backed_document(self): + inner = {"payload": "x" * 8000, "marker": 1} + subdoc = RawBSONDocument(encode({"big": inner}))["big"] + self.assertIsInstance(subdoc.raw, memoryview) + coll = self.client.pymongo_test.test_raw + coll.insert_one(subdoc) + result = coll.find_one({"marker": 1}, {"_id": False}) + self.assertEqual(inner, result) @client_context.require_connection def test_round_trip(self): @@ -101,24 +82,6 @@ def test_round_trip_raw_uuid(self): raw_coll = coll.with_options(codec_options=DEFAULT_RAW_BSON_OPTIONS) self.assertEqual(raw_coll.find_one(), raw) - def test_with_codec_options(self): - # {'date': datetime.datetime(2015, 6, 3, 18, 40, 50, 826000), - # '_id': UUID('026fab8f-975f-4965-9fbf-85ad874c60ff')} - # encoded with JAVA_LEGACY uuid representation. - bson_string = ( - b"-\x00\x00\x00\x05_id\x00\x10\x00\x00\x00\x03eI_\x97\x8f\xabo\x02" - b"\xff`L\x87\xad\x85\xbf\x9f\tdate\x00\x8a\xd6\xb9\xbaM" - b"\x01\x00\x00\x00" - ) - document = RawBSONDocument( - bson_string, - codec_options=CodecOptions( - uuid_representation=JAVA_LEGACY, document_class=RawBSONDocument - ), - ) - - self.assertEqual(uuid.UUID("026fab8f-975f-4965-9fbf-85ad874c60ff"), document["_id"]) - @client_context.require_connection def test_round_trip_codec_options(self): doc = { @@ -188,31 +151,6 @@ def test_write_response_raw_bson(self): coll.update_one(self.document, {"$set": {"a": "b"}}, upsert=True) coll.update_many(self.document, {"$set": {"b": "c"}}) - def test_preserve_key_ordering(self): - keyvaluepairs = [ - ("a", 1), - ("b", 2), - ("c", 3), - ] - rawdoc = RawBSONDocument(encode(SON(keyvaluepairs))) - - for rkey, elt in zip(rawdoc, keyvaluepairs): - self.assertEqual(rkey, elt[0]) - - def test_contains_code_with_scope(self): - doc = RawBSONDocument(encode({"value": Code("x=1", scope={})})) - - self.assertEqual(decode(encode(doc)), {"value": Code("x=1", {})}) - self.assertEqual(doc["value"].scope, RawBSONDocument(encode({}))) - - def test_contains_dbref(self): - doc = RawBSONDocument(encode({"value": DBRef("test", "id")})) - raw = {"$ref": "test", "$id": "id"} - raw_encoded = encode(decode(encode(raw))) - - self.assertEqual(decode(encode(doc)), {"value": DBRef("test", "id")}) - self.assertEqual(doc["value"].raw, raw_encoded) - if __name__ == "__main__": unittest.main() diff --git a/test/test_raw_bson_shared.py b/test/test_raw_bson_shared.py new file mode 100644 index 0000000000..aa8f799636 --- /dev/null +++ b/test/test_raw_bson_shared.py @@ -0,0 +1,271 @@ +# Copyright 2015-present MongoDB, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import copy +import gc +import pickle +import sys +import unittest +import uuid + +from test import UnitTest + +sys.path[0:0] = [""] + +from bson import Code, DBRef, decode, decode_all, encode, has_c +from bson.binary import JAVA_LEGACY +from bson.codec_options import CodecOptions +from bson.errors import InvalidBSON +from bson.raw_bson import DEFAULT_RAW_BSON_OPTIONS, RawBSONDocument +from bson.son import SON + +# {'_id': ObjectId('556df68b6e32ab21a95e0785'), +# 'name': 'Sherlock', +# 'addresses': [{'street': 'Baker Street'}]} +TEST_RAW_BSON = ( + b"Z\x00\x00\x00\x07_id\x00Um\xf6\x8bn2\xab!\xa9^\x07\x85\x02name\x00\t" + b"\x00\x00\x00Sherlock\x00\x04addresses\x00&\x00\x00\x00\x030\x00\x1e" + b"\x00\x00\x00\x02street\x00\r\x00\x00\x00Baker Street\x00\x00\x00\x00" +) + + +class _TaggedRawBSONDocument(RawBSONDocument): + """RawBSONDocument subclass with a different __init__ signature and + extra instance state, stored in __dict__.""" + + def __init__(self, bson_bytes, tag, codec_options=None): + super().__init__(bson_bytes, codec_options) + self.tag = tag + + +class _SlottedRawBSONDocument(RawBSONDocument): + """RawBSONDocument subclass with extra state stored in its own slot.""" + + __slots__ = ("tag",) + + def __init__(self, bson_bytes, tag, codec_options=None): + super().__init__(bson_bytes, codec_options) + self.tag = tag + + +class TestRawBSONDocument(UnitTest): + bson_string = TEST_RAW_BSON + document = RawBSONDocument(bson_string) + + def test_decode(self): + self.assertEqual("Sherlock", self.document["name"]) + first_address = self.document["addresses"][0] + self.assertIsInstance(first_address, RawBSONDocument) + self.assertEqual("Baker Street", first_address["street"]) + + def test_raw(self): + self.assertEqual(self.bson_string, self.document.raw) + + def test_large_subdocument_zero_copy_view(self): + doc = RawBSONDocument(encode({"small": {"n": 1}, "big": {"payload": "x" * 8000}})) + self.assertIsInstance(doc["small"].raw, bytes) + big = doc["big"] + self.assertIsInstance(big.raw, memoryview) + self.assertTrue(big.raw.readonly) + self.assertEqual(encode({"payload": "x" * 8000}), bytes(big.raw)) + self.assertEqual("x" * 8000, big["payload"]) + + def test_large_subdocument_view_keeps_buffer_alive(self): + expected = encode({"payload": "z" * 8000, "n": 42}) + subdoc = RawBSONDocument(encode({"big": {"payload": "z" * 8000, "n": 42}}))["big"] + gc.collect() + churn = [bytearray(8192) for _ in range(100)] + self.assertEqual(42, subdoc["n"]) + self.assertEqual(expected, bytes(subdoc.raw)) + del churn + + def test_decode_whole_buffer_passthrough(self): + data = encode({"payload": "x" * 8000}) + doc = decode(data, DEFAULT_RAW_BSON_OPTIONS) + self.assertIs(data, doc.raw) + + def test_decode_all_zero_copy_views(self): + one = encode({"payload": "w" * 8000}) + docs = decode_all(one * 3, DEFAULT_RAW_BSON_OPTIONS) + self.assertEqual(3, len(docs)) + for doc in docs: + self.assertIsInstance(doc.raw, memoryview) + self.assertEqual(one, bytes(doc.raw)) + self.assertEqual("w" * 8000, docs[0]["payload"]) + (single,) = decode_all(one, DEFAULT_RAW_BSON_OPTIONS) + self.assertIsInstance(single.raw, bytes) + + def test_mutable_buffer_input_copied(self): + one = encode({"payload": "v" * 8000}) + buf = bytearray(one * 2) + docs = decode_all(buf, DEFAULT_RAW_BSON_OPTIONS) + for doc in docs: + self.assertIsInstance(doc.raw, bytes) + buf[:] = bytes(len(buf)) + self.assertEqual("v" * 8000, docs[0]["payload"]) + self.assertEqual(one, docs[1].raw) + + def test_reencode_view_backed_document(self): + inner = {"payload": "x" * 8000} + subdoc = RawBSONDocument(encode({"big": inner}))["big"] + self.assertIsInstance(subdoc.raw, memoryview) + self.assertEqual(encode({"again": inner}), encode({"again": subdoc})) + top = encode(subdoc) + self.assertIsInstance(top, bytes) + self.assertEqual(encode(inner), top) + + @unittest.skipUnless(has_c(), "tests the C extension") + def test_c_encode_rejects_non_bytes_raw(self): + # The C encoder accepts only bytes and memoryview .raw values: + # other buffer types (e.g. a mutable bytearray) raise TypeError. + class _ByteArrayRaw(RawBSONDocument): + @property + def raw(self): + return bytearray(super().raw) + + doc = _ByteArrayRaw(encode({"a": 1})) + with self.assertRaisesRegex(TypeError, "must be bytes or memoryview"): + encode({"sub": doc}) + + def test_pickle_view_backed_document(self): + doc = RawBSONDocument(encode({"big": {"payload": "x" * 8000}})) + subdoc = doc["big"] + self.assertIsInstance(subdoc.raw, memoryview) + for original in (doc, subdoc): + unpickled = pickle.loads(pickle.dumps(original)) + self.assertIsInstance(unpickled.raw, bytes) + self.assertEqual(original, unpickled) + self.assertEqual(dict(original.items()), dict(unpickled.items())) + + def test_decode_mutable_buffer_copied(self): + for decode_one in ( + lambda buf: decode(buf, DEFAULT_RAW_BSON_OPTIONS), + lambda buf: decode_all(buf, DEFAULT_RAW_BSON_OPTIONS)[0], + ): + buf = bytearray(encode({"a": 1})) + doc = decode_one(buf) + self.assertIsInstance(doc.raw, bytes) + buf[-5] = 99 + self.assertEqual(1, doc["a"]) + + def test_buffer_input_copied(self): + big = encode({"payload": "x" * 8000}) + buf = bytearray(big + encode({"a": 1})) + docs = decode_all(memoryview(buf), DEFAULT_RAW_BSON_OPTIONS) + big_raw = docs[0].raw + self.assertIsInstance(big_raw, memoryview) + buf[:] = bytes(len(buf)) + self.assertEqual(big, bytes(big_raw)) + + def test_pickle_deepcopy_subclass(self): + raw_bytes = encode({"payload": "x" * 8000}) + for cls in (_TaggedRawBSONDocument, _SlottedRawBSONDocument): + original = cls(raw_bytes, "tag-value") + for roundtrip in (lambda doc: pickle.loads(pickle.dumps(doc)), copy.deepcopy): + duplicate = roundtrip(original) + self.assertIsInstance(duplicate, cls) + self.assertEqual(original, duplicate) + self.assertEqual("tag-value", duplicate.tag) + self.assertIsInstance(duplicate.raw, bytes) + + def test_repr_view_backed_document(self): + subdoc = RawBSONDocument(encode({"big": {"payload": "x" * 8000}}))["big"] + self.assertIsInstance(subdoc.raw, memoryview) + self.assertIn(repr(bytes(subdoc.raw)), repr(subdoc)) + + @unittest.skipUnless(has_c(), "tests the C extension") + def test_element_to_dict_error_does_not_pin_buffer(self): + from bson import _cbson # type:ignore[attr-defined] + + # An array whose first element is a large subdocument (creates a + # zero-copy view) and whose second element has an invalid type byte. + data = encode({"arr": [{"payload": "x" * 8000}, 1]}) + marker = b"\x101\x00" # type 0x10, key "1" + data = data.replace(marker, b"\xee1\x00") + refcount = sys.getrefcount(data) + for _ in range(5): + with self.assertRaises(InvalidBSON): + _cbson._element_to_dict(data, 4, len(data) - 1, DEFAULT_RAW_BSON_OPTIONS, False) + self.assertEqual(refcount, sys.getrefcount(data)) + + def test_deepcopy_view_backed_document(self): + subdoc = RawBSONDocument(encode({"big": {"payload": "y" * 8000}}))["big"] + self.assertIsInstance(subdoc.raw, memoryview) + copied = copy.deepcopy(subdoc) + self.assertIsInstance(copied.raw, bytes) + self.assertEqual(subdoc, copied) + self.assertEqual("y" * 8000, copied["payload"]) + + def test_empty_doc(self): + doc = RawBSONDocument(encode({})) + with self.assertRaises(KeyError): + doc["does-not-exist"] + + def test_invalid_bson_sequence(self): + bson_byte_sequence = encode({"a": 1}) + encode({}) + with self.assertRaisesRegex(InvalidBSON, "invalid object length"): + RawBSONDocument(bson_byte_sequence) + + def test_invalid_bson_eoo(self): + invalid_bson_eoo = encode({"a": 1})[:-1] + b"\x01" + with self.assertRaisesRegex(InvalidBSON, "bad eoo"): + RawBSONDocument(invalid_bson_eoo) + + def test_with_codec_options(self): + # {'date': datetime.datetime(2015, 6, 3, 18, 40, 50, 826000), + # '_id': UUID('026fab8f-975f-4965-9fbf-85ad874c60ff')} + # encoded with JAVA_LEGACY uuid representation. + bson_string = ( + b"-\x00\x00\x00\x05_id\x00\x10\x00\x00\x00\x03eI_\x97\x8f\xabo\x02" + b"\xff`L\x87\xad\x85\xbf\x9f\tdate\x00\x8a\xd6\xb9\xbaM" + b"\x01\x00\x00\x00" + ) + document = RawBSONDocument( + bson_string, + codec_options=CodecOptions( + uuid_representation=JAVA_LEGACY, document_class=RawBSONDocument + ), + ) + + self.assertEqual(uuid.UUID("026fab8f-975f-4965-9fbf-85ad874c60ff"), document["_id"]) + + def test_preserve_key_ordering(self): + keyvaluepairs = [ + ("a", 1), + ("b", 2), + ("c", 3), + ] + rawdoc = RawBSONDocument(encode(SON(keyvaluepairs))) + + for rkey, elt in zip(rawdoc, keyvaluepairs): + self.assertEqual(rkey, elt[0]) + + def test_contains_code_with_scope(self): + doc = RawBSONDocument(encode({"value": Code("x=1", scope={})})) + + self.assertEqual(decode(encode(doc)), {"value": Code("x=1", {})}) + self.assertEqual(doc["value"].scope, RawBSONDocument(encode({}))) + + def test_contains_dbref(self): + doc = RawBSONDocument(encode({"value": DBRef("test", "id")})) + raw = {"$ref": "test", "$id": "id"} + raw_encoded = encode(decode(encode(raw))) + + self.assertEqual(decode(encode(doc)), {"value": DBRef("test", "id")}) + self.assertEqual(doc["value"].raw, raw_encoded) + + +if __name__ == "__main__": + unittest.main()