feat(firestore): add BSON read deserialization support

Author: ohmayr

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,27 +347,18 @@ 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:
         value (google.cloud.firestore_v1.types.Value): A
             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,37 +385,81 @@ 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:
         value_fields (google.protobuf.pyext._message.MessageMapContainer): A
             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.
         # {"__type__":"__vector__", "value": [1.0, 2.0, 3.0]}.
         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
 
 

packages/google-cloud-firestore/google/cloud/firestore_v1/async_client.py

@@ -105,13 +105,15 @@ def __init__(
         database=None,
         client_info=_CLIENT_INFO,
         client_options=None,
+        decode_bson: bool = False,
     ) -> None:
         super(AsyncClient, self).__init__(
             project=project,
             credentials=credentials,
             database=database,
             client_info=client_info,
             client_options=client_options,
+            decode_bson=decode_bson,
         )
 
     def _to_sync_copy(self):

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.

packages/google-cloud-firestore/google/cloud/firestore_v1/base_document.py

@@ -512,20 +512,34 @@ 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
                 does not exist.
         """
         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)

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,
+}

packages/google-cloud-firestore/google/cloud/firestore_v1/client.py

@@ -94,13 +94,15 @@ def __init__(
         database=None,
         client_info=_CLIENT_INFO,
         client_options=None,
+        decode_bson: bool = False,
     ) -> None:
         super(Client, self).__init__(
             project=project,
             credentials=credentials,
             database=database,
             client_info=client_info,
             client_options=client_options,
+            decode_bson=decode_bson,
         )
 
     @property

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(decode_bson=True) == 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(decode_bson=True) == bson_payload
 
 
 @pytest_asyncio.fixture(scope="module")