Firestore: Implement firestore merge as a boolean

Author: chemelnucfin

firestore/google/cloud/firestore.py

@@ -25,7 +25,6 @@
 from google.cloud.firestore_v1beta1 import ExistsOption
 from google.cloud.firestore_v1beta1 import GeoPoint
 from google.cloud.firestore_v1beta1 import LastUpdateOption
-from google.cloud.firestore_v1beta1 import MergeOption
 from google.cloud.firestore_v1beta1 import Query
 from google.cloud.firestore_v1beta1 import ReadAfterWriteError
 from google.cloud.firestore_v1beta1 import SERVER_TIMESTAMP
@@ -47,7 +46,6 @@
     'ExistsOption',
     'GeoPoint',
     'LastUpdateOption',
-    'MergeOption',
     'Query',
     'ReadAfterWriteError',
     'SERVER_TIMESTAMP',

firestore/google/cloud/firestore_v1beta1/__init__.py

@@ -24,7 +24,6 @@
 from google.cloud.firestore_v1beta1.client import Client
 from google.cloud.firestore_v1beta1.client import ExistsOption
 from google.cloud.firestore_v1beta1.client import LastUpdateOption
-from google.cloud.firestore_v1beta1.client import MergeOption
 from google.cloud.firestore_v1beta1.client import WriteOption
 from google.cloud.firestore_v1beta1.collection import CollectionReference
 from google.cloud.firestore_v1beta1.constants import DELETE_FIELD
@@ -48,7 +47,6 @@
     'ExistsOption',
     'GeoPoint',
     'LastUpdateOption',
-    'MergeOption',
     'Query',
     'ReadAfterWriteError',
     'SERVER_TIMESTAMP',

firestore/google/cloud/firestore_v1beta1/_helpers.py

@@ -517,18 +517,18 @@ def encode_dict(values_dict):
     }
 
 
-def extract_field_paths(actual_data):
-    """Extract field paths from actual data
+def extract_field_paths(document_data):
+    """Extract field paths from document data
 
     Args:
-       actual_data (dict): The dictionary of the actual set data.
+       document_data (dict): The dictionary of the actual set data.
 
     Returns:
        List[~.firestore_v1beta1._helpers.FieldPath]:
            A list of `FieldPath` instances from the actual data.
     """
     field_paths = []
-    for field_name, value in six.iteritems(actual_data):
+    for field_name, value in six.iteritems(document_data):
         if isinstance(value, dict):
             sub_field_paths = extract_field_paths(value)
             for sub_path in sub_field_paths:
@@ -897,7 +897,7 @@ def get_transform_pb(document_path, transform_paths):
     )
 
 
-def pbs_for_set(document_path, document_data, option):
+def pbs_for_set(document_path, document_data, merge=False, exists=None):
     """Make ``Write`` protobufs for ``set()`` methods.
 
     Args:
@@ -920,9 +920,15 @@ def pbs_for_set(document_path, document_data, option):
             fields=encode_dict(actual_data),
         ),
     )
-    if option is not None:
+    if exists is not None:
+        update_pb.current_document.CopyFrom(
+            common_pb2.Precondition(exists=exists))
+
+    if merge:
         field_paths = extract_field_paths(document_data)
-        option.modify_write(update_pb, field_paths=field_paths)
+        field_paths = canonicalize_field_paths(field_paths)
+        mask = common_pb2.DocumentMask(field_paths=sorted(field_paths))
+        update_pb.update_mask.CopyFrom(mask)
 
     write_pbs = [update_pb]
     if transform_paths:

firestore/google/cloud/firestore_v1beta1/batch.py

@@ -57,10 +57,11 @@ def create(self, reference, document_data):
             document_data (dict): Property names and values to use for
                 creating a document.
         """
-        option = self._client.write_option(exists=False)
-        self.set(reference, document_data, option=option)
+        write_pbs = _helpers.pbs_for_set(
+            reference._document_path, document_data, merge=False, exists=False)
+        self._add_write_pbs(write_pbs)
 
-    def set(self, reference, document_data, option=None):
+    def set(self, reference, document_data, merge=False):
         """Add a "change" to replace a document.
 
         See
@@ -69,16 +70,16 @@ def set(self, reference, document_data, option=None):
         applied.
 
         Args:
