add AWS_SESSION_TOKEN_EXPIRES_AT_MS (#1160)

* add AWS_SESSION_TOKEN_EXPIRES_AT_MS

* spotless

* add expiration time to aws integration tests

---------

Co-authored-by: David Lu <dalu@hubspot.com>

Author: wolflex888

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

@@ -23,6 +23,10 @@ public enum PolarisCredentialProperty {
   AWS_KEY_ID(String.class, "s3.access-key-id", "the aws access key id"),
   AWS_SECRET_KEY(String.class, "s3.secret-access-key", "the aws access key secret"),
   AWS_TOKEN(String.class, "s3.session-token", "the aws scoped access token"),
+  AWS_SESSION_TOKEN_EXPIRES_AT_MS(
+      String.class,
+      "s3.session-token-expires-at-ms",
+      "the time the aws session token expires, in milliseconds"),
   CLIENT_REGION(
       String.class, "client.region", "region to configure client for making requests to AWS"),
 

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

@@ -83,9 +83,13 @@ public EnumMap<PolarisCredentialProperty, String> getSubscopedCreds(
     credentialMap.put(PolarisCredentialProperty.AWS_TOKEN, response.credentials().sessionToken());
     Optional.ofNullable(response.credentials().expiration())
         .ifPresent(
-            i ->
-                credentialMap.put(
-                    PolarisCredentialProperty.EXPIRATION_TIME, String.valueOf(i.toEpochMilli())));
+            i -> {
+              credentialMap.put(
+                  PolarisCredentialProperty.EXPIRATION_TIME, String.valueOf(i.toEpochMilli()));
+              credentialMap.put(
+                  PolarisCredentialProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS,
+                  String.valueOf(i.toEpochMilli()));
+            });
 
     if (storageConfig.getRegion() != null) {
       credentialMap.put(PolarisCredentialProperty.CLIENT_REGION, storageConfig.getRegion());

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

@@ -42,6 +42,10 @@ public long getExpirationTime() {
     if (credsMap.containsKey(PolarisCredentialProperty.GCS_ACCESS_TOKEN_EXPIRES_AT)) {
       return Long.parseLong(credsMap.get(PolarisCredentialProperty.GCS_ACCESS_TOKEN_EXPIRES_AT));
     }
+    if (credsMap.containsKey(PolarisCredentialProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS)) {
+      return Long.parseLong(
+          credsMap.get(PolarisCredentialProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS));
+    }
     if (credsMap.containsKey(PolarisCredentialProperty.EXPIRATION_TIME)) {
       return Long.parseLong(credsMap.get(PolarisCredentialProperty.EXPIRATION_TIME));
     }

polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java

@@ -397,6 +397,7 @@ private static List<ScopedCredentialsResult> getFakeScopedCreds(int number, bool
                   ImmutableMap.<PolarisCredentialProperty, String>builder()
                       .put(PolarisCredentialProperty.AWS_KEY_ID, "key_id_" + finalI)
                       .put(PolarisCredentialProperty.AWS_SECRET_KEY, "key_secret_" + finalI)
+                      .put(PolarisCredentialProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS, expireTime)
                       .put(PolarisCredentialProperty.EXPIRATION_TIME, expireTime)
                       .buildOrThrow())));
       if (res.size() == number) return res;

polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java

@@ -21,6 +21,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 import jakarta.annotation.Nonnull;
+import java.time.Instant;
 import java.util.EnumMap;
 import java.util.List;
 import java.util.Set;
@@ -49,13 +50,16 @@
 
 class AwsCredentialsStorageIntegrationTest {
 
+  public static final Instant EXPIRE_TIME = Instant.now().plusMillis(3600_000);
+
   public static final AssumeRoleResponse ASSUME_ROLE_RESPONSE =
       AssumeRoleResponse.builder()
           .credentials(
               Credentials.builder()
                   .accessKeyId("accessKey")
                   .secretAccessKey("secretKey")
                   .sessionToken("sess")
+                  .expiration(EXPIRE_TIME)
                   .build())
           .build();
   public static final String AWS_PARTITION = "aws";
@@ -93,7 +97,10 @@ public void testGetSubscopedCreds() {
         .isNotEmpty()
         .containsEntry(PolarisCredentialProperty.AWS_TOKEN, "sess")
         .containsEntry(PolarisCredentialProperty.AWS_KEY_ID, "accessKey")
-        .containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey");
+        .containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey")
+        .containsEntry(
+            PolarisCredentialProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS,
+            String.valueOf(EXPIRE_TIME.toEpochMilli()));
   }
 
   @ParameterizedTest
@@ -255,7 +262,10 @@ public void testGetSubscopedCredsInlinePolicy(String awsPartition) {
             .isNotEmpty()
             .containsEntry(PolarisCredentialProperty.AWS_TOKEN, "sess")
             .containsEntry(PolarisCredentialProperty.AWS_KEY_ID, "accessKey")
-            .containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey");
+            .containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey")
+            .containsEntry(
+                PolarisCredentialProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS,
+                String.valueOf(EXPIRE_TIME.toEpochMilli()));
         break;
       default:
         throw new IllegalArgumentException("Unknown aws partition: " + awsPartition);
@@ -353,7 +363,10 @@ public void testGetSubscopedCredsInlinePolicyWithoutList() {
         .isNotEmpty()
         .containsEntry(PolarisCredentialProperty.AWS_TOKEN, "sess")
         .containsEntry(PolarisCredentialProperty.AWS_KEY_ID, "accessKey")
-        .containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey");
+        .containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey")
+        .containsEntry(
+            PolarisCredentialProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS,
+            String.valueOf(EXPIRE_TIME.toEpochMilli()));
   }
 
   @Test
@@ -445,7 +458,10 @@ public void testGetSubscopedCredsInlinePolicyWithoutWrites() {
         .isNotEmpty()
         .containsEntry(PolarisCredentialProperty.AWS_TOKEN, "sess")
         .containsEntry(PolarisCredentialProperty.AWS_KEY_ID, "accessKey")
-        .containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey");
+        .containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey")
+        .containsEntry(
+            PolarisCredentialProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS,
+            String.valueOf(EXPIRE_TIME.toEpochMilli()));
   }
 
   @Test
@@ -507,7 +523,10 @@ public void testGetSubscopedCredsInlinePolicyWithEmptyReadAndWrite() {
         .isNotEmpty()
         .containsEntry(PolarisCredentialProperty.AWS_TOKEN, "sess")
         .containsEntry(PolarisCredentialProperty.AWS_KEY_ID, "accessKey")
-        .containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey");
+        .containsEntry(PolarisCredentialProperty.AWS_SECRET_KEY, "secretKey")
+        .containsEntry(
+            PolarisCredentialProperty.AWS_SESSION_TOKEN_EXPIRES_AT_MS,
+            String.valueOf(EXPIRE_TIME.toEpochMilli()));
   }
 
   @ParameterizedTest