diff --git a/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentMask.java b/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FieldMask.java similarity index 64% rename from google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentMask.java rename to google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FieldMask.java index dc2d8f6754bb..89aefce06f10 100644 --- a/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentMask.java +++ b/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()); +/** 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()); private final SortedSet fieldPaths; // Sorted for testing. - DocumentMask(Collection fieldPaths) { + FieldMask(Collection fieldPaths) { this(new TreeSet<>(fieldPaths)); } - private DocumentMask(SortedSet fieldPaths) { + private FieldMask(SortedSet fieldPaths) { this.fieldPaths = fieldPaths; } - static DocumentMask fromObject(Map 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 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 values) { List fieldPaths = extractFromMap(values, FieldPath.empty()); - return new DocumentMask(fieldPaths); + return new FieldMask(fieldPaths); } private static List extractFromMap(Map values, FieldPath path) { diff --git a/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Firestore.java b/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Firestore.java index 261678b6c932..62e2984f1414 100644 --- a/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Firestore.java +++ b/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, AutoCloseable { @@ -93,7 +94,18 @@ ApiFuture runTransaction( * @param documentReferences List of Document References to fetch. */ @Nonnull - ApiFuture> getAll(final DocumentReference... documentReferences); + ApiFuture> 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> getAll( + @Nonnull DocumentReference[] documentReferences, @Nullable FieldMask fieldMask); /** * Gets a Firestore {@link WriteBatch} instance that can be used to combine multiple writes. diff --git a/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java b/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java index dbabafa9ff25..54d479fbda22 100644 --- a/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java +++ b/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java @@ -162,13 +162,23 @@ public Iterable getCollections() { @Nonnull @Override - public ApiFuture> getAll(final DocumentReference... documentReferences) { - return this.getAll(documentReferences, null); + public ApiFuture> getAll( + @Nonnull DocumentReference... documentReferences) { + return this.getAll(documentReferences, null, null); + } + + @Nonnull + @Override + public ApiFuture> getAll( + @Nonnull DocumentReference[] documentReferences, @Nullable FieldMask fieldMask) { + return this.getAll(documentReferences, fieldMask, null); } /** Internal getAll() method that accepts an optional transaction id. */ ApiFuture> getAll( - final DocumentReference[] documentReferences, @Nullable ByteString transactionId) { + final DocumentReference[] documentReferences, + @Nullable FieldMask fieldMask, + @Nullable ByteString transactionId) { final SettableApiFuture> futureList = SettableApiFuture.create(); final Map 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); } diff --git a/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java b/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java index 634118eb9f7b..6954e83ecad1 100644 --- a/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java +++ b/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Transaction.java @@ -134,7 +134,7 @@ public ApiFuture 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, DocumentSnapshot>() { @Override public DocumentSnapshot apply(List snapshots) { @@ -150,10 +150,27 @@ public DocumentSnapshot apply(List snapshots) { * @param documentReferences List of Document References to fetch. */ @Nonnull - public ApiFuture> getAll(final DocumentReference... documentReferences) { + public ApiFuture> 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> getAll( + @Nonnull DocumentReference[] documentReferences, @Nullable FieldMask fieldMask) { + Preconditions.checkState(isEmpty(), READ_BEFORE_WRITE_ERROR_MSG); + + return firestore.getAll(documentReferences, fieldMask, transactionId); } /** diff --git a/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java b/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java index 6b69c40bb4bf..1e937e589b58 100644 --- a/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java +++ b/google-cloud-clients/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java @@ -243,7 +243,7 @@ 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); @@ -251,9 +251,9 @@ private T performSet( if (options.getFieldMask() != null) { List 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()) { diff --git a/google-cloud-clients/google-cloud-firestore/src/test/java/com/google/cloud/firestore/FirestoreTest.java b/google-cloud-clients/google-cloud-firestore/src/test/java/com/google/cloud/firestore/FirestoreTest.java index 68798058a068..bba46c14ae66 100644 --- a/google-cloud-clients/google-cloud-firestore/src/test/java/com/google/cloud/firestore/FirestoreTest.java +++ b/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.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"); diff --git a/google-cloud-clients/google-cloud-firestore/src/test/java/com/google/cloud/firestore/TransactionTest.java b/google-cloud-clients/google-cloud-firestore/src/test/java/com/google/cloud/firestore/TransactionTest.java index 9054d134343e..c37631ee0688 100644 --- a/google-cloud-clients/google-cloud-firestore/src/test/java/com/google/cloud/firestore/TransactionTest.java +++ b/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 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.>any()); + + doAnswer(getAllResponse(SINGLE_FIELD_PROTO)) + .when(firestoreMock) + .streamRequest( + requestCapture.capture(), + streamObserverCapture.capture(), + Matchers.any()); + + final DocumentReference doc1 = firestoreMock.document("coll/doc1"); + final FieldMask fieldMask = FieldMask.of(FieldPath.of("foo", "bar")); + + ApiFuture> transaction = + firestoreMock.runTransaction( + new Transaction.Function>() { + @Override + public List updateCallback(Transaction transaction) + throws ExecutionException, InterruptedException { + return transaction.getAll(new DocumentReference[] {doc1}, fieldMask).get(); + } + }, + options); + transaction.get(); + + List 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()) diff --git a/google-cloud-clients/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITSystemTest.java b/google-cloud-clients/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITSystemTest.java index d05adfe11f23..c935702f544b 100644 --- a/google-cloud-clients/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITSystemTest.java +++ b/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 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();