feat(gax): implement uploadChunk in HttpJsonResumableUploadClient

Author: whowes

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

@@ -34,17 +34,23 @@
 import com.google.api.core.ApiFuture;
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
+import com.google.api.gax.resumable.ChunkUploadRequest;
+import com.google.api.gax.resumable.ChunkUploadResponse;
 import com.google.api.gax.resumable.ResumableUploadClient;
 import com.google.api.gax.resumable.ResumableUploadSession;
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.ApiExceptionFactory;
 import com.google.api.gax.rpc.ClientContext;
 import com.google.api.gax.rpc.StatusCode;
 import com.google.api.gax.rpc.UnaryCallable;
+import com.google.api.pathtemplate.PathTemplate;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import java.io.ByteArrayInputStream;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -67,16 +73,54 @@ public final class HttpJsonResumableUploadClient<RequestT, ResponseT>
 
   private static final String UPLOAD_PROTOCOL_HEADER = "X-Goog-Upload-Protocol";
   private static final String UPLOAD_COMMAND_HEADER = "X-Goog-Upload-Command";
+  private static final String UPLOAD_OFFSET_HEADER = "X-Goog-Upload-Offset";
   private static final String UPLOAD_URL_HEADER = "X-Goog-Upload-URL";
   private static final String UPLOAD_GRANULARITY_HEADER = "X-Goog-Upload-Chunk-Granularity";
+  private static final String UPLOAD_STATUS_HEADER = "X-Goog-Upload-Status";
+  private static final String UPLOAD_SIZE_RECEIVED_HEADER = "X-Goog-Upload-Size-Received";
+  private static final String STATUS_FINAL = "final";
 
   private static final Map<String, List<String>> START_UPLOAD_HEADERS =
       ImmutableMap.of(
           UPLOAD_PROTOCOL_HEADER, ImmutableList.of("resumable"),
           UPLOAD_COMMAND_HEADER, ImmutableList.of("start"));
 
+  private static final PathTemplate PATH_TEMPLATE = PathTemplate.create("{+path}");
+
+  private static final ApiMethodDescriptor<ChunkUploadRequest, String> UPLOAD_CHUNK_DESCRIPTOR =
+      ApiMethodDescriptor.<ChunkUploadRequest, String>newBuilder()
+          .setFullMethodName("ResumableUpload/UploadChunk")
+          .setHttpMethod(HttpMethods.POST)
+          .setType(ApiMethodDescriptor.MethodType.UNARY)
+          .setRequestFormatter(
+              new ResumableUploadChunkRequestFormatter<ChunkUploadRequest>() {
+                @Override
+                public Map<String, List<String>> getQueryParamNames(ChunkUploadRequest request) {
+                  return Collections.emptyMap();
+                }
+
+                @Override
+                public byte[] getBinaryRequestBody(ChunkUploadRequest request) {
+                  return request.getPayload().toByteArray();
+                }
+
+                @Override
+                public String getPath(ChunkUploadRequest request) {
+                  return request.getUploadUrl();
+                }
+
+                @Override
+                public PathTemplate getPathTemplate() {
+                  return PATH_TEMPLATE;
+                }
+              })
+          .setResponseParser(ResumableUploadResponseParser.create())
+          .build();
+
   private final ApiMethodDescriptor<RequestT, String> startUploadDescriptor;
   private final UnaryCallable<RequestT, ResumableUploadSession> startUploadCallable;
