Support WebIdentityTokenFileCredentialsProvider in exchange spooling

trinodb · Mar 22, 2022 · 4ae3121 · 4ae3121
1 parent df66d54
commit 4ae3121
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 0 deletions.
diff --git a/plugin/trino-exchange/pom.xml b/plugin/trino-exchange/pom.xml
@@ -131,6 +131,13 @@
             <artifactId>utils</artifactId>
         </dependency>
 
+        <!-- use of WebIdentityTokenFileCredentialsProvider requires the 'sts' module to be on the classpath -->
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>sts</artifactId>
+            <scope>runtime</scope>
+        </dependency>
+
         <!-- Trino SPI -->
         <dependency>
             <groupId>io.trino</groupId>

diff --git a/plugin/trino-exchange/src/main/java/io/trino/plugin/exchange/s3/ExchangeS3Config.java b/plugin/trino-exchange/src/main/java/io/trino/plugin/exchange/s3/ExchangeS3Config.java
@@ -35,6 +35,7 @@ public class ExchangeS3Config
     private String s3AwsSecretKey;
     private Optional<Region> s3Region = Optional.empty();
     private Optional<String> s3Endpoint = Optional.empty();
+    private boolean s3UseWebIdentityTokenCredentials;
     private int s3MaxErrorRetries = 3;
     // Default to S3 multi-part upload minimum size to avoid excessive memory consumption from buffering
     private DataSize s3UploadPartSize = DataSize.of(5, MEGABYTE);
@@ -91,6 +92,18 @@ public ExchangeS3Config setS3Endpoint(String s3Endpoint)
         return this;
     }
 
+    public boolean isS3UseWebIdentityTokenCredentials()
+    {
+        return s3UseWebIdentityTokenCredentials;
+    }
+
+    @Config("exchange.s3.use-web-identity-token-credentials")
+    public ExchangeS3Config setS3UseWebIdentityTokenCredentials(boolean s3UseWebIdentityTokenCredentials)
+    {
+        this.s3UseWebIdentityTokenCredentials = s3UseWebIdentityTokenCredentials;
+        return this;
+    }
+
     @Min(0)
     public int getS3MaxErrorRetries()
     {

diff --git a/...trino-exchange/src/main/java/io/trino/plugin/exchange/s3/S3FileSystemExchangeStorage.java b/...trino-exchange/src/main/java/io/trino/plugin/exchange/s3/S3FileSystemExchangeStorage.java
@@ -32,6 +32,7 @@
 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.auth.credentials.WebIdentityTokenFileCredentialsProvider;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
 import software.amazon.awssdk.core.retry.RetryPolicy;
@@ -362,6 +363,11 @@ private static AwsCredentialsProvider createAwsCredentialsProvider(ExchangeS3Con
         if (config.getS3AwsAccessKey() != null && config.getS3AwsSecretKey() != null) {
             return StaticCredentialsProvider.create(AwsBasicCredentials.create(config.getS3AwsAccessKey(), config.getS3AwsSecretKey()));
         }
+
+        if (config.isS3UseWebIdentityTokenCredentials()) {
+            return WebIdentityTokenFileCredentialsProvider.create();
+        }
+
         return DefaultCredentialsProvider.create();
     }
 

diff --git a/plugin/trino-exchange/src/test/java/io/trino/plugin/exchange/s3/TestExchangeS3Config.java b/plugin/trino-exchange/src/test/java/io/trino/plugin/exchange/s3/TestExchangeS3Config.java
@@ -34,6 +34,7 @@ public void testDefaults()
                 .setS3AwsSecretKey(null)
                 .setS3Region(null)
                 .setS3Endpoint(null)
+                .setS3UseWebIdentityTokenCredentials(false)
                 .setS3MaxErrorRetries(3)
                 .setS3UploadPartSize(DataSize.of(5, MEGABYTE)));
     }
@@ -46,6 +47,7 @@ public void testExplicitPropertyMappings()
                 .put("exchange.s3.aws-secret-key", "secret")
                 .put("exchange.s3.region", "us-west-1")
                 .put("exchange.s3.endpoint", "https://s3.us-east-1.amazonaws.com")
+                .put("exchange.s3.use-web-identity-token-credentials", "true")
                 .put("exchange.s3.max-error-retries", "8")
                 .put("exchange.s3.upload.part-size", "10MB")
                 .buildOrThrow();
@@ -55,6 +57,7 @@ public void testExplicitPropertyMappings()
                 .setS3AwsSecretKey("secret")
                 .setS3Region("us-west-1")
                 .setS3Endpoint("https://s3.us-east-1.amazonaws.com")
+                .setS3UseWebIdentityTokenCredentials(true)
                 .setS3MaxErrorRetries(8)
                 .setS3UploadPartSize(DataSize.of(10, MEGABYTE));