feat(firestore): add opt-in BSON document read and decoding support (PR 2)

Author: ohmayr

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

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import datetime
+import collections.abc
 import json
 from typing import (
     TYPE_CHECKING,
@@ -50,8 +51,12 @@
 from google.cloud.firestore_v1.bson import (
     BSONBinary,
     BSONDecimal128,
+    BSONInt32,
+    BSONMaxKey,
+    BSONMinKey,
     BSONObjectID,
     BSONRegex,
+    BSONTimestamp,
 )
 from google.cloud.firestore_v1.vector import Vector
 
@@ -415,29 +420,113 @@ def decode_value(
         raise ValueError("Unknown ``value_type``", value_type)
 
 
-def decode_dict(value_fields, client) -> Union[dict, Vector]:
+def _parse_oid(val: Any) -> BSONObjectID:
+    if not isinstance(val, str):
+        raise ValueError(f"Invalid BSONObjectID map value, expected str: {val!r}")
+    return BSONObjectID(val)
+
+
+def _parse_decimal128(val: Any) -> BSONDecimal128:
+    if not isinstance(val, str):
+        raise ValueError(f"Invalid BSONDecimal128 map value, expected str: {val!r}")
+    return BSONDecimal128(val)
+
+
+def _parse_int32(val: Any) -> BSONInt32:
+    if type(val) is not int or isinstance(val, bool):
+        raise ValueError(f"Invalid BSONInt32 map value, expected int: {val!r}")
+    return BSONInt32(val)
+
+
+def _parse_minkey(val: Any) -> BSONMinKey:
+    if type(val) is not int or isinstance(val, bool):
+        raise ValueError(f"Invalid BSONMinKey map value, expected int: {val!r}")
+    return BSONMinKey()
+
+
+def _parse_maxkey(val: Any) -> BSONMaxKey:
+    if type(val) is not int or isinstance(val, bool):
+        raise ValueError(f"Invalid BSONMaxKey map value, expected int: {val!r}")
+    return BSONMaxKey()
+
+
+def _parse_timestamp(val: Any) -> BSONTimestamp:
+    if not isinstance(val, collections.abc.Mapping):
+        raise ValueError(f"Invalid BSONTimestamp map value, expected mapping: {val!r}")
+    sec = val.get("seconds")
+    inc = val.get("increment")
+    if type(sec) is not int or type(inc) is not int or isinstance(sec, bool) or isinstance(inc, bool) or len(val) != 2:
+        raise ValueError(f"Invalid BSONTimestamp fields: {val!r}")
+    return BSONTimestamp(sec, inc)
+
+
+def _parse_regex(val: Any) -> BSONRegex:
+    if not isinstance(val, collections.abc.Mapping):
+        raise ValueError(f"Invalid BSONRegex map value, expected mapping: {val!r}")
+    pat = val.get("pattern")
+    opt = val.get("options", "")
+    if not isinstance(pat, str) or not isinstance(opt, str) or len(val) not in (1, 2):
+        raise ValueError(f"Invalid BSONRegex fields: {val!r}")
+    return BSONRegex(pat, opt)
+
+
+def _parse_binary(val: Any) -> BSONBinary:
+    if not isinstance(val, collections.abc.Mapping):
+        raise ValueError(f"Invalid BSONBinary map value, expected mapping: {val!r}")
+    sub = val.get("sub_type")
+    bdata = val.get("bytes")
+    if type(sub) is not int or isinstance(sub, bool) or not isinstance(bdata, (bytes, bytearray, memoryview)) or len(val) != 2:
+        raise ValueError(f"Invalid BSONBinary fields: {val!r}")
+    return BSONBinary(bdata, subtype=sub)
+
+
+_BSON_MAP_PARSERS = {
+    "__oid__": _parse_oid,
+    "__decimal128__": _parse_decimal128,
+    "__int__": _parse_int32,
+    "__minkey__": _parse_minkey,
+    "__maxkey__": _parse_maxkey,
+    "__timestamp__": _parse_timestamp,
+    "__regex__": _parse_regex,
+    "__binary__": _parse_binary,
+}
+
+
+def _parse_bson_mapping(key: str, val: Any) -> Optional[Any]:
+    """Converts legacy BSON map value representations to native BSON instances."""
+    parser = _BSON_MAP_PARSERS.get(key)
+    if parser is not None:
+        return parser(val)
+    return None
+
+
+def decode_dict(value_fields, client, decode_bson: Optional[bool] = None) -> Union[dict, Vector]:
     """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]): Flag indicating whether to decode BSON map representations.
 
     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``.
+        Dict[str, Any]: A dictionary converted from ``value_fields``.
     """
+    effective_decode = decode_bson if decode_bson is not None else getattr(client, "decode_bson", False)
     value_fields_pb = getattr(value_fields, "_pb", value_fields)
     res = {key: decode_value(value, client) 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)
 
+    if effective_decode and len(res) == 1:
+        single_key = next(iter(res))
+        parsed = _parse_bson_mapping(single_key, res[single_key])
+        if parsed is not None:
+            return parsed
+
     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):
@@ -124,6 +126,7 @@ def _to_sync_copy(self):
                 database=self._database,
                 client_info=self._client_info,
                 client_options=self._client_options,
+                decode_bson=self.decode_bson,
             )
         return self._sync_copy
 

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 = decode_bson
 
     def _firestore_api_helper(self, transport, client_class, client_module) -> Any:
         """Lazy-loading getter GAPIC Firestore API.
@@ -610,14 +612,16 @@ def _parse_batch_get(
     result_type = get_doc_response._pb.WhichOneof("result")
     if result_type == "found":
         reference = _get_reference(get_doc_response.found.name, reference_map)
-        data = _helpers.decode_dict(get_doc_response.found.fields, client)
+        fields = get_doc_response.found.fields
+        data = _helpers.decode_dict(fields, client)
         snapshot = DocumentSnapshot(
             reference,
             data,
             exists=True,
             read_time=get_doc_response.read_time,
             create_time=get_doc_response.found.create_time,
             update_time=get_doc_response.found.update_time,
+            raw_fields=fields,
         )
     elif result_type == "missing":
         reference = _get_reference(get_doc_response.missing, reference_map)

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

@@ -388,21 +388,88 @@ class DocumentSnapshot(object):
     """
 
     def __init__(
-        self, reference, data, exists, read_time, create_time, update_time
+        self,
+        reference,
+        data,
+        exists,
+        read_time,
+        create_time,
+        update_time,
+        raw_fields=None,
+        decode_bson: Optional[bool] = None,
     ) -> None:
         self._reference = reference
