Cluster mode for clickhouse connector

vixns · vixns · commit 4a2c9cea98c0 · 2025-03-28T11:23:46.000+01:00
diff --git a/docs/src/main/sphinx/connector/clickhouse.md b/docs/src/main/sphinx/connector/clickhouse.md
@@ -70,6 +70,15 @@ driver documentation](https://clickhouse.com/docs/en/interfaces/jdbc/)
 ```{include} jdbc-authentication.fragment
 ```
 
+### Cluster mode
+
+If the clickhouse connected by `connection-url` belongs to the cluster mode, it is recommended to add this item, so that the ddl statement will be executed on all nodes of the cluster by default,
+which will be as convenient as the stand-alone clickhouse.
+
+```properties
+cluster-name=examplecluster
+```
+
 ### Multiple ClickHouse servers
 
 If you have multiple ClickHouse servers you need to configure one
diff --git a/plugin/trino-clickhouse/src/main/java/io/trino/plugin/clickhouse/ClickHouseClient.java b/plugin/trino-clickhouse/src/main/java/io/trino/plugin/clickhouse/ClickHouseClient.java
@@ -211,20 +211,23 @@ public class ClickHouseClient
 
     private final ConnectorExpressionRewriter<ParameterizedExpression> connectorExpressionRewriter;
     private final AggregateFunctionRewriter<JdbcExpression, ?> aggregateFunctionRewriter;
+    private final Optional<String> clusterName;
     private final Type uuidType;
     private final Type ipAddressType;
     private final AtomicReference<ClickHouseVersion> clickHouseVersion = new AtomicReference<>();
 
     @Inject
     public ClickHouseClient(
-            BaseJdbcConfig config,
+            ClickHouseConfig config,
+            BaseJdbcConfig baseJdbcConfig,
             ConnectionFactory connectionFactory,
             QueryBuilder queryBuilder,
             TypeManager typeManager,
             IdentifierMapping identifierMapping,
             RemoteQueryModifier queryModifier)
     {
-        super("\"", connectionFactory, queryBuilder, config.getJdbcTypesMappedToVarchar(), identifierMapping, queryModifier, false);
+        super("\"", connectionFactory, queryBuilder, baseJdbcConfig.getJdbcTypesMappedToVarchar(), identifierMapping, queryModifier, false);
+        this.clusterName = config.getClusterName();
         this.uuidType = typeManager.getType(new TypeSignature(StandardTypes.UUID));
         this.ipAddressType = typeManager.getType(new TypeSignature(StandardTypes.IPADDRESS));
         JdbcTypeHandle bigintTypeHandle = new JdbcTypeHandle(Types.BIGINT, Optional.of("bigint"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty());
@@ -253,6 +256,11 @@ public ClickHouseClient(
                         .build());
     }
 
+    private String getClusterInfo()
+    {
+        return clusterName.map(" ON CLUSTER %s "::formatted).orElse("");
+    }
+
     @Override
     public Optional<JdbcExpression> implementAggregation(ConnectorSession session, AggregateFunction aggregate, Map<String, ColumnHandle> assignments)
     {
@@ -340,8 +348,9 @@ protected void copyTableSchema(ConnectorSession session, Connection connection,
         // 1. create table tbl as tbl2
         // 2. create table tbl1 ENGINE=<engine> as select * from tbl2
         String sql = format(
-                "CREATE TABLE %s AS %s ",
+                "CREATE TABLE %s %s AS %s ",
                 quoted(null, schemaName, newTableName),
+                getClusterInfo(),
                 quoted(null, schemaName, tableName));
         try {
             execute(session, connection, sql);
@@ -396,7 +405,7 @@ protected List<String> createTableSqls(RemoteTableName remoteTableName, List<Str
         ClickHouseTableProperties.getSampleBy(tableProperties).ifPresent(value -> tableOptions.add("SAMPLE BY " + quoted(value)));
         tableMetadata.getComment().ifPresent(comment -> tableOptions.add(format("COMMENT %s", clickhouseVarcharLiteral(comment))));
 
-        return ImmutableList.of(format("CREATE TABLE %s (%s) %s", quoted(remoteTableName), join(", ", columns), join(" ", tableOptions.build())));
+        return ImmutableList.of(format("CREATE TABLE %s %s (%s) %s", quoted(remoteTableName), getClusterInfo(), join(", ", columns), join(" ", tableOptions.build())));
     }
 
     @Override
@@ -461,8 +470,9 @@ public void setTableProperties(ConnectorSession session, JdbcTableHandle handle,
 
         try (Connection connection = connectionFactory.openConnection(session)) {
             String sql = format(
-                    "ALTER TABLE %s MODIFY %s",
+                    "ALTER TABLE %s %s MODIFY %s",
                     quoted(handle.asPlainTable().getRemoteTableName()),
+                    getClusterInfo(),
                     join(" ", tableOptions.build()));
             execute(session, connection, sql);
         }
@@ -495,7 +505,7 @@ protected String getColumnDefinitionSql(ConnectorSession session, ColumnMetadata
     protected void createSchema(ConnectorSession session, Connection connection, String remoteSchemaName)
             throws SQLException
     {
-        execute(session, connection, "CREATE DATABASE " + quoted(remoteSchemaName));
+        execute(session, connection, "CREATE DATABASE " + quoted(remoteSchemaName) + getClusterInfo());
     }
 
     @Override
@@ -510,14 +520,14 @@ protected void dropSchema(ConnectorSession session, Connection connection, Strin
                 }
             }
         }
-        execute(session, connection, "DROP DATABASE " + quoted(remoteSchemaName));
+        execute(session, connection, "DROP DATABASE " + quoted(remoteSchemaName) + getClusterInfo());
     }
 
     @Override
     protected void renameSchema(ConnectorSession session, Connection connection, String remoteSchemaName, String newRemoteSchemaName)
             throws SQLException
     {
-        execute(session, connection, "RENAME DATABASE " + quoted(remoteSchemaName) + " TO " + quoted(newRemoteSchemaName));
+        execute(session, connection, "RENAME DATABASE " + quoted(remoteSchemaName) + " TO " + quoted(newRemoteSchemaName) + getClusterInfo());
     }
 
     @Override
@@ -535,8 +545,9 @@ private void addColumn(ConnectorSession session, RemoteTableName table, ColumnMe
         try (Connection connection = connectionFactory.openConnection(session)) {
             String remoteColumnName = getIdentifierMapping().toRemoteColumnName(getRemoteIdentifiers(connection), column.getName());
             String sql = format(
-                    "ALTER TABLE %s ADD COLUMN %s",
+                    "ALTER TABLE %s %s ADD COLUMN %s",
                     quoted(table),
+                    getClusterInfo(),
                     getColumnDefinitionSql(session, column, remoteColumnName));
             execute(session, connection, sql);
         }
@@ -545,12 +556,42 @@ private void addColumn(ConnectorSession session, RemoteTableName table, ColumnMe
         }
     }
 
+    @Override
+    protected void renameColumn(ConnectorSession session, Connection connection, RemoteTableName remoteTableName, String remoteColumnName, String newRemoteColumnName)
+            throws SQLException
+    {
+        execute(session, connection, format(
+                "ALTER TABLE %s RENAME COLUMN %s TO %s %s",
+                quoted(remoteTableName),
+                quoted(remoteColumnName),
+                quoted(newRemoteColumnName),
+                getClusterInfo()));
+    }
+
+    @Override
+    public void dropColumn(ConnectorSession session, JdbcTableHandle handle, JdbcColumnHandle column)
+    {
+        try (Connection connection = connectionFactory.openConnection(session)) {
+            String remoteColumnName = getIdentifierMapping().toRemoteColumnName(getRemoteIdentifiers(connection), column.getColumnName());
+            String sql = format(
+                    "ALTER TABLE %s %s DROP COLUMN %s",
+                    quoted(handle.asPlainTable().getRemoteTableName()),
+                    getClusterInfo(),
+                    getColumnDefinitionSql(session, new ColumnMetadata(column.getColumnName(), column.getColumnType()), remoteColumnName));
+            execute(session, connection, sql);
+        }
+        catch (SQLException e) {
+            throw new TrinoException(JDBC_ERROR, e);
+        }
+    }
+
     @Override
     public void setTableComment(ConnectorSession session, JdbcTableHandle handle, Optional<String> comment)
     {
         String sql = format(
-                "ALTER TABLE %s MODIFY COMMENT %s",
+                "ALTER TABLE %s %s MODIFY COMMENT %s",
                 quoted(handle.asPlainTable().getRemoteTableName()),
+                getClusterInfo(),
                 clickhouseVarcharLiteral(comment.orElse(NO_COMMENT)));
         execute(session, sql);
     }
@@ -559,8 +600,9 @@ public void setTableComment(ConnectorSession session, JdbcTableHandle handle, Op
     public void setColumnComment(ConnectorSession session, JdbcTableHandle handle, JdbcColumnHandle column, Optional<String> comment)
     {
         String sql = format(
-                "ALTER TABLE %s COMMENT COLUMN %s %s",
+                "ALTER TABLE %s %s COMMENT COLUMN %s %s",
                 quoted(handle.asPlainTable().getRemoteTableName()),
+                getClusterInfo(),
                 quoted(column.getColumnName()),
                 clickhouseVarcharLiteral(comment.orElse("")));
         execute(session, sql);
@@ -594,11 +636,19 @@ protected Optional<List<String>> getTableTypes()
     protected void renameTable(ConnectorSession session, Connection connection, String catalogName, String remoteSchemaName, String remoteTableName, String newRemoteSchemaName, String newRemoteTableName)
             throws SQLException
     {
-        execute(session, connection, format("RENAME TABLE %s TO %s",
+        execute(session, connection, format("RENAME TABLE %s %s TO %s",
                 quoted(catalogName, remoteSchemaName, remoteTableName),
+                getClusterInfo(),
                 quoted(catalogName, newRemoteSchemaName, newRemoteTableName)));
     }
 
+    @Override
+    protected void dropTable(ConnectorSession session, RemoteTableName remoteTableName, boolean temporaryTable)
+    {
+        String sql = "DROP TABLE " + quoted(remoteTableName) + getClusterInfo();
+        execute(session, sql);
+    }
+
     @Override
     protected Optional<BiFunction<String, Long, String>> limitFunction()
     {
diff --git a/plugin/trino-clickhouse/src/main/java/io/trino/plugin/clickhouse/ClickHouseConfig.java b/plugin/trino-clickhouse/src/main/java/io/trino/plugin/clickhouse/ClickHouseConfig.java
@@ -13,6 +13,9 @@
  */
 package io.trino.plugin.clickhouse;
 
+import java.util.Optional;
+import jakarta.validation.constraints.NotNull;
+
 import io.airlift.configuration.Config;
 import io.airlift.configuration.ConfigDescription;
 import io.airlift.configuration.DefunctConfig;
@@ -35,4 +38,21 @@ public ClickHouseConfig setMapStringAsVarchar(boolean mapStringAsVarchar)
         this.mapStringAsVarchar = mapStringAsVarchar;
         return this;
     }
+
+    private Optional<String> clusterName = Optional.empty();
+
+    @NotNull
+    public Optional<String> getClusterName()
+    {
+        return clusterName;
+    }
+
+    @Config("clickhouse.cluster-name")
+    @ConfigDescription("ClickHouse cluster name")
+    public ClickHouseConfig setClusterName(String clusterName)
+    {
+        this.clusterName = Optional.ofNullable(clusterName);
+        return this;
+    }
+
 }
diff --git a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseConfig.java b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHouseConfig.java
@@ -42,7 +42,8 @@ public class TestClickHouseConfig
     public void testDefaults()
     {
         assertRecordedDefaults(recordDefaults(ClickHouseConfig.class)
-                .setMapStringAsVarchar(false));
+                .setMapStringAsVarchar(false)
+                .setClusterName(null));
     }
 
     @Test
@@ -53,7 +54,8 @@ public void testExplicitPropertyMappings()
                 .buildOrThrow();
 
         ClickHouseConfig expected = new ClickHouseConfig()
-                .setMapStringAsVarchar(true);
+                .setMapStringAsVarchar(true)
+                .setClusterName("example-cluster");
 
         assertFullMapping(properties, expected);
     }
diff --git a/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHousePlugin.java b/plugin/trino-clickhouse/src/test/java/io/trino/plugin/clickhouse/TestClickHousePlugin.java
@@ -36,4 +36,19 @@ public void testCreateConnector()
                 new TestingConnectorContext())
                 .shutdown();
     }
+
+    @Test
+    public void testCreateClusterConnector()
+    {
+        Plugin plugin = new ClickHousePlugin();
+        ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories());
+        factory.create(
+                "test",
+                ImmutableMap.of(
+                        "connection-url", "jdbc:clickhouse://test",
+                        "cluster-name", "test_cluster",
+                        "bootstrap.quiet", "true"),
+                new TestingConnectorContext())
+                .shutdown();
+    }
 }
diff --git a/testing/trino-product-tests-launcher/src/main/resources/docker/trino-product-tests/conf/environment/multinode-all/clickhouse.properties b/testing/trino-product-tests-launcher/src/main/resources/docker/trino-product-tests/conf/environment/multinode-all/clickhouse.properties
@@ -2,3 +2,4 @@ connector.name=clickhouse
 connection-url=jdbc:clickhouse://host1.invalid:8123/
 connection-user=exampleuser
 connection-password=examplepassword
+cluster-name=test_cluster
diff --git a/testing/trino-product-tests-launcher/src/main/resources/docker/trino-product-tests/conf/environment/multinode-clickhouse/clickhouse.properties b/testing/trino-product-tests-launcher/src/main/resources/docker/trino-product-tests/conf/environment/multinode-clickhouse/clickhouse.properties
@@ -2,3 +2,4 @@ connector.name=clickhouse
 connection-url=jdbc:clickhouse://clickhouse:8124/
 connection-user=test
 connection-password=test
+cluster-name=test_cluster

Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,8 @@ public class TestClickHouseConfig`
`42`	`42`	`public void testDefaults()`
`43`	`43`	`{`
`44`	`44`	`assertRecordedDefaults(recordDefaults(ClickHouseConfig.class)`
`45`		`- .setMapStringAsVarchar(false));`
	`45`	`+ .setMapStringAsVarchar(false)`
	`46`	`+ .setClusterName(null));`
`46`	`47`	`}`
`47`	`48`
`48`	`49`	`@Test`
`@@ -53,7 +54,8 @@ public void testExplicitPropertyMappings()`
`53`	`54`	`.buildOrThrow();`
`54`	`55`
`55`	`56`	`ClickHouseConfig expected = new ClickHouseConfig()`
`56`		`- .setMapStringAsVarchar(true);`
	`57`	`+ .setMapStringAsVarchar(true)`
	`58`	`+ .setClusterName("example-cluster");`
`57`	`59`
`58`	`60`	`assertFullMapping(properties, expected);`
`59`	`61`	`}`