feat(gax): implement queryStatus in HttpJsonResumableUploadClient

Author: whowes

sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonResumableUploadClient.java

@@ -35,6 +35,8 @@
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.resumable.ChunkUploadRequest;
 import com.google.api.gax.resumable.ChunkUploadResponse;
+import com.google.api.gax.resumable.QueryStatusRequest;
+import com.google.api.gax.resumable.QueryStatusResponse;
 import com.google.api.gax.resumable.ResumableUploadClient;
 import com.google.api.gax.resumable.ResumableUploadSession;
 import com.google.api.gax.rpc.ApiCallContext;
@@ -89,6 +91,9 @@ public final class HttpJsonResumableUploadClient<RequestT, ResponseT>
 
   private static final PathTemplate PATH_TEMPLATE = PathTemplate.create("{+path}");
 
+  private static final Map<String, List<String>> QUERY_STATUS_HEADERS =
+      ImmutableMap.of(UPLOAD_COMMAND_HEADER, ImmutableList.of("query"));
+
   private static final ApiMethodDescriptor<ChunkUploadRequest, String> UPLOAD_CHUNK_DESCRIPTOR =
       ApiMethodDescriptor.<ChunkUploadRequest, String>newBuilder()
           .setFullMethodName("ResumableUpload/UploadChunk")
@@ -118,6 +123,36 @@ public PathTemplate getPathTemplate() {
               })
           .setResponseParser(ResumableUploadResponseParser.create())
           .build();
+
+  private static final ApiMethodDescriptor<QueryStatusRequest, String> QUERY_STATUS_DESCRIPTOR =
+      ApiMethodDescriptor.<QueryStatusRequest, String>newBuilder()
+          .setFullMethodName("ResumableUpload/QueryStatus")
+          .setHttpMethod(HttpMethods.POST)
+          .setType(ApiMethodDescriptor.MethodType.UNARY)
+          .setRequestFormatter(
+              new HttpRequestFormatter<QueryStatusRequest>() {
+                @Override
+                public Map<String, List<String>> getQueryParamNames(QueryStatusRequest request) {
+                  return Collections.emptyMap();
+                }
+
+                @Override
+                public String getRequestBody(QueryStatusRequest request) {
+                  return "";
+                }
+
+                @Override
+                public String getPath(QueryStatusRequest request) {
+                  return request.getUploadUrl();
+                }
+
+                @Override
+                public PathTemplate getPathTemplate() {
+                  return PATH_TEMPLATE;
+                }
+              })
+          .setResponseParser(ResumableUploadResponseParser.create())
+          .build();
   private final ClientContext clientContext;
   private final ApiMethodDescriptor<RequestT, String> startUploadDescriptor;
   private final HttpResponseParser<ResponseT> responseParser;
@@ -212,6 +247,35 @@ public ApiFuture<ChunkUploadResponse<ResponseT>> futureCall(
     };
   }
 
