Skip to content

Commit 6054835

Browse files
committed
Allow metastore impersonation in tests
1 parent 2303e30 commit 6054835

File tree

13 files changed

+51
-24
lines changed

13 files changed

+51
-24
lines changed

lib/trino-metastore/src/main/java/io/trino/metastore/HiveMetastoreFactory.java

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,25 +35,27 @@ default boolean hasBuiltInCaching()
3535
*/
3636
HiveMetastore createMetastore(Optional<ConnectorIdentity> identity);
3737

38-
static HiveMetastoreFactory ofInstance(HiveMetastore metastore)
38+
static HiveMetastoreFactory ofInstance(HiveMetastore metastore, boolean impersonationEnabled)
3939
{
40-
return new StaticHiveMetastoreFactory(metastore);
40+
return new StaticHiveMetastoreFactory(metastore, impersonationEnabled);
4141
}
4242

4343
class StaticHiveMetastoreFactory
4444
implements HiveMetastoreFactory
4545
{
4646
private final HiveMetastore metastore;
47+
private final boolean impersonationEnabled;
4748

48-
private StaticHiveMetastoreFactory(HiveMetastore metastore)
49+
private StaticHiveMetastoreFactory(HiveMetastore metastore, boolean impersonationEnabled)
4950
{
5051
this.metastore = requireNonNull(metastore, "metastore is null");
52+
this.impersonationEnabled = impersonationEnabled;
5153
}
5254

5355
@Override
5456
public boolean isImpersonationEnabled()
5557
{
56-
return false;
58+
return impersonationEnabled;
5759
}
5860

5961
@Override

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSplitManager.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,7 @@ public Stream<AddFileEntry> getActiveFiles(
221221
newDirectExecutorService());
222222

223223
TransactionLogReaderFactory transactionLogReaderFactory = new FileSystemTransactionLogReaderFactory(fileSystemFactory);
224-
HiveMetastoreFactory hiveMetastoreFactory = HiveMetastoreFactory.ofInstance(createTestingFileHiveMetastore(new MemoryFileSystemFactory(), Location.of("memory:///")));
224+
HiveMetastoreFactory hiveMetastoreFactory = HiveMetastoreFactory.ofInstance(createTestingFileHiveMetastore(new MemoryFileSystemFactory(), Location.of("memory:///")), false);
225225
DeltaLakeMetadataFactory metadataFactory = new DeltaLakeMetadataFactory(
226226
hiveMetastoreFactory,
227227
new DefaultDeltaLakeFileSystemFactory(hdfsFileSystemFactory, new NoOpVendedCredentialsProvider()),

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/metastore/TestingDeltaLakeMetastoreModule.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ public TestingDeltaLakeMetastoreModule(HiveMetastore metastore)
4141
@Override
4242
public void setup(Binder binder)
4343
{
44-
binder.bind(HiveMetastoreFactory.class).annotatedWith(RawHiveMetastoreFactory.class).toInstance(HiveMetastoreFactory.ofInstance(metastore));
44+
binder.bind(HiveMetastoreFactory.class).annotatedWith(RawHiveMetastoreFactory.class).toInstance(HiveMetastoreFactory.ofInstance(metastore, false));
4545
install(new CachingHiveMetastoreModule());
4646
binder.bind(DeltaLakeTableOperationsProvider.class).to(DeltaLakeFileMetastoreTableOperationsProvider.class).in(Scopes.SINGLETON);
4747

plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConnectorFactory.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ public String getName()
7171
public Connector create(String catalogName, Map<String, String> config, ConnectorContext context)
7272
{
7373
checkStrictSpiVersionMatch(context, this);
74-
return createConnector(catalogName, config, context, EMPTY_MODULE, Optional.empty(), Optional.empty());
74+
return createConnector(catalogName, config, context, EMPTY_MODULE, Optional.empty(), false, Optional.empty());
7575
}
7676

7777
public static Connector createConnector(
@@ -80,6 +80,7 @@ public static Connector createConnector(
8080
ConnectorContext context,
8181
Module module,
8282
Optional<HiveMetastore> metastore,
83+
boolean metastoreImpersonationEnabled,
8384
Optional<TrinoFileSystemFactory> fileSystemFactory)
8485
{
8586
ClassLoader classLoader = HiveConnectorFactory.class.getClassLoader();
@@ -91,7 +92,7 @@ public static Connector createConnector(
9192
new JsonModule(),
9293
new TypeDeserializerModule(),
9394
new HiveModule(),
94-
new HiveMetastoreModule(metastore),
95+
new HiveMetastoreModule(metastore, metastoreImpersonationEnabled),
9596
new HiveSecurityModule(),
9697
fileSystemFactory
9798
.map(factory -> (Module) binder -> binder.bind(TrinoFileSystemFactory.class).toInstance(factory))

plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/HiveMetastoreModule.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,17 +30,19 @@ public class HiveMetastoreModule
3030
extends AbstractConfigurationAwareModule
3131
{
3232
private final Optional<HiveMetastore> metastore;
33+
private final boolean impersonationEnabled;
3334

34-
public HiveMetastoreModule(Optional<HiveMetastore> metastore)
35+
public HiveMetastoreModule(Optional<HiveMetastore> metastore, boolean impersonationEnabled)
3536
{
3637
this.metastore = metastore;
38+
this.impersonationEnabled = impersonationEnabled;
3739
}
3840

3941
@Override
4042
protected void setup(Binder binder)
4143
{
4244
if (metastore.isPresent()) {
43-
binder.bind(HiveMetastoreFactory.class).annotatedWith(RawHiveMetastoreFactory.class).toInstance(HiveMetastoreFactory.ofInstance(metastore.get()));
45+
binder.bind(HiveMetastoreFactory.class).annotatedWith(RawHiveMetastoreFactory.class).toInstance(HiveMetastoreFactory.ofInstance(metastore.get(), impersonationEnabled));
4446
binder.bind(Key.get(boolean.class, AllowHiveTableRename.class)).toInstance(true);
4547
}
4648
else {

plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveQueryRunner.java

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747

4848
import static io.airlift.log.Level.WARN;
4949
import static io.airlift.units.Duration.nanosSince;
50+
import static io.trino.plugin.hive.HiveTestUtils.SESSION;
5051
import static io.trino.plugin.hive.TestingHiveUtils.getConnectorService;
5152
import static io.trino.plugin.tpch.ColumnNaming.SIMPLIFIED;
5253
import static io.trino.plugin.tpch.DecimalTypeMapping.DOUBLE;
@@ -101,6 +102,7 @@ public static class Builder<SELF extends Builder<?>>
101102
private List<TpchTable<?>> initialTables = ImmutableList.of();
102103
private Optional<String> initialSchemasLocationBase = Optional.empty();
103104
private Optional<Function<DistributedQueryRunner, HiveMetastore>> metastore = Optional.empty();
105+
private boolean metastoreImpersonationEnabled;
104106
private boolean tpcdsCatalogEnabled;
105107
private boolean tpchBucketedCatalogEnabled;
106108
private boolean createTpchSchemas = true;
@@ -161,6 +163,13 @@ public SELF setMetastore(Function<DistributedQueryRunner, HiveMetastore> metasto
161163
return self();
162164
}
163165

166+
@CanIgnoreReturnValue
167+
public SELF setMetastoreImpersonationEnabled(boolean metastoreImpersonationEnabled)
168+
{
169+
this.metastoreImpersonationEnabled = metastoreImpersonationEnabled;
170+
return self();
171+
}
172+
164173
@CanIgnoreReturnValue
165174
public SELF setTpcdsCatalogEnabled(boolean tpcdsCatalogEnabled)
166175
{
@@ -236,7 +245,7 @@ public DistributedQueryRunner build()
236245
hiveProperties.put("fs.hadoop.enabled", "true");
237246
}
238247

239-
queryRunner.installPlugin(new TestingHivePlugin(dataDir, metastore, decryptionKeyRetriever));
248+
queryRunner.installPlugin(new TestingHivePlugin(dataDir, metastore, metastoreImpersonationEnabled, decryptionKeyRetriever));
240249

241250
Map<String, String> hiveProperties = new HashMap<>();
242251
if (!skipTimezoneSetup) {
@@ -280,7 +289,7 @@ public DistributedQueryRunner build()
280289
private void populateData(QueryRunner queryRunner)
281290
{
282291
HiveMetastore metastore = getConnectorService(queryRunner, HiveMetastoreFactory.class)
283-
.createMetastore(Optional.empty());
292+
.createMetastore(Optional.of(SESSION.getIdentity()));
284293
if (metastore.getDatabase(TPCH_SCHEMA).isEmpty()) {
285294
metastore.createDatabase(createDatabaseMetastoreObject(TPCH_SCHEMA, initialSchemasLocationBase));
286295
copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, initialTables);

plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHivePageSink.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -394,7 +394,7 @@ private static ConnectorPageSink createPageSink(
394394
getDefaultHiveFileWriterFactories(config, fileSystemFactory),
395395
HDFS_FILE_SYSTEM_FACTORY,
396396
PAGE_SORTER,
397-
HiveMetastoreFactory.ofInstance(metastore),
397+
HiveMetastoreFactory.ofInstance(metastore, false),
398398
new GroupByHashPageIndexerFactory(new FlatHashStrategyCompiler(new TypeOperators())),
399399
TESTING_TYPE_MANAGER,
400400
config,

plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHiveConnectorFactory.java

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -38,17 +38,23 @@ public class TestingHiveConnectorFactory
3838
implements ConnectorFactory
3939
{
4040
private final Optional<HiveMetastore> metastore;
41+
private final boolean metastoreImpersonationEnabled;
4142
private final Module module;
4243

4344
public TestingHiveConnectorFactory(Path localFileSystemRootPath)
4445
{
45-
this(localFileSystemRootPath, Optional.empty(), Optional.empty());
46+
this(localFileSystemRootPath, Optional.empty(), false, Optional.empty());
4647
}
4748

4849
@Deprecated
49-
public TestingHiveConnectorFactory(Path localFileSystemRootPath, Optional<HiveMetastore> metastore, Optional<DecryptionKeyRetriever> decryptionKeyRetriever)
50+
public TestingHiveConnectorFactory(
51+
Path localFileSystemRootPath,
52+
Optional<HiveMetastore> metastore,
53+
boolean metastoreImpersonationEnabled,
54+
Optional<DecryptionKeyRetriever> decryptionKeyRetriever)
5055
{
5156
this.metastore = requireNonNull(metastore, "metastore is null");
57+
this.metastoreImpersonationEnabled = metastoreImpersonationEnabled;
5258

5359
boolean ignored = localFileSystemRootPath.toFile().mkdirs();
5460
this.module = binder -> {
@@ -79,6 +85,6 @@ public Connector create(String catalogName, Map<String, String> config, Connecto
7985
if (metastore.isEmpty() && !config.containsKey("hive.metastore")) {
8086
configBuilder.put("hive.metastore", "file");
8187
}
82-
return createConnector(catalogName, configBuilder.buildOrThrow(), context, module, metastore, Optional.empty());
88+
return createConnector(catalogName, configBuilder.buildOrThrow(), context, module, metastore, metastoreImpersonationEnabled, Optional.empty());
8389
}
8490
}

plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingHivePlugin.java

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,30 +29,36 @@ public class TestingHivePlugin
2929
{
3030
private final Path localFileSystemRootPath;
3131
private final Optional<HiveMetastore> metastore;
32+
private final boolean metastoreImpersonationEnabled;
3233
private final Optional<DecryptionKeyRetriever> decryptionKeyRetriever;
3334

3435
public TestingHivePlugin(Path localFileSystemRootPath)
3536
{
36-
this(localFileSystemRootPath, Optional.empty(), Optional.empty());
37+
this(localFileSystemRootPath, Optional.empty(), false, Optional.empty());
3738
}
3839

3940
@Deprecated
4041
public TestingHivePlugin(Path localFileSystemRootPath, HiveMetastore metastore)
4142
{
42-
this(localFileSystemRootPath, Optional.of(metastore), Optional.empty());
43+
this(localFileSystemRootPath, Optional.of(metastore), false, Optional.empty());
4344
}
4445

4546
@Deprecated
46-
public TestingHivePlugin(Path localFileSystemRootPath, Optional<HiveMetastore> metastore, Optional<DecryptionKeyRetriever> decryptionKeyRetriever)
47+
public TestingHivePlugin(
48+
Path localFileSystemRootPath,
49+
Optional<HiveMetastore> metastore,
50+
boolean metastoreImpersonationEnabled,
51+
Optional<DecryptionKeyRetriever> decryptionKeyRetriever)
4752
{
4853
this.localFileSystemRootPath = requireNonNull(localFileSystemRootPath, "localFileSystemRootPath is null");
4954
this.metastore = requireNonNull(metastore, "metastore is null");
55+
this.metastoreImpersonationEnabled = metastoreImpersonationEnabled;
5056
this.decryptionKeyRetriever = requireNonNull(decryptionKeyRetriever, "decryptionKeyRetriever is null");
5157
}
5258

5359
@Override
5460
public Iterable<ConnectorFactory> getConnectorFactories()
5561
{
56-
return ImmutableList.of(new TestingHiveConnectorFactory(localFileSystemRootPath, metastore, decryptionKeyRetriever));
62+
return ImmutableList.of(new TestingHiveConnectorFactory(localFileSystemRootPath, metastore, metastoreImpersonationEnabled, decryptionKeyRetriever));
5763
}
5864
}

plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestingThriftHiveMetastoreBuilder.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
import static io.trino.plugin.base.security.UserNameProvider.SIMPLE_USER_NAME_PROVIDER;
3636
import static io.trino.plugin.hive.HiveTestUtils.HDFS_ENVIRONMENT;
3737
import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_STATS;
38+
import static io.trino.plugin.hive.HiveTestUtils.SESSION;
3839
import static io.trino.plugin.hive.metastore.thrift.TestingTokenAwareMetastoreClientFactory.TIMEOUT;
3940
import static java.util.Objects.requireNonNull;
4041
import static java.util.concurrent.Executors.newFixedThreadPool;
@@ -108,6 +109,6 @@ public ThriftMetastore build(Consumer<AutoCloseable> registerResource)
108109
thriftMetastoreConfig,
109110
fileSystemFactory,
110111
executorService);
111-
return metastoreFactory.createMetastore(Optional.empty());
112+
return metastoreFactory.createMetastore(Optional.of(SESSION.getIdentity()));
112113
}
113114
}

0 commit comments

Comments
 (0)