Skip to content

Commit 67719c0

Browse files
whoweswhowes
authored
feat(gax): implement queryStatusCallable for resumable uploads (#14155)
This provides the ability for callers to query the server for the committed byte offset so far and detect if the upload has already finalized. This is needed for resumable uploads to be able to recover from occasional interruptions.
1 parent 4d26645 commit 67719c0

6 files changed

Lines changed: 531 additions & 0 deletions

File tree

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@
3434
import com.google.api.core.InternalApi;
3535
import com.google.api.gax.resumable.ChunkUploadRequest;
3636
import com.google.api.gax.resumable.ChunkUploadResponse;
37+
import com.google.api.gax.resumable.QueryStatusRequest;
38+
import com.google.api.gax.resumable.QueryStatusResponse;
3739
import com.google.api.gax.resumable.ResumableUploadClient;
3840
import com.google.api.gax.resumable.ResumableUploadSession;
3941
import com.google.api.gax.rpc.ClientContext;
@@ -58,6 +60,8 @@ public final class HttpJsonResumableUploadClient<RequestT, ResponseT>
5860
private final UnaryCallable<RequestT, ResumableUploadSession> startUploadCallable;
5961
private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
6062
uploadChunkCallable;
63+
private final UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>
64+
queryStatusCallable;
6165

6266
public static <RequestT, ResponseT> HttpJsonResumableUploadClient<RequestT, ResponseT> create(
6367
ClientContext clientContext, ApiMethodDescriptor<RequestT, ResponseT> methodDescriptor) {
@@ -82,6 +86,8 @@ private HttpJsonResumableUploadClient(
8286
this.startUploadCallable =
8387
ResumableUploadStartCallable.create(clientContext, startUploadDescriptor);
8488
this.uploadChunkCallable = ResumableUploadChunkCallable.create(clientContext, responseParser);
89+
this.queryStatusCallable =
90+
ResumableUploadQueryStatusCallable.create(clientContext, responseParser);
8591
}
8692

8793
@Override
@@ -93,4 +99,9 @@ public UnaryCallable<RequestT, ResumableUploadSession> startUploadCallable() {
9399
public UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable() {
94100
return uploadChunkCallable;
95101
}
102+
103+
@Override
104+
public UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> queryStatusCallable() {
105+
return queryStatusCallable;
106+
}
96107
}

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

Lines changed: 262 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,262 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Redistribution and use in source and binary forms, with or without
5+
* modification, are permitted provided that the following conditions are
6+
* met:
7+
*
8+
* * Redistributions of source code must retain the above copyright
9+
* notice, this list of conditions and the following disclaimer.
10+
* * Redistributions in binary form must reproduce the above
11+
* copyright notice, this list of conditions and the following disclaimer
12+
* in the documentation and/or other materials provided with the
13+
* distribution.
14+
* * Neither the name of Google LLC nor the names of its
15+
* contributors may be used to endorse or promote products derived from
16+
* this software without specific prior written permission.
17+
*
18+
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
19+
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
20+
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
21+
* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
22+
* OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
23+
* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
24+
* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
25+
* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
26+
* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
27+
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
28+
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29+
*/
30+
package com.google.api.gax.httpjson;
31+
32+
import com.google.api.client.http.HttpMethods;
33+
import com.google.api.core.ApiFuture;
34+
import com.google.api.gax.resumable.QueryStatusRequest;
35+
import com.google.api.gax.resumable.QueryStatusResponse;
36+
import com.google.api.gax.rpc.ApiCallContext;
37+
import com.google.api.gax.rpc.ApiExceptionFactory;
38+
import com.google.api.gax.rpc.ClientContext;
39+
import com.google.api.gax.rpc.StatusCode;
40+
import com.google.api.gax.rpc.UnaryCallable;
41+
import com.google.api.pathtemplate.PathTemplate;
42+
import com.google.common.base.Preconditions;
43+
import com.google.common.base.Strings;
44+
import com.google.common.collect.ImmutableList;
45+
import com.google.common.collect.ImmutableMap;
46+
import java.io.ByteArrayInputStream;
47+
import java.io.InputStream;
48+
import java.nio.charset.StandardCharsets;
49+
import java.util.Collections;
50+
import java.util.List;
51+
import java.util.Map;
52+
import org.jspecify.annotations.NullMarked;
53+
import org.jspecify.annotations.Nullable;
54+
55+
/**
56+
* A {@link UnaryCallable} that queries the committed status and offset of a resumable upload
57+
* session.
58+
*/
59+
@NullMarked
60+
class ResumableUploadQueryStatusCallable<ResponseT>
61+
extends UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> {
62+
63+
private static final String UPLOAD_COMMAND_HEADER = "X-Goog-Upload-Command";
64+
private static final String UPLOAD_STATUS_HEADER = "X-Goog-Upload-Status";
65+
private static final String UPLOAD_SIZE_RECEIVED_HEADER = "X-Goog-Upload-Size-Received";
66+
private static final String STATUS_FINAL = "final";
67+
private static final String COMMAND_QUERY = "query";
68+
69+
private static final Map<String, List<String>> QUERY_STATUS_HEADERS =
70+
ImmutableMap.of(UPLOAD_COMMAND_HEADER, ImmutableList.of(COMMAND_QUERY));
71+
72+
private static final PathTemplate PATH_TEMPLATE = PathTemplate.create("**");
73+
74+
private static final ApiMethodDescriptor<QueryStatusRequest, String> QUERY_STATUS_DESCRIPTOR =
75+
ApiMethodDescriptor.<QueryStatusRequest, String>newBuilder()
76+
.setFullMethodName("ResumableUpload/QueryStatus")
77+
.setHttpMethod(HttpMethods.POST)
78+
.setType(ApiMethodDescriptor.MethodType.UNARY)
79+
.setRequestFormatter(
80+
new HttpRequestFormatter<QueryStatusRequest>() {
81+
@Override
82+
public Map<String, List<String>> getQueryParamNames(QueryStatusRequest request) {
83+
return Collections.emptyMap();
84+
}
85+
86+
@Override
87+
public String getRequestBody(QueryStatusRequest request) {
88+
return "";
89+
}
90+
91+
@Override
92+
public String getPath(QueryStatusRequest request) {
93+
return request.getUploadUrl();
94+
}
95+
96+
@Override
97+
public PathTemplate getPathTemplate() {
98+
return PATH_TEMPLATE;
99+
}
100+
})
101+
.setResponseParser(ResumableUploadResponseParser.create())
102+
.build();
103+
104+
private final ClientContext clientContext;
105+
private final HttpResponseParser<ResponseT> responseParser;
106+
107+
private ResumableUploadQueryStatusCallable(
108+
ClientContext clientContext, HttpResponseParser<ResponseT> responseParser) {
109+
this.clientContext = Preconditions.checkNotNull(clientContext);
110+
this.responseParser = Preconditions.checkNotNull(responseParser);
111+
}
112+
113+
@Override
114+
public ApiFuture<QueryStatusResponse<ResponseT>> futureCall(
115+
QueryStatusRequest request, @Nullable ApiCallContext inputContext) {
116+
Preconditions.checkNotNull(request);
117+
HttpJsonCallContext context =
118+
(HttpJsonCallContext)
119+
HttpJsonCallContext.createDefault()
120+
.nullToSelf(clientContext.getDefaultCallContext())
121+
.merge(inputContext)
122+
.withExtraHeaders(QUERY_STATUS_HEADERS);
123+
124+
HttpJsonClientCall<QueryStatusRequest, String> clientCall =
125+
HttpJsonClientCalls.newCall(QUERY_STATUS_DESCRIPTOR, context);
126+
127+
ResumableUploadHttpJsonFuture<QueryStatusResponse<ResponseT>> future =
128+
new ResumableUploadHttpJsonFuture<>(clientCall);
129+
HttpJsonClientCalls.startUnaryCall(
130+
clientCall, request, context, new QueryStatusResponseListener<>(future, responseParser));
131+
132+
return future;
133+
}
134+
135+
static <ResponseT> UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> create(
136+
ClientContext clientContext, HttpResponseParser<ResponseT> responseParser) {
137+
UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> rawCallable =
138+
new ResumableUploadQueryStatusCallable<>(clientContext, responseParser);
139+
UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> callable =
140+
new HttpJsonExceptionCallable<>(
141+
rawCallable,
142+
// Wire calls do not retry directly; retries are managed by ResumableUploadCallable.
143+
Collections.emptySet());
144+
return callable.withDefaultCallContext(clientContext.getDefaultCallContext());
145+
}
146+
147+
@Nullable
148+
private static Long parseSizeReceived(HttpJsonMetadata responseHeaders) {
149+
String sizeReceivedStr =
150+
HttpHeadersUtils.getSingleHeader(responseHeaders.getHeaders(), UPLOAD_SIZE_RECEIVED_HEADER);
151+
if (!Strings.isNullOrEmpty(sizeReceivedStr)) {
152+
try {
153+
long parsed = Long.parseLong(sizeReceivedStr);
154+
if (parsed < 0) {
155+
throw ApiExceptionFactory.createException(
156+
"Response contained negative "
157+
+ UPLOAD_SIZE_RECEIVED_HEADER
158+
+ " header: "
159+
+ sizeReceivedStr,
160+
/* cause= */ null,
161+
HttpJsonStatusCode.of(StatusCode.Code.INTERNAL),
162+
/* retryable= */ false);
163+
}
164+
return parsed;
165+
} catch (NumberFormatException e) {
166+
throw ApiExceptionFactory.createException(
167+
"Response contained invalid "
168+
+ UPLOAD_SIZE_RECEIVED_HEADER
169+
+ " header: "
170+
+ sizeReceivedStr,
171+
e,
172+
HttpJsonStatusCode.of(StatusCode.Code.INTERNAL),
173+
/* retryable= */ false);
174+
}
175+
}
176+
return null;
177+
}
178+
179+
/** A listener that parses query response headers to produce the {@link QueryStatusResponse}. */
180+
private static class QueryStatusResponseListener<ResponseT>
181+
extends HttpJsonClientCall.Listener<String> {
182+
183+
private final ResumableUploadHttpJsonFuture<QueryStatusResponse<ResponseT>> future;
184+
private final HttpResponseParser<ResponseT> responseParser;
185+
@Nullable private String uploadStatus = null;
186+
@Nullable private Long committedOffset = null;
187+
@Nullable private Throwable headerParsingException;
188+
private String responseBody = "";
189+
190+
private QueryStatusResponseListener(
191+
ResumableUploadHttpJsonFuture<QueryStatusResponse<ResponseT>> future,
192+
HttpResponseParser<ResponseT> responseParser) {
193+
this.future = future;
194+
this.responseParser = responseParser;
195+
}
196+
197+
@Override
198+
public void onHeaders(HttpJsonMetadata responseHeaders) {
199+
Map<String, Object> headers = responseHeaders.getHeaders();
200+
this.uploadStatus = HttpHeadersUtils.getSingleHeader(headers, UPLOAD_STATUS_HEADER);
201+
try {
202+
this.committedOffset = parseSizeReceived(responseHeaders);
203+
} catch (Throwable t) {
204+
this.headerParsingException = t;
205+
}
206+
}
207+
208+
@Override
209+
public void onMessage(@Nullable String message) {
210+
if (message != null) {
211+
this.responseBody = message;
212+
}
213+
}
214+
215+
@Override
216+
public void onClose(int statusCode, HttpJsonMetadata trailers) {
217+
try {
218+
if (statusCode >= 200 && statusCode < 300) {
219+
if (headerParsingException != null) {
220+
future.setException(headerParsingException);
221+
return;
222+
}
223+
boolean isComplete = STATUS_FINAL.equalsIgnoreCase(uploadStatus);
224+
if (isComplete) {
225+
QueryStatusResponse.Builder<ResponseT> queryResponseBuilder =
226+
QueryStatusResponse.<ResponseT>newBuilder().setComplete(true);
227+
InputStream stream =
228+
new ByteArrayInputStream(responseBody.getBytes(StandardCharsets.UTF_8));
229+
queryResponseBuilder.setResponse(responseParser.parse(stream));
230+
future.set(queryResponseBuilder.build());
231+
} else if (committedOffset != null) {
232+
future.set(
233+
QueryStatusResponse.<ResponseT>newBuilder()
234+
.setComplete(false)
235+
.setCommittedOffset(committedOffset)
236+
.build());
237+
} else {
238+
future.setException(
239+
ApiExceptionFactory.createException(
240+
"Query status response did not contain valid "
241+
+ UPLOAD_SIZE_RECEIVED_HEADER
242+
+ " header",
243+
/* cause= */ null,
244+
HttpJsonStatusCode.of(StatusCode.Code.INTERNAL),
245+
/* retryable= */ false));
246+
}
247+
} else {
248+
Throwable cause = trailers.getException();
249+
future.setException(
250+
cause != null
251+
? cause
252+
: new HttpJsonStatusRuntimeException(
253+
statusCode,
254+
"Failed to query upload status with status code: " + statusCode,
255+
null));
256+
}
257+
} catch (Throwable t) {
258+
future.setException(t);
259+
}
260+
}
261+
}
262+
}

0 commit comments

Comments
 (0)