From f0a94b33b9cca953af130c4d8f50f5315471c6e0 Mon Sep 17 00:00:00 2001 From: Yuya Ebihara Date: Mon, 27 May 2024 17:24:15 +0900 Subject: [PATCH] Convert query runner to builder in ClickHouse --- plugin/trino-clickhouse/pom.xml | 6 + .../clickhouse/ClickHouseQueryRunner.java | 112 +++++++++--------- .../TestAltinityConnectorSmokeTest.java | 10 +- .../TestAltinityLatestConnectorSmokeTest.java | 10 +- .../TestClickHouseConnectorSmokeTest.java | 11 +- .../TestClickHouseConnectorTest.java | 9 +- ...estClickHouseLatestConnectorSmokeTest.java | 10 +- .../TestClickHouseLatestTypeMapping.java | 3 +- .../clickhouse/TestClickHouseTypeMapping.java | 4 +- 9 files changed, 82 insertions(+), 93 deletions(-) diff --git a/plugin/trino-clickhouse/pom.xml b/plugin/trino-clickhouse/pom.xml index d72ad5e67d82..129e45239b72 100644 --- a/plugin/trino-clickhouse/pom.xml +++ b/plugin/trino-clickhouse/pom.xml @@ -90,6 +90,12 @@ provided + + com.google.errorprone + error_prone_annotations + runtime + + io.airlift log diff --git a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/ClickHouseQueryRunner.java b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/ClickHouseQueryRunner.java index d74b3633266e..33aecc1485a2 100644 --- a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/ClickHouseQueryRunner.java +++ b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/ClickHouseQueryRunner.java @@ -14,23 +14,24 @@ package io.trino.plugin.clickhouse; import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; +import com.google.errorprone.annotations.CanIgnoreReturnValue; import io.airlift.log.Level; import io.airlift.log.Logger; import io.airlift.log.Logging; -import io.trino.Session; 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.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 ClickHouseQueryRunner { @@ -43,75 +44,70 @@ public final class ClickHouseQueryRunner private ClickHouseQueryRunner() {} - public static QueryRunner createClickHouseQueryRunner(TestingClickHouseServer server, TpchTable... tables) - throws Exception + public static Builder builder(TestingClickHouseServer server) { - return createClickHouseQueryRunner(server, ImmutableMap.of(), ImmutableList.copyOf(tables)); + return new Builder() + .addConnectorProperty("connection-url", server.getJdbcUrl()); } - // TODO convert to builder - public static QueryRunner createClickHouseQueryRunner( - TestingClickHouseServer server, - Map connectorProperties, - Iterable> tables) - throws Exception + public static final class Builder + extends DistributedQueryRunner.Builder { - return createClickHouseQueryRunner( - server, - ImmutableMap.of(), - connectorProperties, - tables); - } + private final Map connectorProperties = new HashMap<>(); + private List> initialTables = ImmutableList.of(); - // TODO convert to builder - private static QueryRunner createClickHouseQueryRunner( - TestingClickHouseServer server, - Map coordinatorProperties, - Map connectorProperties, - Iterable> tables) - throws Exception - { - QueryRunner queryRunner = null; - try { - queryRunner = DistributedQueryRunner.builder(createSession()) - .setCoordinatorProperties(coordinatorProperties) - .build(); - - queryRunner.installPlugin(new TpchPlugin()); - queryRunner.createCatalog("tpch", "tpch"); - - // note: additional copy via ImmutableList so that if fails on nulls - connectorProperties = new HashMap<>(ImmutableMap.copyOf(connectorProperties)); - connectorProperties.putIfAbsent("connection-url", server.getJdbcUrl()); - - queryRunner.installPlugin(new ClickHousePlugin()); - queryRunner.createCatalog("clickhouse", "clickhouse", connectorProperties); - queryRunner.execute("CREATE SCHEMA " + TPCH_SCHEMA); - copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, tables); - return queryRunner; + private Builder() + { + super(testSessionBuilder() + .setCatalog("clickhouse") + .setSchema(TPCH_SCHEMA) + .build()); } - catch (Throwable e) { - closeAllSuppress(e, queryRunner); - throw e; + + @CanIgnoreReturnValue + public Builder addConnectorProperty(String key, String value) + { + this.connectorProperties.put(key, value); + return this; } - } - public static Session createSession() - { - return testSessionBuilder() - .setCatalog("clickhouse") - .setSchema(TPCH_SCHEMA) - .build(); + @CanIgnoreReturnValue + public Builder setInitialTables(Iterable> initialTables) + { + this.initialTables = ImmutableList.copyOf(requireNonNull(initialTables, "initialTables is null")); + return this; + } + + @Override + public DistributedQueryRunner build() + throws Exception + { + DistributedQueryRunner queryRunner = super.build(); + try { + queryRunner.installPlugin(new TpchPlugin()); + queryRunner.createCatalog("tpch", "tpch"); + + queryRunner.installPlugin(new ClickHousePlugin()); + queryRunner.createCatalog("clickhouse", "clickhouse", connectorProperties); + queryRunner.execute("CREATE SCHEMA " + TPCH_SCHEMA); + copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, initialTables); + + return queryRunner; + } + catch (Throwable e) { + closeAllSuppress(e, queryRunner); + throw e; + } + } } public static void main(String[] args) throws Exception { - QueryRunner queryRunner = createClickHouseQueryRunner( - new TestingClickHouseServer(), - ImmutableMap.of("http-server.http.port", "8080"), - ImmutableMap.of(), - TpchTable.getTables()); + QueryRunner queryRunner = builder(new TestingClickHouseServer()) + .addCoordinatorProperty("http-server.http.port", "8080") + .setInitialTables(TpchTable.getTables()) + .build(); Logger log = Logger.get(ClickHouseQueryRunner.class); log.info("======== SERVER STARTED ========"); diff --git a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestAltinityConnectorSmokeTest.java b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestAltinityConnectorSmokeTest.java index 1b04438f8eb9..7d18ca8c45db 100644 --- a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestAltinityConnectorSmokeTest.java +++ b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestAltinityConnectorSmokeTest.java @@ -13,10 +13,8 @@ */ package io.trino.plugin.clickhouse; -import com.google.common.collect.ImmutableMap; import io.trino.testing.QueryRunner; -import static io.trino.plugin.clickhouse.ClickHouseQueryRunner.createClickHouseQueryRunner; import static io.trino.plugin.clickhouse.TestingClickHouseServer.ALTINITY_DEFAULT_IMAGE; public class TestAltinityConnectorSmokeTest @@ -26,9 +24,9 @@ public class TestAltinityConnectorSmokeTest protected QueryRunner createQueryRunner() throws Exception { - return createClickHouseQueryRunner( - closeAfterClass(new TestingClickHouseServer(ALTINITY_DEFAULT_IMAGE)), - ImmutableMap.of("clickhouse.map-string-as-varchar", "true"), // To handle string types in TPCH tables as varchar instead of varbinary - REQUIRED_TPCH_TABLES); + return ClickHouseQueryRunner.builder(closeAfterClass(new TestingClickHouseServer(ALTINITY_DEFAULT_IMAGE))) + .addConnectorProperty("clickhouse.map-string-as-varchar", "true") // To handle string types in TPCH tables as varchar instead of varbinary + .setInitialTables(REQUIRED_TPCH_TABLES) + .build(); } } diff --git a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestAltinityLatestConnectorSmokeTest.java b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestAltinityLatestConnectorSmokeTest.java index e0581ec9ff84..e00eba74d546 100644 --- a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestAltinityLatestConnectorSmokeTest.java +++ b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestAltinityLatestConnectorSmokeTest.java @@ -13,10 +13,8 @@ */ package io.trino.plugin.clickhouse; -import com.google.common.collect.ImmutableMap; import io.trino.testing.QueryRunner; -import static io.trino.plugin.clickhouse.ClickHouseQueryRunner.createClickHouseQueryRunner; import static io.trino.plugin.clickhouse.TestingClickHouseServer.ALTINITY_LATEST_IMAGE; public class TestAltinityLatestConnectorSmokeTest @@ -26,9 +24,9 @@ public class TestAltinityLatestConnectorSmokeTest protected QueryRunner createQueryRunner() throws Exception { - return createClickHouseQueryRunner( - closeAfterClass(new TestingClickHouseServer(ALTINITY_LATEST_IMAGE)), - ImmutableMap.of("clickhouse.map-string-as-varchar", "true"), // To handle string types in TPCH tables as varchar instead of varbinary - REQUIRED_TPCH_TABLES); + return ClickHouseQueryRunner.builder(closeAfterClass(new TestingClickHouseServer(ALTINITY_LATEST_IMAGE))) + .addConnectorProperty("clickhouse.map-string-as-varchar", "true") // To handle string types in TPCH tables as varchar instead of varbinary + .setInitialTables(REQUIRED_TPCH_TABLES) + .build(); } } diff --git a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseConnectorSmokeTest.java b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseConnectorSmokeTest.java index bc7844b82789..de47922b8e84 100644 --- a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseConnectorSmokeTest.java +++ b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseConnectorSmokeTest.java @@ -13,11 +13,8 @@ */ package io.trino.plugin.clickhouse; -import com.google.common.collect.ImmutableMap; import io.trino.testing.QueryRunner; -import static io.trino.plugin.clickhouse.ClickHouseQueryRunner.createClickHouseQueryRunner; - public class TestClickHouseConnectorSmokeTest extends BaseClickHouseConnectorSmokeTest { @@ -28,9 +25,9 @@ protected QueryRunner createQueryRunner() throws Exception { clickHouseServer = closeAfterClass(new TestingClickHouseServer()); - return createClickHouseQueryRunner( - clickHouseServer, - ImmutableMap.of("clickhouse.map-string-as-varchar", "true"), // To handle string types in TPCH tables as varchar instead of varbinary - REQUIRED_TPCH_TABLES); + return ClickHouseQueryRunner.builder(clickHouseServer) + .addConnectorProperty("clickhouse.map-string-as-varchar", "true") // To handle string types in TPCH tables as varchar instead of varbinary + .setInitialTables(REQUIRED_TPCH_TABLES) + .build(); } } diff --git a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseConnectorTest.java b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseConnectorTest.java index 824029baeb47..e228eeb2bedf 100644 --- a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseConnectorTest.java +++ b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseConnectorTest.java @@ -39,7 +39,6 @@ import java.util.Optional; import java.util.OptionalInt; -import static io.trino.plugin.clickhouse.ClickHouseQueryRunner.createClickHouseQueryRunner; import static io.trino.plugin.clickhouse.ClickHouseTableProperties.ENGINE_PROPERTY; import static io.trino.plugin.clickhouse.ClickHouseTableProperties.ORDER_BY_PROPERTY; import static io.trino.plugin.clickhouse.ClickHouseTableProperties.PARTITION_BY_PROPERTY; @@ -93,10 +92,10 @@ protected QueryRunner createQueryRunner() throws Exception { this.clickhouseServer = closeAfterClass(new TestingClickHouseServer(CLICKHOUSE_LATEST_IMAGE)); - return createClickHouseQueryRunner( - clickhouseServer, - ImmutableMap.of("clickhouse.map-string-as-varchar", "true"), - REQUIRED_TPCH_TABLES); + return ClickHouseQueryRunner.builder(clickhouseServer) + .addConnectorProperty("clickhouse.map-string-as-varchar", "true") + .setInitialTables(REQUIRED_TPCH_TABLES) + .build(); } @Test diff --git a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseLatestConnectorSmokeTest.java b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseLatestConnectorSmokeTest.java index 995615e93782..321d9eec0278 100644 --- a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseLatestConnectorSmokeTest.java +++ b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseLatestConnectorSmokeTest.java @@ -13,10 +13,8 @@ */ package io.trino.plugin.clickhouse; -import com.google.common.collect.ImmutableMap; import io.trino.testing.QueryRunner; -import static io.trino.plugin.clickhouse.ClickHouseQueryRunner.createClickHouseQueryRunner; import static io.trino.plugin.clickhouse.TestingClickHouseServer.CLICKHOUSE_LATEST_IMAGE; public class TestClickHouseLatestConnectorSmokeTest @@ -26,9 +24,9 @@ public class TestClickHouseLatestConnectorSmokeTest protected QueryRunner createQueryRunner() throws Exception { - return createClickHouseQueryRunner( - closeAfterClass(new TestingClickHouseServer(CLICKHOUSE_LATEST_IMAGE)), - ImmutableMap.of("clickhouse.map-string-as-varchar", "true"), // To handle string types in TPCH tables as varchar instead of varbinary - REQUIRED_TPCH_TABLES); + return ClickHouseQueryRunner.builder(closeAfterClass(new TestingClickHouseServer(CLICKHOUSE_LATEST_IMAGE))) + .addConnectorProperty("clickhouse.map-string-as-varchar", "true") // To handle string types in TPCH tables as varchar instead of varbinary + .setInitialTables(REQUIRED_TPCH_TABLES) + .build(); } } diff --git a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseLatestTypeMapping.java b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseLatestTypeMapping.java index 80556e674c53..c1fcd2dc625f 100644 --- a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseLatestTypeMapping.java +++ b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseLatestTypeMapping.java @@ -15,7 +15,6 @@ import io.trino.testing.QueryRunner; -import static io.trino.plugin.clickhouse.ClickHouseQueryRunner.createClickHouseQueryRunner; import static io.trino.plugin.clickhouse.TestingClickHouseServer.CLICKHOUSE_LATEST_IMAGE; public class TestClickHouseLatestTypeMapping @@ -26,6 +25,6 @@ protected QueryRunner createQueryRunner() throws Exception { clickhouseServer = closeAfterClass(new TestingClickHouseServer(CLICKHOUSE_LATEST_IMAGE)); - return createClickHouseQueryRunner(clickhouseServer); + return ClickHouseQueryRunner.builder(clickhouseServer).build(); } } diff --git a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseTypeMapping.java b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseTypeMapping.java index 3e42c49edb0d..13255d9f9fd1 100644 --- a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseTypeMapping.java +++ b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseTypeMapping.java @@ -15,8 +15,6 @@ import io.trino.testing.QueryRunner; -import static io.trino.plugin.clickhouse.ClickHouseQueryRunner.createClickHouseQueryRunner; - public class TestClickHouseTypeMapping extends BaseClickHouseTypeMapping { @@ -25,6 +23,6 @@ protected QueryRunner createQueryRunner() throws Exception { clickhouseServer = closeAfterClass(new TestingClickHouseServer()); - return createClickHouseQueryRunner(clickhouseServer); + return ClickHouseQueryRunner.builder(clickhouseServer).build(); } }