IAM role support for exchange spooling on S3

docs/src/main/sphinx/admin/fault-tolerant-execution.rst

@@ -373,6 +373,14 @@ the property may be configured for:
        be ignored for other S3 storage systems.
      -
      - AWS S3, GCS
+   * - ``exchange.s3.iam-role``
+     - IAM role to assume.
+     -
+     - AWS S3, GCS
+   * - ``exchange.s3.external-id``
+     - External ID for the IAM role trust policy.
+     -
+     - AWS S3, GCS
    * - ``exchange.s3.region``
      - Region of the S3 bucket.
      -

plugin/trino-exchange-filesystem/pom.xml

@@ -297,14 +297,12 @@
 
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
-            <artifactId>utils</artifactId>
+            <artifactId>sts</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>
+            <artifactId>utils</artifactId>
         </dependency>
 
         <!-- Trino SPI -->

plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/s3/ExchangeS3Config.java

@@ -37,6 +37,8 @@ public class ExchangeS3Config
 {
     private String s3AwsAccessKey;
     private String s3AwsSecretKey;
+    private Optional<String> s3IamRole = Optional.empty();
+    private Optional<String> s3ExternalId = Optional.empty();
     private Optional<Region> s3Region = Optional.empty();
     private Optional<String> s3Endpoint = Optional.empty();
     private int s3MaxErrorRetries = 10;
@@ -74,6 +76,32 @@ public ExchangeS3Config setS3AwsSecretKey(String s3AwsSecretKey)
         return this;
     }
 
+    public Optional<String> getS3IamRole()
+    {
+        return s3IamRole;
+    }
+
+    @Config("exchange.s3.iam-role")
+    @ConfigDescription("ARN of an IAM role to assume when connecting to S3")
+    public ExchangeS3Config setS3IamRole(String s3IamRole)
+    {
+        this.s3IamRole = Optional.ofNullable(s3IamRole);
+        return this;
+    }
+
+    public Optional<String> getS3ExternalId()
+    {
+        return s3ExternalId;
+    }
+
+    @Config("exchange.s3.external-id")
+    @ConfigDescription("External ID for the IAM role trust policy when connecting to S3")
+    public ExchangeS3Config setS3ExternalId(String s3ExternalId)
+    {
+        this.s3ExternalId = Optional.ofNullable(s3ExternalId);
+        return this;
+    }
+
     public Optional<Region> getS3Region()
     {
         return s3Region;

plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/s3/S3FileSystemExchangeStorage.java

@@ -75,6 +75,10 @@
 import software.amazon.awssdk.services.s3.model.UploadPartRequest;
 import software.amazon.awssdk.services.s3.paginators.ListObjectsV2Iterable;
 import software.amazon.awssdk.services.s3.paginators.ListObjectsV2Publisher;
+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.model.AssumeRoleRequest;
 
 import javax.annotation.PreDestroy;
 import javax.annotation.concurrent.GuardedBy;
@@ -451,8 +455,46 @@ private static boolean isDirectory(URI uri)
 
     private static AwsCredentialsProvider createAwsCredentialsProvider(ExchangeS3Config config)
     {
-        if (config.getS3AwsAccessKey() != null && config.getS3AwsSecretKey() != null) {
-            return StaticCredentialsProvider.create(AwsBasicCredentials.create(config.getS3AwsAccessKey(), config.getS3AwsSecretKey()));
+        String accessKey = config.getS3AwsAccessKey();
+        String secretKey = config.getS3AwsSecretKey();
+
+        if (accessKey == null && secretKey != null) {
+            throw new IllegalArgumentException("AWS access key set but secret is not set; make sure you set exchange.s3.aws-secret-key config property");
+        }
+
+        if (accessKey != null && secretKey == null) {
+            throw new IllegalArgumentException("AWS secret key set but access is not set; make sure you set exchange.s3.aws-access-key config property");
+        }
+
+        if (accessKey != null) {
+            checkArgument(
+                    config.getS3IamRole().isEmpty(),
+                    "IAM role is not compatible with access key based authentication; make sure you set only one of exchange.s3.aws-access-key, exchange.s3.iam-role config properties");
+            checkArgument(
+                    config.getS3ExternalId().isEmpty(),
+                    "External ID is not compatible with access key based authentication; make sure you set only one of exchange.s3.aws-access-key, exchange.s3.external-id config properties");
+
+            return StaticCredentialsProvider.create(AwsBasicCredentials.create(accessKey, secretKey));
+        }
+
+        if (config.getS3ExternalId().isPresent() && config.getS3IamRole().isEmpty()) {
+            throw new IllegalArgumentException("External ID can only be used with IAM role based authentication; make sure you set exchange.s3.iam-role config property");
+        }
+
+        if (config.getS3IamRole().isPresent()) {
+            AssumeRoleRequest.Builder assumeRoleRequest = AssumeRoleRequest.builder()
+                    .roleArn(config.getS3IamRole().get())
+                    .roleSessionName("trino-exchange");
+            config.getS3ExternalId().ifPresent(assumeRoleRequest::externalId);
+
+            StsClientBuilder stsClientBuilder = StsClient.builder();
+            config.getS3Region().ifPresent(stsClientBuilder::region);
+
+            return StsAssumeRoleCredentialsProvider.builder()
+                    .stsClient(stsClientBuilder.build())
+                    .refreshRequest(assumeRoleRequest.build())
+                    .asyncCredentialUpdateEnabled(true)
+                    .build();
         }
 
         return DefaultCredentialsProvider.create();

plugin/trino-exchange-filesystem/src/test/java/io/trino/plugin/exchange/filesystem/s3/TestExchangeS3Config.java

@@ -36,6 +36,8 @@ public void testDefaults()
         assertRecordedDefaults(recordDefaults(ExchangeS3Config.class)
                 .setS3AwsAccessKey(null)
                 .setS3AwsSecretKey(null)
+                .setS3IamRole(null)
+                .setS3ExternalId(null)
                 .setS3Region(null)
                 .setS3Endpoint(null)
                 .setS3MaxErrorRetries(10)
@@ -54,6 +56,8 @@ public void testExplicitPropertyMappings()
         Map<String, String> properties = ImmutableMap.<String, String>builder()
                 .put("exchange.s3.aws-access-key", "access")
                 .put("exchange.s3.aws-secret-key", "secret")
+                .put("exchange.s3.iam-role", "roleArn")
+                .put("exchange.s3.external-id", "externalId")
                 .put("exchange.s3.region", "us-west-1")
                 .put("exchange.s3.endpoint", "https://s3.us-east-1.amazonaws.com")
                 .put("exchange.s3.max-error-retries", "8")
@@ -69,6 +73,8 @@ public void testExplicitPropertyMappings()
         ExchangeS3Config expected = new ExchangeS3Config()
                 .setS3AwsAccessKey("access")
                 .setS3AwsSecretKey("secret")
+                .setS3IamRole("roleArn")
+                .setS3ExternalId("externalId")
                 .setS3Region("us-west-1")
                 .setS3Endpoint("https://s3.us-east-1.amazonaws.com")
                 .setS3MaxErrorRetries(8)