Skip to content

Commit

Permalink
Allow pinning Glue credentials provider to StsWebIdentityTokenFileCre…
Browse files Browse the repository at this point in the history 
…dentialsProvider

Allow users to only use the StsWebIdentityTokenFileCredentialsProvider instead of
the default credentials provider chain for Glue v2.
  • Loading branch information
@rohanag12
rohanag12 committed Jul 10, 2024
1 parent 34d1bf1 commit a48d6f3
Show file tree
Hide file tree
Showing 3 changed files with 60 additions and 20 deletions.
15 changes: 15 additions & 0 deletions ...trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastoreConfig.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ public class GlueHiveMetastoreConfig
private Optional<String> externalId = Optional.empty();
private Optional<String> awsAccessKey = Optional.empty();
private Optional<String> awsSecretKey = Optional.empty();
private boolean useWebIdentityTokenCredentialsProvider;
private Optional<String> catalogId = Optional.empty();
private int partitionSegments = 5;
private int threads = 40;
Expand Down Expand Up @@ -208,6 +209,20 @@ public GlueHiveMetastoreConfig setAwsSecretKey(String awsSecretKey)
return this;
}

public boolean isUseWebIdentityTokenCredentialsProvider()
{
return useWebIdentityTokenCredentialsProvider;
}

@Config("hive.metastore.glue.use-web-identity-token-credentials-provider")
@ConfigDescription("If true, explicitly use the WebIdentityTokenCredentialsProvider" +
" instead of the default credential provider chain.")
public GlueHiveMetastoreConfig setUseWebIdentityTokenCredentialsProvider(boolean useWebIdentityTokenCredentialsProvider)
{
this.useWebIdentityTokenCredentialsProvider = useWebIdentityTokenCredentialsProvider;
return this;
}

public Optional<String> getCatalogId()
{
return catalogId;
Expand Down
62 changes: 42 additions & 20 deletions plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreModule.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,11 @@
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.EnumSet;
import java.util.Optional;
import java.util.Set;

import static com.google.common.base.Preconditions.checkArgument;
Expand Down Expand Up @@ -123,35 +125,26 @@ public static GlueClient createGlueClient(GlueHiveMetastoreConfig config, OpenTe
.retryPolicy(retry -> retry
.numRetries(config.getMaxGlueErrorRetries())));

if (config.getIamRole().isPresent()) {
StsClientBuilder sts = StsClient.builder();

if (config.getAwsAccessKey().isPresent() && config.getAwsSecretKey().isPresent()) {
sts.credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create(config.getAwsAccessKey().get(), config.getAwsSecretKey().get())));
}

if (config.getGlueStsEndpointUrl().isPresent() && config.getGlueStsRegion().isPresent()) {
sts.endpointOverride(URI.create(config.getGlueStsEndpointUrl().get()))
.region(Region.of(config.getGlueStsRegion().get()));
}
else if (config.getGlueStsRegion().isPresent()) {
sts.region(Region.of(config.getGlueStsRegion().get()));
}
else if (config.getPinGlueClientToCurrentRegion()) {
sts.region(DefaultAwsRegionProviderChain.builder().build().getRegion());
}
Optional<StaticCredentialsProvider> staticCredentialsProvider = getStaticCredentialsProvider(config);

if (config.isUseWebIdentityTokenCredentialsProvider()) {
glue.credentialsProvider(StsWebIdentityTokenFileCredentialsProvider.builder()
.stsClient(getStsClient(config, staticCredentialsProvider))
.asyncCredentialUpdateEnabled(true)
.build());
}
else if (config.getIamRole().isPresent()) {
glue.credentialsProvider(StsAssumeRoleCredentialsProvider.builder()
.refreshRequest(request -> request
.roleArn(config.getIamRole().get())
.roleSessionName("trino-session")
.externalId(config.getExternalId().orElse(null)))
.stsClient(sts.build())
.stsClient(getStsClient(config, staticCredentialsProvider))
.asyncCredentialUpdateEnabled(true)
.build());
}
else if (config.getAwsAccessKey().isPresent() && config.getAwsSecretKey().isPresent()) {
glue.credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create(config.getAwsAccessKey().get(), config.getAwsSecretKey().get())));
else {
staticCredentialsProvider.ifPresent(glue::credentialsProvider);
}

ApacheHttpClient.Builder httpClient = ApacheHttpClient.builder()
Expand All @@ -175,4 +168,33 @@ else if (config.getPinGlueClientToCurrentRegion()) {

return glue.build();
}

private static Optional<StaticCredentialsProvider> getStaticCredentialsProvider(GlueHiveMetastoreConfig config)
{
if (config.getAwsAccessKey().isPresent() && config.getAwsSecretKey().isPresent()) {
return Optional.of(StaticCredentialsProvider.create(
AwsBasicCredentials.create(config.getAwsAccessKey().get(), config.getAwsSecretKey().get())));
}
return Optional.empty();
}

private static StsClient getStsClient(GlueHiveMetastoreConfig config,
Optional<StaticCredentialsProvider> staticCredentialsProvider)
{
StsClientBuilder sts = StsClient.builder();
staticCredentialsProvider.ifPresent(sts::credentialsProvider);

if (config.getGlueStsEndpointUrl().isPresent() && config.getGlueStsRegion().isPresent()) {
sts.endpointOverride(URI.create(config.getGlueStsEndpointUrl().get()))
.region(Region.of(config.getGlueStsRegion().get()));
}
else if (config.getGlueStsRegion().isPresent()) {
sts.region(Region.of(config.getGlueStsRegion().get()));
}
else if (config.getPinGlueClientToCurrentRegion()) {
sts.region(DefaultAwsRegionProviderChain.builder().build().getRegion());
}

return sts.build();
}
}
3 changes: 3 additions & 0 deletions ...o-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueHiveMetastoreConfig.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ void testDefaults()
.setExternalId(null)
.setAwsAccessKey(null)
.setAwsSecretKey(null)
.setUseWebIdentityTokenCredentialsProvider(false)
.setCatalogId(null)
.setPartitionSegments(5)
.setThreads(40)
Expand All @@ -62,6 +63,7 @@ void testExplicitPropertyMapping()
.put("hive.metastore.glue.external-id", "external-id")
.put("hive.metastore.glue.aws-access-key", "ABC")
.put("hive.metastore.glue.aws-secret-key", "DEF")
.put("hive.metastore.glue.use-web-identity-token-credentials-provider", "true")
.put("hive.metastore.glue.catalogid", "0123456789")
.put("hive.metastore.glue.partitions-segments", "10")
.put("hive.metastore.glue.threads", "77")
Expand All @@ -81,6 +83,7 @@ void testExplicitPropertyMapping()
.setExternalId("external-id")
.setAwsAccessKey("ABC")
.setAwsSecretKey("DEF")
.setUseWebIdentityTokenCredentialsProvider(true)
.setCatalogId("0123456789")
.setPartitionSegments(10)
.setThreads(77)
Expand Down

0 comments on commit a48d6f3

Please sign in to comment.