feat(operations): Add WaitOperation API surface [gax-java] (#1284)

* fix(lro): Add Operation name to get, list requests

* fix: add name header to {Cancel,Delete}Operation

* feat(operations): Add WaitOperation API surface

* fix: update OperationsClient comments

* fix: Add override annotations to GrpcOperationsStub

* fix: update comment in OperationsClient

* fix: comment cleanup

* fix: MockOperationsImpl error message

Author: miraleung

gax-grpc/src/main/java/com/google/longrunning/OperationsClient.java

@@ -192,7 +192,7 @@ public final Operation getOperation(String name) {
    * }
    * </code></pre>
    *
-   * @param request The request object containing all of the parameters for the API call.
+   * @param request the request object containing all of the parameters for the API call
    * @throws com.google.api.gax.rpc.ApiException if the remote call fails
    */
   private final Operation getOperation(GetOperationRequest request) {
@@ -502,6 +502,61 @@ public final UnaryCallable<DeleteOperationRequest, Empty> deleteOperationCallabl
     return stub.deleteOperationCallable();
   }
 
+  /**
+   * Waits until the specified long-running operation is done or reaches at most a specified
+   * timeout, returning the latest state. If the operation is already done, the latest state is
+   * immediately returned. If the timeout specified is greater than the default HTTP/RPC timeout,
+   * the HTTP/RPC timeout is used. If the server does not support this method, it returns
+   * `google.rpc.Code.UNIMPLEMENTED`. Note that this method is on a best-effort basis. It may return
+   * the latest state before the specified timeout (including immediately), meaning even an
+   * immediate response is no guarantee that the operation is done.
+   *
+   * <p>Sample code:
+   *
+   * <pre><code>
+   * try (OperationsClient operationsClient = OperationsClient.create()) {
+   *   String name = "";
+   *   WaitOperationRequest request = WaitOperationRequest.newBuilder()
+   *     .setName(name)
+   *     .setTimeout(Duration.ofMillis(100))
+   *     .build();
+   *   Operation response = operationsClient.waitOperation(request);
+   * }
+   * </code></pre>
+   *
+   * @param request the request object containing all of the parameters for the API call.
+   * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+   */
+  public final Operation waitOperation(WaitOperationRequest request) {
+    return waitOperationCallable().call(request);
+  }
+
+  /**
+   * Waits until the specified long-running operation is done or reaches at most a specified
+   * timeout, returning the latest state. If the operation is already done, the latest state is
+   * immediately returned. If the timeout specified is greater than the default HTTP/RPC timeout,
+   * the HTTP/RPC timeout is used. If the server does not support this method, it returns
+   * `google.rpc.Code.UNIMPLEMENTED`. Note that this method is on a best-effort basis. It may return
+   * the latest state before the specified timeout (including immediately), meaning even an
+   * immediate response is no guarantee that the operation is done.
+   *
+   * <p>Sample code:
+   *
+   * <pre><code>
+   * try (OperationsClient operationsClient = OperationsClient.create()) {
+   *   String name = "";
+   *   WaitOperationRequest request = WaitOperationRequest.newBuilder()
+   *     .setName(name)
+   *     .setTimeout(Duration.ofMillis(100))
+   *     .build();
+   *   ApiFuture&lt;Operation&gt; future = operationsClient.waitOperationCallable().futureCall(request);
+   * }
+   * </code></pre>
+   */
+  public final UnaryCallable<WaitOperationRequest, Operation> waitOperationCallable() {
+    return stub.waitOperationCallable();
+  }
+
   @Override
   public final void close() {
     stub.close();

gax-grpc/src/main/java/com/google/longrunning/OperationsSettings.java

@@ -68,6 +68,11 @@ public UnaryCallSettings<DeleteOperationRequest, Empty> deleteOperationSettings(
     return ((OperationsStubSettings) getStubSettings()).deleteOperationSettings();
   }
 
+  /** Returns the object with the settings used for calls to waitOperation. */
+  public UnaryCallSettings<WaitOperationRequest, Operation> waitOperationSettings() {
+    return ((OperationsStubSettings) getStubSettings()).waitOperationSettings();
+  }
+
   public static final OperationsSettings create(OperationsStubSettings stub) throws IOException {
     return new OperationsSettings.Builder(stub.toBuilder()).build();
   }
@@ -166,6 +171,11 @@ public UnaryCallSettings.Builder<DeleteOperationRequest, Empty> deleteOperationS
       return getStubSettingsBuilder().deleteOperationSettings();
     }
 
+    /** Returns the builder for the settings used for calls to deleteOperation. */
+    public UnaryCallSettings.Builder<WaitOperationRequest, Operation> waitOperationSettings() {
+      return getStubSettingsBuilder().waitOperationSettings();
+    }
+
     @Override
     public OperationsSettings build() throws IOException {
       return new OperationsSettings(this);

gax-grpc/src/main/java/com/google/longrunning/stub/GrpcOperationsStub.java

@@ -46,6 +46,7 @@
 import com.google.longrunning.ListOperationsRequest;
 import com.google.longrunning.ListOperationsResponse;
 import com.google.longrunning.Operation;
+import com.google.longrunning.WaitOperationRequest;
 import com.google.protobuf.Empty;
 import io.grpc.MethodDescriptor;
 import io.grpc.protobuf.ProtoUtils;
@@ -97,6 +98,15 @@ public class GrpcOperationsStub extends OperationsStub {
                   ProtoUtils.marshaller(DeleteOperationRequest.getDefaultInstance()))
               .setResponseMarshaller(ProtoUtils.marshaller(Empty.getDefaultInstance()))
               .build();
+  private static final MethodDescriptor<WaitOperationRequest, Operation>
+      waitOperationMethodDescriptor =
+          MethodDescriptor.<WaitOperationRequest, Operation>newBuilder()
+              .setType(MethodDescriptor.MethodType.UNARY)
+              .setFullMethodName("google.longrunning.Operations/WaitOperation")
+              .setRequestMarshaller(
+                  ProtoUtils.marshaller(WaitOperationRequest.getDefaultInstance()))
+              .setResponseMarshaller(ProtoUtils.marshaller(Operation.getDefaultInstance()))
+              .build();
 
   private final BackgroundResource backgroundResources;
 
@@ -106,6 +116,7 @@ public class GrpcOperationsStub extends OperationsStub {
       listOperationsPagedCallable;
   private final UnaryCallable<CancelOperationRequest, Empty> cancelOperationCallable;
   private final UnaryCallable<DeleteOperationRequest, Empty> deleteOperationCallable;
+  private final UnaryCallable<WaitOperationRequest, Operation> waitOperationCallable;
 
   private final GrpcStubCallableFactory callableFactory;
 
@@ -199,6 +210,19 @@ public Map<String, String> extract(DeleteOperationRequest request) {
                   }
                 })
             .build();
+    GrpcCallSettings<WaitOperationRequest, Operation> waitOperationTransportSettings =
+        GrpcCallSettings.<WaitOperationRequest, Operation>newBuilder()
+            .setMethodDescriptor(waitOperationMethodDescriptor)
+            .setParamsExtractor(
+                new RequestParamsExtractor<WaitOperationRequest>() {
+                  @Override
+                  public Map<String, String> extract(WaitOperationRequest request) {
+                    ImmutableMap.Builder<String, String> params = ImmutableMap.builder();
+                    params.put("name", String.valueOf(request.getName()));
+                    return params.build();
+                  }
+                })
+            .build();
 
     this.getOperationCallable =
         callableFactory.createUnaryCallable(
@@ -215,31 +239,44 @@ public Map<String, String> extract(DeleteOperationRequest request) {
     this.deleteOperationCallable =
         callableFactory.createUnaryCallable(
             deleteOperationTransportSettings, settings.deleteOperationSettings(), clientContext);
+    this.waitOperationCallable =
+        callableFactory.createUnaryCallable(
+            waitOperationTransportSettings, settings.waitOperationSettings(), clientContext);
 
     backgroundResources = new BackgroundResourceAggregation(clientContext.getBackgroundResources());
   }
 
+  @Override
   public UnaryCallable<GetOperationRequest, Operation> getOperationCallable() {
     return getOperationCallable;
   }
 
+  @Override
   public UnaryCallable<ListOperationsRequest, ListOperationsPagedResponse>
       listOperationsPagedCallable() {
     return listOperationsPagedCallable;
   }
 
+  @Override
   public UnaryCallable<ListOperationsRequest, ListOperationsResponse> listOperationsCallable() {
     return listOperationsCallable;
   }
 
+  @Override
   public UnaryCallable<CancelOperationRequest, Empty> cancelOperationCallable() {
     return cancelOperationCallable;
   }
 
+  @Override
   public UnaryCallable<DeleteOperationRequest, Empty> deleteOperationCallable() {
     return deleteOperationCallable;
   }
 
+  @Override
+  public UnaryCallable<WaitOperationRequest, Operation> waitOperationCallable() {
+    return waitOperationCallable;
+  }
+
   @Override
   public final void close() {
     shutdown();

gax-grpc/src/main/java/com/google/longrunning/stub/OperationsStub.java

@@ -40,6 +40,7 @@
 import com.google.longrunning.ListOperationsRequest;
 import com.google.longrunning.ListOperationsResponse;
 import com.google.longrunning.Operation;
+import com.google.longrunning.WaitOperationRequest;
 import com.google.protobuf.Empty;
 
 /**
@@ -71,6 +72,10 @@ public UnaryCallable<DeleteOperationRequest, Empty> deleteOperationCallable() {
     throw new UnsupportedOperationException("Not implemented: deleteOperationCallable()");
   }
 
+  public UnaryCallable<WaitOperationRequest, Operation> waitOperationCallable() {
+    throw new UnsupportedOperationException("Not implemented: waitOperationCallable()");
+  }
+
   @Override
   public abstract void close();
 }

gax-grpc/src/main/java/com/google/longrunning/stub/OperationsStubSettings.java

@@ -61,6 +61,7 @@
 import com.google.longrunning.ListOperationsRequest;
 import com.google.longrunning.ListOperationsResponse;
 import com.google.longrunning.Operation;
+import com.google.longrunning.WaitOperationRequest;
 import com.google.protobuf.Empty;
 import java.io.IOException;
 import org.threeten.bp.Duration;
@@ -75,6 +76,7 @@ public class OperationsStubSettings extends StubSettings<OperationsStubSettings>
       listOperationsSettings;
   private final UnaryCallSettings<CancelOperationRequest, Empty> cancelOperationSettings;
   private final UnaryCallSettings<DeleteOperationRequest, Empty> deleteOperationSettings;
+  private final UnaryCallSettings<WaitOperationRequest, Operation> waitOperationSettings;
 
   /** Returns the object with the settings used for calls to getOperation. */
   public UnaryCallSettings<GetOperationRequest, Operation> getOperationSettings() {
@@ -98,6 +100,11 @@ public UnaryCallSettings<DeleteOperationRequest, Empty> deleteOperationSettings(
     return deleteOperationSettings;
   }
 
+  /** Returns the object with the settings used for calls to waitOperation. */
+  public UnaryCallSettings<WaitOperationRequest, Operation> waitOperationSettings() {
+    return waitOperationSettings;
+  }
+
   @BetaApi("A restructuring of stub classes is planned, so this may break in the future")
   public OperationsStub createStub() throws IOException {
     if (getTransportChannelProvider()
@@ -151,6 +158,7 @@ protected OperationsStubSettings(Builder settingsBuilder) throws IOException {
     listOperationsSettings = settingsBuilder.listOperationsSettings().build();
     cancelOperationSettings = settingsBuilder.cancelOperationSettings().build();
     deleteOperationSettings = settingsBuilder.deleteOperationSettings().build();
+    waitOperationSettings = settingsBuilder.waitOperationSettings().build();
   }
 
   private static final PagedListDescriptor<ListOperationsRequest, ListOperationsResponse, Operation>
@@ -215,6 +223,7 @@ public static class Builder extends StubSettings.Builder<OperationsStubSettings,
         listOperationsSettings;
     private final UnaryCallSettings.Builder<CancelOperationRequest, Empty> cancelOperationSettings;
     private final UnaryCallSettings.Builder<DeleteOperationRequest, Empty> deleteOperationSettings;
+    private final UnaryCallSettings.Builder<WaitOperationRequest, Operation> waitOperationSettings;
 
     private static final ImmutableMap<String, ImmutableSet<StatusCode.Code>>
         RETRYABLE_CODE_DEFINITIONS;
@@ -265,6 +274,8 @@ protected Builder(ClientContext clientContext) {
 
       deleteOperationSettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
 
+      waitOperationSettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
+
       unaryMethodSettingsBuilders =
           ImmutableList.<UnaryCallSettings.Builder<?, ?>>of(
               getOperationSettings,
@@ -302,6 +313,11 @@ private static Builder initDefaults(Builder builder) {
           .setRetryableCodes(RETRYABLE_CODE_DEFINITIONS.get("idempotent"))
           .setRetrySettings(RETRY_PARAM_DEFINITIONS.get("default"));
 
+      builder
+          .waitOperationSettings()
+          .setRetryableCodes(RETRYABLE_CODE_DEFINITIONS.get("idempotent"))
+          .setRetrySettings(RETRY_PARAM_DEFINITIONS.get("default"));
+
       return builder;
     }
 
@@ -312,13 +328,15 @@ protected Builder(OperationsStubSettings settings) {
       listOperationsSettings = settings.listOperationsSettings.toBuilder();
       cancelOperationSettings = settings.cancelOperationSettings.toBuilder();
       deleteOperationSettings = settings.deleteOperationSettings.toBuilder();
+      waitOperationSettings = settings.waitOperationSettings.toBuilder();
 
       unaryMethodSettingsBuilders =
           ImmutableList.<UnaryCallSettings.Builder<?, ?>>of(
               getOperationSettings,
               listOperationsSettings,
               cancelOperationSettings,
-              deleteOperationSettings);
+              deleteOperationSettings,
+              waitOperationSettings);
     }
 
     /**
@@ -358,6 +376,11 @@ public UnaryCallSettings.Builder<DeleteOperationRequest, Empty> deleteOperationS
       return deleteOperationSettings;
     }
 
+    /** Returns the builder for the settings used for calls to waitOperation. */
+    public UnaryCallSettings.Builder<WaitOperationRequest, Operation> waitOperationSettings() {
+      return waitOperationSettings;
+    }
+
     @Override
     public OperationsStubSettings build() throws IOException {
       return new OperationsStubSettings(this);

gax-grpc/src/test/java/com/google/longrunning/MockOperationsImpl.java

@@ -130,4 +130,25 @@ public void cancelOperation(
       responseObserver.onError(new IllegalArgumentException("Unrecognized response type"));
     }
   }
+
+  @Override
+  public void waitOperation(
+      WaitOperationRequest request, StreamObserver<Operation> responseObserver) {
+    Object response = responses.remove();
+    if (response instanceof Operation) {
+      requests.add(request);
+      responseObserver.onNext((Operation) response);
+      responseObserver.onCompleted();
+    } else if (response instanceof Exception) {
+      responseObserver.onError((Exception) response);
+    } else {
+      responseObserver.onError(
+          new IllegalArgumentException(
+              String.format(
+                  "Unrecognized response type %s, expected %s or %s",
+                  response.getClass().getName(),
+                  Operation.class.getName(),
+                  Exception.class.getName())));
+    }
+  }
 }

gax-grpc/src/test/java/com/google/longrunning/OperationsClientTest.java

@@ -38,6 +38,7 @@
 import com.google.api.gax.rpc.InvalidArgumentException;
 import com.google.common.collect.Lists;
 import com.google.protobuf.AbstractMessage;
+import com.google.protobuf.Duration;
 import com.google.protobuf.Empty;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
@@ -235,4 +236,46 @@ public void deleteOperationExceptionTest() throws Exception {
       // Expected exception
     }
   }
+
+  @Test
+  @SuppressWarnings("all")
+  public void waitOperationTest() {
+    String name2 = "name2-1052831874";
+    boolean done = true;
+    Operation expectedResponse = Operation.newBuilder().setName(name2).setDone(done).build();
+    mockOperations.addResponse(expectedResponse);
+
+    String name = "name3373707";
+    Duration timeout = Duration.newBuilder().setSeconds(5).build();
+    WaitOperationRequest request =
+        WaitOperationRequest.newBuilder().setName(name).setTimeout(timeout).build();
+
+    Operation actualResponse = client.waitOperation(request);
+    Assert.assertEquals(expectedResponse, actualResponse);
+
+    List<AbstractMessage> actualRequests = mockOperations.getRequests();
+    Assert.assertEquals(1, actualRequests.size());
+    WaitOperationRequest actualRequest = (WaitOperationRequest) actualRequests.get(0);
+
+    Assert.assertEquals(name, actualRequest.getName());
+  }
+
+  @Test
+  @SuppressWarnings("all")
+  public void waitOperationExceptionTest() throws Exception {
+    StatusRuntimeException exception = new StatusRuntimeException(Status.INVALID_ARGUMENT);
+    mockOperations.addException(exception);
+
+    try {
+      String name = "name3373707";
+      Duration timeout = Duration.newBuilder().setSeconds(5).build();
+      WaitOperationRequest request =
+          WaitOperationRequest.newBuilder().setName(name).setTimeout(timeout).build();
+
+      client.waitOperation(request);
+      Assert.fail("No exception raised");
+    } catch (InvalidArgumentException e) {
+      // Expected exception
+    }
+  }
 }