Make OperationFutureImpl support ApiFutures.catching()

gax/src/main/java/com/google/api/gax/longrunning/OperationFuture.java

@@ -55,6 +55,9 @@ public interface OperationFuture<ResponseT, MetadataT> extends ApiFuture<Respons
    */
   ApiFuture<OperationSnapshot> getInitialFuture();
 
+  /** Returns the {@link RetryingFuture} which continues to poll {@link OperationSnapshot}. */
+  RetryingFuture<OperationSnapshot> getPollingFuture();
+
   /**
    * Peeks at the metadata of the operation tracked by this {@link OperationFuture}. If the initial
    * future hasn't completed yet this method returns {@code null}, otherwise it returns the latest

gax/src/main/java/com/google/api/gax/longrunning/OperationFutureImpl.java

@@ -81,6 +81,22 @@ public OperationFutureImpl(
     this.metadataTransformer = checkNotNull(metadataTransformer);
   }
 
+  public OperationFutureImpl(
+      RetryingFuture<OperationSnapshot> pollingFuture,
+      ApiFuture<OperationSnapshot> initialFuture,
+      ApiFunction<OperationSnapshot, ResponseT> responseTransformer,
+      ApiFunction<OperationSnapshot, MetadataT> metadataTransformer,
+      ApiFunction<Exception, ResponseT> exceptionTransformer) {
+    this.pollingFuture = checkNotNull(pollingFuture);
+    this.initialFuture = checkNotNull(initialFuture);
+    this.resultFuture =
+        ApiFutures.catching(
+            ApiFutures.transform(pollingFuture, responseTransformer),
+            Exception.class,
+            exceptionTransformer);
+    this.metadataTransformer = checkNotNull(metadataTransformer);
+  }
+
   @Override
   public void addListener(Runnable listener, Executor executor) {
     pollingFuture.addListener(listener, executor);
@@ -103,15 +119,13 @@ public boolean isDone() {
 
   @Override
   public ResponseT get() throws InterruptedException, ExecutionException {
-    pollingFuture.get();
     return resultFuture.get();
   }
 
   @Override
   public ResponseT get(long timeout, TimeUnit unit)
       throws InterruptedException, ExecutionException, TimeoutException {
-    pollingFuture.get(timeout, unit);
-    return resultFuture.get();
+    return resultFuture.get(timeout, unit);
   }
 
   @Override
@@ -124,6 +138,11 @@ public ApiFuture<OperationSnapshot> getInitialFuture() {
     return initialFuture;
   }
 
+  @Override
+  public RetryingFuture<OperationSnapshot> getPollingFuture() {
+    return pollingFuture;
+  }
+
   // Note, the following two methods are not duplicates of each other even though code checking
   // tools may indicate so. They assign multiple different class fields.
   @Override

gax/src/main/java/com/google/api/gax/longrunning/OperationFutures.java

@@ -32,6 +32,7 @@
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
 import com.google.api.core.BetaApi;
+import com.google.api.gax.retrying.RetryingFuture;
 import com.google.api.gax.rpc.ApiException;
 import com.google.api.gax.rpc.StatusCode;
 import com.google.common.base.Preconditions;
@@ -85,6 +86,11 @@ public ApiFuture<OperationSnapshot> getInitialFuture() {
         return initialFuture;
       }
 
+      @Override
+      public RetryingFuture<OperationSnapshot> getPollingFuture() {
+        throw new UnsupportedOperationException("Not implemented: getPollingFuture().");
+      }
+
       @Override
       public void addListener(Runnable runnable, Executor executor) {
         initialFuture.addListener(runnable, executor);

gax/src/test/java/com/google/api/gax/rpc/testing/FakeOperationApi.java

@@ -35,6 +35,7 @@
 import com.google.api.core.InternalApi;
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.api.gax.longrunning.OperationSnapshot;
+import com.google.api.gax.retrying.RetryingFuture;
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.OperationCallable;
 import java.util.HashMap;
@@ -126,6 +127,11 @@ public ApiFuture<OperationSnapshot> getInitialFuture() {
       throw new UnsupportedOperationException("getInitialFuture() not implemented");
     }
 
+    @Override
+    public RetryingFuture<OperationSnapshot> getPollingFuture() {
+      throw new UnsupportedOperationException("getPollingFuture() not implemented");
+    }
+
     @Override
     public ApiFuture<Long> peekMetadata() {
       return null;