-        # We want immutable data, so callers can't modify this value
-        # out from under us.
-        self._data = copy.deepcopy(data)
         self._exists = exists
         self.read_time = read_time
         self.create_time = create_time
         self.update_time = update_time
+        self._raw_fields = raw_fields
+
+        client = getattr(reference, "_client", None) if reference else None
+        self._decode_bson = (
+            decode_bson
+            if decode_bson is not None
+            else (getattr(client, "decode_bson", False) if client else False)
+        )
+        self._data_raw = None
+        self._data_bson = None
+
+        if raw_fields is not None:
+            if self._decode_bson:
+                self._data_bson = copy.deepcopy(data)
+            else:
+                self._data_raw = copy.deepcopy(data)
+        else:
+            self._data_raw = copy.deepcopy(data) if data is not None else None
+
+    def _get_data(self, decode_bson: Optional[bool] = None) -> Optional[Dict[str, Any]]:
+        effective_decode = (
+            decode_bson
+            if decode_bson is not None
+            else (
+                self._decode_bson
+                if hasattr(self, "_decode_bson") and self._decode_bson is not None
+                else (
+                    self._reference._client.decode_bson
+                    if (self._reference and hasattr(self._reference, "_client") and self._reference._client)
+                    else False
+                )
+            )
+        )
+
+        if effective_decode:
+            if self._data_bson is None:
+                if self._raw_fields is not None:
+                    client = self._reference._client if self._reference else None
+                    self._data_bson = _helpers.decode_dict(
+                        self._raw_fields, client, decode_bson=True
+                    )
+                elif self._data_raw is not None:
+                    self._data_bson = self._data_raw
+            return self._data_bson
+        else:
+            if self._data_raw is None:
+                if self._raw_fields is not None:
+                    client = self._reference._client if self._reference else None
+                    self._data_raw = _helpers.decode_dict(
+                        self._raw_fields, client, decode_bson=False
+                    )
+                elif self._data_bson is not None:
+                    self._data_raw = self._data_bson
+            return self._data_raw
+
+    @property
+    def _data(self) -> Optional[Dict[str, Any]]:
+        return self._get_data()
 
     def __eq__(self, other):
         if not isinstance(other, self.__class__):
             return NotImplemented
