diff --git a/firestore/google/cloud/firestore_v1beta1/_helpers.py b/firestore/google/cloud/firestore_v1beta1/_helpers.py index 8d07bbbf09e1e..9d0cc324f3bbf 100644 --- a/firestore/google/cloud/firestore_v1beta1/_helpers.py +++ b/firestore/google/cloud/firestore_v1beta1/_helpers.py @@ -28,10 +28,9 @@ import grpc import six +from google.cloud import exceptions from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud._helpers import _pb_timestamp_to_datetime -from google.cloud import exceptions - from google.cloud.firestore_v1beta1 import constants from google.cloud.firestore_v1beta1.gapic import enums from google.cloud.firestore_v1beta1.proto import common_pb2 @@ -813,7 +812,6 @@ def pbs_for_set(document_path, document_data, option): or two ``Write`` protobuf instances for ``set()``. """ transform_paths, actual_data = remove_server_timestamp(document_data) - update_pb = write_pb2.Write( update=document_pb2.Document( name=document_path, @@ -821,7 +819,7 @@ def pbs_for_set(document_path, document_data, option): ), ) if option is not None: - option.modify_write(update_pb) + option.modify_write(update_pb, actual_data=actual_data, path=document_path) write_pbs = [update_pb] if transform_paths: @@ -866,7 +864,7 @@ def pbs_for_update(client, document_path, field_updates, option): update_mask=common_pb2.DocumentMask(field_paths=sorted(field_paths)), ) # Due to the default, we don't have to check if ``None``. - option.modify_write(update_pb) + option.modify_write(update_pb, field_paths=field_paths) write_pbs = [update_pb] if transform_paths: diff --git a/firestore/google/cloud/firestore_v1beta1/client.py b/firestore/google/cloud/firestore_v1beta1/client.py index 30adab9ea90c9..ea820856096e9 100644 --- a/firestore/google/cloud/firestore_v1beta1/client.py +++ b/firestore/google/cloud/firestore_v1beta1/client.py @@ -37,6 +37,9 @@ from google.cloud.firestore_v1beta1.document import DocumentSnapshot from google.cloud.firestore_v1beta1.gapic import firestore_client from google.cloud.firestore_v1beta1.transaction import Transaction +from google.cloud.firestore_v1beta1.proto import common_pb2 +from google.cloud.firestore_v1beta1.proto import document_pb2 +from google.cloud.firestore_v1beta1.proto import write_pb2 DEFAULT_DATABASE = '(default)' @@ -251,15 +254,18 @@ def write_option(**kwargs): :meth:`~.DocumentReference.update` and :meth:`~.DocumentReference.delete`. - One of the following two keyword arguments must be provided: + One of the following keyword arguments must be provided: * ``last_update_time`` (:class:`google.protobuf.timestamp_pb2.\ - Timestamp`): A timestamp. When set, the target document must exist - and have been last updated at that time. Protobuf ``update_time`` - timestamps are typically returned from methods that perform write - operations as part of a "write result" protobuf or directly. + Timestamp`): A timestamp. When set, the target document must exist + and have been last updated at that time. Protobuf ``update_time`` + timestamps are typically returned from methods that perform write + operations as part of a "write result" protobuf or directly. * ``exists`` (:class:`bool`): Indicates if the document being modified - should already exist. + should already exist. + * ``merge`` (Any): + Indicates if the document should be merged. + **Note**: argument is ignored Providing no argument would make the option have no effect (so it is not allowed). Providing multiple would be an apparent @@ -283,6 +289,8 @@ def write_option(**kwargs): return LastUpdateOption(value) elif name == 'exists': return ExistsOption(value) + elif name == 'merge': + return MergeOption() else: extra = '{!r} was provided'.format(name) raise TypeError(_BAD_OPTION_ERR, extra) @@ -418,6 +426,39 @@ def modify_write(self, write_pb, **unused_kwargs): write_pb.current_document.CopyFrom(current_doc) +class MergeOption(WriteOption): + """Option used to merge on a write operation. + + This will typically be created by + :meth:`~.firestore_v1beta1.client.Client.write_option`. + """ + def modify_write(self, write_pb, actual_data=None, path=None, **unused_kwargs): + """Modify a ``Write`` protobuf based on the state of this write option. + + Args: + write_pb (google.cloud.firestore_v1beta1.types.Write): A + ``Write`` protobuf instance to be modified with a precondition + determined by the state of this option. + actual_data (dict): + The actual field names and values to use for replacing a + document. + path (str): A fully-qualified document_path + unused_kwargs (Dict[str, Any]): Keyword arguments accepted by + other subclasses that are unused here. + """ + actual_data, field_paths = _helpers.FieldPathHelper.to_field_paths(actual_data) + doc = document_pb2.Document( + name=path, + fields=_helpers.encode_dict(actual_data) + ) + write = write_pb2.Write( + update=doc, + ) + write_pb.CopyFrom(write) + mask = common_pb2.DocumentMask(field_paths=sorted(field_paths)) + write_pb.update_mask.CopyFrom(mask) + + class ExistsOption(WriteOption): """Option used to assert existence on a write operation. diff --git a/firestore/tests/system.py b/firestore/tests/system.py index d3f0d22e02f9a..2c32400d88798 100644 --- a/firestore/tests/system.py +++ b/firestore/tests/system.py @@ -194,6 +194,47 @@ def test_document_set(client, cleanup): assert exc_to_code(exc_info.value.cause) == StatusCode.FAILED_PRECONDITION +def test_document_set_merge(client, cleanup): + document_id = 'for-set' + unique_resource_id('-') + document = client.document('i-did-it', document_id) + # Add to clean-up before API request (in case ``set()`` fails). + cleanup(document) + + # 0. Make sure the document doesn't exist yet using an option. + option0 = client.write_option(exists=True) + with pytest.raises(NotFound) as exc_info: + document.set({'no': 'way'}, option=option0) + + assert exc_info.value.message.startswith(MISSING_DOCUMENT) + assert document_id in exc_info.value.message + + # 1. Use ``set()`` to create the document (using an option). + data1 = {'name': 'Sam', + 'address': {'city': 'SF', + 'state': 'CA'} + } + option1 = client.write_option(exists=False) + write_result1 = document.set(data1, option=option1) + snapshot1 = document.get() + assert snapshot1.to_dict() == data1 + # Make sure the update is what created the document. + assert snapshot1.create_time == snapshot1.update_time + assert snapshot1.update_time == write_result1.update_time + + # 2. Call ``set()`` again to overwrite (no option). + data2 = {'address.city': 'LA'} + option2 = client.write_option(merge=True) + write_result2 = document.set(data2, option=option2) + snapshot2 = document.get() + assert snapshot2.to_dict() == {'name': 'Sam', + 'address': {'city': 'LA', + 'state': 'CA'} + } + # Make sure the create time hasn't changed. + assert snapshot2.create_time == snapshot1.create_time + assert snapshot2.update_time == write_result2.update_time + + def test_update_document(client, cleanup): document_id = 'for-update' + unique_resource_id('-') document = client.document('made', document_id) @@ -207,7 +248,6 @@ def test_update_document(client, cleanup): assert document_id in exc_info.value.message # 1. Try to update before the document exists (now with an option). - option1 = client.write_option(exists=True) with pytest.raises(NotFound) as exc_info: document.update({'still': 'not-there'}, option=option1) diff --git a/firestore/tests/unit/test_batch.py b/firestore/tests/unit/test_batch.py index 8f66b5211d171..59775a5f87c5f 100644 --- a/firestore/tests/unit/test_batch.py +++ b/firestore/tests/unit/test_batch.py @@ -93,6 +93,33 @@ def test_set(self): ) self.assertEqual(batch._write_pbs, [new_write_pb]) + def test_set_merge(self): + from google.cloud.firestore_v1beta1.proto import document_pb2 + from google.cloud.firestore_v1beta1.proto import write_pb2 + from google.cloud.firestore_v1beta1.client import MergeOption + + client = _make_client() + batch = self._make_one(client) + self.assertEqual(batch._write_pbs, []) + + reference = client.document('another', 'one') + field = 'zapzap' + value = u'meadows and flowers' + document_data = {field: value} + option = MergeOption() + ret_val = batch.set(reference, document_data, option) + self.assertIsNone(ret_val) + new_write_pb = write_pb2.Write( + update=document_pb2.Document( + name=reference._document_path, + fields={ + field: _value_pb(string_value=value), + }, + ), + update_mask={'field_paths':[field]} + ) + self.assertEqual(batch._write_pbs, [new_write_pb]) + def test_update(self): from google.cloud.firestore_v1beta1.proto import common_pb2 from google.cloud.firestore_v1beta1.proto import document_pb2 diff --git a/firestore/tests/unit/test_document.py b/firestore/tests/unit/test_document.py index 7e818694f08eb..ac161ff36a2a6 100644 --- a/firestore/tests/unit/test_document.py +++ b/firestore/tests/unit/test_document.py @@ -249,6 +249,7 @@ def _write_pb_for_set(document_path, document_data): ) def _set_helper(self, **option_kwargs): + from google.cloud.firestore_v1beta1._helpers import FieldPathHelper # Create a minimal fake GAPIC with a dummy response. firestore_api = mock.Mock(spec=['commit']) commit_response = mock.Mock( @@ -277,8 +278,10 @@ def _set_helper(self, **option_kwargs): self.assertIs(write_result, mock.sentinel.write_result) write_pb = self._write_pb_for_set( document._document_path, document_data) + update_values, field_paths = FieldPathHelper.to_field_paths(document_data) + if option is not None: - option.modify_write(write_pb) + option.modify_write(write_pb, field_paths=field_paths) firestore_api.commit.assert_called_once_with( client._database_string, [write_pb], transaction=None, options=client._call_options) @@ -286,9 +289,12 @@ def _set_helper(self, **option_kwargs): def test_set(self): self._set_helper() - def test_set_with_option(self): + def test_set_exists(self): self._set_helper(exists=True) + def test_set_merge(self): + self._set_helper(merge='abc') + @staticmethod def _write_pb_for_update(document_path, update_values, field_paths): from google.cloud.firestore_v1beta1.proto import common_pb2