Convert query runners to builder

plugin/trino-accumulo/src/test/java/io/trino/plugin/accumulo/AccumuloQueryRunner.java

@@ -22,6 +22,7 @@
 import io.trino.metadata.QualifiedObjectName;
 import io.trino.plugin.accumulo.conf.AccumuloConfig;
 import io.trino.plugin.accumulo.serializers.LexicoderRowSerializer;
+import io.trino.plugin.base.util.Closables;
 import io.trino.plugin.tpch.TpchPlugin;
 import io.trino.testing.DistributedQueryRunner;
 import io.trino.testing.QueryRunner;
@@ -54,46 +55,59 @@ public final class AccumuloQueryRunner
 
     private AccumuloQueryRunner() {}
 
-    public static synchronized QueryRunner createAccumuloQueryRunner()
-            throws Exception
+    public static Builder builder()
     {
-        return createAccumuloQueryRunner(ImmutableMap.of());
+        return new Builder();
     }
 
-    // TODO convert to builder
-    private static synchronized QueryRunner createAccumuloQueryRunner(Map<String, String> coordinatorProperties)
-            throws Exception
+    public static class Builder
+            extends DistributedQueryRunner.Builder<Builder>
     {
-        QueryRunner queryRunner = DistributedQueryRunner.builder(createSession())
-                .setCoordinatorProperties(coordinatorProperties)
-                .build();
+        protected Builder()
+        {
+            super(testSessionBuilder()
+                    .setCatalog("accumulo")
+                    .setSchema("tpch")
+                    .build());
+        }
 
-        queryRunner.installPlugin(new TpchPlugin());
-        queryRunner.createCatalog("tpch", "tpch");
-
-        TestingAccumuloServer server = TestingAccumuloServer.getInstance();
-        queryRunner.installPlugin(new AccumuloPlugin());
-        Map<String, String> accumuloProperties =
-                ImmutableMap.<String, String>builder()
-                        .put(AccumuloConfig.INSTANCE, server.getInstanceName())
-                        .put(AccumuloConfig.ZOOKEEPERS, server.getZooKeepers())
-                        .put(AccumuloConfig.USERNAME, server.getUser())
-                        .put(AccumuloConfig.PASSWORD, server.getPassword())
-                        .put(AccumuloConfig.ZOOKEEPER_METADATA_ROOT, "/trino-accumulo-test")
-                        .buildOrThrow();
-
-        queryRunner.createCatalog("accumulo", "accumulo", accumuloProperties);
-
-        if (!tpchLoaded) {
-            queryRunner.execute("CREATE SCHEMA accumulo.tpch");
-            copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, createSession(), TpchTable.getTables());
-            try (AccumuloClient client = server.createClient()) {
-                client.tableOperations().addSplits("tpch.orders", ImmutableSortedSet.of(new Text(new LexicoderRowSerializer().encode(BIGINT, 7500L))));
+        @Override
+        public DistributedQueryRunner build()
+                throws Exception
+        {
+            DistributedQueryRunner queryRunner = super.build();
+            try {
+                queryRunner.installPlugin(new TpchPlugin());
+                queryRunner.createCatalog("tpch", "tpch");
+
+                TestingAccumuloServer server = TestingAccumuloServer.getInstance();
+                queryRunner.installPlugin(new AccumuloPlugin());
+                Map<String, String> accumuloProperties =
+                        ImmutableMap.<String, String>builder()
+                                .put(AccumuloConfig.INSTANCE, server.getInstanceName())
+                                .put(AccumuloConfig.ZOOKEEPERS, server.getZooKeepers())
+                                .put(AccumuloConfig.USERNAME, server.getUser())
+                                .put(AccumuloConfig.PASSWORD, server.getPassword())
+                                .put(AccumuloConfig.ZOOKEEPER_METADATA_ROOT, "/trino-accumulo-test")
+                                .buildOrThrow();
+
+                queryRunner.createCatalog("accumulo", "accumulo", accumuloProperties);
+
+                if (!tpchLoaded) {
+                    queryRunner.execute("CREATE SCHEMA accumulo.tpch");
+                    copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, createSession(), TpchTable.getTables());
+                    try (AccumuloClient client = server.createClient()) {
+                        client.tableOperations().addSplits("tpch.orders", ImmutableSortedSet.of(new Text(new LexicoderRowSerializer().encode(BIGINT, 7500L))));
+                    }
+                    tpchLoaded = true;
+                }
             }
-            tpchLoaded = true;
+            catch (Throwable e) {
+                Closables.closeAllSuppress(e, queryRunner);
+                throw e;
+            }
+            return queryRunner;
         }
-
-        return queryRunner;
     }
 
     private static void copyTpchTables(
@@ -161,7 +175,9 @@ public static Session createSession()
     public static void main(String[] args)
             throws Exception
     {
-        QueryRunner queryRunner = createAccumuloQueryRunner(ImmutableMap.of("http-server.http.port", "8080"));
+        QueryRunner queryRunner = builder()
+                .addCoordinatorProperty("http-server.http.port", "8080")
+                .build();
         Logger log = Logger.get(AccumuloQueryRunner.class);
         log.info("======== SERVER STARTED ========");
         log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl());

plugin/trino-accumulo/src/test/java/io/trino/plugin/accumulo/TestAccumuloConnectorTest.java

@@ -25,7 +25,6 @@
 import java.util.OptionalInt;
 import java.util.regex.Pattern;
 
-import static io.trino.plugin.accumulo.AccumuloQueryRunner.createAccumuloQueryRunner;
 import static io.trino.spi.type.VarcharType.VARCHAR;
 import static io.trino.testing.MaterializedResult.resultBuilder;
 import static java.util.regex.Pattern.DOTALL;
@@ -50,7 +49,7 @@ public class TestAccumuloConnectorTest
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return createAccumuloQueryRunner();
+        return AccumuloQueryRunner.builder().build();
     }
 
     @Override

