Skip to content

Add IAM Role support for exchange spooling on S3 #12536

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
losipiuk merged 2 commits into from
May 27, 2022
Merged

Add IAM Role support for exchange spooling on S3 #12536

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fa24173
Remove redundant config exchange.s3.use-web-identity-token-credentials
linzebing May 24, 2022
41f4084
IAM role support for exchange spooling on S3
linzebing May 24, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions docs/src/main/sphinx/admin/fault-tolerant-execution.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
-
Expand Down
6 changes: 2 additions & 4 deletions plugin/trino-exchange-filesystem/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 -->
Expand Down
41 changes: 28 additions & 13 deletions ...nge-filesystem/src/main/java/io/trino/plugin/exchange/filesystem/s3/ExchangeS3Config.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,9 +37,10 @@ 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 boolean s3UseWebIdentityTokenCredentials;
private int s3MaxErrorRetries = 10;
// Default to S3 multi-part upload minimum size to avoid excessive memory consumption from buffering
private DataSize s3UploadPartSize = DataSize.of(5, MEGABYTE);
Expand Down Expand Up @@ -75,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;
Expand Down Expand Up @@ -102,18 +129,6 @@ 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()
{
Expand Down
47 changes: 42 additions & 5 deletions ...tem/src/main/java/io/trino/plugin/exchange/filesystem/s3/S3FileSystemExchangeStorage.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,6 @@
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;
Expand Down Expand Up @@ -76,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;
Expand Down Expand Up @@ -452,12 +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.isS3UseWebIdentityTokenCredentials()) {
return WebIdentityTokenFileCredentialsProvider.create();
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()) {
Copy link
Copy Markdown
Member

@losipiuk losipiuk May 25, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you validate config here that different configuration options are mutually exclusive?

  • both config.getS3AwsAccessKey() and config.getS3AwsSecretKey() must be either set or unset
  • if accessKey and secretKey are set then other auth related config options are not set.
  • externalId can only be set if iaRole is set

Copy link
Copy Markdown
Member Author

@linzebing linzebing May 25, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure if we need to do that because TrinoS3FileSystem didn't perform such checks. Also they are not really mutually exclusive, it's just that priority is different (static credentials > iamRole > default credentials)

Copy link
Copy Markdown
Member

@losipiuk losipiuk May 25, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the fact that TS3FS does not do that is an omission.
Also unless I am missing sth they are mutually exclusive. If you set S3 Access and Secret key it does not matter what you set in IamRole as you would not get to line in code when you use that.
You would return in return StaticCredentialsProvider.create(AwsBasicCredentials.create(config.getS3AwsAccessKey(), config.getS3AwsSecretKey())); right?

Hence if user provides both access/secretKey and IAM role it is fishy, I would throw in such case.

Copy link
Copy Markdown
Member Author

@linzebing linzebing May 25, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK, let me add the validations

Copy link
Copy Markdown
Member

@losipiuk losipiuk May 26, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks. I made the assertions a bit more explicit and used IllegalArgumentException instead of verify (verify should rather be used for consistency assertions - for things that would never happen if system is working correctly).
PTAL if I did not mess anything up.

Copy link
Copy Markdown
Member Author

@linzebing linzebing May 26, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm fine with this, but I think my version is more concise. We can merge it still

Copy link
Copy Markdown
Member

@losipiuk losipiuk May 27, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah it is :) But I think actionable error messages are more important than concise code.
I am also dumb - and reading concise code is hard for me :P

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();
Expand Down
9 changes: 6 additions & 3 deletions ...filesystem/src/test/java/io/trino/plugin/exchange/filesystem/s3/TestExchangeS3Config.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,9 +36,10 @@ public void testDefaults()
assertRecordedDefaults(recordDefaults(ExchangeS3Config.class)
.setS3AwsAccessKey(null)
.setS3AwsSecretKey(null)
.setS3IamRole(null)
.setS3ExternalId(null)
.setS3Region(null)
.setS3Endpoint(null)
.setS3UseWebIdentityTokenCredentials(false)
.setS3MaxErrorRetries(10)
.setS3UploadPartSize(DataSize.of(5, MEGABYTE))
.setStorageClass(StorageClass.STANDARD)
Expand All @@ -55,9 +56,10 @@ 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.use-web-identity-token-credentials", "true")
.put("exchange.s3.max-error-retries", "8")
.put("exchange.s3.upload.part-size", "10MB")
.put("exchange.s3.storage-class", "REDUCED_REDUNDANCY")
Expand All @@ -71,9 +73,10 @@ 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")
.setS3UseWebIdentityTokenCredentials(true)
.setS3MaxErrorRetries(8)
.setS3UploadPartSize(DataSize.of(10, MEGABYTE))
.setStorageClass(StorageClass.REDUCED_REDUNDANCY)
Expand Down