Support Aws STS AssumeRoleWithWebIdentity

.gitignore

@@ -106,10 +106,11 @@ hs_err_pid*
 
 # Python Virtual Env
 .venv
+venv
 
 # Development Experience
 
 # You can create runtime/defaults/src/main/resources/application-local.properties file
 # to override default properties for local development.
 # And then use `./gradlew run -Dquarkus.profile=local` to run Polaris with dev profile.
-application-local.properties
+application-local.properties

polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisPrincipal.java

@@ -18,6 +18,7 @@
  */
 package org.apache.polaris.core.auth;
 
+import jakarta.annotation.Nullable;
 import java.security.Principal;
 import java.util.Map;
 import java.util.Set;
@@ -39,7 +40,23 @@ public interface PolarisPrincipal extends Principal {
    * @param roles the set of roles associated with the principal
    */
   static PolarisPrincipal of(PrincipalEntity principalEntity, Set<String> roles) {
-    return of(principalEntity.getName(), principalEntity.getInternalPropertiesAsMap(), roles);
+    return of(principalEntity.getName(), principalEntity.getInternalPropertiesAsMap(), roles, null);
+  }
+
+  /**
+   * Creates a new instance of {@link PolarisPrincipal} from the given {@link PrincipalEntity} and
+   * roles.
+   *
+   * <p>The created principal will have the same ID and name as the {@link PrincipalEntity}, and its
+   * properties will be derived from the internal properties of the entity.
+   *
+   * @param principalEntity the principal entity representing the user or service
+   * @param roles the set of roles associated with the principal
+   * @param token the access token of the current user
+   */
+  static PolarisPrincipal of(PrincipalEntity principalEntity, Set<String> roles, String token) {
+    return of(
+        principalEntity.getName(), principalEntity.getInternalPropertiesAsMap(), roles, token);
   }
 
   /**
@@ -51,9 +68,24 @@ static PolarisPrincipal of(PrincipalEntity principalEntity, Set<String> roles) {
    * @param roles the set of roles associated with the principal
    */
   static PolarisPrincipal of(String name, Map<String, String> properties, Set<String> roles) {
+    return of(name, properties, roles, null);
+  }
+
+  /**
+   * Creates a new instance of {@link PolarisPrincipal} with the specified ID, name, roles, and
+   * properties.
+   *
+   * @param name the name of the principal
+   * @param properties additional properties associated with the principal
+   * @param roles the set of roles associated with the principal
+   * @param token the access token of the current user
+   */
+  static PolarisPrincipal of(
+      String name, Map<String, String> properties, Set<String> roles, @Nullable String token) {
     return ImmutablePolarisPrincipal.builder()
         .name(name)
         .properties(properties)
+        .token(token)
         .roles(roles)
         .build();
   }
@@ -74,4 +106,8 @@ static PolarisPrincipal of(String name, Map<String, String> properties, Set<Stri
    * as permissions, preferences, or other metadata.
    */
   Map<String, String> getProperties();
+
+  /** Returns the access token of the current user. */
+  @Nullable
+  String getToken();
 }

polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java

@@ -1600,7 +1600,8 @@ public void deletePrincipalSecrets(
       boolean allowListOperation,
       @Nonnull Set<String> allowedReadLocations,
       @Nonnull Set<String> allowedWriteLocations,
-      Optional<String> refreshCredentialsEndpoint) {
+      Optional<String> refreshCredentialsEndpoint,
+      Optional<String> token) {
 
     // get meta store session we should be using
     BasePersistence ms = callCtx.getMetaStore();
@@ -1641,7 +1642,8 @@ public void deletePrincipalSecrets(
               allowListOperation,
               allowedReadLocations,
               allowedWriteLocations,
-              refreshCredentialsEndpoint);
+              refreshCredentialsEndpoint,
+              token);
       return new ScopedCredentialsResult(storageAccessConfig);
     } catch (Exception ex) {
       return new ScopedCredentialsResult(

polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java

@@ -324,7 +324,8 @@ public void deletePrincipalSecrets(
       boolean allowListOperation,
       @Nonnull Set<String> allowedReadLocations,
       @Nonnull Set<String> allowedWriteLocations,
-      Optional<String> refreshCredentialsEndpoint) {
+      Optional<String> refreshCredentialsEndpoint,
+      Optional<String> token) {
     return delegate.getSubscopedCredsForEntity(
         callCtx,
         catalogId,
@@ -333,7 +334,8 @@ public void deletePrincipalSecrets(
         allowListOperation,
         allowedReadLocations,
         allowedWriteLocations,
-        refreshCredentialsEndpoint);
+        refreshCredentialsEndpoint,
+        token);
   }
 
   @Override

polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java

@@ -2098,7 +2098,8 @@ private PolarisEntityResolver resolveSecurableToRoleGrant(
       boolean allowListOperation,
       @Nonnull Set<String> allowedReadLocations,
       @Nonnull Set<String> allowedWriteLocations,
-      Optional<String> refreshCredentialsEndpoint) {
+      Optional<String> refreshCredentialsEndpoint,
+      Optional<String> token) {
 
     // get meta store session we should be using
     TransactionalPersistence ms = ((TransactionalPersistence) callCtx.getMetaStore());
@@ -2134,7 +2135,8 @@ private PolarisEntityResolver resolveSecurableToRoleGrant(
               allowListOperation,
               allowedReadLocations,
               allowedWriteLocations,
-              refreshCredentialsEndpoint);
+              refreshCredentialsEndpoint,
+              token);
       return new ScopedCredentialsResult(storageAccessConfig);
     } catch (Exception ex) {
       return new ScopedCredentialsResult(

polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisCredentialVendor.java

@@ -53,5 +53,6 @@ ScopedCredentialsResult getSubscopedCredsForEntity(
       boolean allowListOperation,
       @Nonnull Set<String> allowedReadLocations,
       @Nonnull Set<String> allowedWriteLocations,
-      Optional<String> refreshCredentialsEndpoint);
+      Optional<String> refreshCredentialsEndpoint,
+      Optional<String> token);
 }

polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisStorageIntegration.java

@@ -67,7 +67,8 @@ public abstract StorageAccessConfig getSubscopedCreds(
       boolean allowListOperation,
       @Nonnull Set<String> allowedReadLocations,
       @Nonnull Set<String> allowedWriteLocations,
-      Optional<String> refreshCredentialsEndpoint);
+      Optional<String> refreshCredentialsEndpoint,
+      Optional<String> token);
 
   /**
    * Validate access for the provided operation actions and locations.

polaris-core/src/main/java/org/apache/polaris/core/storage/StorageCredentialsVendor.java

@@ -67,7 +67,8 @@ public ScopedCredentialsResult getSubscopedCredsForEntity(
       boolean allowListOperation,
       @Nonnull Set<String> allowedReadLocations,
       @Nonnull Set<String> allowedWriteLocations,
-      Optional<String> refreshCredentialsEndpoint) {
+      Optional<String> refreshCredentialsEndpoint,
+      Optional<String> token) {
     return polarisCredentialVendor.getSubscopedCredsForEntity(
         callContext.getPolarisCallContext(),
         entity.getCatalogId(),
@@ -76,6 +77,7 @@ public ScopedCredentialsResult getSubscopedCredsForEntity(
         allowListOperation,
         allowedReadLocations,
         allowedWriteLocations,
-        refreshCredentialsEndpoint);
+        refreshCredentialsEndpoint,
+        token);
   }
 }

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

@@ -44,7 +44,8 @@
 import software.amazon.awssdk.policybuilder.iam.IamStatement;
 import software.amazon.awssdk.services.sts.StsClient;
 import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
-import software.amazon.awssdk.services.sts.model.AssumeRoleResponse;
+import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityRequest;
+import software.amazon.awssdk.services.sts.model.Credentials;
 
 /** Credential vendor that supports generating */
 public class AwsCredentialsStorageIntegration
