Skip to content

Commit

Permalink
feat: add new Options to allow per method header values (#2941)
Browse files Browse the repository at this point in the history
Add new "Option"s for those methods which already have option types to allow providing an ImmutableMap<String, String> to be applied as extra headers to all requests sent as part of that operation.

If an operation has multiple sources of input Options (rewrite) the "first" (i.e. source option) will be the one added to the request.

The following resources do not have "Option"s and therefor do not have extra headers support at this time:
* Acl
* DefaultAcl
* ServiceAccount
* Notification
  • Loading branch information
BenWhitehead authored Feb 25, 2025
1 parent cc4c7f4 commit 297802d
Show file tree
Hide file tree
Showing 21 changed files with 1,913 additions and 238 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
import com.google.cloud.storage.spi.v1.StorageRpc;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.MoreObjects;
import com.google.common.collect.ImmutableMap;
import com.google.common.hash.HashFunction;
import com.google.common.hash.Hashing;
import com.google.common.io.BaseEncoding;
Expand All @@ -51,8 +52,8 @@
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.ScatteringByteChannel;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Map.Entry;
import java.util.function.Function;
import javax.annotation.concurrent.Immutable;
import org.checkerframework.checker.nullness.qual.NonNull;
Expand Down Expand Up @@ -273,6 +274,14 @@ static Get createGetRequest(
"x-goog-encryption-key-sha256",
base64.encode(hashFunction.hashBytes(base64.decode(key)).asBytes()));
});
ifNonNull(
options.get(StorageRpc.Option.EXTRA_HEADERS),
ApiaryUnbufferedReadableByteChannel::cast,
(ImmutableMap<String, String> extraHeaders) -> {
for (Entry<String, String> e : extraHeaders.entrySet()) {
headers.set(e.getKey(), e.getValue());
}
});

