Add builder to MemoryQueryRunner

trinodb · Mar 2, 2022 · 9d67738 · 9d67738
1 parent a0d8caa
commit 9d67738
Showing 1 changed file with 52 additions and 16 deletions.
diff --git a/plugin/trino-memory/src/test/java/io/trino/plugin/memory/MemoryQueryRunner.java b/plugin/trino-memory/src/test/java/io/trino/plugin/memory/MemoryQueryRunner.java
@@ -13,6 +13,7 @@
  */
 package io.trino.plugin.memory;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import io.airlift.log.Logger;
 import io.airlift.log.Logging;
@@ -21,12 +22,14 @@
 import io.trino.testing.DistributedQueryRunner;
 import io.trino.tpch.TpchTable;
 
+import java.util.List;
 import java.util.Map;
 
 import static io.airlift.testing.Closeables.closeAllSuppress;
 import static io.trino.plugin.tpch.TpchMetadata.TINY_SCHEMA_NAME;
 import static io.trino.testing.QueryAssertions.copyTpchTables;
 import static io.trino.testing.TestingSession.testSessionBuilder;
+import static java.util.Objects.requireNonNull;
 
 public final class MemoryQueryRunner
 {
@@ -39,29 +42,62 @@ public static DistributedQueryRunner createMemoryQueryRunner(
             Iterable<TpchTable<?>> tables)
             throws Exception
     {
-        Session session = testSessionBuilder()
-                .setCatalog(CATALOG)
-                .setSchema("default")
-                .build();
-
-        DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(session)
+        return builder()
                 .setExtraProperties(extraProperties)
+                .setInitialTables(tables)
                 .build();
+    }
+
+    public static Builder builder()
+    {
+        return new Builder();
+    }
+
+    public static class Builder
+            extends DistributedQueryRunner.Builder<Builder>
+    {
+        private List<TpchTable<?>> initialTables = ImmutableList.of();
+
+        protected Builder()
+        {
+            super(createSession());
+        }
 
-        try {
-            queryRunner.installPlugin(new MemoryPlugin());
-            queryRunner.createCatalog(CATALOG, "memory", ImmutableMap.of());
+        public Builder setInitialTables(Iterable<TpchTable<?>> initialTables)
+        {
+            this.initialTables = ImmutableList.copyOf(requireNonNull(initialTables, "initialTables is null"));
+            return self();
+        }
 
-            queryRunner.installPlugin(new TpchPlugin());
-            queryRunner.createCatalog("tpch", "tpch", ImmutableMap.of());
+        @Override
+        public DistributedQueryRunner build()
+                throws Exception
+        {
+            DistributedQueryRunner queryRunner = super.build();
 
-            copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, session, tables);
+            try {
+                queryRunner.installPlugin(new MemoryPlugin());
+                queryRunner.createCatalog(CATALOG, "memory", ImmutableMap.of());
 
-            return queryRunner;
+                queryRunner.installPlugin(new TpchPlugin());
+                queryRunner.createCatalog("tpch", "tpch", ImmutableMap.of());
+
+                copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, createSession(), initialTables);
+
+                return queryRunner;
+            }
+            catch (Exception e) {
+                closeAllSuppress(e, queryRunner);
+                throw e;
+            }
         }
-        catch (Exception e) {
-            closeAllSuppress(e, queryRunner);
-            throw e;
+
+        private static Session createSession()
+        {
+            return testSessionBuilder()
+                    .setCatalog(CATALOG)
+                    .setSchema("default")
+                    .build();
         }
     }