Allow pinning to WebIdentityTokenCredentialsProvider

Allow users to only use the WebIdentityTokenCredentialsProvider instead of the default credentials provider chain.
trinodb · Jun 11, 2024 · ad9aa23 · ad9aa23
1 parent 48e6ae4
commit ad9aa23
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 1 deletion.
diff --git a/docs/src/main/sphinx/object-storage/file-system-s3.md b/docs/src/main/sphinx/object-storage/file-system-s3.md
@@ -78,6 +78,12 @@ support:
 * - `s3.max-error-retries`
   - Specifies maximum number of retries the client will make on errors.
     Defaults to `10`.
+* - `s3.use-web-identity-token-credentials-provider`
+  - Set to `true` to only use the web identity token credentials provider,
+    instead of the default providers chain. This can be useful when running
+    Trino on Amazon EKS and using [IAM roles for service accounts
+    (IRSA)](https://docs.aws.amazon.com/eks/latest/userguide/iam-roles-for-service-accounts.html)
+    Defaults to `false`.
 :::
 
 ## Authentication

diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemConfig.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemConfig.java
@@ -88,6 +88,7 @@ public static software.amazon.awssdk.core.retry.RetryMode getRetryMode(RetryMode
     private String stsRegion;
     private S3SseType sseType = S3SseType.NONE;
     private String sseKmsKeyId;
+    private boolean useWebIdentityTokenCredentialsProvider;
     private DataSize streamingPartSize = DataSize.of(16, MEGABYTE);
     private boolean requesterPays;
     private Integer maxConnections;
@@ -294,6 +295,18 @@ public S3FileSystemConfig setSseKmsKeyId(String sseKmsKeyId)
         return this;
     }
 
+    public boolean isUseWebIdentityTokenCredentialsProvider()
+    {
+        return useWebIdentityTokenCredentialsProvider;
+    }
+
+    @Config("s3.use-web-identity-token-credentials-provider")
+    public S3FileSystemConfig setUseWebIdentityTokenCredentialsProvider(boolean useWebIdentityTokenCredentialsProvider)
+    {
+        this.useWebIdentityTokenCredentialsProvider = useWebIdentityTokenCredentialsProvider;
+        return this;
+    }
+
     @NotNull
     @MinDataSize("5MB")
     @MaxDataSize("256MB")

diff --git a/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemFactory.java b/lib/trino-filesystem-s3/src/main/java/io/trino/filesystem/s3/S3FileSystemFactory.java
@@ -34,6 +34,7 @@
 import software.amazon.awssdk.services.sts.StsClient;
 import software.amazon.awssdk.services.sts.StsClientBuilder;
 import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
+import software.amazon.awssdk.services.sts.auth.StsWebIdentityTokenFileCredentialsProvider;
 
 import java.net.URI;
 import java.util.Optional;
@@ -72,7 +73,13 @@ public S3FileSystemFactory(OpenTelemetry openTelemetry, S3FileSystemConfig confi
 
         Optional<StaticCredentialsProvider> staticCredentialsProvider = getStaticCredentialsProvider(config);
 
-        if (config.getIamRole() != null) {
+        if (config.isUseWebIdentityTokenCredentialsProvider()) {
+            s3.credentialsProvider(StsWebIdentityTokenFileCredentialsProvider.builder()
+                    .stsClient(getStsClient(config, staticCredentialsProvider))
+                    .asyncCredentialUpdateEnabled(true)
+                    .build());
+        }
+        else if (config.getIamRole() != null) {
             s3.credentialsProvider(StsAssumeRoleCredentialsProvider.builder()
                     .refreshRequest(request -> request
                             .roleArn(config.getIamRole())

diff --git a/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemConfig.java b/lib/trino-filesystem-s3/src/test/java/io/trino/filesystem/s3/TestS3FileSystemConfig.java
@@ -52,6 +52,7 @@ public void testDefaults()
                 .setRetryMode(LEGACY)
                 .setMaxErrorRetries(10)
                 .setSseKmsKeyId(null)
+                .setUseWebIdentityTokenCredentialsProvider(false)
                 .setStreamingPartSize(DataSize.of(16, MEGABYTE))
                 .setRequesterPays(false)
                 .setMaxConnections(null)
@@ -83,6 +84,7 @@ public void testExplicitPropertyMappings()
                 .put("s3.max-error-retries", "12")
                 .put("s3.sse.type", "KMS")
                 .put("s3.sse.kms-key-id", "mykey")
+                .put("s3.use-web-identity-token-credentials-provider", "true")
                 .put("s3.streaming.part-size", "42MB")
                 .put("s3.requester-pays", "true")
                 .put("s3.max-connections", "42")
@@ -112,6 +114,7 @@ public void testExplicitPropertyMappings()
                 .setMaxErrorRetries(12)
                 .setSseType(S3SseType.KMS)
                 .setSseKmsKeyId("mykey")
+                .setUseWebIdentityTokenCredentialsProvider(true)
                 .setRequesterPays(true)
                 .setMaxConnections(42)
                 .setConnectionTtl(new Duration(1, MINUTES))