Firestore: add support for 'ArrayRemove' / 'ArrayUnion' transforms (#…

…6651)

Closes #6546.

firestore/docs/index.rst

@@ -13,7 +13,7 @@ API Reference
   query
   batch
   transaction
-  constants
+  transforms
   types
 
 

firestore/docs/transforms.rst

@@ -0,0 +1,6 @@
+Transforms
+~~~~~~~~~~
+
+.. automodule:: google.cloud.firestore_v1beta1.transforms
+  :members:
+  :show-inheritance:

firestore/google/cloud/firestore_v1beta1/__init__.py

@@ -26,8 +26,10 @@
 from google.cloud.firestore_v1beta1.batch import WriteBatch
 from google.cloud.firestore_v1beta1.client import Client
 from google.cloud.firestore_v1beta1.collection import CollectionReference
-from google.cloud.firestore_v1beta1.constants import DELETE_FIELD
-from google.cloud.firestore_v1beta1.constants import SERVER_TIMESTAMP
+from google.cloud.firestore_v1beta1.transforms import ArrayRemove
+from google.cloud.firestore_v1beta1.transforms import ArrayUnion
+from google.cloud.firestore_v1beta1.transforms import DELETE_FIELD
+from google.cloud.firestore_v1beta1.transforms import SERVER_TIMESTAMP
 from google.cloud.firestore_v1beta1.document import DocumentReference
 from google.cloud.firestore_v1beta1.document import DocumentSnapshot
 from google.cloud.firestore_v1beta1.gapic import enums
