Skip to content

Commit

Permalink
Convert query runner to builder in ClickHouse
Browse files Browse the repository at this point in the history
  • Loading branch information
@ebyhr
ebyhr committed May 27, 2024
1 parent 483e0b2 commit f0a94b3
Show file tree
Hide file tree
Showing 9 changed files with 82 additions and 93 deletions.
6 changes: 6 additions & 0 deletions plugin/trino-clickhouse/pom.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,12 @@
<scope>provided</scope>
</dependency>

<dependency>
<groupId>com.google.errorprone</groupId>
<artifactId>error_prone_annotations</artifactId>
<scope>runtime</scope>
</dependency>

<dependency>
<groupId>io.airlift</groupId>
<artifactId>log</artifactId>
Expand Down
112 changes: 54 additions & 58 deletions plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/ClickHouseQueryRunner.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
{
Expand All @@ -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<String, String> connectorProperties,
Iterable<TpchTable<?>> tables)
throws Exception
public static final class Builder
extends DistributedQueryRunner.Builder<Builder>
{
return createClickHouseQueryRunner(
server,
ImmutableMap.of(),
connectorProperties,
tables);
}
private final Map<String, String> connectorProperties = new HashMap<>();
private List<TpchTable<?>> initialTables = ImmutableList.of();

// TODO convert to builder
private static QueryRunner createClickHouseQueryRunner(
TestingClickHouseServer server,
Map<String, String> coordinatorProperties,
Map<String, String> connectorProperties,
Iterable<TpchTable<?>> 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<TpchTable<?>> 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 ========");
Expand Down
10 changes: 4 additions & 6 deletions ...o-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestAltinityConnectorSmokeTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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();
}
}
10 changes: 4 additions & 6 deletions ...khouse/src/test/java/io/trino/plugin/clickhouse/TestAltinityLatestConnectorSmokeTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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();
}
}
11 changes: 4 additions & 7 deletions ...clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseConnectorSmokeTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
{
Expand All @@ -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();
}
}
9 changes: 4 additions & 5 deletions ...rino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseConnectorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down
10 changes: 4 additions & 6 deletions ...ouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseLatestConnectorSmokeTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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();
}
}
3 changes: 1 addition & 2 deletions ...-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseLatestTypeMapping.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -26,6 +25,6 @@ protected QueryRunner createQueryRunner()
throws Exception
{
clickhouseServer = closeAfterClass(new TestingClickHouseServer(CLICKHOUSE_LATEST_IMAGE));
return createClickHouseQueryRunner(clickhouseServer);
return ClickHouseQueryRunner.builder(clickhouseServer).build();
}
}
4 changes: 1 addition & 3 deletions .../trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseTypeMapping.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,6 @@

import io.trino.testing.QueryRunner;

import static io.trino.plugin.clickhouse.ClickHouseQueryRunner.createClickHouseQueryRunner;

public class TestClickHouseTypeMapping
extends BaseClickHouseTypeMapping
{
Expand All @@ -25,6 +23,6 @@ protected QueryRunner createQueryRunner()
throws Exception
{
clickhouseServer = closeAfterClass(new TestingClickHouseServer());
return createClickHouseQueryRunner(clickhouseServer);
return ClickHouseQueryRunner.builder(clickhouseServer).build();
}
}

0 comments on commit f0a94b3

Please sign in to comment.