Set user agent in gRPC requests (#1126)

GoogleCloudDataproc · Mar 13, 2024 · e5d7005 · e5d7005
1 parent 4572cab
commit e5d7005
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 1 deletion.
diff --git a/gcs/CHANGES.md b/gcs/CHANGES.md
@@ -1,6 +1,7 @@
 # Release Notes
 
 ## Next
+1. [Bug Fix] Set user agent in gRPC requests
 
 ## 2.2.20 - 2024-02-21
 1. Fix downscoping not working when gRPC is enabled issue

diff --git a/gcsio/src/main/java/com/google/cloud/hadoop/gcsio/GoogleCloudStorageClientImpl.java b/gcsio/src/main/java/com/google/cloud/hadoop/gcsio/GoogleCloudStorageClientImpl.java
@@ -40,8 +40,10 @@
 import com.google.cloud.storage.Storage;
 import com.google.cloud.storage.StorageOptions;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.MoreObjects;
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.GoogleLogger;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
@@ -67,6 +69,7 @@
  */
 @VisibleForTesting
 public class GoogleCloudStorageClientImpl extends ForwardingGoogleCloudStorage {
+  private static final String USER_AGENT = "user-agent";
   private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
 
   private final GoogleCloudStorageOptions storageOptions;
@@ -231,9 +234,10 @@ private Storage createStorage(
       ExecutorService pCUExecutorService,
       Function<List<AccessBoundary>, String> downscopedAccessTokenFn)
       throws IOException {
+    final ImmutableMap<String, String> headers = getUpdatedHeadersWithUserAgent(storageOptions);
     return StorageOptions.grpc()
         .setAttemptDirectPath(storageOptions.isDirectPathPreferred())
-        .setHeaderProvider(() -> storageOptions.getHttpRequestHeaders())
+        .setHeaderProvider(() -> headers)
         .setGrpcInterceptorProvider(
             () -> {
               List<ClientInterceptor> list = new ArrayList<>();
@@ -264,6 +268,22 @@ private Storage createStorage(
         .getService();
   }
 
+  private static ImmutableMap<String, String> getUpdatedHeadersWithUserAgent(
+      GoogleCloudStorageOptions storageOptions) {
+    ImmutableMap<String, String> httpRequestHeaders =
+        MoreObjects.firstNonNull(storageOptions.getHttpRequestHeaders(), ImmutableMap.of());
+    String appName = storageOptions.getAppName();
+    if (!httpRequestHeaders.containsKey(USER_AGENT) && !Strings.isNullOrEmpty(appName)) {
+      logger.atFiner().log("Setting useragent %s", appName);
+      return ImmutableMap.<String, String>builder()
+          .putAll(httpRequestHeaders)
+          .put(USER_AGENT, appName)
+          .build();
+    }
+
+    return httpRequestHeaders;
+  }
+
   private Credentials getNoCredentials(
       Function<List<AccessBoundary>, String> downscopedAccessTokenFn) {
     if (downscopedAccessTokenFn == null) {