@@ -39,6 +41,8 @@
 
 __all__ = [
     '__version__',
+    'ArrayRemove',
+    'ArrayUnion',
     'Client',
     'CollectionReference',
     'DELETE_FIELD',

firestore/google/cloud/firestore_v1beta1/_helpers.py

@@ -30,7 +30,7 @@
 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.firestore_v1beta1 import constants
+from google.cloud.firestore_v1beta1 import transforms
 from google.cloud.firestore_v1beta1 import types
 from google.cloud.firestore_v1beta1.gapic import enums
 from google.cloud.firestore_v1beta1.proto import common_pb2
@@ -654,7 +654,7 @@ def get_doc_id(document_pb, expected_prefix):
     return document_id
 
 
-_EmptyDict = constants.Sentinel("Marker for an empty dict value")
+_EmptyDict = transforms.Sentinel("Marker for an empty dict value")
 
 
 def extract_fields(document_data, prefix_path, expand_dots=False):
@@ -713,6 +713,8 @@ def __init__(self, document_data):
         self.field_paths = []
         self.deleted_fields = []
         self.server_timestamps = []
+        self.array_removes = {}
+        self.array_unions = {}
         self.set_fields = {}
         self.empty_document = False
 
@@ -724,12 +726,18 @@ def __init__(self, document_data):
             if field_path == prefix_path and value is _EmptyDict:
                 self.empty_document = True
 
-            elif value is constants.DELETE_FIELD:
+            elif value is transforms.DELETE_FIELD:
                 self.deleted_fields.append(field_path)
 
-            elif value is constants.SERVER_TIMESTAMP:
+            elif value is transforms.SERVER_TIMESTAMP:
                 self.server_timestamps.append(field_path)
 
+            elif isinstance(value, transforms.ArrayRemove):
+                self.array_removes[field_path] = value.values
+
+            elif isinstance(value, transforms.ArrayUnion):
+                self.array_unions[field_path] = value.values
+
             else:
                 self.field_paths.append(field_path)
                 set_field_value(self.set_fields, field_path, value)
@@ -739,11 +747,18 @@ def _get_document_iterator(self, prefix_path):
 
     @property
     def has_transforms(self):
-        return bool(self.server_timestamps)
+        return bool(
+            self.server_timestamps
+            or self.array_removes
+            or self.array_unions
+        )
 
     @property
     def transform_paths(self):
-        return sorted(self.server_timestamps)
+        return sorted(
+            self.server_timestamps
+            + list(self.array_removes)
+            + list(self.array_unions))
 
     def _get_update_mask(self, allow_empty_mask=False):
         return None
@@ -768,16 +783,34 @@ def get_update_pb(
         return update_pb
 
     def get_transform_pb(self, document_path, exists=None):
+
+        def make_array_value(values):
+            value_list = [encode_value(element) for element in values]
+            return document_pb2.ArrayValue(values=value_list)
+
+        path_field_transforms = [
+            (path, write_pb2.DocumentTransform.FieldTransform(
+                field_path=path.to_api_repr(),
+                set_to_server_value=REQUEST_TIME_ENUM,
+            )) for path in self.server_timestamps
+        ] + [
+            (path, write_pb2.DocumentTransform.FieldTransform(
+                field_path=path.to_api_repr(),
+                remove_all_from_array=make_array_value(values),
+            )) for path, values in self.array_removes.items()
+        ] + [
+            (path, write_pb2.DocumentTransform.FieldTransform(
+                field_path=path.to_api_repr(),
+                append_missing_elements=make_array_value(values),
+            )) for path, values in self.array_unions.items()
+        ]
+        field_transforms = [
+            transform for path, transform in sorted(path_field_transforms)
+        ]
         transform_pb = write_pb2.Write(
             transform=write_pb2.DocumentTransform(
                 document=document_path,
-                field_transforms=[
-                    write_pb2.DocumentTransform.FieldTransform(
-                        field_path=path.to_api_repr(),
-                        set_to_server_value=REQUEST_TIME_ENUM,
-                    )
-                    for path in self.server_timestamps
-                ],
+                field_transforms=field_transforms,
             ),
         )
         if exists is not None:
@@ -953,12 +986,21 @@ def _apply_merge_paths(self, merge):
             ]
             merged_transform_paths.update(tranform_merge_paths)
 
-        # TODO:  other transforms
         self.server_timestamps = [
             path for path in self.server_timestamps
             if path in merged_transform_paths
         ]
 
+        self.array_removes = {
+            path: values for path, values in self.array_removes.items()
+            if path in merged_transform_paths
+        }
+
+        self.array_unions = {
+            path: values for path, values in self.array_unions.items()
+            if path in merged_transform_paths
+        }
+
     def apply_merge(self, merge):
         if merge is True:  # merge all fields
             self._apply_merge_all()

firestore/google/cloud/firestore_v1beta1/document.py

@@ -306,7 +306,7 @@ def update(self, field_updates, option=None):
            ``field_updates``.
 
         To delete / remove a field from an existing document, use the
-        :attr:`~.firestore_v1beta1.constants.DELETE_FIELD` sentinel. So
+        :attr:`~.firestore_v1beta1.transforms.DELETE_FIELD` sentinel. So
         with the example above, sending
 
         .. code-block:: python
@@ -330,7 +330,7 @@ def update(self, field_updates, option=None):
 
         To set a field to the current time on the server when the
         update is received, use the
-        :attr:`~.firestore_v1beta1.constants.SERVER_TIMESTAMP` sentinel.
+        :attr:`~.firestore_v1beta1.transforms.SERVER_TIMESTAMP` sentinel.
         Sending
 
         .. code-block:: python

firestore/google/cloud/firestore_v1beta1/transforms.py

@@ -0,0 +1,82 @@
+# Copyright 2017 Google LLC All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Helpful constants to use for Google Cloud Firestore."""
+
+
+class Sentinel(object):
+    """Sentinel objects used to signal special handling."""
+    __slots__ = ('description',)
+
+    def __init__(self, description):
+        self.description = description
+
+    def __repr__(self):
+        return "Sentinel: {}".format(self.description)
+
+
+DELETE_FIELD = Sentinel("Value used to delete a field in a document.")
+
+
+SERVER_TIMESTAMP = Sentinel(
+    "Value used to set a document field to the server timestamp.")
+
+
+class _ValueList(object):
+    """Read-only list of values.
+
+    Args:
+        values (List | Tuple): values held in the helper.
+    """
+    slots = ('_values',)
+
+    def __init__(self, values):
+        if not isinstance(values, (list, tuple)):
+            raise ValueError("'values' must be a list or tuple.")
+
+        if len(values) == 0:
+            raise ValueError("'values' must be non-empty.")
+
+        self._values = list(values)
+
+    @property
+    def values(self):
+        """Values to append.
+
+        Returns (List):
+            values to be appended by the transform.
+        """
+        return self._values
+
+
+class ArrayUnion(_ValueList):
+    """Field transform: appends missing values to an array field.
+
+    See:
+    https://cloud.google.com/firestore/docs/reference/rpc/google.firestore.v1beta1#google.firestore.v1beta1.DocumentTransform.FieldTransform.FIELDS.google.firestore.v1beta1.ArrayValue.google.firestore.v1beta1.DocumentTransform.FieldTransform.append_missing_elements
+
+    Args:
+        values (List | Tuple): values to append.
+    """
+
+
+class ArrayRemove(_ValueList):
+    """Field transform: remove values from an array field.
+
+    See:
+    https://cloud.google.com/firestore/docs/reference/rpc/google.firestore.v1beta1#google.firestore.v1beta1.DocumentTransform.FieldTransform.FIELDS.google.firestore.v1beta1.ArrayValue.google.firestore.v1beta1.DocumentTransform.FieldTransform.remove_all_from_array
+
+    Args:
+        values (List | Tuple): values to remove.
+    """