fix: For iceberg test failures with unclosed ResolvingFileIO instance (…

…#6217)
deephaven · Oct 18, 2024 · 8f02178 · 8f02178
1 parent 22c20cf
commit 8f02178
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 17 deletions.
diff --git a/extensions/iceberg/s3/src/test/java/io/deephaven/iceberg/util/IcebergToolsTest.java b/extensions/iceberg/s3/src/test/java/io/deephaven/iceberg/util/IcebergToolsTest.java
@@ -14,7 +14,6 @@
 import io.deephaven.iceberg.TestCatalog.IcebergTestCatalog;
 import io.deephaven.test.types.OutOfBandTest;
 import org.apache.iceberg.Snapshot;
-import org.apache.iceberg.catalog.Catalog;
 import org.apache.iceberg.catalog.Namespace;
 import org.apache.iceberg.catalog.TableIdentifier;
 import org.junit.After;
@@ -109,7 +108,7 @@ public abstract class IcebergToolsTest {
     private final List<String> keys = new ArrayList<>();
 
     private String warehousePath;
-    private Catalog resourceCatalog;
+    private IcebergTestCatalog resourceCatalog;
 
     @Rule
     public final EngineCleanup framework = new EngineCleanup();
@@ -134,6 +133,7 @@ public void setUp() throws ExecutionException, InterruptedException {
 
     @After
     public void tearDown() throws ExecutionException, InterruptedException {
+        resourceCatalog.close();
         for (String key : keys) {
             asyncClient.deleteObject(DeleteObjectRequest.builder().bucket(bucket).key(key).build()).get();
         }

diff --git a/extensions/iceberg/src/test/java/io/deephaven/iceberg/TestCatalog/IcebergTestCatalog.java b/extensions/iceberg/src/test/java/io/deephaven/iceberg/TestCatalog/IcebergTestCatalog.java
@@ -3,25 +3,33 @@
 //
 package io.deephaven.iceberg.TestCatalog;
 
+import org.apache.hadoop.conf.Configuration;
 import org.apache.iceberg.*;
 import org.apache.iceberg.catalog.Catalog;
 import org.apache.iceberg.catalog.Namespace;
 import org.apache.iceberg.catalog.SupportsNamespaces;
 import org.apache.iceberg.catalog.TableIdentifier;
 import org.apache.iceberg.exceptions.NamespaceNotEmptyException;
 import org.apache.iceberg.exceptions.NoSuchNamespaceException;
+import org.apache.iceberg.io.ResolvingFileIO;
 import org.jetbrains.annotations.NotNull;
 
 import java.io.File;
 import java.util.*;
 
-public class IcebergTestCatalog implements Catalog, SupportsNamespaces {
+public class IcebergTestCatalog implements Catalog, SupportsNamespaces, AutoCloseable {
     private final Map<Namespace, Map<TableIdentifier, Table>> namespaceTableMap;
     private final Map<TableIdentifier, Table> tableMap;
 
+    private final ResolvingFileIO fileIO;
+
     private IcebergTestCatalog(final String path, @NotNull final Map<String, String> properties) {
         namespaceTableMap = new HashMap<>();
         tableMap = new HashMap<>();
+        final Configuration hadoopConf = new Configuration();
+        fileIO = new ResolvingFileIO();
+        fileIO.setConf(hadoopConf);
+        fileIO.initialize(properties);
 
         // Assume first level is namespace.
         final File root = new File(path);
@@ -33,7 +41,7 @@ private IcebergTestCatalog(final String path, @NotNull final Map<String, String>
                     if (tableFile.isDirectory()) {
                         // Second level is table name.
                         final TableIdentifier tableId = TableIdentifier.of(namespace, tableFile.getName());
-                        final Table table = IcebergTestTable.loadFromMetadata(tableFile.getAbsolutePath(), properties);
+                        final Table table = IcebergTestTable.loadFromMetadata(tableFile.getAbsolutePath(), fileIO);
 
                         // Add it to the maps.
                         namespaceTableMap.get(namespace).put(tableId, table);
@@ -103,4 +111,9 @@ public boolean setProperties(Namespace namespace, Map<String, String> map) throw
     public boolean removeProperties(Namespace namespace, Set<String> set) throws NoSuchNamespaceException {
         return false;
     }
+
+    @Override
+    public void close() {
+        fileIO.close();
+    }
 }
diff --git a/extensions/iceberg/src/test/java/io/deephaven/iceberg/TestCatalog/IcebergTestTable.java b/extensions/iceberg/src/test/java/io/deephaven/iceberg/TestCatalog/IcebergTestTable.java
@@ -3,12 +3,10 @@
 //
 package io.deephaven.iceberg.TestCatalog;
 
-import org.apache.hadoop.conf.Configuration;
 import org.apache.iceberg.*;
 import org.apache.iceberg.encryption.EncryptionManager;
 import org.apache.iceberg.io.FileIO;
 import org.apache.iceberg.io.LocationProvider;
-import org.apache.iceberg.io.ResolvingFileIO;
 import org.jetbrains.annotations.NotNull;
 
 import java.io.File;
@@ -19,12 +17,12 @@
 
 public class IcebergTestTable implements Table {
     private final TableMetadata metadata;
-    private final Map<String, String> properties;
-    private final Configuration hadoopConf;
+    private final FileIO fileIO;
 
-    private IcebergTestTable(@NotNull final String path, @NotNull final Map<String, String> properties) {
-        this.properties = properties;
-        hadoopConf = new Configuration();
+    private IcebergTestTable(
+            @NotNull final String path,
+            @NotNull final FileIO fileIO) {
+        this.fileIO = fileIO;
 
         final File metadataRoot = new File(path, "metadata");
 
@@ -50,8 +48,8 @@ private IcebergTestTable(@NotNull final String path, @NotNull final Map<String,
 
     public static IcebergTestTable loadFromMetadata(
             @NotNull final String path,
-            @NotNull final Map<String, String> properties) {
-        return new IcebergTestTable(path, properties);
+            @NotNull final FileIO fileIO) {
+        return new IcebergTestTable(path, fileIO);
     }
 
     @Override
@@ -220,10 +218,7 @@ public Transaction newTransaction() {
 
     @Override
     public FileIO io() {
-        final ResolvingFileIO io = new ResolvingFileIO();
-        io.setConf(hadoopConf);
-        io.initialize(properties);
-        return io;
+        return fileIO;
     }
 
     @Override