feat: Logical termination for firestore.getAll(...). (#1517)

ehsannas · cherylEnkidu · gcf-owl-bot[bot] · web-flow · commit c6448fc6411d · 2024-01-10T15:01:39.000-08:00
* feat: Logical termination for firestore.getAll(...). * Using existing unit tests to verify the behaviour * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * transaction unit tests are not mocking response correctly, so it cannot run with logical termination --------- Co-authored-by: cherylEnkidu <cheryllin@google.com> Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java
@@ -222,6 +222,7 @@ void getAll(
     ResponseObserver<BatchGetDocumentsResponse> responseObserver =
         new ResponseObserver<BatchGetDocumentsResponse>() {
           int numResponses;
+          boolean hasCompleted = false;
 
           @Override
           public void onStart(StreamController streamController) {}
@@ -265,6 +266,13 @@ public void onResponse(BatchGetDocumentsResponse response) {
                 return;
             }
             apiStreamObserver.onNext(documentSnapshot);
+
+            // Logical termination: if we have already received as many documents as we had
+            // requested, we can
+            // raise the results without waiting for the termination from the server.
+            if (numResponses == documentReferences.length) {
+              onComplete();
+            }
           }
 
           @Override
@@ -277,6 +285,8 @@ public void onError(Throwable throwable) {
 
           @Override
           public void onComplete() {
+            if (hasCompleted) return;
+            hasCompleted = true;
             tracer
                 .getCurrentSpan()
                 .addAnnotation(TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Complete");
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/FirestoreTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/FirestoreTest.java
@@ -23,7 +23,7 @@
 import static com.google.cloud.firestore.LocalFirestoreHelper.arrayUnion;
 import static com.google.cloud.firestore.LocalFirestoreHelper.commit;
 import static com.google.cloud.firestore.LocalFirestoreHelper.commitResponse;
-import static com.google.cloud.firestore.LocalFirestoreHelper.getAllResponse;
+import static com.google.cloud.firestore.LocalFirestoreHelper.getAllResponseWithoutOnComplete;
 import static com.google.cloud.firestore.LocalFirestoreHelper.transform;
 import static com.google.cloud.firestore.LocalFirestoreHelper.update;
 import static org.junit.Assert.assertEquals;
@@ -80,7 +80,7 @@ public void encodeFieldPath() {
 
   @Test
   public void illegalFieldPath() throws Exception {
-    doAnswer(getAllResponse(SINGLE_FIELD_PROTO))
+    doAnswer(getAllResponseWithoutOnComplete(SINGLE_FIELD_PROTO))
         .when(firestoreMock)
         .streamRequest(
             getAllCapture.capture(),
@@ -110,7 +110,7 @@ public void exposesOptions() {
   @Test
   public void getAll() throws Exception {
     doAnswer(
-            getAllResponse(
+            getAllResponseWithoutOnComplete(
                 SINGLE_FIELD_PROTO, SINGLE_FIELD_PROTO, SINGLE_FIELD_PROTO, SINGLE_FIELD_PROTO))
         .when(firestoreMock)
         .streamRequest(
@@ -132,7 +132,7 @@ public void getAll() throws Exception {
 
   @Test
   public void getAllWithFieldMask() throws Exception {
-    doAnswer(getAllResponse(SINGLE_FIELD_PROTO))
+    doAnswer(getAllResponseWithoutOnComplete(SINGLE_FIELD_PROTO))
         .when(firestoreMock)
         .streamRequest(
             getAllCapture.capture(),
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/LocalFirestoreHelper.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/LocalFirestoreHelper.java
@@ -262,6 +262,16 @@ public static Map<String, Object> map() {
 
   public static Answer<BatchGetDocumentsResponse> getAllResponse(
       final Map<String, Value>... fields) {
+    return getAllResponseImpl(true, fields);
+  }
+
+  public static Answer<BatchGetDocumentsResponse> getAllResponseWithoutOnComplete(
+      final Map<String, Value>... fields) {
+    return getAllResponseImpl(false, fields);
+  }
+
+  public static Answer<BatchGetDocumentsResponse> getAllResponseImpl(
+      boolean withOnComplete, final Map<String, Value>... fields) {
     BatchGetDocumentsResponse[] responses = new BatchGetDocumentsResponse[fields.length];
 
     for (int i = 0; i < fields.length; ++i) {
@@ -281,7 +291,13 @@ public static Answer<BatchGetDocumentsResponse> getAllResponse(
       responses[i] = response.build();
     }
 
-    return streamingResponse(responses, null);
+    if (withOnComplete) {
+      return streamingResponse(responses, null);
+    } else {
+      // Verify with logical termination, the return of results no longer depends on calling
+      // OnComplete.
+      return streamingResponseWithoutOnComplete(responses);
+    }
   }
 
   public static ApiFuture<Empty> rollbackResponse() {
