feat: add support for fieldmask to document reference (#245)

suraj-qlogic · web-flow · commit 4a846b1f067a · 2020-06-15T16:04:01.000-04:00
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentReference.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/DocumentReference.java
@@ -354,6 +354,19 @@ public ApiFuture<DocumentSnapshot> get() {
     return extractFirst(firestore.getAll(this));
   }
 
+  /**
+   * Reads the document referenced by this DocumentReference. If the document doesn't exist, the
+   * get(FieldMask fieldMask) will return an empty DocumentSnapshot.
+   *
+   * @param fieldMask A FieldMask object to retrieve the field value
+   * @return An ApiFuture that will be resolved with the contents of the Document at this
+   *     DocumentReference, or a failure if the document does not exist
+   */
+  @Nonnull
+  public ApiFuture<DocumentSnapshot> get(FieldMask fieldMask) {
+    return extractFirst(firestore.getAll(new DocumentReference[] {this}, fieldMask));
+  }
+
   /**
    * Fetches the subcollections that are direct children of this document.
    *
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/DocumentReferenceTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/DocumentReferenceTest.java
@@ -275,6 +275,19 @@ public void deserializeDocumentReference() throws Exception {
     assertEquals(documentReference, snapshot.getReference());
   }
 
+  @Test
+  public void getFieldWithFieldMask() throws Exception {
+    doAnswer(getAllResponse(SINGLE_FIELD_PROTO))
+        .when(firestoreMock)
+        .streamRequest(
+            getAllCapture.capture(),
+            streamObserverCapture.capture(),
+            Matchers.<ServerStreamingCallable>any());
+    DocumentSnapshot snapshot = documentReference.get(FieldMask.of(FieldPath.of("foo"))).get();
+    assertEquals("foo", getAllCapture.getValue().getMask().getFieldPaths(0));
+    assertEquals("bar", snapshot.get("foo"));
+  }
+
   @Test
   public void deserializesDates() throws Exception {
     doAnswer(getAllResponse(ALL_SUPPORTED_TYPES_PROTO))
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITSystemTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITSystemTest.java
@@ -157,6 +157,15 @@ public void getAllWithFieldMask() throws Exception {
     assertEquals(map("foo", "bar"), documentSnapshots.get(0).getData());
   }
 
+  @Test
+  public void getFieldMaskWithDocumentReference() throws Exception {
+    DocumentReference ref = randomColl.document("doc1");
+    ref.set(ALL_SUPPORTED_TYPES_MAP).get();
+    DocumentSnapshot documentSnapshots = ref.get(FieldMask.of("foo", "foobar")).get();
+    assertEquals("bar", documentSnapshots.get("foo"));
+    assertNull(documentSnapshots.get("foobar"));
+  }
+
   @Test
   public void addDocument() throws Exception {
     DocumentReference documentReference = randomColl.add(SINGLE_FIELD_MAP).get();
