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
@@ -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.
 
@@ -402,7 +413,20 @@ def decode_value(
         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_dict(value_fields, client) -> Union[dict, Vector, _BSONType]:
     """Converts a protobuf map of Firestore ``Value``-s.
 
     Args:
@@ -412,9 +436,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,7 +449,7 @@ def decode_dict(value_fields, client) -> Union[dict, Vector]:
         values = cast(Sequence[float], res["value"])
         return Vector(values)
 
-    return res
+    return _decode_bson_dict(res) or res
 
 
 def get_doc_id(document_pb, expected_prefix) -> str:

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",
@@ -531,7 +531,23 @@ def __eq__(self, other: Any) -> bool:
         return NotImplemented
 
     def __hash__(self) -> int:
-        normalized_str = (
-            "NAN" if self._value.upper() == "NAN" else self._value
-        )
+        normalized_str = "NAN" if self._value.upper() == "NAN" else self._value
         return hash((type(self), normalized_str))
+
+
+_BSON_DECODERS: Dict[str, Callable[[Any], Any]] = {
+    "__oid__": BSONObjectId,
+    "__min__": lambda _: BSONMinKey(),
+    "__max__": lambda _: BSONMaxKey(),
+    "__int__": BSONInt32,
+    "__decimal128__": BSONDecimal128,
+    "__binary__": lambda v: BSONBinary(v[1:], subtype=v[0])
+    if isinstance(v, (bytes, bytearray)) and v
+    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)
 
@@ -1307,24 +1307,14 @@ 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"},
+        "user_id": BSONObjectId("507f191e810c19729de860ea"),
+        "min_key": BSONMinKey(),
+        "max_key": BSONMaxKey(),
+        "int32_val": BSONInt32(42),
+        "binary_val_sub128": BSONBinary(b"world", subtype=128),
+        "timestamp_val": BSONTimestamp(1700000000, 1),
+        "regex_val": BSONRegex("^hello.*$", options="i"),
+        "decimal128_val": BSONDecimal128("123.45"),
     }
 
 

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)
 
@@ -1280,24 +1280,14 @@ 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"},
+        "user_id": BSONObjectId("507f191e810c19729de860ea"),
+        "min_key": BSONMinKey(),
+        "max_key": BSONMaxKey(),
+        "int32_val": BSONInt32(42),
+        "binary_val_sub128": BSONBinary(b"world", subtype=128),
+        "timestamp_val": BSONTimestamp(1700000000, 1),
+        "regex_val": BSONRegex("^hello.*$", options="i"),
+        "decimal128_val": BSONDecimal128("123.45"),
     }
 
 

packages/google-cloud-firestore/tests/unit/v1/test__helpers.py

@@ -706,6 +706,36 @@ 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_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