Refactor: Use per-request STS credentials (#1629)

* Refactor: Use per-request STS credentials

No functional changes.

This is mostly to allow more storage integration
flexibility in downstream build.

This might also be useful for non-AWS storage.

Author: dimas-b

polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java

@@ -33,6 +33,7 @@
 import org.apache.polaris.core.storage.InMemoryStorageIntegration;
 import org.apache.polaris.core.storage.StorageAccessProperty;
 import org.apache.polaris.core.storage.StorageUtil;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.policybuilder.iam.IamConditionOperator;
 import software.amazon.awssdk.policybuilder.iam.IamEffect;
 import software.amazon.awssdk.policybuilder.iam.IamPolicy;
@@ -46,10 +47,17 @@
 public class AwsCredentialsStorageIntegration
     extends InMemoryStorageIntegration<AwsStorageConfigurationInfo> {
   private final StsClient stsClient;
+  private final Optional<AwsCredentialsProvider> credentialsProvider;
 
   public AwsCredentialsStorageIntegration(StsClient stsClient) {
+    this(stsClient, Optional.empty());
+  }
+
+  public AwsCredentialsStorageIntegration(
+      StsClient stsClient, Optional<AwsCredentialsProvider> credentialsProvider) {
     super(AwsCredentialsStorageIntegration.class.getName());
     this.stsClient = stsClient;
+    this.credentialsProvider = credentialsProvider;
   }
 
   /** {@inheritDoc} */
@@ -60,21 +68,22 @@ public EnumMap<StorageAccessProperty, String> getSubscopedCreds(
       boolean allowListOperation,
       @Nonnull Set<String> allowedReadLocations,
       @Nonnull Set<String> allowedWriteLocations) {
-    AssumeRoleResponse response =
-        stsClient.assumeRole(
-            AssumeRoleRequest.builder()
-                .externalId(storageConfig.getExternalId())
-                .roleArn(storageConfig.getRoleARN())
-                .roleSessionName("PolarisAwsCredentialsStorageIntegration")
-                .policy(
-                    policyString(
-                            storageConfig.getRoleARN(),
-                            allowListOperation,
-                            allowedReadLocations,
-                            allowedWriteLocations)
-                        .toJson())
-                .durationSeconds(loadConfig(STORAGE_CREDENTIAL_DURATION_SECONDS))
-                .build());
+    AssumeRoleRequest.Builder request =
+        AssumeRoleRequest.builder()
+            .externalId(storageConfig.getExternalId())
+            .roleArn(storageConfig.getRoleARN())
+            .roleSessionName("PolarisAwsCredentialsStorageIntegration")
+            .policy(
+                policyString(
+                        storageConfig.getRoleARN(),
+                        allowListOperation,
+                        allowedReadLocations,
+                        allowedWriteLocations)
+                    .toJson())
+            .durationSeconds(loadConfig(STORAGE_CREDENTIAL_DURATION_SECONDS));
+    credentialsProvider.ifPresent(
+        cp -> request.overrideConfiguration(b -> b.credentialsProvider(cp)));
+    AssumeRoleResponse response = stsClient.assumeRole(request.build());
     EnumMap<StorageAccessProperty, String> credentialMap =
         new EnumMap<>(StorageAccessProperty.class);
     credentialMap.put(StorageAccessProperty.AWS_KEY_ID, response.credentials().accessKeyId());

quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAuthzTestBase.java

@@ -35,6 +35,7 @@
 import java.util.Date;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -219,7 +220,9 @@ public Map<String, String> getConfigOverrides() {
   public static void setUpMocks() {
     PolarisStorageIntegrationProviderImpl mock =
         new PolarisStorageIntegrationProviderImpl(
-            Mockito::mock, () -> GoogleCredentials.create(new AccessToken("abc", new Date())));
+            Mockito::mock,
+            Optional.empty(),
+            () -> GoogleCredentials.create(new AccessToken("abc", new Date())));
     QuarkusMock.installMockForType(mock, PolarisStorageIntegrationProviderImpl.class);
   }
 

service/common/src/main/java/org/apache/polaris/service/storage/PolarisStorageIntegrationProviderImpl.java

@@ -28,6 +28,7 @@
 import jakarta.inject.Inject;
 import java.util.EnumMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Supplier;
 import org.apache.polaris.core.PolarisDiagnostics;
@@ -39,22 +40,30 @@
 import org.apache.polaris.core.storage.aws.AwsCredentialsStorageIntegration;
 import org.apache.polaris.core.storage.azure.AzureCredentialsStorageIntegration;
 import org.apache.polaris.core.storage.gcp.GcpCredentialsStorageIntegration;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.services.sts.StsClient;
 
 @ApplicationScoped
 public class PolarisStorageIntegrationProviderImpl implements PolarisStorageIntegrationProvider {
 
   private final Supplier<StsClient> stsClientSupplier;
+  private final Optional<AwsCredentialsProvider> stsCredentials;
   private final Supplier<GoogleCredentials> gcpCredsProvider;
 
   @Inject
   public PolarisStorageIntegrationProviderImpl(StorageConfiguration storageConfiguration) {
-    this(storageConfiguration.stsClientSupplier(), storageConfiguration.gcpCredentialsSupplier());
+    this(
+        storageConfiguration.stsClientSupplier(false),
+        Optional.ofNullable(storageConfiguration.stsCredentials()),
+        storageConfiguration.gcpCredentialsSupplier());
   }
 
   public PolarisStorageIntegrationProviderImpl(
-      Supplier<StsClient> stsClientSupplier, Supplier<GoogleCredentials> gcpCredsProvider) {
+      Supplier<StsClient> stsClientSupplier,
+      Optional<AwsCredentialsProvider> stsCredentials,
+      Supplier<GoogleCredentials> gcpCredsProvider) {
     this.stsClientSupplier = stsClientSupplier;
+    this.stsCredentials = stsCredentials;
     this.gcpCredsProvider = gcpCredsProvider;
   }
 
@@ -71,7 +80,7 @@ public PolarisStorageIntegrationProviderImpl(
       case S3:
         storageIntegration =
             (PolarisStorageIntegration<T>)
-                new AwsCredentialsStorageIntegration(stsClientSupplier.get());
+                new AwsCredentialsStorageIntegration(stsClientSupplier.get(), stsCredentials);
         break;
       case GCS:
         storageIntegration =

service/common/src/main/java/org/apache/polaris/service/storage/StorageConfiguration.java

@@ -28,6 +28,8 @@
 import java.util.function.Supplier;
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
 import software.amazon.awssdk.services.sts.StsClient;
 import software.amazon.awssdk.services.sts.StsClientBuilder;
@@ -61,21 +63,31 @@ public interface StorageConfiguration {
   Optional<Duration> gcpAccessTokenLifespan();
 
   default Supplier<StsClient> stsClientSupplier() {
+    return stsClientSupplier(true);
+  }
+
+  default Supplier<StsClient> stsClientSupplier(boolean withCredentials) {
     return Suppliers.memoize(
         () -> {
           StsClientBuilder stsClientBuilder = StsClient.builder();
-          if (awsAccessKey().isPresent() && awsSecretKey().isPresent()) {
-            LoggerFactory.getLogger(StorageConfiguration.class)
-                .warn("Using hard-coded AWS credentials - this is not recommended for production");
-            StaticCredentialsProvider awsCredentialsProvider =
-                StaticCredentialsProvider.create(
-                    AwsBasicCredentials.create(awsAccessKey().get(), awsSecretKey().get()));
-            stsClientBuilder.credentialsProvider(awsCredentialsProvider);
+          if (withCredentials) {
+            stsClientBuilder.credentialsProvider(stsCredentials());
           }
           return stsClientBuilder.build();
         });
   }
 
+  default AwsCredentialsProvider stsCredentials() {
+    if (awsAccessKey().isPresent() && awsSecretKey().isPresent()) {
+      LoggerFactory.getLogger(StorageConfiguration.class)
+          .warn("Using hard-coded AWS credentials - this is not recommended for production");
+      return StaticCredentialsProvider.create(
+          AwsBasicCredentials.create(awsAccessKey().get(), awsSecretKey().get()));
+    } else {
+      return DefaultCredentialsProvider.create();
+    }
+  }
+
   default Supplier<GoogleCredentials> gcpCredentialsSupplier() {
     return Suppliers.memoize(
         () -> {

service/common/src/test/java/org/apache/polaris/service/storage/StorageConfigurationTest.java

@@ -31,6 +31,7 @@
 import org.mockito.ArgumentCaptor;
 import org.mockito.MockedStatic;
 import org.mockito.Mockito;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
 import software.amazon.awssdk.services.sts.StsClient;
 import software.amazon.awssdk.services.sts.StsClientBuilder;
@@ -105,7 +106,7 @@ public void testSingletonStsClientWithStaticCredentials() {
       staticMock.when(StsClient::builder).thenReturn(mockBuilder);
 
       StorageConfiguration config = configWithAwsCredentialsAndGcpToken();
-      Supplier<StsClient> supplier = config.stsClientSupplier();
+      Supplier<StsClient> supplier = config.stsClientSupplier(true);
       StsClient client1 = supplier.get();
       StsClient client2 = supplier.get();
 
@@ -119,6 +120,16 @@ public void testSingletonStsClientWithStaticCredentials() {
     }
   }
 
+  @Test
+  public void testStaticStsCredentials() {
+    StorageConfiguration config = configWithAwsCredentialsAndGcpToken();
+    AwsCredentialsProvider credentialsProvider = config.stsCredentials();
+    assertThat(credentialsProvider).isInstanceOf(StaticCredentialsProvider.class);
+    assertThat(credentialsProvider.resolveCredentials().accessKeyId()).isEqualTo(TEST_ACCESS_KEY);
+    assertThat(credentialsProvider.resolveCredentials().secretAccessKey())
+        .isEqualTo(TEST_SECRET_KEY);
+  }
+
   @Test
   public void testCreateGcpCredentialsFromStaticToken() {
     Supplier<GoogleCredentials> supplier =

service/common/src/testFixtures/java/org/apache/polaris/service/TestServices.java

@@ -29,6 +29,7 @@
 import java.util.Date;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import org.apache.polaris.core.PolarisCallContext;
 import org.apache.polaris.core.PolarisDiagnostics;
@@ -150,6 +151,7 @@ public TestServices build() {
       PolarisStorageIntegrationProviderImpl storageIntegrationProvider =
           new PolarisStorageIntegrationProviderImpl(
               () -> stsClient,
+              Optional.empty(),
               () -> GoogleCredentials.create(new AccessToken(GCP_ACCESS_TOKEN, new Date())));
       InMemoryPolarisMetaStoreManagerFactory metaStoreManagerFactory =
           new InMemoryPolarisMetaStoreManagerFactory(