plugin/trino-blackhole/src/test/java/io/trino/plugin/blackhole/BlackHoleQueryRunner.java

@@ -15,58 +15,59 @@
 
 import com.google.common.collect.ImmutableMap;
 import io.airlift.log.Logger;
-import io.trino.Session;
+import io.trino.plugin.base.util.Closables;
 import io.trino.plugin.tpch.TpchPlugin;
 import io.trino.testing.DistributedQueryRunner;
 import io.trino.testing.QueryRunner;
 
-import java.util.Map;
-
-import static io.airlift.testing.Closeables.closeAllSuppress;
 import static io.trino.testing.TestingSession.testSessionBuilder;
 
 public final class BlackHoleQueryRunner
 {
     private BlackHoleQueryRunner() {}
 
-    public static QueryRunner createQueryRunner()
-            throws Exception
+    public static Builder builder()
     {
-        return createQueryRunner(ImmutableMap.of());
+        return new Builder();
     }
 
-    // TODO convert to builder
-    private static QueryRunner createQueryRunner(Map<String, String> coordinatorProperties)
-            throws Exception
+    public static class Builder
+            extends DistributedQueryRunner.Builder<Builder>
     {
-        Session session = testSessionBuilder()
-                .setCatalog("blackhole")
-                .setSchema("default")
-                .build();
-
-        QueryRunner queryRunner = DistributedQueryRunner.builder(session)
-                .setCoordinatorProperties(coordinatorProperties)
-                .build();
-
-        try {
-            queryRunner.installPlugin(new BlackHolePlugin());
-            queryRunner.createCatalog("blackhole", "blackhole", ImmutableMap.of());
+        protected Builder()
+        {
+            super(testSessionBuilder()
+                    .setCatalog("blackhole")
+                    .setSchema("default")
+                    .build());
+        }
 
-            queryRunner.installPlugin(new TpchPlugin());
-            queryRunner.createCatalog("tpch", "tpch", ImmutableMap.of());
+        @Override
+        public DistributedQueryRunner build()
+                throws Exception
+        {
+            DistributedQueryRunner queryRunner = super.build();
+            try {
+                queryRunner.installPlugin(new BlackHolePlugin());
+                queryRunner.createCatalog("blackhole", "blackhole", ImmutableMap.of());
 
+                queryRunner.installPlugin(new TpchPlugin());
+                queryRunner.createCatalog("tpch", "tpch", ImmutableMap.of());
+            }
+            catch (Throwable e) {
+                Closables.closeAllSuppress(e, queryRunner);
+                throw e;
+            }
             return queryRunner;
         }
-        catch (Exception e) {
-            closeAllSuppress(e, queryRunner);
-            throw e;
-        }
     }
 
     public static void main(String[] args)
             throws Exception
     {
-        QueryRunner queryRunner = createQueryRunner(ImmutableMap.of("http-server.http.port", "8080"));
+        QueryRunner queryRunner = builder()
+                .addCoordinatorProperty("http-server.http.port", "8080")
+                .build();
         Logger log = Logger.get(BlackHoleQueryRunner.class);
         log.info("======== SERVER STARTED ========");
         log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl());

plugin/trino-blackhole/src/test/java/io/trino/plugin/blackhole/TestBlackHoleSmoke.java

@@ -59,7 +59,7 @@ public class TestBlackHoleSmoke
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return BlackHoleQueryRunner.createQueryRunner();
+        return BlackHoleQueryRunner.builder().build();
     }
 
     @AfterAll

plugin/trino-cassandra/src/test/java/io/trino/plugin/cassandra/ScyllaQueryRunner.java

@@ -13,92 +13,105 @@
  */
 package io.trino.plugin.cassandra;
 
-import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableList;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import io.airlift.log.Logger;
 import io.airlift.log.Logging;
-import io.trino.Session;
+import io.trino.plugin.base.util.Closables;
 import io.trino.plugin.tpch.TpchPlugin;
 import io.trino.testing.DistributedQueryRunner;
 import io.trino.testing.QueryRunner;
 import io.trino.tpch.TpchTable;
 
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
-import static io.airlift.testing.Closeables.closeAllSuppress;
 import static io.trino.plugin.cassandra.CassandraTestingUtils.createKeyspace;
 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 ScyllaQueryRunner
 {
     private ScyllaQueryRunner() {}
 
-    public static QueryRunner createScyllaQueryRunner(
-            TestingScyllaServer server,
-            Iterable<TpchTable<?>> tables)
-            throws Exception
+    public static Builder builder(TestingScyllaServer server)
     {
-        return createScyllaQueryRunner(server, Map.of(), tables);
+        return new Builder(server)
+                .addConnectorProperty("cassandra.contact-points", server.getHost())
+                .addConnectorProperty("cassandra.native-protocol-port", Integer.toString(server.getPort()))
+                .addConnectorProperty("cassandra.allow-drop-table", "true")
+                .addConnectorProperty("cassandra.load-policy.use-dc-aware", "true")
+                .addConnectorProperty("cassandra.load-policy.dc-aware.local-dc", "datacenter1");
     }
 
-    // TODO convert to builder
-    private static QueryRunner createScyllaQueryRunner(
-            TestingScyllaServer server,
-            Map<String, String> coordinatorProperties,
-            Iterable<TpchTable<?>> tables)
-            throws Exception
+    public static class Builder
+            extends DistributedQueryRunner.Builder<Builder>
     {
-        QueryRunner queryRunner = DistributedQueryRunner.builder(createSession("tpch"))
-                .setCoordinatorProperties(coordinatorProperties)
-                .build();
+        private final TestingScyllaServer server;
+        private final Map<String, String> connectorProperties = new HashMap<>();
+        private List<TpchTable<?>> initialTables = ImmutableList.of();
+
+        private Builder(TestingScyllaServer server)
+        {
+            super(testSessionBuilder()
+                    .setCatalog("cassandra")
+                    .setSchema("tpch")
+                    .build());
+            this.server = requireNonNull(server, "server is null");
+        }
+
+        @CanIgnoreReturnValue
+        public Builder addConnectorProperty(String key, String value)
+        {
+            this.connectorProperties.put(key, value);
+            return this;
+        }
 
-        try {
-            queryRunner.installPlugin(new TpchPlugin());
-            queryRunner.createCatalog("tpch", "tpch");
+        @CanIgnoreReturnValue
+        public Builder setInitialTables(List<TpchTable<?>> initialTables)
+        {
+            this.initialTables = ImmutableList.copyOf(initialTables);
+            return this;
+        }
 
-            // note: additional copy via ImmutableList so that if fails on nulls
-            Map<String, String> connectorProperties = new HashMap<>();
-            connectorProperties.put("cassandra.contact-points", server.getHost());
-            connectorProperties.put("cassandra.native-protocol-port", Integer.toString(server.getPort()));
-            connectorProperties.put("cassandra.allow-drop-table", "true");
-            connectorProperties.put("cassandra.load-policy.use-dc-aware", "true");
-            connectorProperties.put("cassandra.load-policy.dc-aware.local-dc", "datacenter1");
+        @Override
+        public DistributedQueryRunner build()
+                throws Exception
+        {
+            DistributedQueryRunner queryRunner = super.build();
+            try {
+                queryRunner.installPlugin(new TpchPlugin());
+                queryRunner.createCatalog("tpch", "tpch");
 
-            queryRunner.installPlugin(new CassandraPlugin());
-            queryRunner.createCatalog("cassandra", "cassandra", connectorProperties);
+                queryRunner.installPlugin(new CassandraPlugin());
+                queryRunner.createCatalog("cassandra", "cassandra", connectorProperties);
 
-            createKeyspace(server.getSession(), "tpch");
-            copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, tables);
-            for (TpchTable<?> table : tables) {
-                server.refreshSizeEstimates("tpch", table.getTableName());
+                createKeyspace(server.getSession(), "tpch");
+                copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, initialTables);
+                for (TpchTable<?> table : initialTables) {
+                    server.refreshSizeEstimates("tpch", table.getTableName());
+                }
+                return queryRunner;
+            }
+            catch (Throwable e) {
+                Closables.closeAllSuppress(e, queryRunner);
+                throw e;
             }
-            return queryRunner;
-        }
-        catch (Throwable e) {
-            closeAllSuppress(e, queryRunner);
-            throw e;
         }
     }
 
-    public static Session createSession(String schema)
-    {
-        return testSessionBuilder()
-                .setCatalog("cassandra")
-                .setSchema(schema)
-                .build();
-    }
-
     public static void main(String[] args)
             throws Exception
     {
         Logging.initialize();
 
-        QueryRunner queryRunner = createScyllaQueryRunner(
-                new TestingScyllaServer(),
-                ImmutableMap.of("http-server.http.port", "8080"),
-                TpchTable.getTables());
+        QueryRunner queryRunner = builder(new TestingScyllaServer())
+                .addCoordinatorProperty("http-server.http.port", "8080")
+                .setInitialTables(TpchTable.getTables())
+                .build();
 
         Logger log = Logger.get(ScyllaQueryRunner.class);
         log.info("======== SERVER STARTED ========");