From f7e0a0d345086cf59ab69c0e1477c500a267b3c5 Mon Sep 17 00:00:00 2001 From: praveenkrishna Date: Mon, 14 Mar 2022 13:03:24 +0530 Subject: [PATCH] Make AWSCredentialProvider injectable for GlueClient --- .../glue/GlueCredentialsProvider.java | 79 +++++++++++++++++++ .../metastore/glue/GlueHiveMetastore.java | 50 ++---------- .../metastore/glue/GlueMetastoreModule.java | 2 + .../metastore/glue/TestHiveGlueMetastore.java | 2 + .../GlueIcebergTableOperationsProvider.java | 6 +- .../glue/IcebergGlueCatalogModule.java | 3 + .../catalog/glue/TrinoGlueCatalogFactory.java | 5 +- .../iceberg/TestSharedGlueMetastore.java | 2 + .../iceberg/TestTrinoGlueCatalogTest.java | 5 +- 9 files changed, 104 insertions(+), 50 deletions(-) create mode 100644 plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueCredentialsProvider.java diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueCredentialsProvider.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueCredentialsProvider.java new file mode 100644 index 000000000000..2653b8ee8334 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueCredentialsProvider.java @@ -0,0 +1,79 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.hive.metastore.glue; + +import com.amazonaws.auth.AWSCredentialsProvider; +import com.amazonaws.auth.AWSStaticCredentialsProvider; +import com.amazonaws.auth.BasicAWSCredentials; +import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; +import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider; + +import javax.inject.Inject; +import javax.inject.Provider; + +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +public class GlueCredentialsProvider + implements Provider +{ + private final AWSCredentialsProvider credentialsProvider; + + @Inject + public GlueCredentialsProvider(GlueHiveMetastoreConfig config) + { + requireNonNull(config, "config is null"); + if (config.getAwsCredentialsProvider().isPresent()) { + this.credentialsProvider = getCustomAWSCredentialsProvider(config.getAwsCredentialsProvider().get()); + } + else { + AWSCredentialsProvider provider; + if (config.getAwsAccessKey().isPresent() && config.getAwsSecretKey().isPresent()) { + provider = new AWSStaticCredentialsProvider( + new BasicAWSCredentials(config.getAwsAccessKey().get(), config.getAwsSecretKey().get())); + } + else { + provider = DefaultAWSCredentialsProviderChain.getInstance(); + } + if (config.getIamRole().isPresent()) { + provider = new STSAssumeRoleSessionCredentialsProvider + .Builder(config.getIamRole().get(), "trino-session") + .withExternalId(config.getExternalId().orElse(null)) + .withLongLivedCredentialsProvider(provider) + .build(); + } + this.credentialsProvider = provider; + } + } + + @Override + public AWSCredentialsProvider get() + { + return credentialsProvider; + } + + private static AWSCredentialsProvider getCustomAWSCredentialsProvider(String providerClass) + { + try { + Object instance = Class.forName(providerClass).getConstructor().newInstance(); + if (!(instance instanceof AWSCredentialsProvider)) { + throw new RuntimeException("Invalid credentials provider class: " + instance.getClass().getName()); + } + return (AWSCredentialsProvider) instance; + } + catch (ReflectiveOperationException e) { + throw new RuntimeException(format("Error creating an instance of %s", providerClass), e); + } + } +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java index f1fbf517f496..6e362f02fec8 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java @@ -17,10 +17,6 @@ import com.amazonaws.AmazonWebServiceRequest; import com.amazonaws.ClientConfiguration; import com.amazonaws.auth.AWSCredentialsProvider; -import com.amazonaws.auth.AWSStaticCredentialsProvider; -import com.amazonaws.auth.BasicAWSCredentials; -import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; -import com.amazonaws.auth.STSAssumeRoleSessionCredentialsProvider; import com.amazonaws.client.builder.AwsClientBuilder.EndpointConfiguration; import com.amazonaws.handlers.AsyncHandler; import com.amazonaws.handlers.RequestHandler2; @@ -145,7 +141,6 @@ import static io.trino.spi.StandardErrorCode.ALREADY_EXISTS; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static io.trino.spi.security.PrincipalType.USER; -import static java.lang.String.format; import static java.util.Comparator.comparing; import static java.util.Objects.requireNonNull; import static java.util.function.UnaryOperator.identity; @@ -184,15 +179,17 @@ public class GlueHiveMetastore public GlueHiveMetastore( HdfsEnvironment hdfsEnvironment, GlueHiveMetastoreConfig glueConfig, + AWSCredentialsProvider credentialsProvider, @ForGlueHiveMetastore Executor partitionsReadExecutor, GlueColumnStatisticsProviderFactory columnStatisticsProviderFactory, @ForGlueHiveMetastore Optional requestHandler, @ForGlueHiveMetastore Predicate tableFilter) { requireNonNull(glueConfig, "glueConfig is null"); + requireNonNull(credentialsProvider, "credentialsProvider is null"); this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.hdfsContext = new HdfsContext(ConnectorIdentity.ofUser(DEFAULT_METASTORE_USER)); - this.glueClient = createAsyncGlueClient(glueConfig, requestHandler, stats.newRequestMetricsCollector()); + this.glueClient = createAsyncGlueClient(glueConfig, credentialsProvider, requestHandler, stats.newRequestMetricsCollector()); this.defaultDir = glueConfig.getDefaultWarehouseDir(); this.catalogId = glueConfig.getCatalogId().orElse(null); this.partitionSegments = glueConfig.getPartitionSegments(); @@ -202,7 +199,7 @@ public GlueHiveMetastore( this.columnStatisticsProvider = columnStatisticsProviderFactory.createGlueColumnStatisticsProvider(glueClient, stats); } - public static AWSGlueAsync createAsyncGlueClient(GlueHiveMetastoreConfig config, Optional requestHandler, RequestMetricCollector metricsCollector) + public static AWSGlueAsync createAsyncGlueClient(GlueHiveMetastoreConfig config, AWSCredentialsProvider credentialsProvider, Optional requestHandler, RequestMetricCollector metricsCollector) { ClientConfiguration clientConfig = new ClientConfiguration() .withMaxConnections(config.getMaxGlueConnections()) @@ -226,48 +223,11 @@ else if (config.getPinGlueClientToCurrentRegion()) { asyncGlueClientBuilder.setRegion(getCurrentRegionFromEC2Metadata().getName()); } - asyncGlueClientBuilder.setCredentials(getAwsCredentialsProvider(config)); + asyncGlueClientBuilder.setCredentials(credentialsProvider); return asyncGlueClientBuilder.build(); } - private static AWSCredentialsProvider getAwsCredentialsProvider(GlueHiveMetastoreConfig config) - { - if (config.getAwsCredentialsProvider().isPresent()) { - return getCustomAWSCredentialsProvider(config.getAwsCredentialsProvider().get()); - } - AWSCredentialsProvider provider; - if (config.getAwsAccessKey().isPresent() && config.getAwsSecretKey().isPresent()) { - provider = new AWSStaticCredentialsProvider( - new BasicAWSCredentials(config.getAwsAccessKey().get(), config.getAwsSecretKey().get())); - } - else { - provider = DefaultAWSCredentialsProviderChain.getInstance(); - } - if (config.getIamRole().isPresent()) { - provider = new STSAssumeRoleSessionCredentialsProvider - .Builder(config.getIamRole().get(), "trino-session") - .withExternalId(config.getExternalId().orElse(null)) - .withLongLivedCredentialsProvider(provider) - .build(); - } - return provider; - } - - private static AWSCredentialsProvider getCustomAWSCredentialsProvider(String providerClass) - { - try { - Object instance = Class.forName(providerClass).getConstructor().newInstance(); - if (!(instance instanceof AWSCredentialsProvider)) { - throw new RuntimeException("Invalid credentials provider class: " + instance.getClass().getName()); - } - return (AWSCredentialsProvider) instance; - } - catch (ReflectiveOperationException e) { - throw new RuntimeException(format("Error creating an instance of %s", providerClass), e); - } - } - public GlueMetastoreStats getStats() { return stats; diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreModule.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreModule.java index 3ac9d2eba86b..e7c580fab7fe 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreModule.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueMetastoreModule.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.hive.metastore.glue; +import com.amazonaws.auth.AWSCredentialsProvider; import com.amazonaws.handlers.RequestHandler2; import com.amazonaws.services.glue.model.Table; import com.google.inject.Binder; @@ -48,6 +49,7 @@ protected void setup(Binder binder) { configBinder(binder).bindConfig(GlueHiveMetastoreConfig.class); configBinder(binder).bindConfig(HiveConfig.class); + binder.bind(AWSCredentialsProvider.class).toProvider(GlueCredentialsProvider.class).in(Scopes.SINGLETON); newOptionalBinder(binder, Key.get(RequestHandler2.class, ForGlueHiveMetastore.class)); newOptionalBinder(binder, Key.get(new TypeLiteral>() {}, ForGlueHiveMetastore.class)) diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java index e4be75d579b2..e708672d04ec 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.hive.metastore.glue; +import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; import com.amazonaws.services.glue.AWSGlueAsync; import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder; import com.amazonaws.services.glue.model.CreateTableRequest; @@ -217,6 +218,7 @@ protected HiveMetastore createMetastore(File tempDir, HiveIdentity identity) return new GlueHiveMetastore( HDFS_ENVIRONMENT, glueConfig, + DefaultAWSCredentialsProviderChain.getInstance(), executor, new DefaultGlueColumnStatisticsProviderFactory(glueConfig, executor, executor), Optional.empty(), diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergTableOperationsProvider.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergTableOperationsProvider.java index 80241da68763..d603db336050 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergTableOperationsProvider.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/GlueIcebergTableOperationsProvider.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.iceberg.catalog.glue; +import com.amazonaws.auth.AWSCredentialsProvider; import com.amazonaws.services.glue.AWSGlueAsync; import io.trino.plugin.hive.HdfsEnvironment.HdfsContext; import io.trino.plugin.hive.metastore.glue.GlueHiveMetastoreConfig; @@ -38,12 +39,13 @@ public class GlueIcebergTableOperationsProvider private final GlueMetastoreStats stats; @Inject - public GlueIcebergTableOperationsProvider(FileIoProvider fileIoProvider, GlueMetastoreStats stats, GlueHiveMetastoreConfig glueConfig) + public GlueIcebergTableOperationsProvider(FileIoProvider fileIoProvider, GlueMetastoreStats stats, GlueHiveMetastoreConfig glueConfig, AWSCredentialsProvider credentialsProvider) { this.fileIoProvider = requireNonNull(fileIoProvider, "fileIoProvider is null"); this.stats = requireNonNull(stats, "stats is null"); requireNonNull(glueConfig, "glueConfig is null"); - this.glueClient = createAsyncGlueClient(glueConfig, Optional.empty(), stats.newRequestMetricsCollector()); + requireNonNull(credentialsProvider, "credentialsProvider is null"); + this.glueClient = createAsyncGlueClient(glueConfig, credentialsProvider, Optional.empty(), stats.newRequestMetricsCollector()); } @Override diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/IcebergGlueCatalogModule.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/IcebergGlueCatalogModule.java index ba900090e9c6..eb22f5ccdd89 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/IcebergGlueCatalogModule.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/IcebergGlueCatalogModule.java @@ -13,9 +13,11 @@ */ package io.trino.plugin.iceberg.catalog.glue; +import com.amazonaws.auth.AWSCredentialsProvider; import com.google.inject.Binder; import com.google.inject.Scopes; import io.airlift.configuration.AbstractConfigurationAwareModule; +import io.trino.plugin.hive.metastore.glue.GlueCredentialsProvider; import io.trino.plugin.hive.metastore.glue.GlueHiveMetastoreConfig; import io.trino.plugin.hive.metastore.glue.GlueMetastoreStats; import io.trino.plugin.iceberg.catalog.IcebergTableOperationsProvider; @@ -32,6 +34,7 @@ protected void setup(Binder binder) { configBinder(binder).bindConfig(GlueHiveMetastoreConfig.class); binder.bind(GlueMetastoreStats.class).in(Scopes.SINGLETON); + binder.bind(AWSCredentialsProvider.class).toProvider(GlueCredentialsProvider.class).in(Scopes.SINGLETON); binder.bind(IcebergTableOperationsProvider.class).to(GlueIcebergTableOperationsProvider.class).in(Scopes.SINGLETON); binder.bind(TrinoCatalogFactory.class).to(TrinoGlueCatalogFactory.class).in(Scopes.SINGLETON); newExporter(binder).export(TrinoCatalogFactory.class).withGeneratedName(); diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/TrinoGlueCatalogFactory.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/TrinoGlueCatalogFactory.java index ebb03866cd0c..3e8c750a5bb2 100644 --- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/TrinoGlueCatalogFactory.java +++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/catalog/glue/TrinoGlueCatalogFactory.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.iceberg.catalog.glue; +import com.amazonaws.auth.AWSCredentialsProvider; import com.amazonaws.services.glue.AWSGlueAsync; import io.trino.plugin.hive.HdfsEnvironment; import io.trino.plugin.hive.metastore.glue.GlueHiveMetastoreConfig; @@ -48,6 +49,7 @@ public TrinoGlueCatalogFactory( HdfsEnvironment hdfsEnvironment, IcebergTableOperationsProvider tableOperationsProvider, GlueHiveMetastoreConfig glueConfig, + AWSCredentialsProvider credentialsProvider, IcebergConfig icebergConfig, GlueMetastoreStats stats) { @@ -56,7 +58,8 @@ public TrinoGlueCatalogFactory( requireNonNull(glueConfig, "glueConfig is null"); checkArgument(glueConfig.getCatalogId().isEmpty(), "catalogId configuration is not supported"); this.defaultSchemaLocation = glueConfig.getDefaultWarehouseDir(); - this.glueClient = createAsyncGlueClient(glueConfig, Optional.empty(), stats.newRequestMetricsCollector()); + requireNonNull(credentialsProvider, "credentialsProvider is null"); + this.glueClient = createAsyncGlueClient(glueConfig, credentialsProvider, Optional.empty(), stats.newRequestMetricsCollector()); requireNonNull(icebergConfig, "icebergConfig is null"); this.isUniqueTableLocation = icebergConfig.isUniqueTableLocation(); this.stats = requireNonNull(stats, "stats is null"); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestSharedGlueMetastore.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestSharedGlueMetastore.java index 57fdc7dd8e1a..1af0338829d4 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestSharedGlueMetastore.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestSharedGlueMetastore.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.iceberg; +import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; @@ -103,6 +104,7 @@ protected QueryRunner createQueryRunner() this.glueMetastore = new GlueHiveMetastore( hdfsEnvironment, new GlueHiveMetastoreConfig(), + DefaultAWSCredentialsProviderChain.getInstance(), directExecutor(), new DefaultGlueColumnStatisticsProviderFactory(new GlueHiveMetastoreConfig(), directExecutor(), directExecutor()), Optional.empty(), diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestTrinoGlueCatalogTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestTrinoGlueCatalogTest.java index 95c472d3d6e1..9bc4c4f4ad31 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestTrinoGlueCatalogTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestTrinoGlueCatalogTest.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.iceberg; +import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; import com.amazonaws.services.glue.AWSGlueAsyncClientBuilder; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; @@ -59,7 +60,7 @@ protected TrinoCatalog createTrinoCatalog(boolean useUniqueTableLocations) new NoHdfsAuthentication()); return new TrinoGlueCatalog( hdfsEnvironment, - new GlueIcebergTableOperationsProvider(new HdfsFileIoProvider(hdfsEnvironment), new GlueMetastoreStats(), new GlueHiveMetastoreConfig()), + new GlueIcebergTableOperationsProvider(new HdfsFileIoProvider(hdfsEnvironment), new GlueMetastoreStats(), new GlueHiveMetastoreConfig(), DefaultAWSCredentialsProviderChain.getInstance()), AWSGlueAsyncClientBuilder.defaultClient(), new GlueMetastoreStats(), Optional.empty(), @@ -82,7 +83,7 @@ public void testDefaultLocation() new NoHdfsAuthentication()); TrinoCatalog catalogWithDefaultLocation = new TrinoGlueCatalog( hdfsEnvironment, - new GlueIcebergTableOperationsProvider(new HdfsFileIoProvider(hdfsEnvironment), new GlueMetastoreStats(), new GlueHiveMetastoreConfig()), + new GlueIcebergTableOperationsProvider(new HdfsFileIoProvider(hdfsEnvironment), new GlueMetastoreStats(), new GlueHiveMetastoreConfig(), DefaultAWSCredentialsProviderChain.getInstance()), AWSGlueAsyncClientBuilder.defaultClient(), new GlueMetastoreStats(), Optional.of(tmpDirectory.toAbsolutePath().toString()),