-        return self._reference == other._reference and self._data == other._data
+        return (
+            self._reference == other._reference
+            and self.read_time == other.read_time
+            and self._get_data(decode_bson=False) == other._get_data(decode_bson=False)
+        )
 
     def __hash__(self):
         return hash(self._reference) + hash(self.update_time)
@@ -448,84 +515,20 @@ def reference(self) -> BaseDocumentReference:
         """
         return self._reference
 
-    def get(self, field_path: str) -> Any:
-        """Get a value from the snapshot data.
-
-        If the data is nested, for example:
-
-        .. code-block:: python
-
-           >>> snapshot.to_dict()
-           {
-               'top1': {
-                   'middle2': {
-                       'bottom3': 20,
-                       'bottom4': 22,
-                   },
-                   'middle5': True,
-               },
-               'top6': b'\x00\x01 foo',
-           }
-
-        a **field path** can be used to access the nested data. For
-        example:
-
-        .. code-block:: python
-
-           >>> snapshot.get('top1')
-           {
-               'middle2': {
-                   'bottom3': 20,
-                   'bottom4': 22,
-               },
-               'middle5': True,
-           }
-           >>> snapshot.get('top1.middle2')
-           {
-               'bottom3': 20,
-               'bottom4': 22,
-           }
-           >>> snapshot.get('top1.middle2.bottom3')
-           20
-
-        See :meth:`~google.cloud.firestore_v1.client.Client.field_path` for
-        more information on **field paths**.
-
-        A copy is returned since the data may contain mutable values,
-        but the data stored in the snapshot must remain immutable.
-
-        Args:
-            field_path (str): A field path (``.``-delimited list of
-                field names).
-
-        Returns:
-            Any or None:
-                (A copy of) the value stored for the ``field_path`` or
-                None if snapshot document does not exist.
-
-        Raises:
-            KeyError: If the ``field_path`` does not match nested data
-                in the snapshot.
-        """
+    def get(self, field_path: str, decode_bson: Optional[bool] = None) -> Any:
+        """Get a value from the snapshot data."""
         if not self._exists:
             return None
-        nested_data = field_path_module.get_nested_value(field_path, self._data)
+        data = self._get_data(decode_bson=decode_bson)
+        nested_data = field_path_module.get_nested_value(field_path, data)
         return copy.deepcopy(nested_data)
 
-    def to_dict(self) -> 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.
-
-        Returns:
-            Dict[str, Any] or None:
-                The data in the snapshot.  Returns None if reference
-                does not exist.
-        """
+    def to_dict(self, decode_bson: Optional[bool] = None) -> Union[Dict[str, Any], None]:
+        """Retrieve the data contained in this snapshot."""
         if not self._exists:
             return None
-        return copy.deepcopy(self._data)
+        data = self._get_data(decode_bson=decode_bson)
+        return copy.deepcopy(data)
 
     def _to_protobuf(self) -> Optional[Document]:
         return _helpers.document_snapshot_to_protobuf(self)