Properly close Glue client in tests

Fix potential resource leak during tests' execution.
trinodb · May 13, 2024 · d0a716f · d0a716f
1 parent 1c6c7ef
commit d0a716f
Show file tree

Hide file tree

Showing 13 changed files with 24 additions and 17 deletions.
diff --git a/...a-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedGlueMetastoreViews.java b/...a-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeSharedGlueMetastoreViews.java
@@ -29,6 +29,6 @@ public class TestDeltaLakeSharedGlueMetastoreViews
     @Override
     protected HiveMetastore createTestMetastore(Path dataDirectory)
     {
-        return createTestingGlueHiveMetastore(dataDirectory);
+        return createTestingGlueHiveMetastore(dataDirectory, this::closeAfterClass);
     }
 }
diff --git a/...java/io/trino/plugin/deltalake/TestDeltaLakeSharedGlueMetastoreWithTableRedirections.java b/...java/io/trino/plugin/deltalake/TestDeltaLakeSharedGlueMetastoreWithTableRedirections.java
@@ -65,7 +65,7 @@ protected QueryRunner createQueryRunner()
                         .put("delta.hive-catalog-name", "hive_with_redirections")
                         .buildOrThrow());
 
-        this.glueMetastore = createTestingGlueHiveMetastore(dataDirectory);
+        this.glueMetastore = createTestingGlueHiveMetastore(dataDirectory, this::closeAfterClass);
         queryRunner.installPlugin(new TestingHivePlugin(queryRunner.getCoordinator().getBaseDataDir().resolve("hive_data"), glueMetastore));
         queryRunner.createCatalog(
                 "hive_with_redirections",

diff --git a/...ava/io/trino/plugin/deltalake/TestDeltaLakeTableWithCustomLocationUsingGlueMetastore.java b/...ava/io/trino/plugin/deltalake/TestDeltaLakeTableWithCustomLocationUsingGlueMetastore.java
@@ -41,7 +41,7 @@ protected QueryRunner createQueryRunner()
     {
         Path warehouseDir = Files.createTempDirectory("warehouse-dir");
         closeAfterClass(() -> deleteRecursively(warehouseDir, ALLOW_INSECURE));
-        metastore = createTestingGlueHiveMetastore(warehouseDir);
+        metastore = createTestingGlueHiveMetastore(warehouseDir, this::closeAfterClass);
         schema = "test_tables_with_custom_location" + randomNameSuffix();
         return DeltaLakeQueryRunner.builder(schema)
                 .addDeltaProperty("hive.metastore", "glue")

diff --git a/...ino/plugin/deltalake/metastore/glue/TestDeltaLakeConcurrentModificationGlueMetastore.java b/...ino/plugin/deltalake/metastore/glue/TestDeltaLakeConcurrentModificationGlueMetastore.java
@@ -75,7 +75,7 @@ protected QueryRunner createQueryRunner()
         GlueHiveMetastoreConfig glueConfig = new GlueHiveMetastoreConfig()
                 .setDefaultWarehouseDir(dataDirectory.toUri().toString());
 
-        GlueClient glueClient = createGlueClient(new GlueHiveMetastoreConfig(), OpenTelemetry.noop());
+        GlueClient glueClient = closeAfterClass(createGlueClient(new GlueHiveMetastoreConfig(), OpenTelemetry.noop()));
         GlueClient proxiedGlueClient = newProxy(GlueClient.class, (proxy, method, args) -> {
             Object result;
             try {

diff --git a/...io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeRegisterTableProcedureWithGlue.java b/...io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeRegisterTableProcedureWithGlue.java
@@ -40,7 +40,7 @@ protected QueryRunner createQueryRunner()
     {
         Path warehouseDir = Files.createTempDirectory("warehouse-dir");
         closeAfterClass(() -> deleteRecursively(warehouseDir, ALLOW_INSECURE));
-        metastore = createTestingGlueHiveMetastore(warehouseDir);
+        metastore = createTestingGlueHiveMetastore(warehouseDir, this::closeAfterClass);
         schema = "test_delta_lake_register_table" + randomNameSuffix();
         return DeltaLakeQueryRunner.builder(schema)
                 .addDeltaProperty("hive.metastore", "glue")

diff --git a/...c/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeViewsGlueMetastore.java b/...c/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaLakeViewsGlueMetastore.java
@@ -48,7 +48,7 @@ protected QueryRunner createQueryRunner()
     {
         Path warehouseDir = Files.createTempDirectory("warehouse-dir");
         closeAfterClass(() -> deleteRecursively(warehouseDir, ALLOW_INSECURE));
-        metastore = createTestingGlueHiveMetastore(warehouseDir);
+        metastore = createTestingGlueHiveMetastore(warehouseDir, this::closeAfterClass);
         schema = "test_delta_lake_glue_views_" + randomNameSuffix();
         return DeltaLakeQueryRunner.builder(schema)
                 .addDeltaProperty("hive.metastore", "glue")

diff --git a/...c/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaS3AndGlueMetastoreTest.java b/...c/test/java/io/trino/plugin/deltalake/metastore/glue/TestDeltaS3AndGlueMetastoreTest.java
@@ -40,7 +40,7 @@ public TestDeltaS3AndGlueMetastoreTest()
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        metastore = createTestingGlueHiveMetastore(URI.create(schemaPath()));
+        metastore = createTestingGlueHiveMetastore(URI.create(schemaPath()), this::closeAfterClass);
         return DeltaLakeQueryRunner.builder(schemaName)
                 .setDeltaProperties(ImmutableMap.<String, String>builder()
                         .put("hive.metastore", "glue")

diff --git a/...-hive/src/test/java/io/trino/plugin/hive/TestHiveConcurrentModificationGlueMetastore.java b/...-hive/src/test/java/io/trino/plugin/hive/TestHiveConcurrentModificationGlueMetastore.java
@@ -73,7 +73,7 @@ protected QueryRunner createQueryRunner()
         GlueHiveMetastoreConfig glueConfig = new GlueHiveMetastoreConfig()
                 .setDefaultWarehouseDir(dataDirectory.toUri().toString());
 
-        GlueClient glueClient = createGlueClient(new GlueHiveMetastoreConfig(), OpenTelemetry.noop());
+        GlueClient glueClient = closeAfterClass(createGlueClient(new GlueHiveMetastoreConfig(), OpenTelemetry.noop()));
         GlueClient proxiedGlueClient = newProxy(GlueClient.class, (proxy, method, args) -> {
             try {
                 if (method.getName().equals("updateTable")) {

diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestGlueHiveMetastore.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/metastore/TestGlueHiveMetastore.java
@@ -13,6 +13,7 @@
  */
 package io.trino.plugin.hive.metastore;
 
+import io.trino.plugin.base.util.AutoCloseableCloser;
 import io.trino.plugin.hive.metastore.glue.GlueHiveMetastore;
 import org.junit.jupiter.api.AfterAll;
 
@@ -27,22 +28,24 @@
 final class TestGlueHiveMetastore
         extends AbstractTestHiveMetastore
 {
+    private final AutoCloseableCloser closer = AutoCloseableCloser.create();
     private final Path tempDir;
     private final GlueHiveMetastore metastore;
 
     TestGlueHiveMetastore()
             throws IOException
     {
         tempDir = createTempDirectory("test");
-        metastore = createTestingGlueHiveMetastore(tempDir);
+        metastore = createTestingGlueHiveMetastore(tempDir, closer::register);
     }
 
     @AfterAll
     void tearDown()
-            throws IOException
+            throws Exception
     {
         metastore.shutdown();
         deleteRecursively(tempDir, ALLOW_INSECURE);
+        closer.close();
     }
 
     @Override

diff --git a/...rino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestingGlueHiveMetastore.java b/...rino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestingGlueHiveMetastore.java
@@ -15,11 +15,13 @@
 
 import io.opentelemetry.api.OpenTelemetry;
 import io.trino.plugin.hive.metastore.glue.GlueHiveMetastore.TableKind;
+import software.amazon.awssdk.services.glue.GlueClient;
 
 import java.io.IOException;
 import java.net.URI;
 import java.nio.file.Path;
 import java.util.EnumSet;
+import java.util.function.Consumer;
 
 import static com.google.common.base.Verify.verify;
 import static io.trino.plugin.hive.HiveTestUtils.HDFS_FILE_SYSTEM_FACTORY;
@@ -32,7 +34,7 @@ public final class TestingGlueHiveMetastore
 {
     private TestingGlueHiveMetastore() {}
 
-    public static GlueHiveMetastore createTestingGlueHiveMetastore(Path defaultWarehouseDir)
+    public static GlueHiveMetastore createTestingGlueHiveMetastore(Path defaultWarehouseDir, Consumer<AutoCloseable> registerResource)
     {
         if (!exists(defaultWarehouseDir)) {
             try {
@@ -43,15 +45,17 @@ public static GlueHiveMetastore createTestingGlueHiveMetastore(Path defaultWareh
             }
         }
         verify(isDirectory(defaultWarehouseDir), "%s is not a directory", defaultWarehouseDir);
-        return createTestingGlueHiveMetastore(defaultWarehouseDir.toUri());
+        return createTestingGlueHiveMetastore(defaultWarehouseDir.toUri(), registerResource);
     }
 
-    public static GlueHiveMetastore createTestingGlueHiveMetastore(URI warehouseUri)
+    public static GlueHiveMetastore createTestingGlueHiveMetastore(URI warehouseUri, Consumer<AutoCloseable> registerResource)
     {
         GlueHiveMetastoreConfig glueConfig = new GlueHiveMetastoreConfig()
                 .setDefaultWarehouseDir(warehouseUri.toString());
+        GlueClient glueClient = createGlueClient(glueConfig, OpenTelemetry.noop());
+        registerResource.accept(glueClient);
         return new GlueHiveMetastore(
-                createGlueClient(glueConfig, OpenTelemetry.noop()),
+                glueClient,
                 new GlueContext(glueConfig),
                 GlueCache.NOOP,
                 HDFS_FILE_SYSTEM_FACTORY,

diff --git a/...ava/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueTableOperationsInsertFailure.java b/...ava/io/trino/plugin/iceberg/catalog/glue/TestIcebergGlueTableOperationsInsertFailure.java
@@ -91,7 +91,7 @@ protected QueryRunner createQueryRunner()
         queryRunner.installPlugin(new TestingIcebergPlugin(dataDirectory, Optional.of(new TestingIcebergGlueCatalogModule(awsGlueAsyncAdapterProvider))));
         queryRunner.createCatalog(ICEBERG_CATALOG, "iceberg", ImmutableMap.of());
 
-        glueHiveMetastore = createTestingGlueHiveMetastore(dataDirectory);
+        glueHiveMetastore = createTestingGlueHiveMetastore(dataDirectory, this::closeAfterClass);
 
         Database database = Database.builder()
                 .setDatabaseName(schemaName)

diff --git a/...src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergS3AndGlueMetastoreTest.java b/...src/test/java/io/trino/plugin/iceberg/catalog/glue/TestIcebergS3AndGlueMetastoreTest.java
@@ -44,7 +44,7 @@ public TestIcebergS3AndGlueMetastoreTest()
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        metastore = createTestingGlueHiveMetastore(URI.create(schemaPath()));
+        metastore = createTestingGlueHiveMetastore(URI.create(schemaPath()), this::closeAfterClass);
         return IcebergQueryRunner.builder()
                 .setIcebergProperties(ImmutableMap.<String, String>builder()
                         .put("iceberg.catalog.type", "glue")

diff --git a/...o-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestSharedGlueMetastore.java b/...o-iceberg/src/test/java/io/trino/plugin/iceberg/catalog/glue/TestSharedGlueMetastore.java
@@ -93,7 +93,7 @@ protected QueryRunner createQueryRunner()
                         "hive.metastore.glue.default-warehouse-dir", dataDirectory.toString(),
                         "iceberg.hive-catalog-name", "hive"));
 
-        this.glueMetastore = createTestingGlueHiveMetastore(dataDirectory);
+        this.glueMetastore = createTestingGlueHiveMetastore(dataDirectory, this::closeAfterClass);
         queryRunner.installPlugin(new TestingHivePlugin(queryRunner.getCoordinator().getBaseDataDir().resolve("hive_data"), glueMetastore));
         queryRunner.createCatalog(HIVE_CATALOG, "hive");
         queryRunner.createCatalog(