From 6291c61e2366270e3648027f712c0812c7b3a8ac Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 16 Sep 2026 20:34:31 +0000 Subject: [PATCH] feat(firestore): add BSON read deserialization support --- .../google/cloud/firestore_v1/_helpers.py | 75 ++++++++++++++----- .../google/cloud/firestore_v1/async_client.py | 2 + .../google/cloud/firestore_v1/base_client.py | 2 + .../cloud/firestore_v1/base_document.py | 18 ++++- .../google/cloud/firestore_v1/bson.py | 20 ++++- .../google/cloud/firestore_v1/client.py | 2 + .../cloud/firestore_v1/pipeline_result.py | 3 +- .../tests/system/test_system.py | 28 ++----- .../tests/system/test_system_async.py | 28 ++----- .../tests/unit/v1/test__helpers.py | 37 +++++++++ 10 files changed, 145 insertions(+), 70 deletions(-) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py index 9793c0685121..51f288f9a84c 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py @@ -44,7 +44,7 @@ import google from google.cloud import exceptions # type: ignore from google.cloud.firestore_v1 import transforms, types -from google.cloud.firestore_v1.bson import _BSONType +from google.cloud.firestore_v1.bson import _BSON_DECODERS, _BSONType from google.cloud.firestore_v1.field_path import FieldPath, parse_field_path from google.cloud.firestore_v1.types import common, document, write from google.cloud.firestore_v1.types.write import DocumentTransform @@ -347,11 +347,7 @@ def reference_value_to_document(reference_value, client) -> Any: return document -def decode_value( - value, client -) -> Union[ - None, bool, int, float, list, datetime.datetime, str, bytes, dict, GeoPoint, Vector -]: +def decode_value(value, client=None, decode_bson: Optional[bool] = None) -> Any: """Converts a Firestore protobuf ``Value`` to a native Python value. Args: @@ -359,15 +355,10 @@ def decode_value( Firestore protobuf to be decoded / parsed / converted. client (:class:`~google.cloud.firestore_v1.client.Client`): A client that has a document factory. + decode_bson (Optional[bool]): Whether to decode BSON extended types. Returns: - Union[NoneType, bool, int, float, datetime.datetime, \ - str, bytes, dict, ~google.cloud.Firestore.GeoPoint]: A native - Python value converted from the ``value``. - - Raises: - NotImplementedError: If the ``value_type`` is ``reference_value``. - ValueError: If the ``value_type`` is unknown. + Any: A native Python value converted from the ``value``. """ value_pb = getattr(value, "_pb", value) value_type = value_pb.WhichOneof("value_type") @@ -394,15 +385,45 @@ def decode_value( ) elif value_type == "array_value": return [ - decode_value(element, client) for element in value_pb.array_value.values + decode_value(element, client, decode_bson=decode_bson) + for element in value_pb.array_value.values ] elif value_type == "map_value": - return decode_dict(value_pb.map_value.fields, client) + return decode_dict(value_pb.map_value.fields, client, decode_bson=decode_bson) else: raise ValueError("Unknown ``value_type``", value_type) -def decode_dict(value_fields, client) -> Union[dict, Vector]: +def _decode_bson_dict(data: dict) -> Optional[_BSONType]: + """Decode a single-key wire map dictionary if registered.""" + if len(data) == 1: + key, val = next(iter(data.items())) + decoder = _BSON_DECODERS.get(key) + if decoder is not None: + try: + return decoder(val) + except Exception: + pass + return None + + +def _decode_bson_dict_recursive(data: Any) -> Any: + """Recursively decodes BSON wire map dictionaries.""" + if isinstance(data, dict): + decoded = _decode_bson_dict(data) + if decoded is not None: + return decoded + return {k: _decode_bson_dict_recursive(v) for k, v in data.items()} + elif isinstance(data, list): + return [_decode_bson_dict_recursive(item) for item in data] + return data + + +def decode_dict( + value_fields, + client=None, + decode_bson: Optional[bool] = None, +) -> Union[dict, Vector, _BSONType]: """Converts a protobuf map of Firestore ``Value``-s. Args: @@ -410,14 +431,18 @@ def decode_dict(value_fields, client) -> Union[dict, Vector]: protobuf map of Firestore ``Value``-s. client (:class:`~google.cloud.firestore_v1.client.Client`): A client that has a document factory. + decode_bson (Optional[bool]): Whether to decode BSON extended types. Returns: - Dict[str, Union[NoneType, bool, int, float, datetime.datetime, \ - str, bytes, dict, ~google.cloud.Firestore.GeoPoint]]: A dictionary - of native Python values converted from the ``value_fields``. + Union[dict, ~google.cloud.firestore_v1.vector.Vector, \ + ~google.cloud.firestore_v1.bson._BSONType]: A dictionary of native \ + Python values, Vector, or BSON object converted from ``value_fields``. """ value_fields_pb = getattr(value_fields, "_pb", value_fields) - res = {key: decode_value(value, client) for key, value in value_fields_pb.items()} + res = { + key: decode_value(value, client, decode_bson=decode_bson) + for key, value in value_fields_pb.items() + } if res.get("__type__", None) == "__vector__": # Vector data type is represented as mapping. @@ -425,6 +450,16 @@ def decode_dict(value_fields, client) -> Union[dict, Vector]: values = cast(Sequence[float], res["value"]) return Vector(values) + should_decode = ( + decode_bson + if decode_bson is not None + else getattr(client, "_decode_bson", False) + ) + if should_decode: + decoded = _decode_bson_dict(res) + if decoded is not None: + return decoded + return res diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_client.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_client.py index 3167335e0385..4cd625387048 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_client.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_client.py @@ -105,6 +105,7 @@ def __init__( database=None, client_info=_CLIENT_INFO, client_options=None, + decode_bson: bool = False, ) -> None: super(AsyncClient, self).__init__( project=project, @@ -112,6 +113,7 @@ def __init__( database=database, client_info=client_info, client_options=client_options, + decode_bson=decode_bson, ) def _to_sync_copy(self): diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/base_client.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/base_client.py index 95166266bef2..5cbdf47c1b8d 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/base_client.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/base_client.py @@ -132,6 +132,7 @@ def __init__( database=None, client_info=_CLIENT_INFO, client_options=None, + decode_bson: bool = False, ) -> None: database = database or DEFAULT_DATABASE # NOTE: This API has no use for the _http argument, but sending it @@ -165,6 +166,7 @@ def __init__( self._client_options = client_options self._database = database + self._decode_bson: bool = decode_bson def _firestore_api_helper(self, transport, client_class, client_module) -> Any: """Lazy-loading getter GAPIC Firestore API. diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/base_document.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/base_document.py index 92d8daa21fd6..a2a423d36de8 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/base_document.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/base_document.py @@ -512,12 +512,17 @@ def get(self, field_path: str) -> Any: nested_data = field_path_module.get_nested_value(field_path, self._data) return copy.deepcopy(nested_data) - def to_dict(self) -> Union[Dict[str, Any], None]: + def to_dict( + self, decode_bson: Optional[bool] = None + ) -> Union[Dict[str, Any], None]: """Retrieve the data contained in this snapshot. A copy is returned since the data may contain mutable values, but the data stored in the snapshot must remain immutable. + Args: + decode_bson (Optional[bool]): Whether to decode BSON extended types. + Returns: Dict[str, Any] or None: The data in the snapshot. Returns None if reference @@ -525,7 +530,16 @@ def to_dict(self) -> Union[Dict[str, Any], None]: """ if not self._exists: return None - return copy.deepcopy(self._data) + data = copy.deepcopy(self._data) + client = self._reference._client if self._reference is not None else None + should_decode = ( + decode_bson + if decode_bson is not None + else getattr(client, "_decode_bson", False) + ) + if should_decode: + return _helpers._decode_bson_dict_recursive(data) + return data def _to_protobuf(self) -> Optional[Document]: return _helpers.document_snapshot_to_protobuf(self) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/bson.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/bson.py index 40e1d77f09fb..da663d14ad63 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/bson.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/bson.py @@ -27,7 +27,7 @@ import abc import decimal import re -from typing import Any, Dict, Union +from typing import Any, Callable, Dict, Union __all__ = [ "BSONObjectId", @@ -508,3 +508,21 @@ def __hash__(self) -> int: return hash(self.to_decimal) except decimal.InvalidOperation: return hash((type(self), self._value)) + + +_BSON_DECODERS: Dict[str, Callable[[Any], Any]] = { + "__oid__": BSONObjectId, + "__min__": lambda _: BSONMinKey(), + "__max__": lambda _: BSONMaxKey(), + "__int__": BSONInt32, + "__decimal128__": BSONDecimal128, + "__binary__": lambda v: (v[1:] if v[0] == 0 else BSONBinary(v[1:], subtype=v[0])) + if isinstance(v, (bytes, bytearray)) and len(v) >= 1 + else None, + "__request_timestamp__": lambda v: BSONTimestamp(v["seconds"], v["increment"]) + if isinstance(v, dict) and "seconds" in v and "increment" in v + else None, + "__regex__": lambda v: BSONRegex(v["pattern"], v.get("options", "")) + if isinstance(v, dict) and "pattern" in v + else None, +} diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/client.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/client.py index e29d07cb09ac..7b97f3ac13a5 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/client.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/client.py @@ -94,6 +94,7 @@ def __init__( database=None, client_info=_CLIENT_INFO, client_options=None, + decode_bson: bool = False, ) -> None: super(Client, self).__init__( project=project, @@ -101,6 +102,7 @@ def __init__( database=database, client_info=client_info, client_options=client_options, + decode_bson=decode_bson, ) @property diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py index e3fd74677a1e..7edb9808d292 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py @@ -44,6 +44,7 @@ from google.cloud.firestore_v1.async_transaction import AsyncTransaction from google.cloud.firestore_v1.base_client import BaseClient from google.cloud.firestore_v1.base_document import BaseDocumentReference + from google.cloud.firestore_v1.bson import _BSONType from google.cloud.firestore_v1.client import Client from google.cloud.firestore_v1.pipeline import Pipeline from google.cloud.firestore_v1.pipeline_expressions import Constant @@ -138,7 +139,7 @@ def __eq__(self, other: object) -> bool: return NotImplemented return (self._ref == other._ref) and (self._fields_pb == other._fields_pb) - def data(self) -> dict | "Vector" | None: + def data(self) -> dict | "Vector" | "_BSONType" | None: """ Retrieves all fields in the result. diff --git a/packages/google-cloud-firestore/tests/system/test_system.py b/packages/google-cloud-firestore/tests/system/test_system.py index 85a1b3546ba8..874adb1d6f8b 100644 --- a/packages/google-cloud-firestore/tests/system/test_system.py +++ b/packages/google-cloud-firestore/tests/system/test_system.py @@ -1285,9 +1285,9 @@ def test_unicode_doc(client, cleanup, database): @pytest.mark.parametrize("database", [FIRESTORE_ENTERPRISE_DB], indirect=True) -def test_bson_document_writes(client, cleanup, database): - """Test write operations for BSON types on Enterprise DB.""" - collection_id = "bson_type_writes_" + UNIQUE_RESOURCE_ID +def test_bson_document_read_and_write(client, cleanup, database): + """Test read and write operations for BSON types on Enterprise DB.""" + collection_id = "bson_type_read_write_" + UNIQUE_RESOURCE_ID doc_ref = client.collection(collection_id).document("bson_doc") cleanup(doc_ref.delete) @@ -1296,6 +1296,7 @@ def test_bson_document_writes(client, cleanup, database): "min_key": BSONMinKey(), "max_key": BSONMaxKey(), "int32_val": BSONInt32(42), + "binary_val_sub0": b"hello", "binary_val_sub128": BSONBinary(b"world", subtype=128), "timestamp_val": BSONTimestamp(1700000000, 1), "regex_val": BSONRegex("^hello.*$", options="i"), @@ -1306,26 +1307,7 @@ def test_bson_document_writes(client, cleanup, database): snapshot = doc_ref.get() assert snapshot.exists - assert snapshot.to_dict() == { - "user_id": {"__oid__": "507f191e810c19729de860ea"}, - "min_key": {"__min__": None}, - "max_key": {"__max__": None}, - "int32_val": {"__int__": 42}, - "binary_val_sub128": {"__binary__": b"\x80world"}, - "timestamp_val": { - "__request_timestamp__": { - "seconds": 1700000000, - "increment": 1, - } - }, - "regex_val": { - "__regex__": { - "pattern": "^hello.*$", - "options": "i", - } - }, - "decimal128_val": {"__decimal128__": "123.45"}, - } + assert snapshot.to_dict(decode_bson=True) == bson_payload @pytest.fixture(scope="module") diff --git a/packages/google-cloud-firestore/tests/system/test_system_async.py b/packages/google-cloud-firestore/tests/system/test_system_async.py index fac9aef81fe4..f7704ba010e1 100644 --- a/packages/google-cloud-firestore/tests/system/test_system_async.py +++ b/packages/google-cloud-firestore/tests/system/test_system_async.py @@ -1258,9 +1258,9 @@ async def test_list_collections_with_read_time(client, cleanup, database): @pytest.mark.asyncio @pytest.mark.parametrize("database", [FIRESTORE_ENTERPRISE_DB], indirect=True) -async def test_async_bson_document_writes(client, cleanup, database): - """Test async write operations for BSON types on Enterprise DB.""" - collection_id = "async_bson_type_writes_" + UNIQUE_RESOURCE_ID +async def test_async_bson_document_read_and_write(client, cleanup, database): + """Test async read and write operations for BSON types on Enterprise DB.""" + collection_id = "async_bson_type_read_write_" + UNIQUE_RESOURCE_ID doc_ref = client.collection(collection_id).document("bson_doc") cleanup(doc_ref.delete) @@ -1269,6 +1269,7 @@ async def test_async_bson_document_writes(client, cleanup, database): "min_key": BSONMinKey(), "max_key": BSONMaxKey(), "int32_val": BSONInt32(42), + "binary_val_sub0": b"hello", "binary_val_sub128": BSONBinary(b"world", subtype=128), "timestamp_val": BSONTimestamp(1700000000, 1), "regex_val": BSONRegex("^hello.*$", options="i"), @@ -1279,26 +1280,7 @@ async def test_async_bson_document_writes(client, cleanup, database): snapshot = await doc_ref.get() assert snapshot.exists - assert snapshot.to_dict() == { - "user_id": {"__oid__": "507f191e810c19729de860ea"}, - "min_key": {"__min__": None}, - "max_key": {"__max__": None}, - "int32_val": {"__int__": 42}, - "binary_val_sub128": {"__binary__": b"\x80world"}, - "timestamp_val": { - "__request_timestamp__": { - "seconds": 1700000000, - "increment": 1, - } - }, - "regex_val": { - "__regex__": { - "pattern": "^hello.*$", - "options": "i", - } - }, - "decimal128_val": {"__decimal128__": "123.45"}, - } + assert snapshot.to_dict(decode_bson=True) == bson_payload @pytest_asyncio.fixture(scope="module") diff --git a/packages/google-cloud-firestore/tests/unit/v1/test__helpers.py b/packages/google-cloud-firestore/tests/unit/v1/test__helpers.py index 4ce48424d3c4..b0b81eb98e51 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test__helpers.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test__helpers.py @@ -706,6 +706,43 @@ def test_decode_dict_w_many_types(): assert decode_dict(value_fields, mock.sentinel.client) == expected +def test_decode_dict_w_bson_types(): + from google.cloud.firestore_v1._helpers import decode_dict, encode_dict + from google.cloud.firestore_v1.bson import ( + BSONBinary, + BSONDecimal128, + BSONInt32, + BSONMaxKey, + BSONMinKey, + BSONObjectId, + BSONRegex, + BSONTimestamp, + ) + + original_dict = { + "oid": BSONObjectId("507f191e810c19729de860ea"), + "min_k": BSONMinKey(), + "max_k": BSONMaxKey(), + "int32_v": BSONInt32(42), + "bin_sub0": b"hello", + "bin_sub0_empty": b"", + "bin_sub128": BSONBinary(b"world", subtype=128), + "ts_v": BSONTimestamp(1700000000, 1), + "regex_v": BSONRegex("^hello.*$", options="i"), + "dec_v": BSONDecimal128("123.45"), + } + + pb_fields = encode_dict(original_dict) + # Default (decode_bson=False) returns raw dict + raw_decoded = decode_dict(pb_fields, mock.sentinel.client) + assert raw_decoded != original_dict + assert raw_decoded["oid"] == {"__oid__": "507f191e810c19729de860ea"} + + # decode_bson=True returns deserialized BSON objects + decoded = decode_dict(pb_fields, mock.sentinel.client, decode_bson=True) + assert decoded == original_dict + + def _dummy_ref_string(collection_id): from google.cloud.firestore_v1.base_client import DEFAULT_DATABASE