Allow injecting authentication components for ThriftMetastore

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));

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<ConnectorIdentity> identity)
+            throws TException;
+}

plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/StaticMetastoreLocator.java → 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<Backoff> 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<URI> metastoreUris, @Nullable String metastoreUsername, ThriftMetastoreClientFactory clientFactory)
+    public StaticTokenAwareMetastoreClientFactory(List<URI> metastoreUris, @Nullable String metastoreUsername, ThriftMetastoreClientFactory clientFactory)
     {
         this(metastoreUris, metastoreUsername, clientFactory, Ticker.systemTicker());
     }
 
-    private StaticMetastoreLocator(List<URI> metastoreUris, @Nullable String metastoreUsername, ThriftMetastoreClientFactory clientFactory, Ticker ticker)
+    private StaticTokenAwareMetastoreClientFactory(List<URI> 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());

plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftHiveMetastore.java

@@ -143,7 +143,7 @@ public class ThriftHiveMetastore
 
     private final Optional<ConnectorIdentity> 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<ConnectorIdentity> 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<PrivilegeGrantInfo> requestedPri
     private ThriftMetastoreClient createMetastoreClient()
             throws TException
     {
-        return metastoreFactory.createMetastoreClient(identity);
+        return metastoreClientFactory.createMetastoreClientFor(identity);
     }
 
     private RetryDriver retry()

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<ConnectorIdentity> identity)
         return new ThriftHiveMetastore(
                 identity,
                 hdfsEnvironment,
-                metastoreFactory,
+                metastoreClientFactory,
                 backoffScaleFactor,
                 minBackoffDelay,
                 maxBackoffDelay,

plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/ThriftMetastoreAuthenticationModule.java

@@ -24,33 +24,24 @@
 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
 {
     @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);
         }
 

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);
 

plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/MetastoreLocator.java → 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}

plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/thrift/TokenDelegationThriftMetastoreFactory.java → 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<String, String> 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<ConnectorIdentity> identity)
+    @Override
+    public ThriftMetastoreClient createMetastoreClientFor(Optional<ConnectorIdentity> identity)
             throws TException
     {
         if (!impersonationEnabled) {
@@ -75,21 +69,16 @@ public ThriftMetastoreClient createMetastoreClient(Optional<ConnectorIdentity> 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
-            }
-        }
-    }
 }