-            reference (~.firestore_v1beta1.document.DocumentReference): A
-               document reference that will have values set in this batch.
-            document_data (dict): Property names and values to use for
-                replacing a document.
-            option (Optional[~.firestore_v1beta1.client.WriteOption]): A
-               write option to make assertions / preconditions on the server
-               state of the document before applying changes.
+            reference (~.firestore_v1beta1.document.DocumentReference):
+                A document reference that will have values set in this batch.
+            document_data (dict):
+                Property names and values to use for replacing a document.
+            merge (Optional[bool]):
+                If True, apply merging instead of overwriting the state
+                of the document.
         """
         write_pbs = _helpers.pbs_for_set(
-            reference._document_path, document_data, option)
+            reference._document_path, document_data, merge=merge)
         self._add_write_pbs(write_pbs)
 
     def update(self, reference, field_updates, option=None):

firestore/google/cloud/firestore_v1beta1/client.py

@@ -33,14 +33,13 @@
 from google.cloud.firestore_v1beta1.document import DocumentReference
 from google.cloud.firestore_v1beta1.document import DocumentSnapshot
 from google.cloud.firestore_v1beta1.gapic import firestore_client
-from google.cloud.firestore_v1beta1.proto import common_pb2
 from google.cloud.firestore_v1beta1.transaction import Transaction
 
 
 DEFAULT_DATABASE = '(default)'
 """str: The default database used in a :class:`~.firestore.client.Client`."""
 _BAD_OPTION_ERR = (
-    'Exactly one of ``last_update_time``, ``merge`` or ``exists`` '
+    'Exactly one of ``last_update_time``, ``exists`` '
     'must be provided.')
 _BAD_DOC_TEMPLATE = (
     'Document {!r} appeared in response but was not present among references')
@@ -261,9 +260,6 @@ def write_option(**kwargs):
                protobuf or directly.
         * ``exists`` (:class:`bool`): Indicates if the document being modified
               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
@@ -287,8 +283,6 @@ 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)
@@ -424,32 +418,6 @@ 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, field_paths=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.
-            field_paths (Sequence[str]):
-                The actual field names to use for replacing a document.
-            unused_kwargs (Dict[str, Any]): Keyword arguments accepted by
-                other subclasses that are unused here.
-        """
-        if field_paths is None:
-            field_paths = []
-        field_paths = _helpers.canonicalize_field_paths(field_paths)
-        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/google/cloud/firestore_v1beta1/document.py

@@ -192,7 +192,7 @@ def create(self, document_data):
         write_results = batch.commit()
         return _first_write_result(write_results)
 
-    def set(self, document_data, option=None):
+    def set(self, document_data, merge=False):
         """Replace the current document in the Firestore database.
 
         A write ``option`` can be specified to indicate preconditions of
@@ -219,7 +219,7 @@ def set(self, document_data, option=None):
             result contains an ``update_time`` field.
         """
         batch = self._client.batch()
-        batch.set(self, document_data, option=option)
+        batch.set(self, document_data, merge=merge)
         write_results = batch.commit()
         return _first_write_result(write_results)
 

firestore/tests/system.py

@@ -143,10 +143,9 @@ def test_document_set(client, cleanup):
     snapshot = document.get()
     assert snapshot.to_dict() is None
 
-    # 1. Use ``set()`` to create the document (using an option).
+    # 1. Use ``create()`` to create the document (using an option).
     data1 = {'foo': 88}
-    option1 = client.write_option(exists=False)
-    write_result1 = document.set(data1, option=option1)
+    write_result1 = document.create(data1)
     snapshot1 = document.get()
     assert snapshot1.to_dict() == data1
     # Make sure the update is what created the document.
@@ -162,30 +161,6 @@ def test_document_set(client, cleanup):
     assert snapshot2.create_time == snapshot1.create_time
     assert snapshot2.update_time == write_result2.update_time
 
-    # 3. Call ``set()`` with a valid "last timestamp" option.
-    data3 = {'skates': 88}
-    option3 = client.write_option(last_update_time=snapshot2.update_time)
-    write_result3 = document.set(data3, option=option3)
-    snapshot3 = document.get()
-    assert snapshot3.to_dict() == data3
-    # Make sure the create time hasn't changed.
-    assert snapshot3.create_time == snapshot1.create_time
-    assert snapshot3.update_time == write_result3.update_time
-
-    # 4. Call ``set()`` with invalid (in the past) "last timestamp" option.
-    assert_timestamp_less(option3._last_update_time, snapshot3.update_time)
-    with pytest.raises(FailedPrecondition):
-        document.set({'bad': 'time-past'}, option=option3)
-
-    # 5. Call ``set()`` with invalid (in the future) "last timestamp" option.
-    timestamp_pb = timestamp_pb2.Timestamp(
-        seconds=snapshot3.update_time.nanos + 120,
-        nanos=snapshot3.update_time.nanos,
-    )
-    option5 = client.write_option(last_update_time=timestamp_pb)
-    with pytest.raises(FailedPrecondition):
-        document.set({'bad': 'time-future'}, option=option5)
-
 
 def test_document_integer_field(client, cleanup):
     document_id = 'for-set' + unique_resource_id('-')
