Adding FieldMask support to GetAll() (#4017)

google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentMask.java → google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FieldMask.java

@@ -17,29 +17,57 @@
 package com.google.cloud.firestore;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.SortedSet;
 import java.util.TreeSet;
+import javax.annotation.Nonnull;
 
-/** A DocumentMask contains the field paths affected by an update. */
-final class DocumentMask {
-  static final DocumentMask EMPTY_MASK = new DocumentMask(new TreeSet<FieldPath>());
+/** A FieldMask can be used to limit the number of fields returned by a `getAll()` call. */
+public final class FieldMask {
+  static final FieldMask EMPTY_MASK = new FieldMask(new TreeSet<FieldPath>());
 
   private final SortedSet<FieldPath> fieldPaths; // Sorted for testing.
 
-  DocumentMask(Collection<FieldPath> fieldPaths) {
+  FieldMask(Collection<FieldPath> fieldPaths) {
     this(new TreeSet<>(fieldPaths));
   }
 
-  private DocumentMask(SortedSet<FieldPath> fieldPaths) {
+  private FieldMask(SortedSet<FieldPath> fieldPaths) {
     this.fieldPaths = fieldPaths;
   }
 
-  static DocumentMask fromObject(Map<String, Object> values) {
+  /**
+   * Creates a FieldMask from the provided field paths.
+   *
+   * @param fieldPaths A list of field paths.
+   * @return A {@code FieldMask} that describes a subset of fields.
+   */
+  @Nonnull
+  public static FieldMask of(String... fieldPaths) {
+    List<FieldPath> paths = new ArrayList<>();
+    for (String fieldPath : fieldPaths) {
+      paths.add(FieldPath.fromDotSeparatedString(fieldPath));
+    }
+    return new FieldMask(paths);
+  }
+
+  /**
+   * Creates a FieldMask from the provided field paths.
+   *
+   * @param fieldPaths A list of field paths.
+   * @return A {@code FieldMask} that describes a subset of fields.
+   */
+  @Nonnull
+  public static FieldMask of(FieldPath... fieldPaths) {
+    return new FieldMask(Arrays.asList(fieldPaths));
+  }
+
+  static FieldMask fromObject(Map<String, Object> values) {
     List<FieldPath> fieldPaths = extractFromMap(values, FieldPath.empty());
-    return new DocumentMask(fieldPaths);
+    return new FieldMask(fieldPaths);
   }
 
   private static List<FieldPath> extractFromMap(Map<String, Object> values, FieldPath path) {

google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Firestore.java

@@ -20,6 +20,7 @@
 import com.google.cloud.Service;
 import java.util.List;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
 /** Represents a Firestore Database and is the entry point for all Firestore operations */
 public interface Firestore extends Service<FirestoreOptions>, AutoCloseable {
@@ -93,7 +94,18 @@ <T> ApiFuture<T> runTransaction(
    * @param documentReferences List of Document References to fetch.
    */
   @Nonnull
-  ApiFuture<List<DocumentSnapshot>> getAll(final DocumentReference... documentReferences);
+  ApiFuture<List<DocumentSnapshot>> getAll(@Nonnull DocumentReference... documentReferences);
+
+  /**
+   * Retrieves multiple documents from Firestore, while optionally applying a field mask to reduce
+   * the amount of data transmitted.
+   *
+   * @param documentReferences Array with Document References to fetch.
+   * @param fieldMask If set, specifies the subset of fields to return.
+   */
+  @Nonnull
+  ApiFuture<List<DocumentSnapshot>> getAll(
+      @Nonnull DocumentReference[] documentReferences, @Nullable FieldMask fieldMask);
 
   /**
    * Gets a Firestore {@link WriteBatch} instance that can be used to combine multiple writes.

google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java

@@ -162,13 +162,23 @@ public Iterable<CollectionReference> getCollections() {
 
   @Nonnull
   @Override
-  public ApiFuture<List<DocumentSnapshot>> getAll(final DocumentReference... documentReferences) {
-    return this.getAll(documentReferences, null);
+  public ApiFuture<List<DocumentSnapshot>> getAll(
+      @Nonnull DocumentReference... documentReferences) {
+    return this.getAll(documentReferences, null, null);
+  }
+
+  @Nonnull
+  @Override
+  public ApiFuture<List<DocumentSnapshot>> getAll(
+      @Nonnull DocumentReference[] documentReferences, @Nullable FieldMask fieldMask) {
+    return this.getAll(documentReferences, fieldMask, null);
   }
 
   /** Internal getAll() method that accepts an optional transaction id. */
   ApiFuture<List<DocumentSnapshot>> getAll(
-      final DocumentReference[] documentReferences, @Nullable ByteString transactionId) {
+      final DocumentReference[] documentReferences,
+      @Nullable FieldMask fieldMask,
+      @Nullable ByteString transactionId) {
     final SettableApiFuture<List<DocumentSnapshot>> futureList = SettableApiFuture.create();
     final Map<DocumentReference, DocumentSnapshot> resultMap = new HashMap<>();
 
@@ -238,6 +248,10 @@ public void onCompleted() {
     BatchGetDocumentsRequest.Builder request = BatchGetDocumentsRequest.newBuilder();
     request.setDatabase(getDatabaseName());
 
+    if (fieldMask != null) {
+      request.setMask(fieldMask.toPb());
+    }
+
     if (transactionId != null) {
       request.setTransaction(transactionId);
     }

google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java

@@ -134,7 +134,7 @@ public ApiFuture<DocumentSnapshot> get(@Nonnull DocumentReference documentRef) {
     Preconditions.checkState(isEmpty(), READ_BEFORE_WRITE_ERROR_MSG);
 
     return ApiFutures.transform(
-        firestore.getAll(new DocumentReference[] {documentRef}, transactionId),
+        firestore.getAll(new DocumentReference[] {documentRef}, /*fieldMask=*/ null, transactionId),
         new ApiFunction<List<DocumentSnapshot>, DocumentSnapshot>() {
           @Override
           public DocumentSnapshot apply(List<DocumentSnapshot> snapshots) {
@@ -150,10 +150,27 @@ public DocumentSnapshot apply(List<DocumentSnapshot> snapshots) {
    * @param documentReferences List of Document References to fetch.
    */
   @Nonnull
-  public ApiFuture<List<DocumentSnapshot>> getAll(final DocumentReference... documentReferences) {
+  public ApiFuture<List<DocumentSnapshot>> getAll(
+      @Nonnull DocumentReference... documentReferences) {
     Preconditions.checkState(isEmpty(), READ_BEFORE_WRITE_ERROR_MSG);
 
-    return firestore.getAll(documentReferences, transactionId);
+    return firestore.getAll(documentReferences, /*fieldMask=*/ null, transactionId);
+  }
+
+  /**
+   * Retrieves multiple documents from Firestore, while optionally applying a field mask to reduce
+   * the amount of data transmitted from the backend. Holds a pessimistic lock on all returned
+   * documents.
+   *
+   * @param documentReferences Array with Document References to fetch.
+   * @param fieldMask If set, specifies the subset of fields to return.
+   */
+  @Nonnull
+  public ApiFuture<List<DocumentSnapshot>> getAll(
+      @Nonnull DocumentReference[] documentReferences, @Nullable FieldMask fieldMask) {
+    Preconditions.checkState(isEmpty(), READ_BEFORE_WRITE_ERROR_MSG);
+
+    return firestore.getAll(documentReferences, fieldMask, transactionId);
   }
 
   /**

google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java

@@ -243,17 +243,17 @@ private T performSet(
     DocumentSnapshot documentSnapshot =
         DocumentSnapshot.fromObject(
             firestore, documentReference, expandObject(documentData), options.getEncodingOptions());
-    DocumentMask documentMask = DocumentMask.EMPTY_MASK;
+    FieldMask documentMask = FieldMask.EMPTY_MASK;
     DocumentTransform documentTransform =
         DocumentTransform.fromFieldPathMap(documentReference, documentData);
 
     if (options.isMerge()) {
       if (options.getFieldMask() != null) {
         List<FieldPath> fieldMask = new ArrayList<>(options.getFieldMask());
         fieldMask.removeAll(documentTransform.getFields());
-        documentMask = new DocumentMask(fieldMask);
+        documentMask = new FieldMask(fieldMask);
       } else {
-        documentMask = DocumentMask.fromObject(fields);
+        documentMask = FieldMask.fromObject(fields);
       }
     }
 
@@ -528,14 +528,14 @@ public boolean allowTransform() {
     DocumentTransform documentTransform =
         DocumentTransform.fromFieldPathMap(documentReference, fields);
     fieldPaths.removeAll(documentTransform.getFields());
-    DocumentMask documentMask = new DocumentMask(fieldPaths);
+    FieldMask fieldMask = new FieldMask(fieldPaths);
 
     Mutation mutation = addMutation();
     mutation.precondition = precondition.toPb();
 
-    if (!documentSnapshot.isEmpty() || !documentMask.isEmpty()) {
+    if (!documentSnapshot.isEmpty() || !fieldMask.isEmpty()) {
       mutation.document = documentSnapshot.toPb();
-      mutation.document.setUpdateMask(documentMask.toPb());
+      mutation.document.setUpdateMask(fieldMask.toPb());
     }
 
     if (!documentTransform.isEmpty()) {

google-cloud-clients/google-cloud-firestore/src/test/java/com/google/cloud/firestore/FirestoreTest.java

@@ -118,6 +118,25 @@ public void getAll() throws Exception {
     assertEquals("doc3", snapshot.get(3).getId());
   }
 
+  @Test
+  public void getAllWithFieldMask() throws Exception {
+    doAnswer(getAllResponse(SINGLE_FIELD_PROTO))
+        .when(firestoreMock)
+        .streamRequest(
+            getAllCapture.capture(),
+            streamObserverCapture.capture(),
+            Matchers.<ServerStreamingCallable>any());
+
+    DocumentReference doc1 = firestoreMock.document("coll/doc1");
+    FieldMask fieldMask = FieldMask.of(FieldPath.of("foo", "bar"));
+
+    firestoreMock.getAll(new DocumentReference[]{doc1}, fieldMask).get();
+
+    BatchGetDocumentsRequest request = getAllCapture.getValue();
+    assertEquals(1, request.getMask().getFieldPathsCount());
+    assertEquals("foo.bar", request.getMask().getFieldPaths(0));
+  }
+
   @Test
   public void arrayUnionEquals() {
     FieldValue arrayUnion1 = FieldValue.arrayUnion("foo", "bar");

google-cloud-clients/google-cloud-firestore/src/test/java/com/google/cloud/firestore/TransactionTest.java

@@ -46,6 +46,8 @@
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.cloud.Timestamp;
 import com.google.cloud.firestore.spi.v1beta1.FirestoreRpc;
+import com.google.firestore.v1beta1.BatchGetDocumentsRequest;
+import com.google.firestore.v1beta1.DocumentMask;
 import com.google.firestore.v1beta1.Write;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Message;
@@ -386,6 +388,50 @@ public List<DocumentSnapshot> updateCallback(Transaction transaction)
     assertEquals(commit(TRANSACTION_ID), requests.get(2));
   }
 
+  @Test
+  public void getMultipleDocumentsWithFieldMask() throws Exception {
+    doReturn(beginResponse())
+        .doReturn(commitResponse(0, 0))
+        .when(firestoreMock)
+        .sendRequest(requestCapture.capture(), Matchers.<UnaryCallable<Message, Message>>any());
+
+    doAnswer(getAllResponse(SINGLE_FIELD_PROTO))
+        .when(firestoreMock)
+        .streamRequest(
+            requestCapture.capture(),
+            streamObserverCapture.capture(),
+            Matchers.<ServerStreamingCallable>any());
+
+    final DocumentReference doc1 = firestoreMock.document("coll/doc1");
+    final FieldMask fieldMask = FieldMask.of(FieldPath.of("foo", "bar"));
+
+    ApiFuture<List<DocumentSnapshot>> transaction =
+        firestoreMock.runTransaction(
+            new Transaction.Function<List<DocumentSnapshot>>() {
+              @Override
+              public List<DocumentSnapshot> updateCallback(Transaction transaction)
+                  throws ExecutionException, InterruptedException {
+                return transaction.getAll(new DocumentReference[] {doc1}, fieldMask).get();
+              }
+            },
+            options);
+    transaction.get();
+
+    List<Message> requests = requestCapture.getAllValues();
+    assertEquals(3, requests.size());
+
+    assertEquals(begin(), requests.get(0));
+    BatchGetDocumentsRequest expectedGetAll =
+        getAll(TRANSACTION_ID, doc1.getResourcePath().toString());
+    expectedGetAll =
+        expectedGetAll
+            .toBuilder()
+            .setMask(DocumentMask.newBuilder().addFieldPaths("foo.bar"))
+            .build();
+    assertEquals(expectedGetAll, requests.get(1));
+    assertEquals(commit(TRANSACTION_ID), requests.get(2));
+  }
+
   @Test
   public void getQuery() throws Exception {
     doReturn(beginResponse())

google-cloud-clients/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITSystemTest.java

@@ -32,6 +32,7 @@
 import com.google.cloud.firestore.DocumentReference;
 import com.google.cloud.firestore.DocumentSnapshot;
 import com.google.cloud.firestore.EventListener;
+import com.google.cloud.firestore.FieldMask;
 import com.google.cloud.firestore.FieldValue;
 import com.google.cloud.firestore.Firestore;
 import com.google.cloud.firestore.FirestoreException;
@@ -122,6 +123,15 @@ public void getAll() throws Exception {
     assertEquals(SINGLE_FIELD_OBJECT, documentSnapshots.get(1).toObject(SingleField.class));
   }
 
+  @Test
+  public void getAllWithFieldMask() throws Exception {
+    DocumentReference ref = randomColl.document("doc1");
+    ref.set(ALL_SUPPORTED_TYPES_MAP).get();
+    List<DocumentSnapshot> documentSnapshots =
+        firestore.getAll(new DocumentReference[] {ref}, FieldMask.of("foo")).get();
+    assertEquals(map("foo", "bar"), documentSnapshots.get(0).getData());
+  }
+
   @Test
   public void addDocument() throws Exception {
     DocumentReference documentReference = randomColl.add(SINGLE_FIELD_MAP).get();