From dfbda86c6f0b631b520a2252b569d0dedc525878 Mon Sep 17 00:00:00 2001 From: Piotr Findeisen Date: Wed, 15 May 2024 12:32:42 +0200 Subject: [PATCH 1/2] Close thrift metastore client in tests --- .../BaseDeltaLakeConnectorSmokeTest.java | 2 +- ...tDeltaLakeFlushMetadataCacheProcedure.java | 2 +- .../io/trino/plugin/hive/HiveQueryRunner.java | 4 +-- .../plugin/hive/TestHive3OnDataLake.java | 2 +- .../TestingThriftHiveMetastoreBuilder.java | 8 ++++-- .../metastore/TestBridgingHiveMetastore.java | 10 ++++--- .../cache/TestCachingHiveMetastore.java | 12 ++++++--- .../thrift/TestThriftHttpMetastoreClient.java | 12 ++++++++- .../metastore/thrift/TestUnityMetastore.java | 27 ++++++++++--------- .../plugin/hive/s3/S3HiveQueryRunner.java | 2 +- .../BaseIcebergMinioConnectorSmokeTest.java | 4 +-- .../TestIcebergAbfsConnectorSmokeTest.java | 4 +-- .../TestIcebergGcsConnectorSmokeTest.java | 4 +-- ...TestTrinoHiveCatalogWithHiveMetastore.java | 12 ++++----- .../trino/testing/DistributedQueryRunner.java | 19 +++++++++++++ 15 files changed, 84 insertions(+), 40 deletions(-) diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeConnectorSmokeTest.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeConnectorSmokeTest.java index 19e93546b034..b3ef29bd386d 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeConnectorSmokeTest.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/BaseDeltaLakeConnectorSmokeTest.java @@ -170,7 +170,7 @@ protected QueryRunner createQueryRunner() this.metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveHadoop.getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); QueryRunner queryRunner = createDeltaLakeQueryRunner(); try { diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeFlushMetadataCacheProcedure.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeFlushMetadataCacheProcedure.java index c1d8c4716951..2eaa59bed38e 100644 --- a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeFlushMetadataCacheProcedure.java +++ b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeFlushMetadataCacheProcedure.java @@ -44,7 +44,7 @@ protected QueryRunner createQueryRunner() metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveMinioDataLake.getHiveHadoop().getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); return DeltaLakeQueryRunner.builder("default") .addMetastoreProperties(hiveMinioDataLake.getHiveHadoop()) diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveQueryRunner.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveQueryRunner.java index 14dd29a7ffb2..0ba9a2824da6 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveQueryRunner.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/HiveQueryRunner.java @@ -99,7 +99,7 @@ public static class Builder> private ImmutableMap.Builder hiveProperties = ImmutableMap.builder(); private List> initialTables = ImmutableList.of(); private Optional initialSchemasLocationBase = Optional.empty(); - private Optional> metastore = Optional.empty(); + private Optional> metastore = Optional.empty(); private boolean tpcdsCatalogEnabled; private boolean tpchBucketedCatalogEnabled; private boolean createTpchSchemas = true; @@ -153,7 +153,7 @@ public SELF setInitialSchemasLocationBase(String initialSchemasLocationBase) } @CanIgnoreReturnValue - public SELF setMetastore(Function metastore) + public SELF setMetastore(Function metastore) { this.metastore = Optional.of(metastore); return self(); diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHive3OnDataLake.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHive3OnDataLake.java index 0ef43c35bff8..b9fb17ab55db 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHive3OnDataLake.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHive3OnDataLake.java @@ -93,7 +93,7 @@ protected QueryRunner createQueryRunner() this.metastoreClient = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(this.hiveMinioDataLake.getHiveHadoop().getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); return S3HiveQueryRunner.builder(hiveMinioDataLake) .addExtraProperty("sql.path", "hive.functions") .addExtraProperty("sql.default-function-catalog", "hive") 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 bc7aaf8727ad..c344df637391 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 @@ -28,6 +28,8 @@ import java.net.URI; import java.util.Optional; +import java.util.concurrent.ExecutorService; +import java.util.function.Consumer; import static com.google.common.base.Preconditions.checkState; import static io.trino.plugin.base.security.UserNameProvider.SIMPLE_USER_NAME_PROVIDER; @@ -95,15 +97,17 @@ public TestingThriftHiveMetastoreBuilder fileSystemFactory(TrinoFileSystemFactor return this; } - public ThriftMetastore build() + public ThriftMetastore build(Consumer registerResource) { checkState(tokenAwareMetastoreClientFactory != null, "metastore client not set"); + ExecutorService executorService = newFixedThreadPool(thriftMetastoreConfig.getWriteStatisticsThreads()); + registerResource.accept(executorService); ThriftHiveMetastoreFactory metastoreFactory = new ThriftHiveMetastoreFactory( new UgiBasedMetastoreClientFactory(tokenAwareMetastoreClientFactory, SIMPLE_USER_NAME_PROVIDER, thriftMetastoreConfig), new HiveMetastoreConfig().isHideDeltaLakeTables(), thriftMetastoreConfig, fileSystemFactory, - newFixedThreadPool(thriftMetastoreConfig.getWriteStatisticsThreads())); + executorService); return metastoreFactory.createMetastore(Optional.empty()); } } diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestBridgingHiveMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestBridgingHiveMetastore.java index 9e241568d409..86719b75d1ab 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestBridgingHiveMetastore.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestBridgingHiveMetastore.java @@ -13,6 +13,7 @@ */ package io.trino.plugin.hive.metastore; +import io.trino.plugin.base.util.AutoCloseableCloser; import io.trino.plugin.hive.SchemaAlreadyExistsException; import io.trino.plugin.hive.TableAlreadyExistsException; import io.trino.plugin.hive.containers.HiveHadoop; @@ -47,12 +48,12 @@ final class TestBridgingHiveMetastore extends AbstractTestHiveMetastore { - private final HiveHadoop hiveHadoop; + private final AutoCloseableCloser closer = AutoCloseableCloser.create(); private final HiveMetastore metastore; TestBridgingHiveMetastore() { - hiveHadoop = HiveHadoop.builder().build(); + HiveHadoop hiveHadoop = closer.register(HiveHadoop.builder().build()); hiveHadoop.start(); MetastoreClientAdapterProvider metastoreClientAdapterProvider = delegate -> newProxy(ThriftMetastoreClient.class, (proxy, method, methodArgs) -> { @@ -72,13 +73,14 @@ final class TestBridgingHiveMetastore metastore = new BridgingHiveMetastore(testingThriftHiveMetastoreBuilder() .metastoreClient(hiveHadoop.getHiveMetastoreEndpoint(), metastoreClientAdapterProvider) .thriftMetastoreConfig(new ThriftMetastoreConfig().setDeleteFilesOnDrop(true)) - .build()); + .build(closer::register)); } @AfterAll void afterAll() + throws Exception { - hiveHadoop.stop(); + closer.close(); } @Override diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastore.java index 858776e1e6f8..6c902b1e6439 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastore.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/cache/TestCachingHiveMetastore.java @@ -22,6 +22,7 @@ import io.trino.hive.thrift.metastore.ColumnStatisticsData; import io.trino.hive.thrift.metastore.ColumnStatisticsObj; import io.trino.hive.thrift.metastore.LongColumnStatsData; +import io.trino.plugin.base.util.AutoCloseableCloser; import io.trino.plugin.hive.HiveBasicStatistics; import io.trino.plugin.hive.HiveColumnHandle; import io.trino.plugin.hive.HiveMetastoreClosure; @@ -118,6 +119,7 @@ public class TestCachingHiveMetastore private static final TableInfo TEST_TABLE_INFO = new TableInfo(TEST_SCHEMA_TABLE, TableInfo.ExtendedRelationType.TABLE); private static final Duration CACHE_TTL = new Duration(5, TimeUnit.MINUTES); + private AutoCloseableCloser closer; private MockThriftMetastoreClient mockClient; private ThriftMetastore thriftHiveMetastore; private ListeningExecutorService executor; @@ -128,6 +130,7 @@ public class TestCachingHiveMetastore @BeforeEach public void setUp() { + closer = AutoCloseableCloser.create(); mockClient = new MockThriftMetastoreClient(); thriftHiveMetastore = createThriftHiveMetastore(); executor = listeningDecorator(newCachedThreadPool(daemonThreadsNamed(getClass().getSimpleName() + "-%s"))); @@ -140,10 +143,13 @@ public void setUp() @AfterAll public void tearDown() + throws Exception { executor.shutdownNow(); executor = null; metastore = null; + closer.close(); + closer = null; } private ThriftMetastore createThriftHiveMetastore() @@ -151,11 +157,11 @@ private ThriftMetastore createThriftHiveMetastore() return createThriftHiveMetastore(mockClient); } - private static ThriftMetastore createThriftHiveMetastore(ThriftMetastoreClient client) + private ThriftMetastore createThriftHiveMetastore(ThriftMetastoreClient client) { return testingThriftHiveMetastoreBuilder() .metastoreClient(client) - .build(); + .build(closer::register); } @Test @@ -1050,7 +1056,7 @@ private PartitionCachingAssertions assertThatCachingWithDisabledPartitionCache() return new PartitionCachingAssertions(executor); } - static class PartitionCachingAssertions + class PartitionCachingAssertions { private final CachingHiveMetastore cachingHiveMetastore; private final MockThriftMetastoreClient thriftClient; diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftHttpMetastoreClient.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftHttpMetastoreClient.java index 5a6e3411024a..1cbe6d01e5b4 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftHttpMetastoreClient.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftHttpMetastoreClient.java @@ -17,9 +17,11 @@ import io.opentelemetry.api.OpenTelemetry; import io.trino.hive.thrift.metastore.Database; import io.trino.hive.thrift.metastore.NoSuchObjectException; +import io.trino.plugin.base.util.AutoCloseableCloser; import io.trino.testing.TestingNodeManager; import jakarta.servlet.http.HttpServletRequest; import org.apache.http.HttpHeaders; +import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -38,6 +40,7 @@ public class TestThriftHttpMetastoreClient { + private static final AutoCloseableCloser closer = AutoCloseableCloser.create(); private static ThriftMetastore delegate; @BeforeAll @@ -46,7 +49,14 @@ public static void setup() { File tempDir = Files.createTempDirectory(null).toFile(); tempDir.deleteOnExit(); - delegate = testingThriftHiveMetastoreBuilder().metastoreClient(createFakeMetastoreClient()).build(); + delegate = testingThriftHiveMetastoreBuilder().metastoreClient(createFakeMetastoreClient()).build(closer::register); + } + + @AfterAll + public static void tearDown() + throws Exception + { + closer.close(); } private static ThriftMetastoreClient createFakeMetastoreClient() diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestUnityMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestUnityMetastore.java index b864a94e0fbb..eaeba00ed43a 100644 --- a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestUnityMetastore.java +++ b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/thrift/TestUnityMetastore.java @@ -14,6 +14,7 @@ package io.trino.plugin.hive.metastore.thrift; import io.opentelemetry.api.OpenTelemetry; +import io.trino.plugin.base.util.AutoCloseableCloser; import io.trino.plugin.hive.metastore.HiveMetastore; import io.trino.plugin.hive.metastore.TableInfo; import io.trino.spi.connector.SchemaTableName; @@ -38,7 +39,7 @@ void test() String databricksHost = requireNonNull(System.getenv("DATABRICKS_HOST"), "Environment variable not set: DATABRICKS_HOST"); String databricksToken = requireNonNull(System.getenv("DATABRICKS_TOKEN"), "Environment variable not set: DATABRICKS_TOKEN"); String databricksCatalogName = requireNonNull(System.getenv("DATABRICKS_UNITY_CATALOG_NAME"), "Environment variable not set: DATABRICKS_UNITY_CATALOG_NAME"); - URI metastoreUri = URI.create("https://%s:443/api/2.0/unity-hms-proxy/metadata".formatted(databricksHost)); + URI metastoreUri = URI.create("https://%s:443/api/2.0/unity-hms-proxy/metadata" .formatted(databricksHost)); ThriftHttpMetastoreConfig config = new ThriftHttpMetastoreConfig() .setAuthenticationMode(BEARER) @@ -47,18 +48,20 @@ void test() ThriftMetastoreClient client = ((ThriftMetastoreClientFactory) new HttpThriftMetastoreClientFactory(config, new TestingNodeManager(), OpenTelemetry.noop())) .create(metastoreUri, Optional.empty()); - HiveMetastore metastore = new BridgingHiveMetastore(testingThriftHiveMetastoreBuilder() - .metastoreClient(client) - .thriftMetastoreConfig(new ThriftMetastoreConfig().setDeleteFilesOnDrop(true)) - .build()); + try (var closer = AutoCloseableCloser.create()) { + HiveMetastore metastore = new BridgingHiveMetastore(testingThriftHiveMetastoreBuilder() + .metastoreClient(client) + .thriftMetastoreConfig(new ThriftMetastoreConfig().setDeleteFilesOnDrop(true)) + .build(closer::register)); - List tables = metastore.getAllDatabases().stream() - .map(metastore::getTables) - .flatMap(List::stream) - .toList(); - assertThat(tables).isNotEmpty(); + List tables = metastore.getAllDatabases().stream() + .map(metastore::getTables) + .flatMap(List::stream) + .toList(); + assertThat(tables).isNotEmpty(); - SchemaTableName schemaTableName = tables.getFirst().tableName(); - assertThat(metastore.getTable(schemaTableName.getSchemaName(), schemaTableName.getTableName())).isPresent(); + SchemaTableName schemaTableName = tables.getFirst().tableName(); + assertThat(metastore.getTable(schemaTableName.getSchemaName(), schemaTableName.getTableName())).isPresent(); + } } } 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 eb614990b6f6..9df3a85934b3 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 @@ -164,7 +164,7 @@ public DistributedQueryRunner build() testingThriftHiveMetastoreBuilder() .metastoreClient(hiveMetastoreEndpoint, thriftMetastoreTimeout) .thriftMetastoreConfig(thriftMetastoreConfig) - .build())); + .build(distributedQueryRunner::registerResource))); setInitialSchemasLocationBase("s3a://" + bucketName); // cannot use s3:// as Hive metastore is not configured to accept it return super.build(); } diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergMinioConnectorSmokeTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergMinioConnectorSmokeTest.java index 84428f3a1d13..3d54196f5d8f 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergMinioConnectorSmokeTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergMinioConnectorSmokeTest.java @@ -233,7 +233,7 @@ protected void dropTableFromMetastore(String tableName) HiveMetastore metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveMinioDataLake.getHiveHadoop().getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); metastore.dropTable(schemaName, tableName, false); assertThat(metastore.getTable(schemaName, tableName)).isEmpty(); } @@ -244,7 +244,7 @@ protected String getMetadataLocation(String tableName) HiveMetastore metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveMinioDataLake.getHiveHadoop().getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); return metastore .getTable(schemaName, tableName).orElseThrow() .getParameters().get("metadata_location"); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergAbfsConnectorSmokeTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergAbfsConnectorSmokeTest.java index 63a5d80d4d5c..66ec24a21d9e 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergAbfsConnectorSmokeTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergAbfsConnectorSmokeTest.java @@ -123,7 +123,7 @@ protected void dropTableFromMetastore(String tableName) HiveMetastore metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveHadoop.getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); metastore.dropTable(schemaName, tableName, false); assertThat(metastore.getTable(schemaName, tableName)).isEmpty(); } @@ -134,7 +134,7 @@ protected String getMetadataLocation(String tableName) HiveMetastore metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveHadoop.getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); return metastore .getTable(schemaName, tableName).orElseThrow() .getParameters().get("metadata_location"); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergGcsConnectorSmokeTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergGcsConnectorSmokeTest.java index df1241e9a20d..857d3c8f6f83 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergGcsConnectorSmokeTest.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergGcsConnectorSmokeTest.java @@ -175,7 +175,7 @@ protected void dropTableFromMetastore(String tableName) HiveMetastore metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveHadoop.getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); metastore.dropTable(schema, tableName, false); assertThat(metastore.getTable(schema, tableName)).isEmpty(); } @@ -186,7 +186,7 @@ protected String getMetadataLocation(String tableName) HiveMetastore metastore = new BridgingHiveMetastore( testingThriftHiveMetastoreBuilder() .metastoreClient(hiveHadoop.getHiveMetastoreEndpoint()) - .build()); + .build(this::closeAfterClass)); return metastore .getTable(schema, tableName).orElseThrow() .getParameters().get("metadata_location"); diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/hms/TestTrinoHiveCatalogWithHiveMetastore.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/hms/TestTrinoHiveCatalogWithHiveMetastore.java index 330b28e5f7cc..b13affd3017d 100644 --- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/hms/TestTrinoHiveCatalogWithHiveMetastore.java +++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/hms/TestTrinoHiveCatalogWithHiveMetastore.java @@ -25,6 +25,7 @@ import io.trino.hdfs.authentication.NoHdfsAuthentication; import io.trino.hdfs.s3.HiveS3Config; import io.trino.hdfs.s3.TrinoS3ConfigurationInitializer; +import io.trino.plugin.base.util.AutoCloseableCloser; import io.trino.plugin.hive.TrinoViewHiveMetastore; import io.trino.plugin.hive.containers.HiveMinioDataLake; import io.trino.plugin.hive.metastore.cache.CachingHiveMetastore; @@ -66,13 +67,14 @@ public class TestTrinoHiveCatalogWithHiveMetastore { private static final String bucketName = "test-hive-catalog-with-hms-" + randomNameSuffix(); + private AutoCloseableCloser closer = AutoCloseableCloser.create(); // Use MinIO for storage, since HDFS is hard to get working in a unit test private HiveMinioDataLake dataLake; @BeforeAll public void setUp() { - dataLake = new HiveMinioDataLake(bucketName, HIVE3_IMAGE); + dataLake = closer.register(new HiveMinioDataLake(bucketName, HIVE3_IMAGE)); dataLake.start(); } @@ -80,10 +82,8 @@ public void setUp() public void tearDown() throws Exception { - if (dataLake != null) { - dataLake.stop(); - dataLake = null; - } + dataLake = null; + closer.close(); } @Override @@ -108,7 +108,7 @@ protected TrinoCatalog createTrinoCatalog(boolean useUniqueTableLocations) // Read timed out sometimes happens with the default timeout .setReadTimeout(new Duration(1, MINUTES))) .metastoreClient(dataLake.getHiveHadoop().getHiveMetastoreEndpoint()) - .build(); + .build(closer::register); CachingHiveMetastore metastore = createPerTransactionCache(new BridgingHiveMetastore(thriftMetastore), 1000); return new TrinoHiveCatalog( new CatalogName("catalog"), diff --git a/testing/trino-testing/src/main/java/io/trino/testing/DistributedQueryRunner.java b/testing/trino-testing/src/main/java/io/trino/testing/DistributedQueryRunner.java index 471bd91f6e2d..b8d19736a726 100644 --- a/testing/trino-testing/src/main/java/io/trino/testing/DistributedQueryRunner.java +++ b/testing/trino-testing/src/main/java/io/trino/testing/DistributedQueryRunner.java @@ -81,6 +81,7 @@ import static com.google.common.base.MoreObjects.firstNonNull; import static com.google.common.base.Preconditions.checkState; +import static com.google.common.base.Throwables.throwIfInstanceOf; import static com.google.common.base.Throwables.throwIfUnchecked; import static com.google.common.base.Verify.verify; import static com.google.common.collect.Iterables.getOnlyElement; @@ -231,6 +232,22 @@ private DistributedQueryRunner( log.info("Created DistributedQueryRunner in %s (unclosed instances = %s)", nanosSince(start), unclosedInstances.incrementAndGet()); } + public void registerResource(AutoCloseable resource) + { + requireNonNull(resource, "resource is null"); + checkState(!closed, "already closed"); + closer.register(() -> { + try { + resource.close(); + } + catch (Exception e) { + throwIfUnchecked(e); + throwIfInstanceOf(e, IOException.class); + throw new RuntimeException(e); + } + }); + } + private TestingTrinoServer createServer( boolean coordinator, Map extraCoordinatorProperties, @@ -835,6 +852,7 @@ public SELF enableBackupCoordinator() @CanIgnoreReturnValue public SELF registerResource(AutoCloseable closeable) { + checkState(extraCloseables != null, "query runner is already built"); extraCloseables.add(requireNonNull(closeable, "closeable is null")); return self(); } @@ -899,6 +917,7 @@ public void queryCompleted(QueryCompletedEvent queryCompletedEvent) eventListeners, extraCloseables.build(), testingTrinoClientFactory); + extraCloseables = null; try { additionalSetup.accept(queryRunner); From a7846c59005982072cb7fcdab5167e495f0064c7 Mon Sep 17 00:00:00 2001 From: Piotr Findeisen Date: Wed, 15 May 2024 14:51:49 +0200 Subject: [PATCH 2/2] =?UTF-8?q?empty:=20roll=20the=20dice=20=F0=9F=8E=B2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit