[wip] firestore bson comprehensive

packages/google-cloud-firestore/google/cloud/firestore/__init__.py

@@ -35,6 +35,14 @@
     AsyncQuery,
     AsyncTransaction,
     AsyncWriteBatch,
+    BSONBinary,
+    BSONDecimal128,
+    BSONInt32,
+    BSONMaxKey,
+    BSONMinKey,
+    BSONObjectID,
+    BSONRegex,
+    BSONTimestamp,
     Client,
     CollectionGroup,
     CollectionReference,
@@ -92,6 +100,14 @@
     "async_transactional",
     "AsyncTransaction",
     "AsyncWriteBatch",
+    "BSONBinary",
+    "BSONDecimal128",
+    "BSONInt32",
+    "BSONMaxKey",
+    "BSONMinKey",
+    "BSONObjectID",
+    "BSONRegex",
+    "BSONTimestamp",
     "Client",
     "CountAggregation",
     "CollectionGroup",

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

@@ -31,6 +31,16 @@
     ReadAfterWriteError,
     WriteOption,
 )
+from google.cloud.firestore_v1.bson import (
+    BSONBinary,
+    BSONDecimal128,
+    BSONInt32,
+    BSONMaxKey,
+    BSONMinKey,
+    BSONObjectID,
+    BSONRegex,
+    BSONTimestamp,
+)
 from google.cloud.firestore_v1.async_batch import AsyncWriteBatch
 from google.cloud.firestore_v1.async_client import AsyncClient
 from google.cloud.firestore_v1.async_collection import AsyncCollectionReference
@@ -147,6 +157,14 @@
     "async_transactional",
     "AsyncTransaction",
     "AsyncWriteBatch",
+    "BSONBinary",
+    "BSONDecimal128",
+    "BSONInt32",
+    "BSONMaxKey",
+    "BSONMinKey",
+    "BSONObjectID",
+    "BSONRegex",
+    "BSONTimestamp",
     "Client",
     "CountAggregation",
     "CollectionGroup",

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

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import collections.abc
 import datetime
 import json
 from typing import (
@@ -44,6 +45,16 @@
 import google
 from google.cloud import exceptions  # type: ignore
 from google.cloud.firestore_v1 import transforms, types
+from google.cloud.firestore_v1.bson import (
+    BSONBinary,
+    BSONDecimal128,
+    BSONInt32,
+    BSONMaxKey,
+    BSONMinKey,
+    BSONObjectID,
+    BSONRegex,
+    BSONTimestamp,
+)
 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
@@ -182,6 +193,9 @@ def encode_value(value) -> types.document.Value:
     if value is None:
         return document.Value(null_value=struct_pb2.NULL_VALUE)
 
+    if hasattr(value, "to_map_value"):
+        return encode_value(value.to_map_value())
+
     # Must come before int since ``bool`` is an integer subtype.
     if isinstance(value, bool):
         return document.Value(boolean_value=value)
@@ -218,9 +232,6 @@ def encode_value(value) -> types.document.Value:
         value_pb = document.ArrayValue(values=value_list)
         return document.Value(array_value=value_pb)
 
-    if isinstance(value, Vector):
-        return encode_value(value.to_map_value())
-
     if isinstance(value, dict):
         value_dict = encode_dict(value)
         value_pb = document.MapValue(fields=value_dict)
@@ -344,22 +355,19 @@ 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
-]:
+    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 types.
 
     Returns:
-        Union[NoneType, bool, int, float, datetime.datetime, \
-            str, bytes, dict, ~google.cloud.Firestore.GeoPoint]: A native
-        Python value converted from the ``value``.
+        Any: A native Python value converted from the ``value``.
 
     Raises:
         NotImplementedError: If the ``value_type`` is ``reference_value``.
@@ -390,40 +398,92 @@ 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=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=client, decode_bson=decode_bson
+        )
     else:
         raise ValueError("Unknown ``value_type``", value_type)
 
 
-def decode_dict(value_fields, client) -> Union[dict, Vector]:
+def decode_dict(
+    value_fields, client=None, decode_bson: Optional[bool] = None
+) -> Any:
     """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 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``.
+        Any: A dictionary or BSON object/Vector converted from ``value_fields``.
     """
