diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py index fdb12a36e741..2d45140bd96f 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py @@ -18,6 +18,7 @@ import datetime import json +import re from typing import ( TYPE_CHECKING, Any, @@ -44,7 +45,13 @@ 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 ( + BSONBinary, + BSONDecimal128, + BSONObjectID, + BSONRegex, + 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 @@ -164,13 +171,28 @@ def verify_path(path, is_collection) -> None: raise ValueError(msg) -def encode_value(value) -> types.document.Value: - """Converts a native Python value into a Firestore protobuf ``Value``. +_REGEX_FLAG_MAP = ( + (re.IGNORECASE, "i"), + (re.MULTILINE, "m"), + (re.DOTALL, "s"), + (re.VERBOSE, "x"), + (re.LOCALE, "l"), +) + + +def _extract_regex_options(flags: Union[int, str]) -> str: + if isinstance(flags, str): + return flags + if isinstance(flags, int): + return "".join(char for bit, char in _REGEX_FLAG_MAP if flags & bit) + return "" + + +def encode_value(value: Any) -> document.Value: + """Convert a Python value into a Value protobuf. Args: - value (Union[NoneType, bool, int, float, datetime.datetime, \ - str, bytes, dict, ~google.cloud.Firestore.GeoPoint, \ - ~google.cloud.firestore_v1.vector.Vector]): A native + value (Any): The Python value to convert to a protobuf field. Returns: @@ -186,6 +208,23 @@ def encode_value(value) -> types.document.Value: if isinstance(value, BSONType): return encode_value(value.to_map_value()) + # Duck-typing input bridge for external PyMongo / bson package objects (zero dependency) + binary_attr = getattr(value, "binary", None) + if binary_attr is not None and not isinstance( + value, (bytes, bytearray, BSONBinary) + ): + return encode_value(BSONObjectID(binary_attr)) + + to_decimal_fn = getattr(value, "to_decimal", None) + if callable(to_decimal_fn) and not isinstance(value, BSONDecimal128): + return encode_value(BSONDecimal128(to_decimal_fn())) + + pattern_attr = getattr(value, "pattern", None) + if pattern_attr is not None and not isinstance(value, (str, BSONRegex)): + flags_attr = getattr(value, "flags", "") + options_str = _extract_regex_options(flags_attr) + return encode_value(BSONRegex(pattern_attr, options_str)) + # Must come before int since ``bool`` is an integer subtype. if isinstance(value, bool): return document.Value(boolean_value=value) diff --git a/packages/google-cloud-firestore/tests/system/test_bson.py b/packages/google-cloud-firestore/tests/system/test_bson.py index b74a87bb3e02..747786580b5a 100644 --- a/packages/google-cloud-firestore/tests/system/test_bson.py +++ b/packages/google-cloud-firestore/tests/system/test_bson.py @@ -66,3 +66,40 @@ def test_bson_document_writes(client, cleanup, database): snapshot = doc_ref.get() assert snapshot.exists + + +@pytest.mark.parametrize("database", [FIRESTORE_ENTERPRISE_DB], indirect=True) +def test_pymongo_document_writes(client, cleanup, database): + """Test write operations using native duck-typed PyMongo objects on Enterprise DB.""" + import decimal + + class DummyPyMongoObjectId: + def __init__(self, raw: bytes): + self.binary = raw + + class DummyPyMongoDecimal128: + def __init__(self, d: decimal.Decimal): + self._d = d + + def to_decimal(self): + return self._d + + class DummyPyMongoRegex: + def __init__(self, pat: str, flags: str): + self.pattern = pat + self.flags = flags + + collection_id = "pymongo_docs_write_" + UNIQUE_RESOURCE_ID + doc_ref = client.collection(collection_id).document("pymongo_doc") + cleanup(doc_ref.delete) + + payload = { + "_id": DummyPyMongoObjectId(bytes.fromhex("507f191e810c19729de860ea")), + "price": DummyPyMongoDecimal128(decimal.Decimal("99.99")), + "pattern": DummyPyMongoRegex("^test.*", "i"), + } + + doc_ref.set(payload) + + snapshot = doc_ref.get() + assert snapshot.exists diff --git a/packages/google-cloud-firestore/tests/unit/v1/test__helpers.py b/packages/google-cloud-firestore/tests/unit/v1/test__helpers.py index 4ce48424d3c4..8322089664d6 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test__helpers.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test__helpers.py @@ -354,6 +354,45 @@ def test_encode_dict_w_many_types(): assert encoded_dict == expected_dict +def test_encode_value_duck_typed_pymongo(): + import decimal + + from google.cloud.firestore_v1._helpers import encode_value + + class DummyPyMongoObjectId: + def __init__(self, raw: bytes): + self.binary = raw + + class DummyPyMongoDecimal128: + def __init__(self, d: decimal.Decimal): + self._d = d + + def to_decimal(self): + return self._d + + class DummyPyMongoRegex: + def __init__(self, pat: str, flags: str): + self.pattern = pat + self.flags = flags + + dummy_oid = DummyPyMongoObjectId(bytes.fromhex("507f1f77bcf86cd799439011")) + res_oid = encode_value(dummy_oid) + assert ( + res_oid.map_value.fields["__oid__"].string_value == "507f1f77bcf86cd799439011" + ) + + dummy_dec = DummyPyMongoDecimal128(decimal.Decimal("99.99")) + res_dec = encode_value(dummy_dec) + assert res_dec.map_value.fields["__decimal128__"].string_value == "99.99" + + dummy_reg = DummyPyMongoRegex("^test$", "i") + res_reg = encode_value(dummy_reg) + assert ( + res_reg.map_value.fields["__regex__"].map_value.fields["pattern"].string_value + == "^test$" + ) + + def test_reference_value_to_document_w_bad_format(): from google.cloud.firestore_v1._helpers import ( BAD_REFERENCE_ERROR, @@ -2572,3 +2611,17 @@ def _make_field_path(*fields): from google.cloud.firestore_v1 import field_path return field_path.FieldPath(*fields) + + +def test_encode_value_w_compiled_regex_flags(): + import re + + from google.cloud.firestore_v1._helpers import encode_value + + compiled_re = re.compile("abc", re.I | re.M) + encoded = encode_value(compiled_re) + # Checks that compiled regex integer flags translate to 'im' options + fields = encoded.map_value.fields + assert fields["__regex__"].map_value.fields["pattern"].string_value == "abc" + assert "i" in fields["__regex__"].map_value.fields["options"].string_value + assert "m" in fields["__regex__"].map_value.fields["options"].string_value