Firestore: MergeOption

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,15 +812,14 @@ 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,
             fields=encode_dict(actual_data),
         ),
     )
     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:

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.
 

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)

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

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,18 +278,27 @@ 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)
+
         if option is not None:
-            option.modify_write(write_pb)
+            from google.cloud.firestore_v1beta1.client import MergeOption
+            if isinstance(option, MergeOption):
+
+                option.modify_write(write_pb, actual_data=document_data, path=document._document_path)
+            else:
+                option.modify_write(write_pb)
         firestore_api.commit.assert_called_once_with(
             client._database_string, [write_pb], transaction=None,
             options=client._call_options)
 
     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