MergeOption

firestore/google/cloud/firestore_v1beta1/_helpers.py

@@ -24,10 +24,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
@@ -764,6 +763,22 @@ def get_doc_id(document_pb, expected_prefix):
     return document_id
 
 
+def get_field_paths(update_data):
+    field_paths = []
+    for field_name, value in six.iteritems(update_data):
+        match = re.match(FieldPath.simple_field_name, field_name)
+        if not (match and match.group(0) == field_name):
+            field_name = field_name.replace('\\', '\\\\').replace('`', '\\`')
+            field_name = '`' + field_name + '`'
+        if isinstance(value, dict):
+            sub_field_paths = get_field_paths(value)
+            field_paths.extend(
+                [field_name + "." + sub_path for sub_path in sub_field_paths])
+        else:
+            field_paths.append(field_name)
+    return field_paths
+
+
 def remove_server_timestamp(document_data):
     """Remove all server timestamp sentinel values from data.
 
@@ -876,15 +891,16 @@ 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)
+        field_paths = get_field_paths(actual_data)
+        option.modify_write(
+            update_pb, field_paths=field_paths, path=document_path)
 
     write_pbs = [update_pb]
     if transform_paths:
@@ -949,7 +965,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

@@ -34,6 +34,7 @@
 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
 
 
 DEFAULT_DATABASE = '(default)'
@@ -252,12 +253,15 @@ def write_option(**kwargs):
         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
@@ -281,6 +285,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)
@@ -416,46 +422,28 @@ def modify_write(self, write_pb, **unused_kwargs):
         write_pb.current_document.CopyFrom(current_doc)
 
 
-class CreateIfMissingOption(WriteOption):
-    """Option used to assert "create if missing" on a write operation.
+class MergeOption(WriteOption):
+    """Option used to merge on a write operation.
 
     This will typically be created by
     :meth:`~.firestore_v1beta1.client.Client.write_option`.
-
-    Args:
-        create_if_missing (bool): Indicates if the document should be created
-            if it doesn't already exist.
     """
-    def __init__(self, create_if_missing):
-        self._create_if_missing = create_if_missing
-
-    def modify_write(self, write_pb, no_create_msg=None):
+    def modify_write(self, write_pb, field_paths=None, path=None, **unused_kwargs):
         """Modify a ``Write`` protobuf based on the state of this write option.
 
-        If:
-
-        * ``create_if_missing=False``, adds a precondition that requires
-          existence
-        * ``create_if_missing=True``, does not add any precondition
-        * ``no_create_msg`` is passed, raises an exception. For example, in a
-          :meth:`~.DocumentReference.delete`, no "create" can occur, so it
-          wouldn't make sense to "create if missing".
-
         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.
-            no_create_msg (Optional[str]): A message to use to indicate that
-                a create operation is not allowed.
-
-        Raises:
-            ValueError: If ``no_create_msg`` is passed.
+            field_paths (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.
         """
-        if no_create_msg is not None:
-            raise ValueError(no_create_msg)
-        elif not self._create_if_missing:
-            current_doc = types.Precondition(exists=True)
-            write_pb.current_document.CopyFrom(current_doc)
+        mask = common_pb2.DocumentMask(field_paths=sorted(field_paths))
+        write_pb.update_mask.CopyFrom(mask)
 
 
 class ExistsOption(WriteOption):

firestore/tests/system.py

@@ -132,7 +132,7 @@ def assert_timestamp_less(timestamp_pb1, timestamp_pb2):
 def test_no_document(client, cleanup):
     document_id = 'no_document' + unique_resource_id('-')
     document = client.document('abcde', document_id)
-    option0 = client.write_option(create_if_missing=False)
+    option0 = client.write_option(exists=True)
     with pytest.raises(NotFound):
         document.set({'no': 'way'}, option=option0)
     snapshot = document.get()
@@ -214,9 +214,9 @@ def test_document_integer_field(client, cleanup):
     option1 = client.write_option(exists=False)
     document.set(data1, option=option1)
 
-    data2 = {'1a.ab': '4d', '6f.7g': '9h'}
-    option2 = client.write_option(create_if_missing=True)
-    document.update(data2, option=option2)
+    data2 = {'1a': {'ab': '4d'}, '6f': {'7g': '9h'}}
+    option2 = client.write_option(merge=True)
+    document.set(data2, option=option2)
     snapshot = document.get()
     expected = {
         '1a': {
@@ -229,6 +229,45 @@ def test_document_integer_field(client, cleanup):
     assert snapshot.to_dict() == expected
 
 
+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)

firestore/tests/unit/test__helpers.py

@@ -1506,7 +1506,7 @@ def _helper(self, option=None, do_transform=False, **write_kwargs):
             'yum': _value_pb(bytes_value=value),
         })
         if isinstance(option, ExistsOption):
-            write_kwargs.update({'current_document' : {'exists': option._exists}})
+            write_kwargs.update({'current_document': {'exists': option._exists}})
 
         expected_update_pb = write_pb2.Write(
             update=document_pb2.Document(
@@ -1541,7 +1541,6 @@ def test_without_option(self):
         self._helper(current_document=precondition)
 
     def test_with_option(self):
-        from google.cloud.firestore_v1beta1.proto import common_pb2
         from google.cloud.firestore_v1beta1.client import ExistsOption
 
         option = ExistsOption(False)

firestore/tests/unit/test_batch.py

@@ -90,6 +90,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_client.py

@@ -480,6 +480,31 @@ def test_modify_write(self):
             self.assertEqual(write_pb.current_document, expected_doc)
 
 
+class TestMergeOption(unittest.TestCase):
+
+    @staticmethod
+    def _get_target_class():
+        from google.cloud.firestore_v1beta1.client import MergeOption
+        return MergeOption
+
+    def _make_one(self, *args, **kwargs):
+        klass = self._get_target_class()
+        return klass()
+
+    def test_modify_write(self):
+        from google.cloud.firestore_v1beta1.proto import common_pb2
+        from google.cloud.firestore_v1beta1.proto import write_pb2
+
+        for merge in (True, False):
+            option = self._make_one(merge)
+            write_pb = write_pb2.Write()
+            field_paths = ['a', 'b', 'c']
+            ret_val = option.modify_write(write_pb, field_paths=field_paths)
+            mask = common_pb2.DocumentMask(field_paths=sorted(field_paths))
+            self.assertIsNone(ret_val)
+            self.assertEqual(write_pb.update_mask, mask)
+
+
 class Test__reference_info(unittest.TestCase):
 
     @staticmethod

firestore/tests/unit/test_document.py

@@ -280,17 +280,25 @@ def _set_helper(self, **option_kwargs):
         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, field_paths=document_data.keys())
+            else:
+                option.modify_write(write_pb)
+
         firestore_api.commit.assert_called_once_with(
             client._database_string, [write_pb], transaction=None,
             metadata=client._rpc_metadata)
 
     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