@@ -201,8 +176,7 @@ def test_document_integer_field(client, cleanup):
             '7g': '8h',
             'cd': '0j'}
     }
-    option1 = client.write_option(exists=False)
-    document.set(data1, option=option1)
+    document.create(data1)
 
     data2 = {'1a.ab': '4d', '6f.7g': '9h'}
     option2 = client.write_option(exists=True)
@@ -225,30 +199,24 @@ def test_document_set_merge(client, cleanup):
     # 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
+    # 0. Make sure the document doesn't exist yet
+    snapshot = document.get()
+    assert not snapshot.exists
 
     # 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)
+    write_result1 = document.create(data1)
     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).
+    # 2. Call ``set()`` to merge
     data2 = {'address': {'city': 'LA'}}
-    option2 = client.write_option(merge=True)
-    write_result2 = document.set(data2, option=option2)
+    write_result2 = document.set(data2, merge=True)
     snapshot2 = document.get()
     assert snapshot2.to_dict() == {'name': 'Sam',
                                    'address': {'city': 'LA',
@@ -333,7 +301,7 @@ def test_update_document(client, cleanup):
     )
     option6 = client.write_option(last_update_time=timestamp_pb)
     with pytest.raises(FailedPrecondition) as exc_info:
-        document.set({'bad': 'time-future'}, option=option6)
+        document.update({'bad': 'time-future'}, option=option6)
 
 
 def check_snapshot(snapshot, document, data, write_result):

firestore/tests/unit/test__helpers.py

@@ -1397,8 +1397,7 @@ def _call_fut(document_path, document_data, option):
 
         return pbs_for_set(document_path, document_data, option)
 
-    def _helper(self, option=None, do_transform=False, **write_kwargs):
-        from google.cloud.firestore_v1beta1.client import MergeOption
+    def _helper(self, merge=False, do_transform=False, **write_kwargs):
         from google.cloud.firestore_v1beta1.constants import SERVER_TIMESTAMP
         from google.cloud.firestore_v1beta1.gapic import enums
         from google.cloud.firestore_v1beta1.proto import common_pb2
@@ -1420,7 +1419,7 @@ def _helper(self, option=None, do_transform=False, **write_kwargs):
         if do_transform:
             document_data[field_name3] = SERVER_TIMESTAMP
 
-        write_pbs = self._call_fut(document_path, document_data, option)
+        write_pbs = self._call_fut(document_path, document_data, merge)
 
         expected_update_pb = write_pb2.Write(
             update=document_pb2.Document(
@@ -1434,7 +1433,7 @@ def _helper(self, option=None, do_transform=False, **write_kwargs):
         )
         expected_pbs = [expected_update_pb]
 
-        if isinstance(option, MergeOption):
+        if merge:
             field_paths = [field_name1, field_name2]
             mask = common_pb2.DocumentMask(field_paths=sorted(field_paths))
             expected_pbs[0].update_mask.CopyFrom(mask)
@@ -1459,19 +1458,8 @@ def _helper(self, option=None, do_transform=False, **write_kwargs):
     def test_without_option(self):
         self._helper()
 
-    def test_with_exists_option(self):
-        from google.cloud.firestore_v1beta1.proto import common_pb2
-        from google.cloud.firestore_v1beta1.client import ExistsOption
-
-        option = ExistsOption(True)
-        precondition = common_pb2.Precondition(exists=True)
-        self._helper(option=option, current_document=precondition)
-
     def test_with_merge_option(self):
-        from google.cloud.firestore_v1beta1.client import MergeOption
-
-        option = MergeOption()
-        self._helper(option=option)
+        self._helper(merge=True)
 
     def test_update_and_transform(self):
         self._helper(do_transform=True)

firestore/tests/unit/test_batch.py

@@ -93,7 +93,6 @@ def test_set(self):
     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)
@@ -103,8 +102,7 @@ def test_set_merge(self):
         field = 'zapzap'
         value = u'meadows and flowers'
         document_data = {field: value}
-        option = MergeOption()
-        ret_val = batch.set(reference, document_data, option)
+        ret_val = batch.set(reference, document_data, merge=True)
         self.assertIsNone(ret_val)
         new_write_pb = write_pb2.Write(
             update=document_pb2.Document(