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..9403dcfc999a 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py @@ -350,7 +350,18 @@ def reference_value_to_document(reference_value, client) -> Any: def decode_value( value, client ) -> Union[ - None, bool, int, float, list, datetime.datetime, str, bytes, dict, GeoPoint, Vector + None, + bool, + int, + float, + list, + datetime.datetime, + str, + bytes, + dict, + GeoPoint, + Vector, + _BSONType, ]: """Converts a Firestore protobuf ``Value`` to a native Python value. @@ -362,7 +373,9 @@ def decode_value( Returns: Union[NoneType, bool, int, float, datetime.datetime, \ - str, bytes, dict, ~google.cloud.Firestore.GeoPoint]: A native + str, bytes, dict, ~google.cloud.Firestore.GeoPoint, \ + ~google.cloud.firestore_v1.vector.Vector, \ + ~google.cloud.firestore_v1.bson._BSONType]: A native \ Python value converted from the ``value``. Raises: @@ -402,7 +415,22 @@ def decode_value( raise ValueError("Unknown ``value_type``", value_type) -def decode_dict(value_fields, client) -> Union[dict, Vector]: +def _decode_bson_dict_recursive(data: Any) -> Any: + """Recursively decodes BSON wire map dictionaries.""" + if isinstance(data, dict): + decoded = _BSONType._from_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, +) -> Union[dict, Vector, _BSONType]: """Converts a protobuf map of Firestore ``Value``-s. Args: @@ -412,9 +440,9 @@ def decode_dict(value_fields, client) -> Union[dict, Vector]: A client that has a document factory. 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()} @@ -425,6 +453,10 @@ def decode_dict(value_fields, client) -> Union[dict, Vector]: values = cast(Sequence[float], res["value"]) return Vector(values) + decoded = _BSONType._from_dict(res) + if decoded is not None: + return decoded + return res 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..95f515659a56 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 @@ -525,7 +525,8 @@ 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) + return _helpers._decode_bson_dict_recursive(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 fab546a09b72..453154f05398 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, FrozenSet, Union +from typing import Any, Callable, Dict, FrozenSet, Union __all__ = [ "BSONObjectId", @@ -65,6 +65,24 @@ def __eq__(self, other: Any) -> bool: def __hash__(self) -> int: """Hash representation contract for set and dictionary keys.""" + @classmethod + def _from_dict(cls, data: Any) -> Any: + """Deserializes a BSON wire map dictionary into a BSON instance or bytes. + + Args: + data (Any): Potential BSON wire map dictionary. + + Returns: + Any: Deserialized BSON container instance/bytes, or None if not a BSON wire map. + """ + if not isinstance(data, dict) or len(data) != 1: + return None + key, val = next(iter(data.items())) + decoder = _BSON_DECODERS.get(key) + if decoder is None: + return None + return decoder(val) + def __repr__(self) -> str: return f"{self.__class__.__name__}()" @@ -513,3 +531,21 @@ def __hash__(self) -> int: return hash(d) 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/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..f15fb4ec6b14 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() == 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..a79db558b574 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() == 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..88fe361eee31 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,37 @@ 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) + decoded = decode_dict(pb_fields, mock.sentinel.client) + assert decoded == original_dict + + def _dummy_ref_string(collection_id): from google.cloud.firestore_v1.base_client import DEFAULT_DATABASE