// gzip handling is performed upstream of here. Ensure we always get the raw input stream from
// the request
Expand Down Expand Up @@ -302,7 +311,7 @@ private static String getHeaderValue(@NonNull HttpHeaders headers, @NonNull Stri
if (list.isEmpty()) {
return null;
} else {
return list.get(0).trim().toLowerCase(Locale.ENGLISH);
return Utils.headerNameToLowerCase(list.get(0).trim());
}
} else if (o instanceof String) {
return (String) o;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -650,8 +650,10 @@ public CopyWriter copy(CopyRequest copyRequest) {
Opts<ObjectTargetOpt> dstOpts =
Opts.unwrap(copyRequest.getTargetOptions()).resolveFrom(dst).prepend(defaultOpts);

Mapper<RewriteObjectRequest.Builder> mapper =
Mapper<RewriteObjectRequest.Builder> requestBuilderMapper =
srcOpts.rewriteObjectsRequest().andThen(dstOpts.rewriteObjectsRequest());
Mapper<GrpcCallContext> grpcCallContextMapper =
srcOpts.grpcMetadataMapper().andThen(dstOpts.grpcMetadataMapper());

Object srcProto = codecs.blobId().encode(src);
Object dstProto = codecs.blobInfo().encode(dst);
Expand Down Expand Up @@ -686,9 +688,8 @@ public CopyWriter copy(CopyRequest copyRequest) {
b.setMaxBytesRewrittenPerCall(copyRequest.getMegabytesCopiedPerChunk() * _1MiB);
}

RewriteObjectRequest req = mapper.apply(b).build();
GrpcCallContext grpcCallContext =
srcOpts.grpcMetadataMapper().apply(GrpcCallContext.createDefault());
RewriteObjectRequest req = requestBuilderMapper.apply(b).build();
GrpcCallContext grpcCallContext = grpcCallContextMapper.apply(GrpcCallContext.createDefault());
UnaryCallable<RewriteObjectRequest, RewriteResponse> callable =
storageClient.rewriteObjectCallable().withDefaultCallContext(grpcCallContext);
GrpcCallContext retryContext = Retrying.newCallContext();
Expand Down Expand Up @@ -733,7 +734,7 @@ public GrpcBlobReadChannel reader(String bucket, String blob, BlobSourceOption..
public GrpcBlobReadChannel reader(BlobId blob, BlobSourceOption... options) {
Opts<ObjectSourceOpt> opts = Opts.unwrap(options).resolveFrom(blob).prepend(defaultOpts);
ReadObjectRequest request = getReadObjectRequest(blob, opts);
GrpcCallContext grpcCallContext = Retrying.newCallContext();
GrpcCallContext grpcCallContext = opts.grpcMetadataMapper().apply(Retrying.newCallContext());

return new GrpcBlobReadChannel(
storageClient.readObjectCallable().withDefaultCallContext(grpcCallContext),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,6 @@
import java.util.IdentityHashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
Expand Down Expand Up @@ -161,6 +160,11 @@ StorageSettings getStorageSettings() throws IOException {
return resolveSettingsAndOpts().x();
}

@InternalApi
GrpcInterceptorProvider getGrpcInterceptorProvider() {
return grpcInterceptorProvider;
}

private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
in.defaultReadObject();
this.openTelemetry = HttpStorageOptions.getDefaultInstance().getOpenTelemetry();
Expand Down Expand Up @@ -225,7 +229,7 @@ private Tuple<StorageSettings, Opts<UserProject>> resolveSettingsAndOpts() throw
Map<String, List<String>> requestMetadata = credentials.getRequestMetadata(uri);
for (Entry<String, List<String>> e : requestMetadata.entrySet()) {
String key = e.getKey();
if ("x-goog-user-project".equals(key.trim().toLowerCase(Locale.ENGLISH))) {
if ("x-goog-user-project".equals(Utils.headerNameToLowerCase(key.trim()))) {
List<String> value = e.getValue();
if (!value.isEmpty()) {
foundQuotaProject = true;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -54,14 +54,13 @@ final class JsonResumableSession {
* have the concept of nested retry handling.
*/
ResumableOperationResult<@Nullable StorageObject> query() {
return new JsonResumableSessionQueryTask(context, resumableWrite.getUploadId()).call();
return new JsonResumableSessionQueryTask(context, resumableWrite).call();
}

ResumableOperationResult<@Nullable StorageObject> put(
RewindableContent content, HttpContentRange contentRange) {
JsonResumableSessionPutTask task =
new JsonResumableSessionPutTask(
context, resumableWrite.getUploadId(), content, contentRange);
new JsonResumableSessionPutTask(context, resumableWrite, content, contentRange);
HttpRpcContext httpRpcContext = HttpRpcContext.getInstance();
try {
httpRpcContext.newInvocationId();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
package com.google.cloud.storage;

import com.google.api.client.http.GenericUrl;
import com.google.api.client.http.HttpHeaders;
import com.google.api.client.http.HttpRequest;
import com.google.api.client.http.HttpResponse;
import com.google.api.client.http.HttpResponseException;
Expand All @@ -31,14 +32,15 @@
import java.io.IOException;
import java.math.BigInteger;
import java.util.Locale;
import java.util.Map.Entry;
import java.util.concurrent.Callable;
import org.checkerframework.checker.nullness.qual.Nullable;

final class JsonResumableSessionPutTask
implements Callable<ResumableOperationResult<@Nullable StorageObject>> {

private final HttpClientContext context;
private final String uploadId;
private final JsonResumableWrite jsonResumableWrite;
private final RewindableContent content;
private final HttpContentRange originalContentRange;

Expand All @@ -47,11 +49,11 @@ final class JsonResumableSessionPutTask
@VisibleForTesting
JsonResumableSessionPutTask(
HttpClientContext httpClientContext,
String uploadId,
JsonResumableWrite jsonResumableWrite,
RewindableContent content,
HttpContentRange originalContentRange) {
this.context = httpClientContext;
this.uploadId = uploadId;
this.jsonResumableWrite = jsonResumableWrite;
this.content = content;
this.originalContentRange = originalContentRange;
this.contentRange = originalContentRange;
Expand Down Expand Up @@ -87,13 +89,18 @@ public void rewindTo(long offset) {
boolean success = false;
boolean finalizing = originalContentRange.isFinalizing();

String uploadId = jsonResumableWrite.getUploadId();
HttpRequest req =
context
.getRequestFactory()
.buildPutRequest(new GenericUrl(uploadId), content)
.setParser(context.getObjectParser());
req.setThrowExceptionOnExecuteError(false);
req.getHeaders().setContentRange(contentRange.getHeaderValue());
HttpHeaders headers = req.getHeaders();
headers.setContentRange(contentRange.getHeaderValue());
for (Entry<String, String> e : jsonResumableWrite.getExtraHeaders().entrySet()) {
headers.set(e.getKey(), e.getValue());
}

HttpResponse response = null;
try {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,37 +20,44 @@

import com.google.api.client.http.EmptyContent;
import com.google.api.client.http.GenericUrl;
import com.google.api.client.http.HttpHeaders;
import com.google.api.client.http.HttpRequest;
import com.google.api.client.http.HttpResponse;
import com.google.api.client.http.HttpResponseException;
import com.google.api.services.storage.model.StorageObject;
import java.io.IOException;
import java.math.BigInteger;
import java.util.Locale;
import java.util.Map.Entry;
import java.util.concurrent.Callable;
import org.checkerframework.checker.nullness.qual.Nullable;

final class JsonResumableSessionQueryTask
implements Callable<ResumableOperationResult<@Nullable StorageObject>> {

private final HttpClientContext context;
private final String uploadId;
private final JsonResumableWrite jsonResumableWrite;

JsonResumableSessionQueryTask(HttpClientContext context, String uploadId) {
JsonResumableSessionQueryTask(HttpClientContext context, JsonResumableWrite jsonResumableWrite) {
this.context = context;
this.uploadId = uploadId;
this.jsonResumableWrite = jsonResumableWrite;
}

public ResumableOperationResult<@Nullable StorageObject> call() {
HttpResponse response = null;
String uploadId = jsonResumableWrite.getUploadId();
try {
HttpRequest req =
context
.getRequestFactory()
.buildPutRequest(new GenericUrl(uploadId), new EmptyContent())
.setParser(context.getObjectParser());
req.setThrowExceptionOnExecuteError(false);
req.getHeaders().setContentRange(HttpContentRange.query().getHeaderValue());
HttpHeaders headers = req.getHeaders();
headers.setContentRange(HttpContentRange.query().getHeaderValue());
for (Entry<String, String> e : jsonResumableWrite.getExtraHeaders().entrySet()) {
headers.set(e.getKey(), e.getValue());
}

response = req.execute();

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import com.google.api.services.storage.model.StorageObject;
import com.google.cloud.storage.spi.v1.StorageRpc;
import com.google.common.base.MoreObjects;
import com.google.common.collect.ImmutableMap;
import com.google.gson.Gson;
import com.google.gson.stream.JsonReader;
import java.io.IOException;
Expand Down Expand Up @@ -60,6 +61,16 @@ private JsonResumableWrite(
this.beginOffset = beginOffset;
}

ImmutableMap<String, String> getExtraHeaders() {
if (options != null) {
Object tmp = options.get(StorageRpc.Option.EXTRA_HEADERS);
if (tmp != null) {
return (ImmutableMap<String, String>) tmp;
}
}
return ImmutableMap.of();
}

public @NonNull String getUploadId() {
return uploadId;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.function.Consumer;
import java.util.function.Predicate;
Expand Down Expand Up @@ -295,12 +294,12 @@ static boolean isContinue(int code) {
// The header names from HttpHeaders are lower cased, define some utility methods to create
// predicates where we can specify values ignoring case
private static Predicate<String> matches(String expected) {
String lower = expected.toLowerCase(Locale.US);
String lower = Utils.headerNameToLowerCase(expected);
return lower::equals;
}

private static Predicate<String> startsWith(String prefix) {
String lower = prefix.toLowerCase(Locale.US);
String lower = Utils.headerNameToLowerCase(prefix);
return s -> s.startsWith(lower);
}

Expand Down
Loading

0 comments on commit 297802d

Please sign in to comment.