AWS: Make sure Signer + User Agent config are both applied (#10198)

apache · Apr 22, 2024 · e3b78be · e3b78be
1 parent a23021d
commit e3b78be
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 9 deletions.
diff --git a/aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIOProperties.java b/aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIOProperties.java
@@ -36,6 +36,7 @@
 import org.apache.iceberg.util.PropertyUtil;
 import org.apache.iceberg.util.SerializableMap;
 import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider;
+import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
 import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption;
 import software.amazon.awssdk.services.s3.S3ClientBuilder;
 import software.amazon.awssdk.services.s3.S3Configuration;
@@ -788,10 +789,15 @@ public <T extends S3ClientBuilder> void applyServiceConfigurations(T builder) {
    */
   public <T extends S3ClientBuilder> void applySignerConfiguration(T builder) {
     if (isRemoteSigningEnabled) {
+      ClientOverrideConfiguration.Builder configBuilder =
+          null != builder.overrideConfiguration()
+              ? builder.overrideConfiguration().toBuilder()
+              : ClientOverrideConfiguration.builder();
       builder.overrideConfiguration(
-          c ->
-              c.putAdvancedOption(
-                  SdkAdvancedClientOption.SIGNER, S3V4RestSignerClient.create(allProperties)));
+          configBuilder
+              .putAdvancedOption(
+                  SdkAdvancedClientOption.SIGNER, S3V4RestSignerClient.create(allProperties))
+              .build());
     }
   }
 
@@ -829,8 +835,14 @@ public <T extends S3ClientBuilder> void applyS3AccessGrantsConfigurations(T buil
   }
 
   public <T extends S3ClientBuilder> void applyUserAgentConfigurations(T builder) {
+    ClientOverrideConfiguration.Builder configBuilder =
+        null != builder.overrideConfiguration()
+            ? builder.overrideConfiguration().toBuilder()
+            : ClientOverrideConfiguration.builder();
     builder.overrideConfiguration(
-        c -> c.putAdvancedOption(SdkAdvancedClientOption.USER_AGENT_PREFIX, S3_FILE_IO_USER_AGENT));
+        configBuilder
+            .putAdvancedOption(SdkAdvancedClientOption.USER_AGENT_PREFIX, S3_FILE_IO_USER_AGENT)
+            .build());
   }
 
   /**

diff --git a/aws/src/test/java/org/apache/iceberg/aws/TestS3FileIOProperties.java b/aws/src/test/java/org/apache/iceberg/aws/TestS3FileIOProperties.java
@@ -236,6 +236,33 @@ public void testS3RemoteSigningEnabled() {
     Assertions.assertThat(signerClient.properties()).isEqualTo(properties);
   }
 
+  @Test
+  public void s3RemoteSigningEnabledWithUserAgent() {
+    String uri = "http://localhost:12345";
+    Map<String, String> properties =
+        ImmutableMap.of(
+            S3FileIOProperties.REMOTE_SIGNING_ENABLED, "true", CatalogProperties.URI, uri);
+    S3FileIOProperties s3Properties = new S3FileIOProperties(properties);
+    S3ClientBuilder builder = S3Client.builder();
+
+    s3Properties.applySignerConfiguration(builder);
+    s3Properties.applyUserAgentConfigurations(builder);
+
+    Optional<String> userAgent =
+        builder.overrideConfiguration().advancedOption(SdkAdvancedClientOption.USER_AGENT_PREFIX);
+    Assertions.assertThat(userAgent)
+        .isPresent()
+        .get()
+        .satisfies(x -> Assertions.assertThat(x).startsWith("s3fileio"));
+
+    Optional<Signer> signer =
+        builder.overrideConfiguration().advancedOption(SdkAdvancedClientOption.SIGNER);
+    Assertions.assertThat(signer).isPresent().get().isInstanceOf(S3V4RestSignerClient.class);
+    S3V4RestSignerClient signerClient = (S3V4RestSignerClient) signer.get();
+    Assertions.assertThat(signerClient.baseSignerUri()).isEqualTo(uri);
+    Assertions.assertThat(signerClient.properties()).isEqualTo(properties);
+  }
+
   @Test
   public void testS3RemoteSigningDisabled() {
     Map<String, String> properties =

diff --git a/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3FileIOProperties.java b/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3FileIOProperties.java
@@ -22,16 +22,18 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Consumer;
 import java.util.stream.Collectors;
+import org.apache.iceberg.CatalogProperties;
 import org.apache.iceberg.aws.AwsClientProperties;
 import org.apache.iceberg.exceptions.ValidationException;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.relocated.com.google.common.collect.Maps;
 import org.apache.iceberg.relocated.com.google.common.collect.Sets;
 import org.assertj.core.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
+import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
 import software.amazon.awssdk.services.s3.S3ClientBuilder;
 import software.amazon.awssdk.services.s3.S3Configuration;
 import software.amazon.awssdk.services.s3.model.ObjectCannedACL;
@@ -459,13 +461,18 @@ public void testApplyS3ServiceConfigurations() {
 
   @Test
   public void testApplySignerConfiguration() {
-    Map<String, String> properties = Maps.newHashMap();
-    properties.put(S3FileIOProperties.REMOTE_SIGNING_ENABLED, "true");
+    Map<String, String> properties =
+        ImmutableMap.of(
+            S3FileIOProperties.REMOTE_SIGNING_ENABLED,
+            "true",
+            CatalogProperties.URI,
+            "http://localhost:12345");
     S3FileIOProperties s3FileIOProperties = new S3FileIOProperties(properties);
     S3ClientBuilder mockS3ClientBuilder = Mockito.mock(S3ClientBuilder.class);
     s3FileIOProperties.applySignerConfiguration(mockS3ClientBuilder);
 
-    Mockito.verify(mockS3ClientBuilder).overrideConfiguration(Mockito.any(Consumer.class));
+    Mockito.verify(mockS3ClientBuilder)
+        .overrideConfiguration(Mockito.any(ClientOverrideConfiguration.class));
   }
 
   @Test
@@ -486,6 +493,7 @@ public void testApplyUserAgentConfigurations() {
     S3ClientBuilder mockS3ClientBuilder = Mockito.mock(S3ClientBuilder.class);
     s3FileIOProperties.applyUserAgentConfigurations(mockS3ClientBuilder);
 
-    Mockito.verify(mockS3ClientBuilder).overrideConfiguration(Mockito.any(Consumer.class));
+    Mockito.verify(mockS3ClientBuilder)
+        .overrideConfiguration(Mockito.any(ClientOverrideConfiguration.class));
   }
 }