feat: Add REST AIP-151 LRO support (#1484)

* feat: Add REST AIP-151 LRO support

* reformat code

gax-grpc/src/main/java/com/google/api/gax/grpc/ProtoOperationTransformers.java

@@ -65,6 +65,11 @@ public ResponseT apply(OperationSnapshot operationSnapshot) {
             operationSnapshot.getErrorCode(),
             false);
       }
+
+      if (!(operationSnapshot.getResponse() instanceof Any)) {
+        return (ResponseT) operationSnapshot.getResponse();
+      }
+
       try {
         return transformer.apply((Any) operationSnapshot.getResponse());
       } catch (RuntimeException e) {
@@ -94,9 +99,11 @@ private MetadataTransformer(Class<MetadataT> packedClass) {
 
     @Override
     public MetadataT apply(OperationSnapshot operationSnapshot) {
+      if (!(operationSnapshot.getMetadata() instanceof Any)) {
+        return (MetadataT) operationSnapshot.getMetadata();
+      }
       try {
-        return transformer.apply(
-            operationSnapshot.getMetadata() != null ? (Any) operationSnapshot.getMetadata() : null);
+        return transformer.apply((Any) operationSnapshot.getMetadata());
       } catch (RuntimeException e) {
         throw ApiExceptionFactory.createException(
             "Polling operation with name \""

gax-grpc/src/test/java/com/google/api/gax/grpc/ProtoOperationTransformersTest.java

@@ -48,9 +48,8 @@
 
 @RunWith(JUnit4.class)
 public class ProtoOperationTransformersTest {
-
   @Test
-  public void testResponseTransformer() {
+  public void testAnyResponseTransformer() {
     ResponseTransformer<Money> transformer = ResponseTransformer.create(Money.class);
     Money inputMoney = Money.newBuilder().setCurrencyCode("USD").build();
     OperationSnapshot operationSnapshot =
@@ -60,7 +59,7 @@ public void testResponseTransformer() {
   }
 
   @Test
-  public void testResponseTransformer_exception() {
+  public void testAnyResponseTransformer_exception() {
     ResponseTransformer<Money> transformer = ResponseTransformer.create(Money.class);
     Money inputMoney = Money.newBuilder().setCurrencyCode("USD").build();
     Status status = Status.newBuilder().setCode(Code.UNAVAILABLE.value()).build();
@@ -78,7 +77,7 @@ public void testResponseTransformer_exception() {
   }
 
   @Test
-  public void testResponseTransformer_mismatchedTypes() {
+  public void testAnyResponseTransformer_mismatchedTypes() {
     ResponseTransformer<Money> transformer = ResponseTransformer.create(Money.class);
     Status status = Status.newBuilder().setCode(Code.OK.value()).build();
     OperationSnapshot operationSnapshot =
@@ -96,7 +95,7 @@ public void testResponseTransformer_mismatchedTypes() {
   }
 
   @Test
-  public void testMetadataTransformer() {
+  public void testAnyMetadataTransformer() {
     MetadataTransformer<Money> transformer = MetadataTransformer.create(Money.class);
     Money inputMoney = Money.newBuilder().setCurrencyCode("USD").build();
     OperationSnapshot operationSnapshot =
@@ -106,7 +105,7 @@ public void testMetadataTransformer() {
   }
 
   @Test
-  public void testMetadataTransformer_mismatchedTypes() {
+  public void testAnyMetadataTransformer_mismatchedTypes() {
     MetadataTransformer<Money> transformer = MetadataTransformer.create(Money.class);
     Status status = Status.newBuilder().setCode(Code.OK.value()).build();
     OperationSnapshot operationSnapshot =

gax-httpjson/src/main/java/com/google/api/gax/httpjson/ApiMessageHttpResponseParser.java

@@ -37,6 +37,7 @@
 import com.google.gson.TypeAdapter;
 import com.google.gson.stream.JsonReader;
 import com.google.gson.stream.JsonWriter;
+import com.google.protobuf.TypeRegistry;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.lang.reflect.Type;
@@ -104,6 +105,11 @@ public ResponseT parse(InputStream httpResponseBody) {
     }
   }
 
+  @Override
+  public ResponseT parse(InputStream httpResponseBody, TypeRegistry registry) {
+    return parse(httpResponseBody);
+  }
+
   @Override
   public String serialize(ResponseT response) {
     return getResponseMarshaller().toJson(response);

gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonCallOptions.java

@@ -32,6 +32,7 @@
 import com.google.api.core.BetaApi;
 import com.google.auth.Credentials;
 import com.google.auto.value.AutoValue;
+import com.google.protobuf.TypeRegistry;
 import javax.annotation.Nullable;
 import org.threeten.bp.Instant;
 
@@ -45,6 +46,9 @@ public abstract class HttpJsonCallOptions {
   @Nullable
   public abstract Credentials getCredentials();
 
+  @Nullable
+  public abstract TypeRegistry getTypeRegistry();
+
   public static Builder newBuilder() {
     return new AutoValue_HttpJsonCallOptions.Builder();
   }
@@ -55,6 +59,8 @@ public abstract static class Builder {
 
     public abstract Builder setCredentials(Credentials value);
 
+    public abstract Builder setTypeRegistry(TypeRegistry value);
+
     public abstract HttpJsonCallOptions build();
   }
 }

gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonCallSettings.java

@@ -29,18 +29,27 @@
  */
 package com.google.api.gax.httpjson;
 
+import com.google.protobuf.TypeRegistry;
+
 /** HTTP-specific settings for creating callables. */
 public class HttpJsonCallSettings<RequestT, ResponseT> {
   private final ApiMethodDescriptor<RequestT, ResponseT> methodDescriptor;
+  private final TypeRegistry typeRegistry;
 
-  private HttpJsonCallSettings(ApiMethodDescriptor<RequestT, ResponseT> methodDescriptor) {
+  private HttpJsonCallSettings(
+      ApiMethodDescriptor<RequestT, ResponseT> methodDescriptor, TypeRegistry typeRegistry) {
     this.methodDescriptor = methodDescriptor;
+    this.typeRegistry = typeRegistry;
   }
 
   public ApiMethodDescriptor<RequestT, ResponseT> getMethodDescriptor() {
     return methodDescriptor;
   }
 
+  public TypeRegistry getTypeRegistry() {
+    return typeRegistry;
+  }
+
   public static <RequestT, ResponseT> Builder<RequestT, ResponseT> newBuilder() {
     return new Builder<>();
   }
@@ -58,6 +67,7 @@ public Builder toBuilder() {
 
   public static class Builder<RequestT, ResponseT> {
     private ApiMethodDescriptor<RequestT, ResponseT> methodDescriptor;
+    private TypeRegistry typeRegistry;
 
     private Builder() {}
 
@@ -71,8 +81,13 @@ public Builder<RequestT, ResponseT> setMethodDescriptor(
       return this;
     }
 
+    public Builder<RequestT, ResponseT> setTypeRegistry(TypeRegistry typeRegistry) {
+      this.typeRegistry = typeRegistry;
+      return this;
+    }
+
     public HttpJsonCallSettings<RequestT, ResponseT> build() {
-      return new HttpJsonCallSettings<>(methodDescriptor);
+      return new HttpJsonCallSettings<>(methodDescriptor, typeRegistry);
     }
   }
 }

gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonCallableFactory.java

@@ -61,7 +61,7 @@ private HttpJsonCallableFactory() {}
   private static <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createDirectUnaryCallable(
       HttpJsonCallSettings<RequestT, ResponseT> httpJsonCallSettings) {
     return new HttpJsonDirectCallable<RequestT, ResponseT>(
-        httpJsonCallSettings.getMethodDescriptor());
+        httpJsonCallSettings.getMethodDescriptor(), httpJsonCallSettings.getTypeRegistry());
   }
 
   static <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createUnaryCallable(

gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonDirectCallable.java

@@ -33,6 +33,7 @@
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.common.base.Preconditions;
+import com.google.protobuf.TypeRegistry;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import org.threeten.bp.Instant;
@@ -44,9 +45,16 @@
  */
 class HttpJsonDirectCallable<RequestT, ResponseT> extends UnaryCallable<RequestT, ResponseT> {
   private final ApiMethodDescriptor<RequestT, ResponseT> descriptor;
+  private final TypeRegistry typeRegistry;
 
   HttpJsonDirectCallable(ApiMethodDescriptor<RequestT, ResponseT> descriptor) {
+    this(descriptor, null);
+  }
+
+  HttpJsonDirectCallable(
+      ApiMethodDescriptor<RequestT, ResponseT> descriptor, TypeRegistry typeRegistry) {
     this.descriptor = descriptor;
+    this.typeRegistry = typeRegistry;
   }
 
   @Override
@@ -68,6 +76,7 @@ public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext inputCon
         HttpJsonCallOptions.newBuilder()
             .setDeadline(deadline)
             .setCredentials(context.getCredentials())
+            .setTypeRegistry(typeRegistry)
             .build();
     return context.getChannel().issueFutureUnaryCall(callOptions, request, descriptor);
   }

gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonOperationSnapshot.java

@@ -33,6 +33,8 @@
 import com.google.api.core.InternalApi;
 import com.google.api.gax.longrunning.OperationSnapshot;
 import com.google.api.gax.rpc.StatusCode;
+import com.google.api.gax.rpc.StatusCode.Code;
+import com.google.longrunning.Operation;
 
 /**
  * Implementation of OperationSnapshot based on REST transport.
@@ -42,32 +44,26 @@
 @BetaApi("The surface for long-running operations is not stable yet and may change in the future.")
 @InternalApi
 public class HttpJsonOperationSnapshot implements OperationSnapshot {
-
   private final String name;
   private final Object metadata;
   private final boolean done;
   private final Object response;
   private final StatusCode errorCode;
   private final String errorMessage;
 
-  public HttpJsonOperationSnapshot(
+  private HttpJsonOperationSnapshot(
       String name,
       Object metadata,
       boolean done,
       Object response,
-      int errorCode,
+      StatusCode errorCode,
       String errorMessage) {
     this.name = name;
     this.metadata = metadata;
     this.done = done;
-    this.response = done ? response : null;
-    if (done && errorCode != 0) {
-      this.errorCode = HttpJsonStatusCode.of(errorCode, errorMessage);
-      this.errorMessage = errorMessage;
-    } else {
-      this.errorCode = null;
-      this.errorMessage = null;
-    }
+    this.response = response;
+    this.errorCode = errorCode;
+    this.errorMessage = errorMessage;
   }
 
   /** {@inheritDoc} */
@@ -106,6 +102,10 @@ public String getErrorMessage() {
     return this.errorMessage;
   }
 
+  public static HttpJsonOperationSnapshot create(Operation operation) {
+    return newBuilder().setOperation(operation).build();
+  }
+
   public static Builder newBuilder() {
     return new HttpJsonOperationSnapshot.Builder();
   }
@@ -115,7 +115,7 @@ public static class Builder {
     private Object metadata;
     private boolean done;
     private Object response;
-    private int errorCode;
+    private StatusCode errorCode;
     private String errorMessage;
 
     public Builder setName(String name) {
@@ -139,11 +139,24 @@ public Builder setResponse(Object response) {
     }
 
     public Builder setError(int errorCode, String errorMessage) {
-      this.errorCode = errorCode;
+      this.errorCode =
+          HttpJsonStatusCode.of(
+              errorCode == 0 ? Code.OK.getHttpStatusCode() : errorCode, errorMessage);
       this.errorMessage = errorMessage;
       return this;
     }
 
+    private Builder setOperation(Operation operation) {
+      this.name = operation.getName();
+      this.done = operation.getDone();
+      this.response = operation.getResponse();
+      this.metadata = operation.getMetadata();
+      this.errorCode =
+          HttpJsonStatusCode.of(com.google.rpc.Code.forNumber(operation.getError().getCode()));
+      this.errorMessage = operation.getError().getMessage();
+      return this;
+    }
+
     public HttpJsonOperationSnapshot build() {
       return new HttpJsonOperationSnapshot(name, metadata, done, response, errorCode, errorMessage);
     }

gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonStatusCode.java

@@ -57,6 +57,51 @@ public static HttpJsonStatusCode of(StatusCode.Code statusCode) {
     return new HttpJsonStatusCode(statusCode.getHttpStatusCode(), statusCode);
   }
 
+  public static HttpJsonStatusCode of(com.google.rpc.Code rpcCode) {
+    return new HttpJsonStatusCode(rpcCode.getNumber(), rpcCodeToStatusCode(rpcCode));
+  }
+
+  static StatusCode.Code rpcCodeToStatusCode(com.google.rpc.Code rpcCode) {
+    switch (rpcCode) {
+      case OK:
+        return Code.OK;
+      case CANCELLED:
+        return Code.CANCELLED;
+      case UNKNOWN:
+        return Code.UNKNOWN;
+      case INVALID_ARGUMENT:
+        return Code.INVALID_ARGUMENT;
+      case DEADLINE_EXCEEDED:
+        return Code.DEADLINE_EXCEEDED;
+      case NOT_FOUND:
+        return Code.DEADLINE_EXCEEDED;
+      case ALREADY_EXISTS:
+        return Code.ALREADY_EXISTS;
+      case PERMISSION_DENIED:
+        return Code.PERMISSION_DENIED;
+      case RESOURCE_EXHAUSTED:
+        return Code.RESOURCE_EXHAUSTED;
+      case FAILED_PRECONDITION:
+        return Code.FAILED_PRECONDITION;
+      case ABORTED:
+        return Code.ABORTED;
+      case OUT_OF_RANGE:
+        return Code.OUT_OF_RANGE;
+      case UNIMPLEMENTED:
+        return Code.UNIMPLEMENTED;
+      case INTERNAL:
+        return Code.INTERNAL;
+      case UNAVAILABLE:
+        return Code.UNAVAILABLE;
+      case DATA_LOSS:
+        return Code.DATA_LOSS;
+      case UNAUTHENTICATED:
+        return Code.UNAUTHENTICATED;
+      default:
+        throw new IllegalArgumentException("Unrecognized rpc code: " + rpcCode);
+    }
+  }
+
   static StatusCode.Code httpStatusToStatusCode(int httpStatus, String errorMessage) {
     String causeMessage = Strings.nullToEmpty(errorMessage).toUpperCase();
     switch (httpStatus) {
@@ -143,4 +188,9 @@ public boolean equals(Object o) {
   public int hashCode() {
     return Objects.hash(statusCode);
   }
+
+  @Override
+  public String toString() {
+    return "HttpJsonStatusCode{" + "statusCode=" + statusCode + "}";
+  }
 }

gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpRequestRunnable.java

@@ -195,9 +195,13 @@ public void run() {
             HttpJsonStatusCode.of(httpResponse.getStatusCode(), httpResponse.getStatusMessage()),
             false);
       }
+
       if (getApiMethodDescriptor().getResponseParser() != null) {
         ResponseT response =
-            getApiMethodDescriptor().getResponseParser().parse(httpResponse.getContent());
+            getApiMethodDescriptor()
+                .getResponseParser()
+                .parse(httpResponse.getContent(), getHttpJsonCallOptions().getTypeRegistry());
+
         getResponseFuture().set(response);
       } else {
         getResponseFuture().set(null);