From 7e566b9ea078645cea88ed1158222cd2690fdc50 Mon Sep 17 00:00:00 2001 From: Zebing Lin Date: Tue, 24 May 2022 15:20:27 -0700 Subject: [PATCH] IAM role support for exchange spooling on S3 --- .../sphinx/admin/fault-tolerant-execution.rst | 8 ++++ plugin/trino-exchange-filesystem/pom.xml | 6 +-- .../filesystem/s3/ExchangeS3Config.java | 28 +++++++++++ .../s3/S3FileSystemExchangeStorage.java | 46 ++++++++++++++++++- .../filesystem/s3/TestExchangeS3Config.java | 6 +++ 5 files changed, 88 insertions(+), 6 deletions(-) diff --git a/docs/src/main/sphinx/admin/fault-tolerant-execution.rst b/docs/src/main/sphinx/admin/fault-tolerant-execution.rst index 43efd2cc7ed2..d61db346bee3 100644 --- a/docs/src/main/sphinx/admin/fault-tolerant-execution.rst +++ b/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. - diff --git a/plugin/trino-exchange-filesystem/pom.xml b/plugin/trino-exchange-filesystem/pom.xml index 03f3bdcdd57a..572b3c8d0b91 100644 --- a/plugin/trino-exchange-filesystem/pom.xml +++ b/plugin/trino-exchange-filesystem/pom.xml @@ -297,14 +297,12 @@ software.amazon.awssdk - utils + sts - software.amazon.awssdk - sts - runtime + utils diff --git a/plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/s3/ExchangeS3Config.java b/plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/s3/ExchangeS3Config.java index 51e2e1c7747d..44c698eaaa9d 100644 --- a/plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/s3/ExchangeS3Config.java +++ b/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 s3IamRole = Optional.empty(); + private Optional s3ExternalId = Optional.empty(); private Optional s3Region = Optional.empty(); private Optional s3Endpoint = Optional.empty(); private int s3MaxErrorRetries = 10; @@ -74,6 +76,32 @@ public ExchangeS3Config setS3AwsSecretKey(String s3AwsSecretKey) return this; } + public Optional 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 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 getS3Region() { return s3Region; diff --git a/plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/s3/S3FileSystemExchangeStorage.java b/plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/s3/S3FileSystemExchangeStorage.java index 188e6b243190..fa8c272972da 100644 --- a/plugin/trino-exchange-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/s3/S3FileSystemExchangeStorage.java +++ b/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(); diff --git a/plugin/trino-exchange-filesystem/src/test/java/io/trino/plugin/exchange/filesystem/s3/TestExchangeS3Config.java b/plugin/trino-exchange-filesystem/src/test/java/io/trino/plugin/exchange/filesystem/s3/TestExchangeS3Config.java index fe1716f7c75d..5dc520c876dd 100644 --- a/plugin/trino-exchange-filesystem/src/test/java/io/trino/plugin/exchange/filesystem/s3/TestExchangeS3Config.java +++ b/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 properties = ImmutableMap.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)