feat: add logical termination for RunQueryResponse (#956)

* feat: add logical termination for RunQueryResponse * Add one test and comments * Change code based on feedback * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * add comments * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: cherylEnkidu

google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java

@@ -1330,6 +1330,8 @@ public void stream(@Nonnull final ApiStreamObserver<DocumentSnapshot> responseOb
 
  internalStream(
  new QuerySnapshotObserver() {
+ boolean hasCompleted = false;
+
  @Override
  public void onNext(QueryDocumentSnapshot documentSnapshot) {
  responseObserver.onNext(documentSnapshot);
@@ -1342,6 +1344,8 @@ public void onError(Throwable throwable) {
 
  @Override
  public void onCompleted() {
+ if (hasCompleted) return;
+ hasCompleted = true;
  responseObserver.onCompleted();
  }
  },
@@ -1537,6 +1541,16 @@ public void onResponse(RunQueryResponse response) {
  if (readTime == null) {
  readTime = Timestamp.fromProto(response.getReadTime());
  }
+
+ if (response.hasDone() && response.getDone()) {
+ Tracing.getTracer()
+ .getCurrentSpan()
+ .addAnnotation(
+ "Firestore.Query: Completed",
+ ImmutableMap.of(
+ "numDocuments", AttributeValue.longAttributeValue(numDocuments)));
+ documentObserver.onCompleted(readTime);
+ }
  }
 
  @Override
@@ -1640,6 +1654,9 @@ ApiFuture<QuerySnapshot> get(@Nullable ByteString transactionId) {
  internalStream(
  new QuerySnapshotObserver() {
  final List<QueryDocumentSnapshot> documentSnapshots = new ArrayList<>();
+ // The stream's onCompleted could be called more than once,
+ // this flag makes sure only the first one is actually processed.
+ boolean hasCompleted = false;
 
  @Override
  public void onNext(QueryDocumentSnapshot documentSnapshot) {
@@ -1653,6 +1670,9 @@ public void onError(Throwable throwable) {
 
  @Override
  public void onCompleted() {
+ if (hasCompleted) return;
+ hasCompleted = true;
+
  // The results for limitToLast queries need to be flipped since we reversed the
  // ordering constraints before sending the query to the backend.
  List<QueryDocumentSnapshot> resultView =

google-cloud-firestore/src/test/java/com/google/cloud/firestore/LocalFirestoreHelper.java

@@ -311,6 +311,29 @@ public static Answer<RunQueryResponse> queryResponse(
  return streamingResponse(responses, throwable);
  }
 
+ /** Returns a stream of responses when RunQueryResponse.done set to true */
+ public static Answer<RunQueryResponse> queryResponseWithDone(
+ boolean callWithoutOnComplete, String... documentNames) {
+ RunQueryResponse[] responses = new RunQueryResponse[documentNames.length];
+
+ for (int i = 0; i < documentNames.length; ++i) {
+ final RunQueryResponse.Builder runQueryResponse = RunQueryResponse.newBuilder();
+ runQueryResponse.setDocument(
+ Document.newBuilder().setName(documentNames[i]).putAllFields(SINGLE_FIELD_PROTO));
+ runQueryResponse.setReadTime(
+ com.google.protobuf.Timestamp.newBuilder().setSeconds(1).setNanos(2));
+ if (i == (documentNames.length - 1)) {
+ runQueryResponse.setDone(true);
+ }
+ responses[i] = runQueryResponse.build();
+ }
+ if (callWithoutOnComplete) {
+ return streamingResponseWithoutOnComplete(responses);
+ } else {
+ return streamingResponse(responses, null);
+ }
+ }
+
  /** Returns a stream of responses followed by an optional exception. */
  public static <T> Answer<T> streamingResponse(
  final T[] response, @Nullable final Throwable throwable) {
@@ -328,6 +351,18 @@ public static <T> Answer<T> streamingResponse(
  };
  }
 
+ /** Returns a stream of responses even though onComplete() wasn't triggered */
+ public static <T> Answer<T> streamingResponseWithoutOnComplete(final T[] response) {
+ return invocation -> {
+ Object[] args = invocation.getArguments();
+ ResponseObserver<T> observer = (ResponseObserver<T>) args[1];
+ for (T resp : response) {
+ observer.onResponse(resp);
+ }
+ return null;
+ };
+ }
+
  public static ApiFuture<CommitResponse> commitResponse(int adds, int deletes) {
  CommitResponse.Builder commitResponse = CommitResponse.newBuilder();
  commitResponse.getCommitTimeBuilder().setSeconds(0).setNanos(0);

google-cloud-firestore/src/test/java/com/google/cloud/firestore/QueryTest.java

@@ -27,6 +27,7 @@
 import static com.google.cloud.firestore.LocalFirestoreHelper.order;
 import static com.google.cloud.firestore.LocalFirestoreHelper.query;
 import static com.google.cloud.firestore.LocalFirestoreHelper.queryResponse;
+import static com.google.cloud.firestore.LocalFirestoreHelper.queryResponseWithDone;
 import static com.google.cloud.firestore.LocalFirestoreHelper.reference;
 import static com.google.cloud.firestore.LocalFirestoreHelper.select;
 import static com.google.cloud.firestore.LocalFirestoreHelper.startAt;
@@ -952,6 +953,86 @@ public void onCompleted() {
  semaphore.acquire();
  }
 
+ @Test
+ public void successfulReturnWithoutOnComplete() throws Exception {
+ doAnswer(
+ queryResponseWithDone(
+ /* callWithoutOnComplete */ true, DOCUMENT_NAME + "1", DOCUMENT_NAME + "2"))
+ .when(firestoreMock)
+ .streamRequest(
+ runQuery.capture(),
+ streamObserverCapture.capture(),
+ Matchers.<ServerStreamingCallable>any());
+
+ final Semaphore semaphore = new Semaphore(0);
+ final Iterator<String> iterator = Arrays.asList("doc1", "doc2").iterator();
+
+ query.stream(
+ new ApiStreamObserver<DocumentSnapshot>() {
+ @Override
+ public void onNext(DocumentSnapshot documentSnapshot) {
+ assertEquals(iterator.next(), documentSnapshot.getId());
+ }
+
+ @Override
+ public void onError(Throwable throwable) {
+ fail();
+ }
+
+ @Override
+ public void onCompleted() {
+ semaphore.release();
+ }
+ });
+
+ semaphore.acquire();
+ }
+
+ @Test
+ /**
+ * onComplete() will be called twice. The first time is when it detects RunQueryResponse.done set
+ * to true. The second time is when it receives half close
+ */
+ public void successfulReturnCallsOnCompleteTwice() throws Exception {
+ doAnswer(
+ queryResponseWithDone(
+ /* callWithoutOnComplete */ false, DOCUMENT_NAME + "1", DOCUMENT_NAME + "2"))
+ .when(firestoreMock)
+ .streamRequest(
+ runQuery.capture(),
+ streamObserverCapture.capture(),
+ Matchers.<ServerStreamingCallable>any());
+
+ final Semaphore semaphore = new Semaphore(0);
+ final Iterator<String> iterator = Arrays.asList("doc1", "doc2").iterator();
+ final int[] counter = {0};
+
+ query.stream(
+ new ApiStreamObserver<DocumentSnapshot>() {
+ @Override
+ public void onNext(DocumentSnapshot documentSnapshot) {
+ assertEquals(iterator.next(), documentSnapshot.getId());
+ }
+
+ @Override
+ public void onError(Throwable throwable) {
+ fail();
+ }
+
+ @Override
+ public void onCompleted() {
+ counter[0]++;
+ semaphore.release();
+ }
+ });
+
+ semaphore.acquire();
+
+ // Wait for some time to see whether onCompleted() has been called more than once
+ Thread.sleep(200);
+ assertEquals(1, counter[0]);
+ }
+
  @Test
  public void retriesAfterRetryableError() throws Exception {
  final boolean[] returnError = new boolean[] {true};