Use builder pattern for IcebergQueryRunner

Additionally, add support for tracking catalog properties added to iceberg. This can be useful for tests that may require their own catalogs
prestodb · Feb 7, 2025 · bdd370e · bdd370e
1 parent 8789cd9
commit bdd370e
Show file tree

Hide file tree

Showing 27 changed files with 308 additions and 313 deletions.
diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/BenchmarkIcebergHadoopCatalog.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/BenchmarkIcebergHadoopCatalog.java
@@ -14,7 +14,6 @@
 package com.facebook.presto.iceberg;
 
 import com.facebook.presto.tests.DistributedQueryRunner;
-import com.google.common.collect.ImmutableMap;
 import org.openjdk.jmh.annotations.BenchmarkMode;
 import org.openjdk.jmh.annotations.Fork;
 import org.openjdk.jmh.annotations.Measurement;
@@ -27,12 +26,9 @@
 import org.openjdk.jmh.runner.options.OptionsBuilder;
 import org.openjdk.jmh.runner.options.VerboseMode;
 
-import java.util.Optional;
 import java.util.OptionalInt;
 
 import static com.facebook.presto.iceberg.CatalogType.HADOOP;
-import static com.facebook.presto.iceberg.FileFormat.PARQUET;
-import static com.facebook.presto.iceberg.IcebergQueryRunner.createIcebergQueryRunner;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static org.openjdk.jmh.annotations.Mode.AverageTime;
 import static org.openjdk.jmh.annotations.Scope.Benchmark;
@@ -50,14 +46,11 @@ public class BenchmarkIcebergHadoopCatalog
     public DistributedQueryRunner getQueryRunner()
     {
         try {
-            return createIcebergQueryRunner(
-                    ImmutableMap.of(),
-                    ImmutableMap.of("iceberg.catalog.type", HADOOP.name()),
-                    PARQUET,
-                    false,
-                    true,
-                    OptionalInt.of(1),
-                    Optional.empty());
+            return IcebergQueryRunner.builder()
+                    .setCatalogType(HADOOP)
+                    .setNodeCount(OptionalInt.of(4))
+                    .build()
+                    .getQueryRunner();
         }
         catch (Exception e) {
             e.printStackTrace();

diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/BenchmarkIcebergHiveCatalog.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/BenchmarkIcebergHiveCatalog.java
@@ -14,7 +14,6 @@
 package com.facebook.presto.iceberg;
 
 import com.facebook.presto.tests.DistributedQueryRunner;
-import com.google.common.collect.ImmutableMap;
 import org.openjdk.jmh.annotations.BenchmarkMode;
 import org.openjdk.jmh.annotations.Fork;
 import org.openjdk.jmh.annotations.Measurement;
@@ -27,11 +26,9 @@
 import org.openjdk.jmh.runner.options.OptionsBuilder;
 import org.openjdk.jmh.runner.options.VerboseMode;
 
-import java.util.Optional;
 import java.util.OptionalInt;
 
-import static com.facebook.presto.iceberg.FileFormat.PARQUET;
-import static com.facebook.presto.iceberg.IcebergQueryRunner.createIcebergQueryRunner;
+import static com.facebook.presto.iceberg.CatalogType.HIVE;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static org.openjdk.jmh.annotations.Mode.AverageTime;
 import static org.openjdk.jmh.annotations.Scope.Benchmark;
@@ -49,14 +46,10 @@ public class BenchmarkIcebergHiveCatalog
     public DistributedQueryRunner getQueryRunner()
     {
         try {
-            return createIcebergQueryRunner(
-                    ImmutableMap.of(),
-                    ImmutableMap.of(),
-                    PARQUET,
-                    false,
-                    true,
-                    OptionalInt.of(1),
-                    Optional.empty());
+            return IcebergQueryRunner.builder()
+                    .setCatalogType(HIVE)
+                    .setNodeCount(OptionalInt.of(1))
+                    .build().getQueryRunner();
         }
         catch (Exception e) {
             e.printStackTrace();

diff --git a/...to-iceberg/src/test/java/com/facebook/presto/iceberg/IcebergDistributedSmokeTestBase.java b/...to-iceberg/src/test/java/com/facebook/presto/iceberg/IcebergDistributedSmokeTestBase.java
@@ -23,7 +23,6 @@
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.testing.assertions.Assert;
 import com.facebook.presto.tests.AbstractTestIntegrationSmokeTest;
-import com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.Table;
 import org.apache.iceberg.UpdateProperties;
 import org.intellij.lang.annotations.Language;
@@ -74,7 +73,7 @@ protected IcebergDistributedSmokeTestBase(CatalogType catalogType)
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return IcebergQueryRunner.createIcebergQueryRunner(ImmutableMap.of(), catalogType);
+        return IcebergQueryRunner.builder().setCatalogType(catalogType).build().getQueryRunner();
     }
 
     @Test
@@ -123,6 +122,7 @@ public void testDescribeTable()
         MaterializedResult actualColumns = computeActual("DESCRIBE orders");
         Assert.assertEquals(actualColumns, expectedColumns);
     }
+
     @Test
     public void testShowCreateTable()
     {
@@ -771,15 +771,15 @@ private void testCreateTableLike()
 
             assertUpdate(session, "CREATE TABLE test_create_table_like_copy4 (LIKE test_create_table_like_original INCLUDING PROPERTIES) WITH (format = 'ORC')");
             assertEquals(getTablePropertiesString("test_create_table_like_copy4"), format("WITH (\n" +
-                    "   delete_mode = 'merge-on-read',\n" +
-                    "   format = 'ORC',\n" +
-                    "   format_version = '2',\n" +
-                    "   location = '%s',\n" +
-                    "   metadata_delete_after_commit = false,\n" +
-                    "   metadata_previous_versions_max = 100,\n" +
-                    "   metrics_max_inferred_column = 100,\n" +
-                    "   partitioning = ARRAY['adate']\n" +
-                    ")",
+                            "   delete_mode = 'merge-on-read',\n" +
+                            "   format = 'ORC',\n" +
+                            "   format_version = '2',\n" +
+                            "   location = '%s',\n" +
+                            "   metadata_delete_after_commit = false,\n" +
+                            "   metadata_previous_versions_max = 100,\n" +
+                            "   metrics_max_inferred_column = 100,\n" +
+                            "   partitioning = ARRAY['adate']\n" +
+                            ")",
                     getLocation(schemaName, "test_create_table_like_original")));
             dropTable(session, "test_create_table_like_copy4");
         }

diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/IcebergDistributedTestBase.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/IcebergDistributedTestBase.java
@@ -186,7 +186,10 @@ protected IcebergDistributedTestBase(CatalogType catalogType)
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return IcebergQueryRunner.createIcebergQueryRunner(ImmutableMap.of(), catalogType, extraConnectorProperties);
+        return IcebergQueryRunner.builder()
+                .setCatalogType(catalogType)
+                .setExtraConnectorProperties(extraConnectorProperties)
+                .build().getQueryRunner();
     }
 
     @Test