You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by yh...@apache.org on 2025/03/10 20:40:03 UTC
(beam) branch master updated: Implement ClickHouse Integration Testing Resource Manager (#34230)
This is an automated email from the ASF dual-hosted git repository.
yhu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new d89ccc96ca3 Implement ClickHouse Integration Testing Resource Manager (#34230)
d89ccc96ca3 is described below
commit d89ccc96ca39d5f72120cc8ff38354355b8dc966
Author: Bentsi Leviav <be...@clickhouse.com>
AuthorDate: Mon Mar 10 22:39:57 2025 +0200
Implement ClickHouse Integration Testing Resource Manager (#34230)
* implement a clickhouse resource manager
* create integration test for ClickHouseResourceManager
* create unit tests for ClickHouseResourceManager
* fix warnings
* address cr comments
* fix mixed junit4 and 5 imports
* arrange import
---
it/clickhouse/build.gradle | 53 ++++
.../it/clickhouse/ClickHouseResourceManager.java | 353 +++++++++++++++++++++
.../ClickHouseResourceManagerException.java | 25 ++
.../apache/beam/it/clickhouse/ClickHouseUtils.java | 58 ++++
.../apache/beam/it/clickhouse/package-info.java | 20 ++
.../it/clickhouse/ClickHouseResourceManagerIT.java | 85 +++++
.../clickhouse/ClickHouseResourceManagerTest.java | 186 +++++++++++
.../beam/it/clickhouse/ClickHouseUtilsTest.java | 96 ++++++
settings.gradle.kts | 4 +-
9 files changed, 879 insertions(+), 1 deletion(-)
diff --git a/it/clickhouse/build.gradle b/it/clickhouse/build.gradle
new file mode 100644
index 00000000000..063ae767a3e
--- /dev/null
+++ b/it/clickhouse/build.gradle
@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * License); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an AS IS BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+plugins {
+ id 'org.apache.beam.module'
+}
+applyJavaNature(
+ automaticModuleName: 'org.apache.beam.it.clickhouse',
+ exportJavadoc: false,
+)
+
+description = "Apache Beam :: IT :: ClickHouse"
+ext.summary = "Integration test utilities for ClickHouse."
+
+def clickhouse_jdbc_version = "0.6.4"
+
+dependencies {
+ implementation project(path: ":it:common")
+ implementation project(path: ":it:testcontainers")
+ testImplementation platform('org.junit:junit-bom:5.10.0')
+ implementation library.java.slf4j_api
+ testImplementation 'org.junit.jupiter:junit-jupiter'
+ implementation library.java.testcontainers_base
+ implementation library.java.testcontainers_clickhouse
+ implementation "com.clickhouse:clickhouse-jdbc:$clickhouse_jdbc_version:all"
+
+ testImplementation(library.java.truth)
+ testImplementation 'junit:junit:4.13.2'
+
+ testImplementation library.java.guava
+ permitUsedUndeclared library.java.guava
+ testImplementation library.java.mockito_inline
+ testImplementation library.java.commons_lang3
+ testRuntimeOnly library.java.slf4j_simple
+
+ permitUnusedDeclared "com.clickhouse:clickhouse-jdbc:$clickhouse_jdbc_version:all"
+
+}
\ No newline at end of file
diff --git a/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/ClickHouseResourceManager.java b/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/ClickHouseResourceManager.java
new file mode 100644
index 00000000000..200fb82905e
--- /dev/null
+++ b/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/ClickHouseResourceManager.java
@@ -0,0 +1,353 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.it.clickhouse;
+
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import org.apache.beam.it.common.ResourceManager;
+import org.apache.beam.it.testcontainers.TestContainerResourceManager;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testcontainers.clickhouse.ClickHouseContainer;
+import org.testcontainers.containers.GenericContainer;
+import org.testcontainers.containers.wait.strategy.LogMessageWaitStrategy;
+import org.testcontainers.utility.DockerImageName;
+
+/**
+ * Client for managing ClickHouse resources.
+ *
+ * <p>The class is thread-safe.
+ */
+public class ClickHouseResourceManager extends TestContainerResourceManager<GenericContainer<?>>
+ implements ResourceManager {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ClickHouseResourceManager.class);
+
+ private final Connection connection;
+
+ // 8123 is the default port that ClickHouse is configured to listen on
+ private static final int CLICKHOUSE_INTERNAL_PORT = 8123;
+
+ private static final String DEFAULT_CLICKHOUSE_CONTAINER_NAME = "clickhouse/clickhouse-server";
+
+ private static final String DEFAULT_CLICKHOUSE_CONTAINER_TAG = "23.8";
+
+ private final String jdbcConnectionString;
+
+ final List<String> managedTableNames = new ArrayList<>();
+
+ ClickHouseResourceManager(Builder builder) throws SQLException {
+ this(/* clickHouseConnection*/ null, buildContainer(builder), builder);
+ }
+
+ private static ClickHouseContainer buildContainer(Builder builder) {
+ ClickHouseContainer container =
+ new ClickHouseContainer(
+ DockerImageName.parse(builder.containerImageName)
+ .withTag(builder.containerImageTag))
+ .withStartupAttempts(10);
+
+ Duration startupTimeout = Duration.ofMinutes(2);
+ container.setWaitStrategy(new LogMessageWaitStrategy().withStartupTimeout(startupTimeout));
+
+ return container;
+ }
+
+ // @VisibleForTesting
+ @SuppressWarnings("nullness")
+ ClickHouseResourceManager(
+ @Nullable Connection clickHosueConnection, ClickHouseContainer container, Builder builder)
+ throws SQLException {
+ super(container, builder);
+
+ this.jdbcConnectionString =
+ "http://" + this.getHost() + ":" + this.getPort(CLICKHOUSE_INTERNAL_PORT);
+
+ this.connection =
+ clickHosueConnection != null ? clickHosueConnection : container.createConnection("");
+ }
+
+ /** Builder for {@link ClickHouseResourceManager}. */
+ public static final class Builder
+ extends TestContainerResourceManager.Builder<ClickHouseResourceManager> {
+
+ Builder(String testId) {
+ super(testId, DEFAULT_CLICKHOUSE_CONTAINER_NAME, DEFAULT_CLICKHOUSE_CONTAINER_TAG);
+ }
+
+ @Override
+ public ClickHouseResourceManager build() {
+ try {
+ return new ClickHouseResourceManager(this);
+ } catch (SQLException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ }
+
+ public static Builder builder(String testId) {
+ return new Builder(testId);
+ }
+
+ /** Returns the JDBC connection string to the ClickHouse service. */
+ public synchronized String getUri() {
+ return jdbcConnectionString;
+ }
+
+ synchronized boolean tableExists(String tableName) throws SQLException {
+
+ ClickHouseUtils.checkValidTableName(tableName);
+ String query = String.format("EXISTS TABLE %s", tableName);
+
+ try (Statement statement = connection.createStatement();
+ ResultSet res = statement.executeQuery(query)) {
+
+ if (res.next()) {
+ return res.getInt(1) == 1; // ClickHouse returns 1 if the table exists, 0 otherwise
+ }
+
+ return false;
+ }
+ }
+
+ /**
+ * Creates a table on ClickHouse.
+ *
+ * @param tableName Name of the table to create.
+ * @return A boolean indicating whether the resource was created.
+ * @throws ClickHouseResourceManagerException if there is an error creating the table in
+ * ClickHouse.
+ */
+ public synchronized boolean createTable(
+ String tableName, Map<String, String> columns, String engine, List<String> orderBy)
+ throws ClickHouseResourceManagerException {
+
+ LOG.info("Creating table '{}' in ClickHouse.", tableName);
+
+ try {
+
+ ClickHouseUtils.checkValidTableName(tableName);
+
+ if (tableExists(tableName)) {
+ LOG.warn("Table '{}' already exists.", tableName);
+ return false;
+ }
+
+ if (columns == null || columns.isEmpty()) {
+ throw new ClickHouseResourceManagerException(
+ "Cannot create table with no columns.",
+ new Exception("Cannot create table with no columns."));
+ }
+
+ if (engine == null || engine.isEmpty()) {
+ engine = "MergeTree()";
+ }
+
+ if (orderBy == null || orderBy.isEmpty()) {
+ orderBy = Collections.singletonList(columns.keySet().iterator().next()); // First column
+ }
+
+ String columnDefinitions =
+ columns.entrySet().stream()
+ .map(entry -> entry.getKey() + " " + entry.getValue())
+ .collect(Collectors.joining(", "));
+
+ // Construct CREATE TABLE query
+ String createTableQuery =
+ String.format(
+ "CREATE TABLE IF NOT EXISTS %s (%s) ENGINE = %s ORDER BY (%s)",
+ tableName, columnDefinitions, engine, String.join(", ", orderBy));
+
+ LOG.info("Executing query: {}", createTableQuery);
+
+ try (Statement statement = connection.createStatement()) {
+ statement.execute(createTableQuery);
+ return tableExists(tableName);
+ }
+
+ } catch (Exception e) {
+ throw new ClickHouseResourceManagerException("Error creating table: " + tableName, e);
+ }
+ }
+
+ /**
+ * Inserts the given rows into a ClickHouse table.
+ *
+ * <p>Note: Implementations may create the table if it does not already exist.
+ *
+ * @param tableName The name of the table to insert the rows into.
+ * @param rows A list of rows, where each row is a map of (column name -> value).
+ * @return A boolean indicating whether the rows were inserted successfully.
+ * @throws ClickHouseResourceManagerException if there is an error inserting the rows.
+ */
+ public synchronized boolean insertRows(String tableName, List<Map<String, Object>> rows)
+ throws ClickHouseResourceManagerException {
+ LOG.info("Attempting to write {} rows to {}.", rows.size(), tableName);
+
+ if (rows.isEmpty()) {
+ LOG.warn("No rows provided for insertion.");
+ return false;
+ }
+
+ try {
+ // Extract column names from the first row
+ List<String> columns = new ArrayList<>(rows.get(0).keySet());
+
+ String columnNames = String.join(", ", columns);
+ String placeholders = String.join(", ", Collections.nCopies(columns.size(), "?"));
+ String insertQuery =
+ String.format("INSERT INTO %s (%s) VALUES (%s)", tableName, columnNames, placeholders);
+
+ try (PreparedStatement preparedStatement = connection.prepareStatement(insertQuery)) {
+ // Add batch inserts
+ for (Map<String, Object> row : rows) {
+ int index = 1;
+ for (String column : columns) {
+ preparedStatement.setObject(index++, row.get(column));
+ }
+ preparedStatement.addBatch();
+ }
+
+ // Execute batch insert
+ int[] result = preparedStatement.executeBatch();
+
+ boolean allSuccessful =
+ Arrays.stream(result)
+ .allMatch(resRows -> resRows > 0 || resRows == Statement.SUCCESS_NO_INFO);
+ if (!allSuccessful) {
+ LOG.warn("Some rows might have failed to insert into {}", tableName);
+ return false;
+ }
+
+ int affectedRows = Arrays.stream(result).sum();
+ LOG.info("Successfully wrote {} rows to {}", affectedRows, tableName);
+ return true;
+ }
+ } catch (Exception e) {
+ throw new ClickHouseResourceManagerException("Error inserting rows into ClickHouse. ", e);
+ }
+ }
+
+ /**
+ * Reads all the rows in a table.
+ *
+ * @param tableName The name of the table to read from.
+ * @return An iterable of all the Rows in the collection.
+ * @throws ClickHouseResourceManagerException if there is an error reading the data.
+ */
+ public synchronized List<Map<String, Object>> fetchAll(String tableName)
+ throws ClickHouseResourceManagerException {
+ LOG.info("Reading all rows from {}.", tableName);
+
+ String query = String.format("SELECT * FROM %s", tableName);
+
+ try (Statement statement = connection.createStatement();
+ ResultSet resultSet = statement.executeQuery(query)) {
+
+ List<Map<String, Object>> rows = new ArrayList<>();
+ ResultSetMetaData metaData = resultSet.getMetaData();
+ int columnCount = metaData.getColumnCount();
+
+ while (resultSet.next()) {
+ Map<String, Object> row = new HashMap<>();
+ for (int i = 1; i <= columnCount; i++) {
+ Object value = resultSet.getObject(i);
+ row.put(metaData.getColumnName(i), value != null ? value : "NULL");
+ }
+ rows.add(row);
+ }
+
+ return rows;
+
+ } catch (SQLException e) {
+ LOG.error("Error during reading table {}.", tableName, e);
+ throw new ClickHouseResourceManagerException(
+ "Error during reading table " + tableName + ".", e);
+ }
+ }
+
+ /**
+ * Gets the count of rows in a table.
+ *
+ * @param tableName The name of the table to read from.
+ * @return The number of rows for the given table
+ * @throws ClickHouseResourceManagerException if there is an error reading the data.
+ */
+ public long count(String tableName) throws ClickHouseResourceManagerException {
+ LOG.info("Fetching count from {}.", tableName);
+
+ String query = String.format("SELECT count(*) FROM %s", tableName);
+
+ try (Statement statement = connection.createStatement();
+ ResultSet countRes = statement.executeQuery(query)) {
+ if (countRes.next()) {
+ return countRes.getLong(1);
+ } else {
+ throw new ClickHouseResourceManagerException(
+ "No result returned for count query.", new Error());
+ }
+
+ } catch (Exception e) {
+ throw new ClickHouseResourceManagerException(
+ "Error fetching count from " + tableName + ".", e);
+ }
+ }
+
+ /**
+ * Deletes all created resources and cleans up the ClickHouse client, making the manager object
+ * unusable.
+ *
+ * @throws ClickHouseResourceManagerException if there is an error deleting the ClickHouse
+ * resources.
+ */
+ @Override
+ public synchronized void cleanupAll() throws ClickHouseResourceManagerException {
+ LOG.info("Attempting to clean up ClickHouse manager.");
+
+ if (!managedTableNames.isEmpty()) {
+ try (Statement statement = connection.createStatement()) {
+ for (String tableName : managedTableNames) {
+ String dropTableQuery = String.format("DROP TABLE IF EXISTS %s", tableName);
+ statement.execute(dropTableQuery);
+ LOG.info("Dropped table: {}", tableName);
+ }
+ managedTableNames.clear();
+ } catch (Exception e) {
+ LOG.error("Failed to delete ClickHouse tables: {}", managedTableNames, e);
+ throw new ClickHouseResourceManagerException("Failed deleting ClickHouse tables", e);
+ }
+ }
+
+ super.cleanupAll();
+
+ LOG.info("ClickHouse manager successfully cleaned up.");
+ }
+}
diff --git a/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/ClickHouseResourceManagerException.java b/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/ClickHouseResourceManagerException.java
new file mode 100644
index 00000000000..7a7faf4a7b5
--- /dev/null
+++ b/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/ClickHouseResourceManagerException.java
@@ -0,0 +1,25 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.it.clickhouse;
+
+/** Custom exception for {@link ClickHouseResourceManager} implementations. */
+public class ClickHouseResourceManagerException extends RuntimeException {
+ public ClickHouseResourceManagerException(String message, Throwable err) {
+ super(message, err);
+ }
+}
diff --git a/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/ClickHouseUtils.java b/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/ClickHouseUtils.java
new file mode 100644
index 00000000000..03a6d4c7a53
--- /dev/null
+++ b/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/ClickHouseUtils.java
@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.it.clickhouse;
+
+import java.util.regex.Pattern;
+
+/** Utilities for {@link ClickHouseResourceManager} implementations. */
+public class ClickHouseUtils {
+
+ // ClickHouse max table name length
+ private static final int MAX_TABLE_NAME_LENGTH = 206;
+
+ // ClickHouse table name restrictions:
+ // - Must not contain '.', '/', '\\', ' ', ',', '`'
+ // - Must start with a letter or an underscore
+ private static final Pattern VALID_TABLE_NAME_PATTERN =
+ Pattern.compile("^[a-zA-Z_][a-zA-Z0-9_]*$");
+
+ private ClickHouseUtils() {}
+
+ /**
+ * Checks whether the given table name is valid according to ClickHouse constraints.
+ *
+ * @param tableName the table name to check.
+ * @throws ClickHouseResourceManagerException if the table name is invalid.
+ */
+ static void checkValidTableName(String tableName) {
+ if (tableName.length() > MAX_TABLE_NAME_LENGTH) {
+ throw new IllegalArgumentException(
+ "Table name "
+ + tableName
+ + " cannot be longer than "
+ + MAX_TABLE_NAME_LENGTH
+ + " characters.");
+ }
+ if (!VALID_TABLE_NAME_PATTERN.matcher(tableName).matches()) {
+ throw new IllegalArgumentException(
+ "Table name "
+ + tableName
+ + " is not valid. It must contain only letters, numbers, and underscores and must start with a letter or an underscore.");
+ }
+ }
+}
diff --git a/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/package-info.java b/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/package-info.java
new file mode 100644
index 00000000000..28e5f596f30
--- /dev/null
+++ b/it/clickhouse/src/main/java/org/apache/beam/it/clickhouse/package-info.java
@@ -0,0 +1,20 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/** Package for managing ClickHouse resources within integration tests. */
+package org.apache.beam.it.clickhouse;
diff --git a/it/clickhouse/src/test/java/org/apache/beam/it/clickhouse/ClickHouseResourceManagerIT.java b/it/clickhouse/src/test/java/org/apache/beam/it/clickhouse/ClickHouseResourceManagerIT.java
new file mode 100644
index 00000000000..d9990579a80
--- /dev/null
+++ b/it/clickhouse/src/test/java/org/apache/beam/it/clickhouse/ClickHouseResourceManagerIT.java
@@ -0,0 +1,85 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.it.clickhouse;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.it.testcontainers.TestContainersIntegrationTest;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Integration tests for {@link ClickHouseResourceManager}. */
+@Category(TestContainersIntegrationTest.class)
+@RunWith(JUnit4.class)
+public class ClickHouseResourceManagerIT {
+
+ private ClickHouseResourceManager clickHouseResourceManager;
+
+ @Before
+ public void setUp() {
+ clickHouseResourceManager = ClickHouseResourceManager.builder("dummy").build();
+ }
+
+ @Test
+ public void testResourceManagerE2E() {
+
+ String testTableName = "dummy_insert";
+ Map<String, String> columns = new HashMap<>();
+ columns.put("name", "String");
+ columns.put("age", "Integer");
+ boolean createTable = clickHouseResourceManager.createTable(testTableName, columns, null, null);
+
+ assertThat(createTable).isTrue();
+
+ List<Map<String, Object>> rows = new ArrayList<>();
+
+ Map<String, Object> row1 = new HashMap<>();
+ row1.put("name", "Alice");
+ row1.put("age", 25);
+ rows.add(row1);
+
+ Map<String, Object> row2 = new HashMap<>();
+ row2.put("name", "Bob");
+ row2.put("age", 30);
+ rows.add(row2);
+
+ Map<String, Object> row3 = new HashMap<>();
+ row3.put("name", "Charlie");
+ row3.put("age", 28);
+ rows.add(row3);
+
+ boolean insertRows = clickHouseResourceManager.insertRows(testTableName, rows);
+
+ assertThat(insertRows).isTrue();
+
+ long count = clickHouseResourceManager.count(testTableName);
+
+ assertThat(count).isEqualTo(3L);
+
+ List<Map<String, Object>> fetchedRows = clickHouseResourceManager.fetchAll(testTableName);
+ assertThat(fetchedRows).hasSize(3);
+ assertThat(fetchedRows).containsExactlyElementsIn(rows);
+ }
+}
diff --git a/it/clickhouse/src/test/java/org/apache/beam/it/clickhouse/ClickHouseResourceManagerTest.java b/it/clickhouse/src/test/java/org/apache/beam/it/clickhouse/ClickHouseResourceManagerTest.java
new file mode 100644
index 00000000000..74e1bfd3b1e
--- /dev/null
+++ b/it/clickhouse/src/test/java/org/apache/beam/it/clickhouse/ClickHouseResourceManagerTest.java
@@ -0,0 +1,186 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.it.clickhouse;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyInt;
+import static org.mockito.Mockito.anyString;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.ResultSetMetaData;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Answers;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+import org.testcontainers.clickhouse.ClickHouseContainer;
+
+@RunWith(JUnit4.class)
+public class ClickHouseResourceManagerTest {
+
+ @Rule public final MockitoRule mockito = MockitoJUnit.rule();
+
+ @Mock(answer = Answers.RETURNS_DEEP_STUBS)
+ private Connection mockConnection;
+
+ @Mock private Statement mockStatement;
+
+ @Mock private PreparedStatement mockPreparedStatement;
+
+ @Mock private ResultSet mockResultSet;
+
+ @Mock private ClickHouseContainer mockContainer;
+
+ private ClickHouseResourceManager resourceManager;
+
+ @Before
+ public void setUp() throws SQLException {
+ MockitoAnnotations.initMocks(this); // This initializes the mocks
+ when(mockContainer.createConnection("")).thenReturn(mockConnection);
+ doReturn(mockContainer).when(mockContainer).withLogConsumer(any());
+
+ resourceManager =
+ new ClickHouseResourceManager(
+ mockConnection, mockContainer, ClickHouseResourceManager.builder("test-id"));
+ }
+
+ @Test
+ public void testTableExists_WhenTableExists() throws SQLException {
+ when(mockConnection.createStatement()).thenReturn(mockStatement);
+ when(mockStatement.executeQuery("EXISTS TABLE test_table")).thenReturn(mockResultSet);
+ when(mockResultSet.next()).thenReturn(true);
+ when(mockResultSet.getInt(1)).thenReturn(1);
+
+ assertTrue(resourceManager.tableExists("test_table"));
+ }
+
+ @Test
+ public void testTableExists_WhenTableDoesNotExist() throws SQLException {
+ when(mockConnection.createStatement()).thenReturn(mockStatement);
+ when(mockStatement.executeQuery("EXISTS TABLE test_table")).thenReturn(mockResultSet);
+ when(mockResultSet.next()).thenReturn(false);
+
+ assertFalse(resourceManager.tableExists("test_table"));
+ }
+
+ @Test
+ public void testCreateTable_Success() throws Exception {
+ when(mockConnection.createStatement()).thenReturn(mockStatement);
+ when(mockStatement.execute(anyString())).thenReturn(true);
+ when(mockStatement.executeQuery(anyString())).thenReturn(mockResultSet);
+ when(mockResultSet.next()).thenReturn(true);
+ // 0 for the first does-exists call, and 1 for the table creation validation call
+ when(mockResultSet.getInt(1)).thenReturn(0, 1);
+
+ Map<String, String> columns = new HashMap<>();
+ columns.put("id", "Int32");
+ columns.put("name", "String");
+
+ assertTrue(
+ resourceManager.createTable(
+ "test_table", columns, "MergeTree()", Collections.singletonList("id")));
+ }
+
+ @Test
+ public void testInsertRows_Success() throws Exception {
+ when(mockConnection.prepareStatement(anyString())).thenReturn(mockPreparedStatement);
+ when(mockPreparedStatement.executeBatch()).thenReturn(new int[] {1, 1});
+
+ List<Map<String, Object>> rows = new ArrayList<>();
+ Map<String, Object> row1 = new HashMap<>();
+ row1.put("id", 1);
+ row1.put("name", "Alice");
+ rows.add(row1);
+
+ Map<String, Object> row2 = new HashMap<>();
+ row2.put("id", 2);
+ row2.put("name", "Bob");
+ rows.add(row2);
+
+ assertTrue(resourceManager.insertRows("test_table", rows));
+ assertEquals(2, rows.size());
+ }
+
+ @Test
+ public void testFetchAll_Success() throws Exception {
+ when(mockConnection.createStatement()).thenReturn(mockStatement);
+ when(mockStatement.executeQuery(anyString())).thenReturn(mockResultSet);
+ when(mockResultSet.getMetaData()).thenReturn(mock(ResultSetMetaData.class));
+ when(mockResultSet.next()).thenReturn(true, false);
+ when(mockResultSet.getObject(anyInt())).thenReturn("Alice");
+ when(mockResultSet.getMetaData().getColumnCount()).thenReturn(1);
+ when(mockResultSet.getMetaData().getColumnName(1)).thenReturn("name");
+
+ List<Map<String, Object>> result = resourceManager.fetchAll("test_table");
+ assertEquals(1, result.size());
+ assertEquals("Alice", result.get(0).get("name"));
+ }
+
+ @Test
+ public void testCountRowsInTable() throws SQLException {
+ String tableName = "test_table";
+ ResultSet countRes = mock(ResultSet.class);
+ when(countRes.next()).thenReturn(true);
+ when(countRes.getLong(1)).thenReturn(10L);
+
+ Statement statement = mock(Statement.class);
+ when(statement.executeQuery("SELECT count(*) FROM " + tableName)).thenReturn(countRes);
+ when(mockConnection.createStatement()).thenReturn(statement);
+
+ long rowCount = resourceManager.count(tableName);
+
+ assertEquals(10L, rowCount, "Row count should be 10.");
+ }
+
+ @Test
+ public void testInsertRowsWithInvalidData() {
+ // Arrange
+ String tableName = "test_table";
+ List<Map<String, Object>> rows = new ArrayList<>();
+ Map<String, Object> row = new HashMap<>();
+ row.put("id", "string"); // Invalid data
+ rows.add(row);
+
+ // Act & Assert
+ assertThrows(
+ ClickHouseResourceManagerException.class,
+ () -> resourceManager.insertRows(tableName, rows),
+ "Should throw exception when inserting invalid data.");
+ }
+}
diff --git a/it/clickhouse/src/test/java/org/apache/beam/it/clickhouse/ClickHouseUtilsTest.java b/it/clickhouse/src/test/java/org/apache/beam/it/clickhouse/ClickHouseUtilsTest.java
new file mode 100644
index 00000000000..99bf2c1f5a3
--- /dev/null
+++ b/it/clickhouse/src/test/java/org/apache/beam/it/clickhouse/ClickHouseUtilsTest.java
@@ -0,0 +1,96 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.it.clickhouse;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import org.junit.jupiter.api.Test;
+
+class ClickHouseUtilsTest {
+
+ @Test
+ void testValidTableName() {
+ String validTableName = "valid_table_name";
+
+ // Test should pass without any exception
+ assertDoesNotThrow(() -> ClickHouseUtils.checkValidTableName(validTableName));
+ }
+
+ @Test
+ void testTableNameTooLong() {
+ // Create a table name of 207 characters manually (using StringBuilder for Java 8)
+ StringBuilder tableNameTooLong = new StringBuilder();
+ for (int i = 0; i < 207; i++) {
+ tableNameTooLong.append("t");
+ }
+
+ // Test should throw IllegalArgumentException due to length constraint
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> ClickHouseUtils.checkValidTableName(tableNameTooLong.toString()));
+
+ assertEquals(
+ "Table name " + tableNameTooLong + " cannot be longer than 206 characters.",
+ exception.getMessage());
+ }
+
+ @Test
+ void testTableNameWithInvalidCharacters() {
+ String invalidTableName = "invalid.table,name";
+
+ // Test should throw IllegalArgumentException due to invalid characters in the name
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> ClickHouseUtils.checkValidTableName(invalidTableName));
+
+ assertEquals(
+ "Table name "
+ + invalidTableName
+ + " is not valid. It must contain only letters, numbers, and underscores and must start with a letter or an underscore.",
+ exception.getMessage());
+ }
+
+ @Test
+ void testTableNameStartingWithNonLetter() {
+ String invalidTableName = "1invalid_table_name";
+
+ // Test should throw IllegalArgumentException because table name starts with a non-letter
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> ClickHouseUtils.checkValidTableName(invalidTableName));
+
+ assertEquals(
+ "Table name "
+ + invalidTableName
+ + " is not valid. It must contain only letters, numbers, and underscores and must start with a letter or an underscore.",
+ exception.getMessage());
+ }
+
+ @Test
+ void testTableNameWithUnderscoreAtStart() {
+ String validTableName = "_valid_table_name";
+
+ // Test should pass without any exception
+ assertDoesNotThrow(() -> ClickHouseUtils.checkValidTableName(validTableName));
+ }
+}
diff --git a/settings.gradle.kts b/settings.gradle.kts
index 470d6c020a5..dca142c9eef 100644
--- a/settings.gradle.kts
+++ b/settings.gradle.kts
@@ -360,4 +360,6 @@ findProject(":sdks:java:extensions:combiners")?.name = "combiners"
include("sdks:java:io:iceberg:hive")
findProject(":sdks:java:io:iceberg:hive")?.name = "hive"
include("sdks:java:io:iceberg:bqms")
-findProject(":sdks:java:io:iceberg:bqms")?.name = "bqms"
\ No newline at end of file
+findProject(":sdks:java:io:iceberg:bqms")?.name = "bqms"
+include("it:clickhouse")
+findProject(":it:clickhouse")?.name = "clickhouse"