+  @Override
+  public UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> queryStatusCallable() {
+    return new UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>() {
+      @Override
+      public ApiFuture<QueryStatusResponse<ResponseT>> futureCall(
+          QueryStatusRequest request, @Nullable ApiCallContext inputContext) {
+        Preconditions.checkNotNull(request);
+        HttpJsonCallContext context =
+            (HttpJsonCallContext)
+                HttpJsonCallContext.createDefault()
+                    .nullToSelf(clientContext.getDefaultCallContext())
+                    .merge(inputContext)
+                    .withExtraHeaders(QUERY_STATUS_HEADERS);
+
+        HttpJsonClientCall<QueryStatusRequest, String> clientCall =
+            HttpJsonClientCalls.newCall(QUERY_STATUS_DESCRIPTOR, context);
+
+        SettableApiFuture<QueryStatusResponse<ResponseT>> future = SettableApiFuture.create();
+        HttpJsonClientCalls.startUnaryCall(
+            clientCall,
+            request,
+            context,
+            new QueryStatusResponseListener<>(future, responseParser));
+
+        return future;
+      }
+    };
+  }
+
   private static class StartUploadResponseListener extends HttpJsonClientCall.Listener<String> {
 
     private final SettableApiFuture<ResumableUploadSession> future;
@@ -300,25 +364,15 @@ private static class ChunkUploadResponseListener<ResponseT>
 
     @Override
     public void onHeaders(HttpJsonMetadata responseHeaders) {
-      Map<String, Object> headers = responseHeaders.getHeaders();
-
-      String statusStr = HttpHeadersUtils.getSingleHeader(headers, UPLOAD_STATUS_HEADER);
+      String statusStr =
+          HttpHeadersUtils.getSingleHeader(responseHeaders.getHeaders(), UPLOAD_STATUS_HEADER);
       if (statusStr != null) {
         this.hasUploadStatusHeader = true;
-        if (STATUS_FINAL.equalsIgnoreCase(statusStr)) {
-          this.isComplete = true;
-        }
+        this.isComplete = STATUS_FINAL.equalsIgnoreCase(statusStr);
       }
-
-      String sizeReceivedStr =
-          HttpHeadersUtils.getSingleHeader(headers, UPLOAD_SIZE_RECEIVED_HEADER);
-      if (!Strings.isNullOrEmpty(sizeReceivedStr)) {
-        try {
-          this.committedOffset = Long.parseLong(sizeReceivedStr);
-        } catch (NumberFormatException ignored) {
-          // Ignore invalid/malformed size received header and fall back to local offset
-          // calculation.
-        }
+      Long sizeReceived = parseSizeReceived(responseHeaders);
+      if (sizeReceived != null) {
+        this.committedOffset = sizeReceived;
       }
     }
 
@@ -367,12 +421,93 @@ public void onClose(int statusCode, HttpJsonMetadata trailers) {
     }
   }
 
+  private static class QueryStatusResponseListener<ResponseT>
+      extends HttpJsonClientCall.Listener<String> {
+
+    private final SettableApiFuture<QueryStatusResponse<ResponseT>> future;
+    private final HttpResponseParser<ResponseT> responseParser;
+    private boolean isComplete = false;
+    @Nullable private Long committedOffset = null;
+    private String responseBody = "";
+
+    QueryStatusResponseListener(
+        SettableApiFuture<QueryStatusResponse<ResponseT>> future,
+        HttpResponseParser<ResponseT> responseParser) {
+      this.future = future;
+      this.responseParser = responseParser;
+    }
+
+    @Override
+    public void onHeaders(HttpJsonMetadata responseHeaders) {
+      this.isComplete = isUploadFinal(responseHeaders);
+      this.committedOffset = parseSizeReceived(responseHeaders);
+    }
+
+    @Override
+    public void onMessage(@Nullable String message) {
+      if (message != null) {
+        this.responseBody = message;
+      }
+    }
+
+    @Override
+    public void onClose(int statusCode, HttpJsonMetadata trailers) {
+      try {
+        if (statusCode >= 200 && statusCode < 300) {
+          if (isComplete || committedOffset != null) {
+            ResponseT response = null;
+            if (isComplete) {
+              InputStream stream =
+                  new ByteArrayInputStream(responseBody.getBytes(StandardCharsets.UTF_8));
+              response = responseParser.parse(stream);
+            }
+            future.set(
+                QueryStatusResponse.create(
+                    committedOffset != null ? committedOffset : 0L, isComplete, response));
+          } else {
+            future.setException(
+                ApiExceptionFactory.createException(
+                    "Query status response did not contain valid X-Goog-Upload-Size-Received"
+                        + " header",
+                    /* cause= */ null,
+                    HttpJsonStatusCode.of(StatusCode.Code.INTERNAL),
+                    /* retryable= */ false));
+          }
+        } else {
+          future.setException(
+              createApiException(statusCode, trailers, "Failed to query upload status"));
+        }
+      } catch (Exception e) {
+        future.setException(
+            ApiExceptionFactory.createException(
+                "Internal error processing query status response",
+                e,
+                HttpJsonStatusCode.of(StatusCode.Code.INTERNAL),
+                /* retryable= */ false));
+      }
+    }
+  }
+
   private static boolean isUploadFinal(HttpJsonMetadata responseHeaders) {
     String statusStr =
         HttpHeadersUtils.getSingleHeader(responseHeaders.getHeaders(), UPLOAD_STATUS_HEADER);
     return STATUS_FINAL.equalsIgnoreCase(statusStr);
   }
 
+  @Nullable
+  private static Long parseSizeReceived(HttpJsonMetadata responseHeaders) {
+    String sizeReceivedStr =
+        HttpHeadersUtils.getSingleHeader(responseHeaders.getHeaders(), UPLOAD_SIZE_RECEIVED_HEADER);
+    if (!Strings.isNullOrEmpty(sizeReceivedStr)) {
+      try {
+        return Long.parseLong(sizeReceivedStr);
+      } catch (NumberFormatException ignored) {
+        // Unparseable header; return null and let the listener decide how to handle it.
+      }
+    }
+    return null;
+  }
+
   private static ApiException createApiException(
       int statusCode, @Nullable HttpJsonMetadata trailers, String actionDescription) {
     Throwable cause = trailers != null ? trailers.getException() : null;