diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakePerTransactionMetastoreCache.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakePerTransactionMetastoreCache.java index 7049e24dfbdb..bd80b88fe018 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakePerTransactionMetastoreCache.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakePerTransactionMetastoreCache.java @@ -26,16 +26,15 @@ import io.trino.plugin.hive.metastore.RawHiveMetastoreFactory; import io.trino.plugin.hive.metastore.thrift.BridgingHiveMetastoreFactory; import io.trino.plugin.hive.metastore.thrift.DefaultThriftMetastoreClientFactory; -import io.trino.plugin.hive.metastore.thrift.MetastoreLocator; import io.trino.plugin.hive.metastore.thrift.StaticMetastoreConfig; -import io.trino.plugin.hive.metastore.thrift.StaticMetastoreLocator; +import io.trino.plugin.hive.metastore.thrift.StaticTokenAwareMetastoreClientFactory; import io.trino.plugin.hive.metastore.thrift.ThriftHiveMetastore; import io.trino.plugin.hive.metastore.thrift.ThriftHiveMetastoreFactory; import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreAuthenticationModule; import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreClientFactory; import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreConfig; import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreFactory; -import io.trino.plugin.hive.metastore.thrift.TokenDelegationThriftMetastoreFactory; +import io.trino.plugin.hive.metastore.thrift.TokenAwareMetastoreClientFactory; import io.trino.spi.security.ConnectorIdentity; import io.trino.testing.DistributedQueryRunner; import io.trino.tpch.TpchEntity; @@ -97,10 +96,9 @@ private DistributedQueryRunner createQueryRunner(boolean enablePerTransactionHiv protected void setup(Binder binder) { newOptionalBinder(binder, ThriftMetastoreClientFactory.class).setDefault().to(DefaultThriftMetastoreClientFactory.class).in(Scopes.SINGLETON); - binder.bind(MetastoreLocator.class).to(StaticMetastoreLocator.class).in(Scopes.SINGLETON); + binder.bind(TokenAwareMetastoreClientFactory.class).to(StaticTokenAwareMetastoreClientFactory.class).in(Scopes.SINGLETON); configBinder(binder).bindConfig(StaticMetastoreConfig.class); configBinder(binder).bindConfig(ThriftMetastoreConfig.class); - binder.bind(TokenDelegationThriftMetastoreFactory.class); binder.bind(ThriftMetastoreFactory.class).to(ThriftHiveMetastoreFactory.class).in(Scopes.SINGLETON); newExporter(binder).export(ThriftMetastoreFactory.class) .as(generator -> generator.generatedNameOf(ThriftHiveMetastore.class)); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/IdentityAwareMetastoreClientFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/IdentityAwareMetastoreClientFactory.java new file mode 100644 index 000000000000..6f371ace1132 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/IdentityAwareMetastoreClientFactory.java @@ -0,0 +1,25 @@ +/* + * 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.thrift; + +import io.trino.spi.security.ConnectorIdentity; +import org.apache.thrift.TException; + +import java.util.Optional; + +public interface IdentityAwareMetastoreClientFactory +{ + ThriftMetastoreClient createMetastoreClientFor(Optional identity) + throws TException; +} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/StaticMetastoreLocator.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/StaticTokenAwareMetastoreClientFactory.java similarity index 89% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/StaticMetastoreLocator.java rename to plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/StaticTokenAwareMetastoreClientFactory.java index 28f6b9c3901a..dda309091ef5 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/StaticMetastoreLocator.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/StaticTokenAwareMetastoreClientFactory.java @@ -41,21 +41,21 @@ import static java.util.concurrent.TimeUnit.NANOSECONDS; import static java.util.concurrent.TimeUnit.SECONDS; -public class StaticMetastoreLocator - implements MetastoreLocator +public class StaticTokenAwareMetastoreClientFactory + implements TokenAwareMetastoreClientFactory { private final List backoffs; private final ThriftMetastoreClientFactory clientFactory; private final String metastoreUsername; @Inject - public StaticMetastoreLocator(StaticMetastoreConfig config, ThriftMetastoreAuthenticationConfig authenticationConfig, ThriftMetastoreClientFactory clientFactory) + public StaticTokenAwareMetastoreClientFactory(StaticMetastoreConfig config, ThriftMetastoreAuthenticationConfig authenticationConfig, ThriftMetastoreClientFactory clientFactory) { this(config, authenticationConfig, clientFactory, Ticker.systemTicker()); } @VisibleForTesting - StaticMetastoreLocator(StaticMetastoreConfig config, ThriftMetastoreAuthenticationConfig authenticationConfig, ThriftMetastoreClientFactory clientFactory, Ticker ticker) + StaticTokenAwareMetastoreClientFactory(StaticMetastoreConfig config, ThriftMetastoreAuthenticationConfig authenticationConfig, ThriftMetastoreClientFactory clientFactory, Ticker ticker) { this(config.getMetastoreUris(), config.getMetastoreUsername(), clientFactory, ticker); @@ -66,17 +66,17 @@ public StaticMetastoreLocator(StaticMetastoreConfig config, ThriftMetastoreAuthe authenticationConfig.getAuthenticationType()); } - public StaticMetastoreLocator(List metastoreUris, @Nullable String metastoreUsername, ThriftMetastoreClientFactory clientFactory) + public StaticTokenAwareMetastoreClientFactory(List metastoreUris, @Nullable String metastoreUsername, ThriftMetastoreClientFactory clientFactory) { this(metastoreUris, metastoreUsername, clientFactory, Ticker.systemTicker()); } - private StaticMetastoreLocator(List metastoreUris, @Nullable String metastoreUsername, ThriftMetastoreClientFactory clientFactory, Ticker ticker) + private StaticTokenAwareMetastoreClientFactory(List metastoreUris, @Nullable String metastoreUsername, ThriftMetastoreClientFactory clientFactory, Ticker ticker) { requireNonNull(metastoreUris, "metastoreUris is null"); checkArgument(!metastoreUris.isEmpty(), "metastoreUris must specify at least one URI"); this.backoffs = metastoreUris.stream() - .map(StaticMetastoreLocator::checkMetastoreUri) + .map(StaticTokenAwareMetastoreClientFactory::checkMetastoreUri) .map(uri -> HostAndPort.fromParts(uri.getHost(), uri.getPort())) .map(address -> new Backoff(address, ticker)) .collect(toImmutableList()); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastore.java index 48d89bbe5058..9160e8b5e9f8 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastore.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastore.java @@ -143,7 +143,7 @@ public class ThriftHiveMetastore private final Optional identity; private final HdfsEnvironment hdfsEnvironment; - private final TokenDelegationThriftMetastoreFactory metastoreFactory; + private final IdentityAwareMetastoreClientFactory metastoreClientFactory; private final double backoffScaleFactor; private final Duration minBackoffDelay; private final Duration maxBackoffDelay; @@ -158,7 +158,7 @@ public class ThriftHiveMetastore public ThriftHiveMetastore( Optional identity, HdfsEnvironment hdfsEnvironment, - TokenDelegationThriftMetastoreFactory metastoreFactory, + IdentityAwareMetastoreClientFactory metastoreClientFactory, double backoffScaleFactor, Duration minBackoffDelay, Duration maxBackoffDelay, @@ -172,7 +172,7 @@ public ThriftHiveMetastore( { this.identity = requireNonNull(identity, "identity is null"); this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); - this.metastoreFactory = requireNonNull(metastoreFactory, "metastoreFactory is null"); + this.metastoreClientFactory = requireNonNull(metastoreClientFactory, "metastoreClientFactory is null"); this.backoffScaleFactor = backoffScaleFactor; this.minBackoffDelay = requireNonNull(minBackoffDelay, "minBackoffDelay is null"); this.maxBackoffDelay = requireNonNull(maxBackoffDelay, "maxBackoffDelay is null"); @@ -1848,7 +1848,7 @@ private static boolean containsAllPrivilege(Set requestedPri private ThriftMetastoreClient createMetastoreClient() throws TException { - return metastoreFactory.createMetastoreClient(identity); + return metastoreClientFactory.createMetastoreClientFor(identity); } private RetryDriver retry() diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastoreFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastoreFactory.java index cefb6431a5b4..edd02394023f 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastoreFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastoreFactory.java @@ -31,7 +31,7 @@ public class ThriftHiveMetastoreFactory implements ThriftMetastoreFactory { private final HdfsEnvironment hdfsEnvironment; - private final TokenDelegationThriftMetastoreFactory metastoreFactory; + private final IdentityAwareMetastoreClientFactory metastoreClientFactory; private final double backoffScaleFactor; private final Duration minBackoffDelay; private final Duration maxBackoffDelay; @@ -46,13 +46,13 @@ public class ThriftHiveMetastoreFactory @Inject public ThriftHiveMetastoreFactory( - TokenDelegationThriftMetastoreFactory metastoreFactory, + IdentityAwareMetastoreClientFactory metastoreClientFactory, @HideDeltaLakeTables boolean hideDeltaLakeTables, @TranslateHiveViews boolean translateHiveViews, ThriftMetastoreConfig thriftConfig, HdfsEnvironment hdfsEnvironment) { - this.metastoreFactory = requireNonNull(metastoreFactory, "metastoreFactory is null"); + this.metastoreClientFactory = requireNonNull(metastoreClientFactory, "metastoreClientFactory is null"); this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.backoffScaleFactor = thriftConfig.getBackoffScaleFactor(); this.minBackoffDelay = thriftConfig.getMinBackoffDelay(); @@ -87,7 +87,7 @@ public ThriftMetastore createMetastore(Optional identity) return new ThriftHiveMetastore( identity, hdfsEnvironment, - metastoreFactory, + metastoreClientFactory, backoffScaleFactor, minBackoffDelay, maxBackoffDelay, diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreAuthenticationModule.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreAuthenticationModule.java index fe8114888cbf..d48422cdb8ea 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreAuthenticationModule.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreAuthenticationModule.java @@ -24,8 +24,10 @@ import io.trino.plugin.hive.ForHiveMetastore; import static com.google.inject.Scopes.SINGLETON; +import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; import static io.airlift.configuration.ConfigBinder.configBinder; import static io.trino.hdfs.authentication.AuthenticationModules.createCachingKerberosHadoopAuthentication; +import static io.trino.plugin.hive.metastore.thrift.ThriftMetastoreAuthenticationConfig.ThriftMetastoreAuthenticationType.KERBEROS; public class ThriftMetastoreAuthenticationModule extends AbstractConfigurationAwareModule @@ -33,24 +35,13 @@ public class ThriftMetastoreAuthenticationModule @Override protected void setup(Binder binder) { - install(getAuthenticationModule()); - } + newOptionalBinder(binder, IdentityAwareMetastoreClientFactory.class) + .setDefault().to(UgiBasedMetastoreClientFactory.class).in(SINGLETON); + newOptionalBinder(binder, HiveMetastoreAuthentication.class) + .setDefault().to(NoHiveMetastoreAuthentication.class).in(SINGLETON); - private Module getAuthenticationModule() - { - return switch (buildConfigObject(ThriftMetastoreAuthenticationConfig.class).getAuthenticationType()) { - case NONE -> new NoHiveMetastoreAuthenticationModule(); - case KERBEROS -> new KerberosHiveMetastoreAuthenticationModule(); - }; - } - - public static class NoHiveMetastoreAuthenticationModule - implements Module - { - @Override - public void configure(Binder binder) - { - binder.bind(HiveMetastoreAuthentication.class).to(NoHiveMetastoreAuthentication.class).in(SINGLETON); + if (buildConfigObject(ThriftMetastoreAuthenticationConfig.class).getAuthenticationType() == KERBEROS) { + install(new KerberosHiveMetastoreAuthenticationModule()); } } @@ -60,7 +51,10 @@ public static class KerberosHiveMetastoreAuthenticationModule @Override public void configure(Binder binder) { - binder.bind(HiveMetastoreAuthentication.class).to(KerberosHiveMetastoreAuthentication.class).in(SINGLETON); + newOptionalBinder(binder, IdentityAwareMetastoreClientFactory.class) + .setBinding().to(TokenFetchingMetastoreClientFactory.class).in(SINGLETON); + newOptionalBinder(binder, HiveMetastoreAuthentication.class) + .setBinding().to(KerberosHiveMetastoreAuthentication.class).in(SINGLETON); configBinder(binder).bindConfig(MetastoreKerberosConfig.class); } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreModule.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreModule.java index d553f1a65600..7c2e3d4a23e6 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreModule.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreModule.java @@ -31,8 +31,7 @@ protected void setup(Binder binder) { OptionalBinder.newOptionalBinder(binder, ThriftMetastoreClientFactory.class) .setDefault().to(DefaultThriftMetastoreClientFactory.class).in(Scopes.SINGLETON); - binder.bind(MetastoreLocator.class).to(StaticMetastoreLocator.class).in(Scopes.SINGLETON); - binder.bind(TokenDelegationThriftMetastoreFactory.class); + binder.bind(TokenAwareMetastoreClientFactory.class).to(StaticTokenAwareMetastoreClientFactory.class).in(Scopes.SINGLETON); configBinder(binder).bindConfig(StaticMetastoreConfig.class); configBinder(binder).bindConfig(ThriftMetastoreConfig.class); diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/MetastoreLocator.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/TokenAwareMetastoreClientFactory.java similarity index 94% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/MetastoreLocator.java rename to plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/TokenAwareMetastoreClientFactory.java index 82ebd390dc4d..fc32062398fd 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/MetastoreLocator.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/TokenAwareMetastoreClientFactory.java @@ -17,7 +17,7 @@ import java.util.Optional; -public interface MetastoreLocator +public interface TokenAwareMetastoreClientFactory { /** * Create a connected {@link ThriftMetastoreClient} diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/TokenDelegationThriftMetastoreFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/TokenFetchingMetastoreClientFactory.java similarity index 58% rename from plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/TokenDelegationThriftMetastoreFactory.java rename to plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/TokenFetchingMetastoreClientFactory.java index b48dc22d50e3..e589fe993963 100644 --- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/TokenDelegationThriftMetastoreFactory.java +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/TokenFetchingMetastoreClientFactory.java @@ -17,16 +17,12 @@ import com.google.common.cache.CacheLoader; import com.google.common.util.concurrent.UncheckedExecutionException; import io.trino.collect.cache.NonEvictableLoadingCache; -import io.trino.hdfs.HdfsEnvironment; -import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreAuthenticationConfig.ThriftMetastoreAuthenticationType; import io.trino.spi.TrinoException; import io.trino.spi.security.ConnectorIdentity; import org.apache.thrift.TException; import javax.inject.Inject; -import java.io.Closeable; -import java.io.IOException; import java.util.Optional; import static com.google.common.base.Throwables.throwIfInstanceOf; @@ -35,23 +31,20 @@ import static java.util.Objects.requireNonNull; import static java.util.concurrent.TimeUnit.MILLISECONDS; -public class TokenDelegationThriftMetastoreFactory +public class TokenFetchingMetastoreClientFactory + implements IdentityAwareMetastoreClientFactory { - private final MetastoreLocator clientProvider; + private final TokenAwareMetastoreClientFactory clientProvider; private final boolean impersonationEnabled; - private final boolean authenticationEnabled; private final NonEvictableLoadingCache delegationTokenCache; @Inject - public TokenDelegationThriftMetastoreFactory( - MetastoreLocator metastoreLocator, - ThriftMetastoreConfig thriftConfig, - ThriftMetastoreAuthenticationConfig authenticationConfig, - HdfsEnvironment hdfsEnvironment) + public TokenFetchingMetastoreClientFactory( + TokenAwareMetastoreClientFactory tokenAwareMetastoreClientFactory, + ThriftMetastoreConfig thriftConfig) { - this.clientProvider = requireNonNull(metastoreLocator, "metastoreLocator is null"); + this.clientProvider = requireNonNull(tokenAwareMetastoreClientFactory, "tokenAwareMetastoreClientFactory is null"); this.impersonationEnabled = thriftConfig.isImpersonationEnabled(); - this.authenticationEnabled = authenticationConfig.getAuthenticationType() != ThriftMetastoreAuthenticationType.NONE; this.delegationTokenCache = buildNonEvictableCache( CacheBuilder.newBuilder() @@ -66,7 +59,8 @@ private ThriftMetastoreClient createMetastoreClient() return clientProvider.createMetastoreClient(Optional.empty()); } - public ThriftMetastoreClient createMetastoreClient(Optional identity) + @Override + public ThriftMetastoreClient createMetastoreClientFor(Optional identity) throws TException { if (!impersonationEnabled) { @@ -75,21 +69,16 @@ public ThriftMetastoreClient createMetastoreClient(Optional i String username = identity.map(ConnectorIdentity::getUser) .orElseThrow(() -> new IllegalStateException("End-user name should exist when metastore impersonation is enabled")); - if (authenticationEnabled) { - String delegationToken; - try { - delegationToken = delegationTokenCache.getUnchecked(username); - } - catch (UncheckedExecutionException e) { - throwIfInstanceOf(e.getCause(), TrinoException.class); - throw e; - } - return clientProvider.createMetastoreClient(Optional.of(delegationToken)); - } - ThriftMetastoreClient client = createMetastoreClient(); - setMetastoreUserOrClose(client, username); - return client; + String delegationToken; + try { + delegationToken = delegationTokenCache.getUnchecked(username); + } + catch (UncheckedExecutionException e) { + throwIfInstanceOf(e.getCause(), TrinoException.class); + throw e; + } + return clientProvider.createMetastoreClient(Optional.of(delegationToken)); } private String loadDelegationToken(String username) @@ -101,21 +90,4 @@ private String loadDelegationToken(String username) throw new TrinoException(HIVE_METASTORE_ERROR, e); } } - - private static void setMetastoreUserOrClose(ThriftMetastoreClient client, String username) - throws TException - { - try { - client.setUGI(username); - } - catch (Throwable t) { - // close client and suppress any error from close - try (Closeable ignored = client) { - throw t; - } - catch (IOException e) { - // impossible; will be suppressed - } - } - } } diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/UgiBasedMetastoreClientFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/UgiBasedMetastoreClientFactory.java new file mode 100644 index 000000000000..fd3396e80ef3 --- /dev/null +++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/UgiBasedMetastoreClientFactory.java @@ -0,0 +1,73 @@ +/* + * 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.thrift; + +import io.trino.spi.security.ConnectorIdentity; +import org.apache.thrift.TException; + +import javax.inject.Inject; + +import java.io.Closeable; +import java.io.IOException; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +public class UgiBasedMetastoreClientFactory + implements IdentityAwareMetastoreClientFactory +{ + private final TokenAwareMetastoreClientFactory clientProvider; + private final boolean impersonationEnabled; + + @Inject + public UgiBasedMetastoreClientFactory( + TokenAwareMetastoreClientFactory clientProvider, + ThriftMetastoreConfig thriftConfig) + { + this.clientProvider = requireNonNull(clientProvider, "clientProvider is null"); + this.impersonationEnabled = thriftConfig.isImpersonationEnabled(); + } + + @Override + public ThriftMetastoreClient createMetastoreClientFor(Optional identity) + throws TException + { + ThriftMetastoreClient client = clientProvider.createMetastoreClient(Optional.empty()); + + if (impersonationEnabled) { + String username = identity.map(ConnectorIdentity::getUser) + .orElseThrow(() -> new IllegalStateException("End-user name should exist when metastore impersonation is enabled")); + setMetastoreUserOrClose(client, username); + } + + return client; + } + + private static void setMetastoreUserOrClose(ThriftMetastoreClient client, String username) + throws TException + { + try { + client.setUGI(username); + } + catch (Throwable t) { + // close client and suppress any error from close + try (Closeable ignored = client) { + throw t; + } + catch (IOException e) { + // impossible; will be suppressed + } + } + } +} diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingThriftHiveMetastoreBuilder.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingThriftHiveMetastoreBuilder.java index 6e485a4c00c2..e76ced158586 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingThriftHiveMetastoreBuilder.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingThriftHiveMetastoreBuilder.java @@ -26,14 +26,13 @@ import io.trino.plugin.hive.gcs.GoogleGcsConfigurationInitializer; import io.trino.plugin.hive.gcs.HiveGcsConfig; import io.trino.plugin.hive.metastore.HiveMetastoreConfig; -import io.trino.plugin.hive.metastore.thrift.MetastoreLocator; -import io.trino.plugin.hive.metastore.thrift.TestingMetastoreLocator; +import io.trino.plugin.hive.metastore.thrift.TestingTokenAwareMetastoreClientFactory; import io.trino.plugin.hive.metastore.thrift.ThriftHiveMetastoreFactory; import io.trino.plugin.hive.metastore.thrift.ThriftMetastore; -import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreAuthenticationConfig; import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreClient; import io.trino.plugin.hive.metastore.thrift.ThriftMetastoreConfig; -import io.trino.plugin.hive.metastore.thrift.TokenDelegationThriftMetastoreFactory; +import io.trino.plugin.hive.metastore.thrift.TokenAwareMetastoreClientFactory; +import io.trino.plugin.hive.metastore.thrift.UgiBasedMetastoreClientFactory; import io.trino.plugin.hive.s3.HiveS3Config; import io.trino.plugin.hive.s3.TrinoS3ConfigurationInitializer; @@ -57,7 +56,7 @@ public final class TestingThriftHiveMetastoreBuilder new HdfsConfig(), new NoHdfsAuthentication()); - private MetastoreLocator metastoreLocator; + private TokenAwareMetastoreClientFactory tokenAwareMetastoreClientFactory; private HiveConfig hiveConfig = new HiveConfig(); private ThriftMetastoreConfig thriftMetastoreConfig = new ThriftMetastoreConfig(); private HdfsEnvironment hdfsEnvironment = HDFS_ENVIRONMENT; @@ -73,24 +72,24 @@ public TestingThriftHiveMetastoreBuilder metastoreClient(HostAndPort address, Du { requireNonNull(address, "address is null"); requireNonNull(timeout, "timeout is null"); - checkState(metastoreLocator == null, "Metastore client already set"); - metastoreLocator = new TestingMetastoreLocator(HiveTestUtils.SOCKS_PROXY, address, timeout); + checkState(tokenAwareMetastoreClientFactory == null, "Metastore client already set"); + tokenAwareMetastoreClientFactory = new TestingTokenAwareMetastoreClientFactory(HiveTestUtils.SOCKS_PROXY, address, timeout); return this; } public TestingThriftHiveMetastoreBuilder metastoreClient(HostAndPort address) { requireNonNull(address, "address is null"); - checkState(metastoreLocator == null, "Metastore client already set"); - metastoreLocator = new TestingMetastoreLocator(HiveTestUtils.SOCKS_PROXY, address); + checkState(tokenAwareMetastoreClientFactory == null, "Metastore client already set"); + tokenAwareMetastoreClientFactory = new TestingTokenAwareMetastoreClientFactory(HiveTestUtils.SOCKS_PROXY, address); return this; } public TestingThriftHiveMetastoreBuilder metastoreClient(ThriftMetastoreClient client) { requireNonNull(client, "client is null"); - checkState(metastoreLocator == null, "Metastore client already set"); - metastoreLocator = token -> client; + checkState(tokenAwareMetastoreClientFactory == null, "Metastore client already set"); + tokenAwareMetastoreClientFactory = token -> client; return this; } @@ -114,13 +113,9 @@ public TestingThriftHiveMetastoreBuilder hdfsEnvironment(HdfsEnvironment hdfsEnv public ThriftMetastore build() { - checkState(metastoreLocator != null, "metastore client not set"); + checkState(tokenAwareMetastoreClientFactory != null, "metastore client not set"); ThriftHiveMetastoreFactory metastoreFactory = new ThriftHiveMetastoreFactory( - new TokenDelegationThriftMetastoreFactory( - metastoreLocator, - thriftMetastoreConfig, - new ThriftMetastoreAuthenticationConfig(), - hdfsEnvironment), + new UgiBasedMetastoreClientFactory(tokenAwareMetastoreClientFactory, thriftMetastoreConfig), new HiveMetastoreConfig().isHideDeltaLakeTables(), hiveConfig.isTranslateHiveViews(), thriftMetastoreConfig, diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestStaticMetastoreLocator.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestStaticTokenAwareMetastoreClientFactory.java similarity index 60% rename from plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestStaticMetastoreLocator.java rename to plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestStaticTokenAwareMetastoreClientFactory.java index 9ff4b96cf9c0..7e211bee0695 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestStaticMetastoreLocator.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestStaticTokenAwareMetastoreClientFactory.java @@ -29,7 +29,7 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.testng.Assert.assertEquals; -public class TestStaticMetastoreLocator +public class TestStaticTokenAwareMetastoreClientFactory { private static final ThriftMetastoreClient DEFAULT_CLIENT = createFakeMetastoreClient(); private static final ThriftMetastoreClient FALLBACK_CLIENT = createFakeMetastoreClient(); @@ -58,59 +58,59 @@ public class TestStaticMetastoreLocator public void testDefaultHiveMetastore() throws TException { - MetastoreLocator locator = createMetastoreLocator(CONFIG_WITH_FALLBACK, ImmutableMap.of(DEFAULT_URI, Optional.of(DEFAULT_CLIENT))); - assertEqualHiveClient(locator.createMetastoreClient(Optional.empty()), DEFAULT_CLIENT); + TokenAwareMetastoreClientFactory clientFactory = createMetastoreClientFactory(CONFIG_WITH_FALLBACK, ImmutableMap.of(DEFAULT_URI, Optional.of(DEFAULT_CLIENT))); + assertEqualHiveClient(clientFactory.createMetastoreClient(Optional.empty()), DEFAULT_CLIENT); } @Test public void testFallbackHiveMetastore() throws TException { - MetastoreLocator locator = createMetastoreLocator(CONFIG_WITH_FALLBACK, ImmutableMap.of(DEFAULT_URI, Optional.empty(), FALLBACK_URI, Optional.of(FALLBACK_CLIENT))); - assertEqualHiveClient(locator.createMetastoreClient(Optional.empty()), FALLBACK_CLIENT); + TokenAwareMetastoreClientFactory clientFactory = createMetastoreClientFactory(CONFIG_WITH_FALLBACK, ImmutableMap.of(DEFAULT_URI, Optional.empty(), FALLBACK_URI, Optional.of(FALLBACK_CLIENT))); + assertEqualHiveClient(clientFactory.createMetastoreClient(Optional.empty()), FALLBACK_CLIENT); } @Test public void testFallbackHiveMetastoreFails() { - MetastoreLocator locator = createMetastoreLocator(CONFIG_WITH_FALLBACK, ImmutableMap.of()); - assertCreateClientFails(locator, "Failed connecting to Hive metastore: [default:8080, fallback:8090, fallback2:8090]"); + TokenAwareMetastoreClientFactory clientFactory = createMetastoreClientFactory(CONFIG_WITH_FALLBACK, ImmutableMap.of()); + assertCreateClientFails(clientFactory, "Failed connecting to Hive metastore: [default:8080, fallback:8090, fallback2:8090]"); } @Test public void testMetastoreFailedWithoutFallback() { - MetastoreLocator locator = createMetastoreLocator(CONFIG_WITHOUT_FALLBACK, ImmutableMap.of(DEFAULT_URI, Optional.empty())); - assertCreateClientFails(locator, "Failed connecting to Hive metastore: [default:8080]"); + TokenAwareMetastoreClientFactory clientFactory = createMetastoreClientFactory(CONFIG_WITHOUT_FALLBACK, ImmutableMap.of(DEFAULT_URI, Optional.empty())); + assertCreateClientFails(clientFactory, "Failed connecting to Hive metastore: [default:8080]"); } @Test public void testFallbackHiveMetastoreWithHiveUser() throws TException { - MetastoreLocator locator = createMetastoreLocator(CONFIG_WITH_FALLBACK_WITH_USER, ImmutableMap.of(DEFAULT_URI, Optional.empty(), FALLBACK_URI, Optional.empty(), FALLBACK2_URI, Optional.of(FALLBACK_CLIENT))); - assertEqualHiveClient(locator.createMetastoreClient(Optional.empty()), FALLBACK_CLIENT); + TokenAwareMetastoreClientFactory clientFactory = createMetastoreClientFactory(CONFIG_WITH_FALLBACK_WITH_USER, ImmutableMap.of(DEFAULT_URI, Optional.empty(), FALLBACK_URI, Optional.empty(), FALLBACK2_URI, Optional.of(FALLBACK_CLIENT))); + assertEqualHiveClient(clientFactory.createMetastoreClient(Optional.empty()), FALLBACK_CLIENT); } @Test public void testMetastoreFailedWithoutFallbackWithHiveUser() { - MetastoreLocator locator = createMetastoreLocator(CONFIG_WITHOUT_FALLBACK_WITH_USER, ImmutableMap.of(DEFAULT_URI, Optional.empty())); - assertCreateClientFails(locator, "Failed connecting to Hive metastore: [default:8080]"); + TokenAwareMetastoreClientFactory clientFactory = createMetastoreClientFactory(CONFIG_WITHOUT_FALLBACK_WITH_USER, ImmutableMap.of(DEFAULT_URI, Optional.empty())); + assertCreateClientFails(clientFactory, "Failed connecting to Hive metastore: [default:8080]"); } @Test public void testFallbackHiveMetastoreOnTimeOut() throws TException { - MetastoreLocator cluster = createMetastoreLocator(CONFIG_WITH_FALLBACK, CLIENTS); + TokenAwareMetastoreClientFactory clientFactory = createMetastoreClientFactory(CONFIG_WITH_FALLBACK, CLIENTS); - ThriftMetastoreClient metastoreClient1 = cluster.createMetastoreClient(Optional.empty()); + ThriftMetastoreClient metastoreClient1 = clientFactory.createMetastoreClient(Optional.empty()); assertEqualHiveClient(metastoreClient1, DEFAULT_CLIENT); assertGetTableException(metastoreClient1); - ThriftMetastoreClient metastoreClient2 = cluster.createMetastoreClient(Optional.empty()); + ThriftMetastoreClient metastoreClient2 = clientFactory.createMetastoreClient(Optional.empty()); assertEqualHiveClient(metastoreClient2, FALLBACK_CLIENT); assertGetTableException(metastoreClient2); @@ -120,22 +120,22 @@ public void testFallbackHiveMetastoreOnTimeOut() public void testFallbackHiveMetastoreOnAllTimeOut() throws TException { - MetastoreLocator cluster = createMetastoreLocator(CONFIG_WITH_FALLBACK, CLIENTS); + TokenAwareMetastoreClientFactory clientFactory = createMetastoreClientFactory(CONFIG_WITH_FALLBACK, CLIENTS); - ThriftMetastoreClient metastoreClient1 = cluster.createMetastoreClient(Optional.empty()); + ThriftMetastoreClient metastoreClient1 = clientFactory.createMetastoreClient(Optional.empty()); assertEqualHiveClient(metastoreClient1, DEFAULT_CLIENT); for (int i = 0; i < 20; ++i) { assertGetTableException(metastoreClient1); } - ThriftMetastoreClient metastoreClient2 = cluster.createMetastoreClient(Optional.empty()); + ThriftMetastoreClient metastoreClient2 = clientFactory.createMetastoreClient(Optional.empty()); assertEqualHiveClient(metastoreClient2, FALLBACK_CLIENT); assertGetTableException(metastoreClient2); // Still get FALLBACK_CLIENT because DEFAULT_CLIENT failed more times before and therefore longer backoff - ThriftMetastoreClient metastoreClient3 = cluster.createMetastoreClient(Optional.empty()); + ThriftMetastoreClient metastoreClient3 = clientFactory.createMetastoreClient(Optional.empty()); assertEqualHiveClient(metastoreClient3, FALLBACK_CLIENT); } @@ -144,20 +144,20 @@ public void testStickToFallbackAfterBackoff() throws TException { TestingTicker ticker = new TestingTicker(); - MetastoreLocator cluster = createMetastoreLocator(CONFIG_WITH_FALLBACK, CLIENTS, ticker); + TokenAwareMetastoreClientFactory clientFactory = createMetastoreClientFactory(CONFIG_WITH_FALLBACK, CLIENTS, ticker); ticker.increment(10, NANOSECONDS); - ThriftMetastoreClient metastoreClient1 = cluster.createMetastoreClient(Optional.empty()); + ThriftMetastoreClient metastoreClient1 = clientFactory.createMetastoreClient(Optional.empty()); assertEqualHiveClient(metastoreClient1, DEFAULT_CLIENT); assertGetTableException(metastoreClient1); ticker.increment(10, NANOSECONDS); - ThriftMetastoreClient metastoreClient2 = cluster.createMetastoreClient(Optional.empty()); + ThriftMetastoreClient metastoreClient2 = clientFactory.createMetastoreClient(Optional.empty()); assertEqualHiveClient(metastoreClient2, FALLBACK_CLIENT); // even after backoff for DEFAULT_CLIENT passes we should stick to client which we saw working correctly most recently - ticker.increment(StaticMetastoreLocator.Backoff.MAX_BACKOFF, NANOSECONDS); - ThriftMetastoreClient metastoreClient3 = cluster.createMetastoreClient(Optional.empty()); + ticker.increment(StaticTokenAwareMetastoreClientFactory.Backoff.MAX_BACKOFF, NANOSECONDS); + ThriftMetastoreClient metastoreClient3 = clientFactory.createMetastoreClient(Optional.empty()); assertEqualHiveClient(metastoreClient3, FALLBACK_CLIENT); } @@ -166,20 +166,20 @@ public void testReturnsToDefaultClientAfterErrorOnFallback() throws TException { TestingTicker ticker = new TestingTicker(); - MetastoreLocator cluster = createMetastoreLocator(CONFIG_WITH_FALLBACK, CLIENTS, ticker); + TokenAwareMetastoreClientFactory clientFactory = createMetastoreClientFactory(CONFIG_WITH_FALLBACK, CLIENTS, ticker); ticker.increment(10, NANOSECONDS); - ThriftMetastoreClient metastoreClient1 = cluster.createMetastoreClient(Optional.empty()); + ThriftMetastoreClient metastoreClient1 = clientFactory.createMetastoreClient(Optional.empty()); assertEqualHiveClient(metastoreClient1, DEFAULT_CLIENT); assertGetTableException(metastoreClient1); ticker.increment(10, NANOSECONDS); - ThriftMetastoreClient metastoreClient2 = cluster.createMetastoreClient(Optional.empty()); + ThriftMetastoreClient metastoreClient2 = clientFactory.createMetastoreClient(Optional.empty()); assertEqualHiveClient(metastoreClient2, FALLBACK_CLIENT); assertGetTableException(metastoreClient2); ticker.increment(10, NANOSECONDS); - ThriftMetastoreClient metastoreClient3 = cluster.createMetastoreClient(Optional.empty()); + ThriftMetastoreClient metastoreClient3 = clientFactory.createMetastoreClient(Optional.empty()); assertEqualHiveClient(metastoreClient3, DEFAULT_CLIENT); } @@ -190,21 +190,21 @@ private static void assertGetTableException(ThriftMetastoreClient client) .hasMessageContaining("Read timeout"); } - private static void assertCreateClientFails(MetastoreLocator locator, String message) + private static void assertCreateClientFails(TokenAwareMetastoreClientFactory clientFactory, String message) { - assertThatThrownBy(() -> locator.createMetastoreClient(Optional.empty())) + assertThatThrownBy(() -> clientFactory.createMetastoreClient(Optional.empty())) .hasCauseInstanceOf(TException.class) .hasMessage(message); } - private static MetastoreLocator createMetastoreLocator(StaticMetastoreConfig config, Map> clients) + private static TokenAwareMetastoreClientFactory createMetastoreClientFactory(StaticMetastoreConfig config, Map> clients) { - return createMetastoreLocator(config, clients, Ticker.systemTicker()); + return createMetastoreClientFactory(config, clients, Ticker.systemTicker()); } - private static MetastoreLocator createMetastoreLocator(StaticMetastoreConfig config, Map> clients, Ticker ticker) + private static TokenAwareMetastoreClientFactory createMetastoreClientFactory(StaticMetastoreConfig config, Map> clients, Ticker ticker) { - return new StaticMetastoreLocator(config, new ThriftMetastoreAuthenticationConfig(), new MockThriftMetastoreClientFactory(clients), ticker); + return new StaticTokenAwareMetastoreClientFactory(config, new ThriftMetastoreAuthenticationConfig(), new MockThriftMetastoreClientFactory(clients), ticker); } private static ThriftMetastoreClient createFakeMetastoreClient() diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestingMetastoreLocator.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestingTokenAwareMetastoreClientFactory.java similarity index 82% rename from plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestingMetastoreLocator.java rename to plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestingTokenAwareMetastoreClientFactory.java index e355f598f583..4f65cb429a6c 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestingMetastoreLocator.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestingTokenAwareMetastoreClientFactory.java @@ -22,8 +22,8 @@ import static java.util.Objects.requireNonNull; import static java.util.concurrent.TimeUnit.SECONDS; -public class TestingMetastoreLocator - implements MetastoreLocator +public class TestingTokenAwareMetastoreClientFactory + implements TokenAwareMetastoreClientFactory { private static final HiveMetastoreAuthentication AUTHENTICATION = new NoHiveMetastoreAuthentication(); public static final Duration TIMEOUT = new Duration(20, SECONDS); @@ -31,12 +31,12 @@ public class TestingMetastoreLocator private final DefaultThriftMetastoreClientFactory factory; private final HostAndPort address; - public TestingMetastoreLocator(Optional socksProxy, HostAndPort address) + public TestingTokenAwareMetastoreClientFactory(Optional socksProxy, HostAndPort address) { this(socksProxy, address, TIMEOUT); } - public TestingMetastoreLocator(Optional socksProxy, HostAndPort address, Duration timeout) + public TestingTokenAwareMetastoreClientFactory(Optional socksProxy, HostAndPort address, Duration timeout) { this.factory = new DefaultThriftMetastoreClientFactory(Optional.empty(), socksProxy, timeout, AUTHENTICATION, "localhost"); this.address = requireNonNull(address, "address is null"); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/S3HiveQueryRunner.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/S3HiveQueryRunner.java index 2e676c70a683..992ced70ad1b 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/S3HiveQueryRunner.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/s3/S3HiveQueryRunner.java @@ -20,7 +20,7 @@ import io.trino.plugin.hive.HiveQueryRunner; import io.trino.plugin.hive.containers.HiveMinioDataLake; import io.trino.plugin.hive.metastore.thrift.BridgingHiveMetastore; -import io.trino.plugin.hive.metastore.thrift.TestingMetastoreLocator; +import io.trino.plugin.hive.metastore.thrift.TestingTokenAwareMetastoreClientFactory; import io.trino.testing.DistributedQueryRunner; import io.trino.tpch.TpchTable; @@ -84,7 +84,7 @@ public static class Builder extends HiveQueryRunner.Builder { private HostAndPort hiveMetastoreEndpoint; - private Duration thriftMetastoreTimeout = TestingMetastoreLocator.TIMEOUT; + private Duration thriftMetastoreTimeout = TestingTokenAwareMetastoreClientFactory.TIMEOUT; private String s3Endpoint; private String s3AccessKey; private String s3SecretKey;