feat(firestore): add BSON read deserialization support

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
 
 

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)

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",
@@ -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__}()"
 
@@ -503,3 +521,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,
+}

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.
 

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")

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")

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