+  private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
+      uploadChunkCallable;
 
   public static <RequestT, ResponseT> HttpJsonResumableUploadClient<RequestT, ResponseT> create(
       ClientContext clientContext, ApiMethodDescriptor<RequestT, ResponseT> methodDescriptor) {
@@ -87,6 +131,8 @@ private HttpJsonResumableUploadClient(
       ClientContext clientContext, ApiMethodDescriptor<RequestT, ResponseT> methodDescriptor) {
     Preconditions.checkNotNull(clientContext);
     Preconditions.checkNotNull(methodDescriptor);
+    HttpResponseParser<ResponseT> responseParser =
+        Preconditions.checkNotNull(methodDescriptor.getResponseParser());
 
     this.startUploadDescriptor =
         ApiMethodDescriptor.<RequestT, String>newBuilder()
@@ -97,13 +143,19 @@ private HttpJsonResumableUploadClient(
             .setResponseParser(ResumableUploadResponseParser.create())
             .build();
     this.startUploadCallable = createStartUploadCallable(clientContext);
+    this.uploadChunkCallable = createUploadChunkCallable(clientContext, responseParser);
   }
 
   @Override
   public UnaryCallable<RequestT, ResumableUploadSession> startUploadCallable() {
     return startUploadCallable;
   }
 
+  @Override
+  public UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable() {
+    return uploadChunkCallable;
+  }
+
   private UnaryCallable<RequestT, ResumableUploadSession> createStartUploadCallable(
       ClientContext clientContext) {
     UnaryCallable<RequestT, ResumableUploadSession> rawCallable =
@@ -129,6 +181,51 @@ public ApiFuture<ResumableUploadSession> futureCall(
     return createClientCallable(rawCallable, clientContext);
   }
 
+  private UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
+      createUploadChunkCallable(
+          ClientContext clientContext, HttpResponseParser<ResponseT> responseParser) {
+    UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> rawCallable =
+        new UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>() {
+          @Override
+          public ApiFuture<ChunkUploadResponse<ResponseT>> futureCall(
+              ChunkUploadRequest request, @Nullable ApiCallContext inputContext) {
+            Preconditions.checkNotNull(request);
+            boolean isPayloadEmpty = request.getPayload().isEmpty();
+            String command;
+            if (request.isFinal()) {
+              command = !isPayloadEmpty ? "upload, finalize" : "finalize";
+            } else {
+              command = "upload";
+            }
+            ImmutableMap.Builder<String, List<String>> chunkHeadersBuilder =
+                ImmutableMap.<String, List<String>>builder()
+                    .put(UPLOAD_COMMAND_HEADER, ImmutableList.of(command));
+            if (!"finalize".equals(command)) {
+              chunkHeadersBuilder.put(
+                  UPLOAD_OFFSET_HEADER, ImmutableList.of(String.valueOf(request.getOffset())));
+            }
+            Map<String, List<String>> chunkHeaders = chunkHeadersBuilder.build();
+
+            HttpJsonCallContext context =
+                createCallContext(clientContext, inputContext, chunkHeaders);
+
+            HttpJsonClientCall<ChunkUploadRequest, String> clientCall =
+                HttpJsonClientCalls.newCall(UPLOAD_CHUNK_DESCRIPTOR, context);
+
+            HttpJsonCallFuture<ChunkUploadResponse<ResponseT>> future =
+                new HttpJsonCallFuture<>(clientCall);
+            HttpJsonClientCalls.startUnaryCall(
+                clientCall,
+                request,
+                context,
+                new ChunkUploadResponseListener<>(request, future, responseParser));
+
+            return future;
+          }
+        };
+    return createClientCallable(rawCallable, clientContext);
+  }
+
   private static HttpJsonCallContext createCallContext(
       ClientContext clientContext,
       @Nullable ApiCallContext inputContext,
@@ -150,6 +247,20 @@ private static <CallReqT, CallRespT> UnaryCallable<CallReqT, CallRespT> createCl
     return callable.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
+  @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;
+  }
+
   /**
    * An {@link ApiFuture} that cancels the underlying {@link HttpJsonClientCall} to prevent
    * connection leaks.
@@ -167,12 +278,10 @@ protected void interruptTask() {
       call.cancel("Call was cancelled", null);
     }
 
-    @Override
     public boolean set(T value) {
       return super.set(value);
     }
 
-    @Override
     public boolean setException(Throwable throwable) {
       return super.setException(throwable);
     }
@@ -268,4 +377,88 @@ public void onClose(int statusCode, HttpJsonMetadata trailers) {
       }
     }
   }
+
+  /**
+   * A listener that processes chunk upload response headers and bodies to produce the {@link
+   * ChunkUploadResponse}.
+   */
+  private static class ChunkUploadResponseListener<ResponseT>
+      extends HttpJsonClientCall.Listener<String> {
+
+    private final ChunkUploadRequest request;
+    private final HttpJsonCallFuture<ChunkUploadResponse<ResponseT>> future;
+    private final HttpResponseParser<ResponseT> responseParser;
+    private boolean hasUploadStatusHeader = false;
+    private boolean isComplete = false;
+    @Nullable private Long committedOffset = null;
+    private String responseBody = "";
+
+    ChunkUploadResponseListener(
+        ChunkUploadRequest request,
+        HttpJsonCallFuture<ChunkUploadResponse<ResponseT>> future,
+        HttpResponseParser<ResponseT> responseParser) {
+      this.request = request;
+      this.future = future;
+      this.responseParser = responseParser;
+    }
+
+    @Override
+    public void onHeaders(HttpJsonMetadata responseHeaders) {
+      Map<String, Object> headers = responseHeaders.getHeaders();
+
+      String statusStr = HttpHeadersUtils.getSingleHeader(headers, UPLOAD_STATUS_HEADER);
+      if (statusStr != null) {
+        this.hasUploadStatusHeader = true;
+        if (STATUS_FINAL.equalsIgnoreCase(statusStr)) {
+          this.isComplete = true;
+        }
+      }
+
+      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 (!hasUploadStatusHeader) {
+            future.setException(
+                ApiExceptionFactory.createException(
+                    "Upload chunk response did not contain valid X-Goog-Upload-Status header",
+                    /* cause= */ null,
+                    HttpJsonStatusCode.of(StatusCode.Code.INTERNAL),
+                    /* retryable= */ false));
+            return;
+          }
+          long confirmedOffset =
+              committedOffset != null
+                  ? committedOffset
+                  : request.getOffset() + request.getPayload().size();
+          ResponseT response = null;
+          if (isComplete) {
+            InputStream stream =
+                new ByteArrayInputStream(responseBody.getBytes(StandardCharsets.UTF_8));
+            response = responseParser.parse(stream);
+          }
+          future.set(ChunkUploadResponse.create(confirmedOffset, isComplete, response));
+        } else {
+          Throwable cause = trailers.getException();
+          future.setException(
+              cause != null
+                  ? cause
+                  : new HttpJsonStatusRuntimeException(
+                      statusCode, "Failed to upload chunk with status code: " + statusCode, null));
+        }
+      } catch (Throwable t) {
+        future.setException(t);
+      }
+    }
+  }
 }