+    effective_decode = (
+        decode_bson
+        if decode_bson is not None
+        else (client.decode_bson if (client is not None and hasattr(client, "decode_bson")) else False)
+    )
+
     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=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)
 
+    if effective_decode and len(res) == 1:
+        key, val = next(iter(res.items()))
+        bson_obj = _parse_bson_mapping(key, val)
+        if bson_obj is not None:
+            return bson_obj
+
     return res
 
 
+def _parse_bson_mapping(key: str, val: Any) -> Optional[Any]:
+    """Converts legacy BSON map value representations to native BSON instances."""
+    try:
+        if key == "__oid__" and isinstance(val, str):
+            return BSONObjectID(val)
+        elif key == "__decimal128__" and isinstance(val, str):
+            return BSONDecimal128(val)
+        elif key == "__int__" and type(val) is int:
+            return BSONInt32(val)
+        elif key == "__minkey__" and type(val) is int:
+            return BSONMinKey()
+        elif key == "__maxkey__" and type(val) is int:
+            return BSONMaxKey()
+        elif key == "__timestamp__" and isinstance(val, collections.abc.Mapping):
+            sec = val.get("seconds")
+            inc = val.get("increment")
+            if type(sec) is int and type(inc) is int and len(val) == 2:
+                return BSONTimestamp(sec, inc)
+        elif key == "__regex__" and isinstance(val, collections.abc.Mapping):
+            pat = val.get("pattern")
+            opt = val.get("options", "")
+            if isinstance(pat, str) and isinstance(opt, str) and len(val) in (1, 2):
+                return BSONRegex(pat, opt)
+        elif key == "__binary__" and isinstance(val, collections.abc.Mapping):
+            sub = val.get("sub_type")
+            bdata = val.get("bytes")
+            if type(sub) is int and isinstance(bdata, (bytes, bytearray, memoryview)) and len(val) == 2:
+                return BSONBinary(bdata, subtype=sub)
+    except (ValueError, TypeError):
+        pass
+    return None
+
+
 def get_doc_id(document_pb, expected_prefix) -> str:
     """Parse a document ID from a document protobuf.
 
@@ -524,32 +584,29 @@ def __init__(self, document_data) -> None:
         prefix_path = FieldPath()
         iterator = self._get_document_iterator(prefix_path)
 
+        handlers = (
+            (lambda v: v is transforms.DELETE_FIELD, lambda fp, v: self.deleted_fields.append(fp)),
+            (lambda v: v is transforms.SERVER_TIMESTAMP, lambda fp, v: self.server_timestamps.append(fp)),
+            (lambda v: isinstance(v, transforms.ArrayRemove), lambda fp, v: self.array_removes.update({fp: v.values})),
+            (lambda v: isinstance(v, transforms.ArrayUnion), lambda fp, v: self.array_unions.update({fp: v.values})),
+            (lambda v: isinstance(v, transforms.Increment), lambda fp, v: self.increments.update({fp: v.value})),
+            (lambda v: isinstance(v, transforms.Maximum), lambda fp, v: self.maximums.update({fp: v.value})),
+            (lambda v: isinstance(v, transforms.Minimum), lambda fp, v: self.minimums.update({fp: v.value})),
+        )
+
         for field_path, value in iterator:
             if field_path == prefix_path and value is _EmptyDict:
                 self.empty_document = True
+                continue
 
-            elif value is transforms.DELETE_FIELD:
-                self.deleted_fields.append(field_path)
-
-            elif value is transforms.SERVER_TIMESTAMP:
-                self.server_timestamps.append(field_path)
-
-            elif isinstance(value, transforms.ArrayRemove):
-                self.array_removes[field_path] = value.values
-
-            elif isinstance(value, transforms.ArrayUnion):
-                self.array_unions[field_path] = value.values
-
-            elif isinstance(value, transforms.Increment):
-                self.increments[field_path] = value.value
-
-            elif isinstance(value, transforms.Maximum):
-                self.maximums[field_path] = value.value
-
-            elif isinstance(value, transforms.Minimum):
-                self.minimums[field_path] = value.value
+            handled = False
+            for match, action in handlers:
+                if match(value):
+                    action(field_path, value)
+                    handled = True
+                    break
 
-            else:
+            if not handled:
                 self.field_paths.append(field_path)
                 set_field_value(self.set_fields, field_path, value)
 

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)