fix: retry watch on InternalException (#875)

schmidt-sebastian · web-flow · commit a76a0fddbcd8 · 2022-02-22T10:18:54.000-07:00
Fixes #866
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Watch.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Watch.java
@@ -570,6 +570,9 @@ private static Status getStatus(Throwable throwable) {
       return ((StatusRuntimeException) throwable).getStatus();
     } else if (throwable instanceof StatusException) {
       return ((StatusException) throwable).getStatus();
+    } else if (throwable instanceof ApiException
+        && ((ApiException) throwable).getStatusCode().getTransportCode() instanceof Code) {
+      return ((Code) ((ApiException) throwable).getStatusCode().getTransportCode()).toStatus();
     }
     return null;
   }
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/WatchTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/WatchTest.java
@@ -32,9 +32,11 @@
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 
+import com.google.api.gax.grpc.GrpcStatusCode;
 import com.google.api.gax.rpc.BidiStreamObserver;
 import com.google.api.gax.rpc.BidiStreamingCallable;
 import com.google.api.gax.rpc.ClientStream;
+import com.google.api.gax.rpc.InternalException;
 import com.google.cloud.firestore.Query.Direction;
 import com.google.cloud.firestore.WatchTest.SnapshotDocument.ChangeType;
 import com.google.cloud.firestore.spi.v1.FirestoreRpc;
@@ -393,6 +395,19 @@ public void queryWatchRetriesBasedOnErrorCode() throws InterruptedException {
     }
   }
 
+  @Test
+  public void queryWatchRetriesOnInternalException() throws InterruptedException {
+    addQueryListener();
+    awaitAddTarget();
+    send(addTarget());
+    destroy(new InternalException(null, GrpcStatusCode.of(Code.INTERNAL), true));
+    awaitAddTarget();
+    send(addTarget());
+    send(current());
+    send(snapshot());
+    awaitQuerySnapshot();
+  }
+
   @Test
   public void queryWatchHandlesDocumentChange() throws InterruptedException {
     addQueryListener();
@@ -991,7 +1006,11 @@ private void send(ListenResponse response) {
   }
 
   private void destroy(Code code) {
-    streamObserverCapture.getValue().onError(new StatusException(io.grpc.Status.fromCode(code)));
+    destroy(new StatusException(io.grpc.Status.fromCode(code)));
+  }
+
+  private void destroy(Exception e) {
+    streamObserverCapture.getValue().onError(e);
   }
 
   private void close() {