@@ -81,7 +82,8 @@ public StorageAccessConfig getSubscopedCreds(
       boolean allowListOperation,
       @Nonnull Set<String> allowedReadLocations,
       @Nonnull Set<String> allowedWriteLocations,
-      Optional<String> refreshCredentialsEndpoint) {
+      Optional<String> refreshCredentialsEndpoint,
+      Optional<String> token) {
     int storageCredentialDurationSeconds =
         realmConfig.getConfig(STORAGE_CREDENTIAL_DURATION_SECONDS);
     AwsStorageConfigurationInfo storageConfig = config();
@@ -90,35 +92,58 @@ public StorageAccessConfig getSubscopedCreds(
     StorageAccessConfig.Builder accessConfig = StorageAccessConfig.builder();
 
     if (shouldUseSts(storageConfig)) {
-      AssumeRoleRequest.Builder request =
-          AssumeRoleRequest.builder()
-              .externalId(storageConfig.getExternalId())
-              .roleArn(storageConfig.getRoleARN())
-              .roleSessionName("PolarisAwsCredentialsStorageIntegration")
-              .policy(
-                  policyString(
-                          storageConfig,
-                          allowListOperation,
-                          allowedReadLocations,
-                          allowedWriteLocations,
-                          region,
-                          accountId)
-                      .toJson())
-              .durationSeconds(storageCredentialDurationSeconds);
-      credentialsProvider.ifPresent(
-          cp -> request.overrideConfiguration(b -> b.credentialsProvider(cp)));
 
       @SuppressWarnings("resource")
       // Note: stsClientProvider returns "thin" clients that do not need closing
       StsClient stsClient =
           stsClientProvider.stsClient(StsDestination.of(storageConfig.getStsEndpointUri(), region));
-
-      AssumeRoleResponse response = stsClient.assumeRole(request.build());
-      accessConfig.put(StorageAccessProperty.AWS_KEY_ID, response.credentials().accessKeyId());
-      accessConfig.put(
-          StorageAccessProperty.AWS_SECRET_KEY, response.credentials().secretAccessKey());
-      accessConfig.put(StorageAccessProperty.AWS_TOKEN, response.credentials().sessionToken());
-      Optional.ofNullable(response.credentials().expiration())
+      Credentials credentials;
+      if (Boolean.TRUE.equals(storageConfig.getPropagateApiUserIdentity())) {
+        AssumeRoleWithWebIdentityRequest.Builder request =
+            AssumeRoleWithWebIdentityRequest.builder()
+                .webIdentityToken(
+                    token.orElseThrow(
+                        () ->
+                            new IllegalArgumentException(
+                                "Token must be provided when PROPAGATE_API_USER_IDENTITY is true")))
+                .roleArn(storageConfig.getRoleARN())
+                .roleSessionName("PolarisAwsCredentialsStorageIntegration")
+                .policy(
+                    policyString(
+                            storageConfig,
+                            allowListOperation,
+                            allowedReadLocations,
+                            allowedWriteLocations,
+                            region,
+                            accountId)
+                        .toJson())
+                .durationSeconds(storageCredentialDurationSeconds);
+
+        credentials = stsClient.assumeRoleWithWebIdentity(request.build()).credentials();
+      } else {
+        AssumeRoleRequest.Builder request =
+            AssumeRoleRequest.builder()
+                .externalId(storageConfig.getExternalId())
+                .roleArn(storageConfig.getRoleARN())
+                .roleSessionName("PolarisAwsCredentialsStorageIntegration")
+                .policy(
+                    policyString(
+                            storageConfig,
+                            allowListOperation,
+                            allowedReadLocations,
+                            allowedWriteLocations,
+                            region,
+                            accountId)
+                        .toJson())
+                .durationSeconds(storageCredentialDurationSeconds);
+        credentialsProvider.ifPresent(
+            cp -> request.overrideConfiguration(b -> b.credentialsProvider(cp)));
+        credentials = stsClient.assumeRole(request.build()).credentials();
+      }
+      accessConfig.put(StorageAccessProperty.AWS_KEY_ID, credentials.accessKeyId());
+      accessConfig.put(StorageAccessProperty.AWS_SECRET_KEY, credentials.secretAccessKey());
+      accessConfig.put(StorageAccessProperty.AWS_TOKEN, credentials.sessionToken());
+      Optional.ofNullable(credentials.expiration())
           .ifPresent(
               i -> {
                 accessConfig.put(

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

@@ -64,6 +64,8 @@ public String getFileIoImplClassName() {
   @Nullable
   public abstract String getRoleARN();
 
+  public abstract @Nullable Boolean getPropagateApiUserIdentity();
+
   /** KMS Key ARN for server-side encryption,used for writes, optional */
   @Nullable
   public abstract String getCurrentKmsKey();

polaris-core/src/main/java/org/apache/polaris/core/storage/azure/AzureCredentialsStorageIntegration.java

@@ -78,7 +78,8 @@ public StorageAccessConfig getSubscopedCreds(
       boolean allowListOperation,
       @Nonnull Set<String> allowedReadLocations,
       @Nonnull Set<String> allowedWriteLocations,
-      Optional<String> refreshCredentialsEndpoint) {
+      Optional<String> refreshCredentialsEndpoint,
+      Optional<String> token) {
     String loc =
         !allowedWriteLocations.isEmpty()
             ? allowedWriteLocations.stream().findAny().orElse(null)

polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCache.java

@@ -109,7 +109,8 @@ public StorageAccessConfig getOrGenerateSubScopeCreds(
       boolean allowListOperation,
       @Nonnull Set<String> allowedReadLocations,
       @Nonnull Set<String> allowedWriteLocations,
-      Optional<String> refreshCredentialsEndpoint) {
+      Optional<String> refreshCredentialsEndpoint,
+      Optional<String> token) {
     RealmContext realmContext = storageCredentialsVendor.getRealmContext();
     RealmConfig realmConfig = storageCredentialsVendor.getRealmConfig();
     if (!isTypeSupported(polarisEntity.getType())) {
@@ -134,7 +135,8 @@ public StorageAccessConfig getOrGenerateSubScopeCreds(
                   allowListOperation,
                   allowedReadLocations,
                   allowedWriteLocations,
-                  refreshCredentialsEndpoint);
+                  refreshCredentialsEndpoint,
+                  token);
           if (scopedCredentialsResult.isSuccess()) {
             long maxCacheDurationMs = maxCacheDurationMs(realmConfig);
             return new StorageCredentialCacheEntry(

polaris-core/src/main/java/org/apache/polaris/core/storage/gcp/GcpCredentialsStorageIntegration.java

@@ -77,7 +77,8 @@ public StorageAccessConfig getSubscopedCreds(
       boolean allowListOperation,
       @Nonnull Set<String> allowedReadLocations,
       @Nonnull Set<String> allowedWriteLocations,
-      Optional<String> refreshCredentialsEndpoint) {
+      Optional<String> refreshCredentialsEndpoint,
+      Optional<String> token) {
     try {
       sourceCredentials.refresh();
     } catch (IOException e) {
@@ -93,9 +94,9 @@ public StorageAccessConfig getSubscopedCreds(
             .setSourceCredential(sourceCredentials)
             .setCredentialAccessBoundary(accessBoundary)
             .build();
-    AccessToken token;
+    AccessToken accessToken;
     try {
-      token = credentials.refreshAccessToken();
+      accessToken = credentials.refreshAccessToken();
     } catch (IOException e) {
       LOGGER
           .atError()
@@ -110,10 +111,10 @@ public StorageAccessConfig getSubscopedCreds(
     // If expires_in missing, use source credential's expire time, which require another api call to
     // get.
     StorageAccessConfig.Builder accessConfig = StorageAccessConfig.builder();
-    accessConfig.put(StorageAccessProperty.GCS_ACCESS_TOKEN, token.getTokenValue());
+    accessConfig.put(StorageAccessProperty.GCS_ACCESS_TOKEN, accessToken.getTokenValue());
     accessConfig.put(
         StorageAccessProperty.GCS_ACCESS_TOKEN_EXPIRES_AT,
-        String.valueOf(token.getExpirationTime().getTime()));
+        String.valueOf(accessToken.getExpirationTime().getTime()));
 
     refreshCredentialsEndpoint.ifPresent(
         endpoint -> {

polaris-core/src/test/java/org/apache/polaris/core/storage/InMemoryStorageIntegrationTest.java

@@ -199,7 +199,8 @@ public StorageAccessConfig getSubscopedCreds(
         boolean allowListOperation,
         @Nonnull Set<String> allowedReadLocations,
         @Nonnull Set<String> allowedWriteLocations,
-        Optional<String> refreshCredentialsEndpoint) {
+        Optional<String> refreshCredentialsEndpoint,
+        Optional<String> token) {
       return null;
     }
   }