Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ public class ColumnInfo implements Serializable {
private String comment;
private Object defaultValue;
private String properties;
private Boolean nullable;

@Builder
public ColumnInfo(String name, String type) {
Expand All @@ -58,6 +59,12 @@ public ColumnInfo(String name,
this.defaultValue = defaultValue;
}

public ColumnInfo(String name, String type, Boolean nullable) {
this.name = name;
this.type = type;
this.nullable = nullable;
}

public String getComment() {
return comment;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,4 +18,6 @@

public class ClickhouseConstants {
public static String CLICKHOUSE_CONNECTOR_NAME = "clickhouse";

public static String CLICKHOUSE_DECIMAL_INPUT_TYPE = "Decimal(76, 76)";
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
/*
* Copyright 2022 Bytedance Ltd. and/or its affiliates.
*
* Licensed 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 com.bytedance.bitsail.connector.clickhouse.option;

import com.bytedance.bitsail.common.annotation.Essential;
import com.bytedance.bitsail.common.option.ConfigOption;
import com.bytedance.bitsail.common.option.WriterOptions;

import com.alibaba.fastjson.TypeReference;

import java.util.Map;

import static com.bytedance.bitsail.common.option.ConfigOptions.key;
import static com.bytedance.bitsail.common.option.WriterOptions.WRITER_PREFIX;

public interface ClickhouseWriterOptions extends WriterOptions.BaseWriterOptions {
/**
* Standard format:
* jdbc:(ch|clickhouse)[:protocol]://endpoint[,endpoint][/database][?parameters][#tags]<br/>
* - endpoint: [protocol://]host[:port][/database][?parameters][#tags]<br/>
* - protocol: (grpc|grpcs|http|https|tcp|tcps)
*/
@Essential
ConfigOption<String> JDBC_URL =
key(WRITER_PREFIX + "jdbc_url")
.noDefaultValue(String.class);

/* ConfigOption<String> WRITE_MODE =
key(WRITER_PREFIX + "write_mode")
.noDefaultValue(String.class);

ConfigOption<String> WRITER_PARALLELISM_NUM =
key(WRITER_PREFIX + "writer_parallelism_num")
.noDefaultValue(String.class);*/

// Connection properties.
ConfigOption<Map<String, String>> CUSTOMIZED_CONNECTION_PROPERTIES =
key(WRITER_PREFIX + "customized_connection_properties")
.onlyReference(new TypeReference<Map<String, String>>() {});

ConfigOption<Integer> BATCH_SIZE =
key(WRITER_PREFIX + "batch_size")
.defaultValue(10);
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
/*
* Copyright 2022 Bytedance Ltd. and/or its affiliates.
*
* Licensed 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 com.bytedance.bitsail.connector.clickhouse.sink;

import com.bytedance.bitsail.base.connector.writer.v1.Sink;
import com.bytedance.bitsail.base.connector.writer.v1.Writer;
import com.bytedance.bitsail.base.connector.writer.v1.state.EmptyState;
import com.bytedance.bitsail.common.configuration.BitSailConfiguration;
import com.bytedance.bitsail.common.row.Row;
import com.bytedance.bitsail.common.type.TypeInfoConverter;
import com.bytedance.bitsail.common.type.filemapping.FileMappingTypeInfoConverter;
import com.bytedance.bitsail.connector.clickhouse.constant.ClickhouseConstants;

import java.io.IOException;
import java.io.Serializable;

public class ClickhouseSink<CommitT extends Serializable> implements Sink<Row, CommitT, EmptyState> {
private BitSailConfiguration jobConf;

@Override
public String getWriterName() {
return ClickhouseConstants.CLICKHOUSE_CONNECTOR_NAME;
}

@Override
public void configure(BitSailConfiguration commonConfiguration, BitSailConfiguration writerConfiguration) throws Exception {
this.jobConf = writerConfiguration;
}

@Override
public Writer<Row, CommitT, EmptyState> createWriter(Writer.Context<EmptyState> context) throws IOException {
return new ClickhouseWriter<>(this.jobConf, context);
}

@Override
public TypeInfoConverter createTypeInfoConverter() {
return new FileMappingTypeInfoConverter(getWriterName());
}
}

Original file line number Diff line number Diff line change
@@ -0,0 +1,196 @@
/*
* Copyright 2022 Bytedance Ltd. and/or its affiliates.
*
* Licensed 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 com.bytedance.bitsail.connector.clickhouse.sink;

import com.bytedance.bitsail.base.connector.writer.v1.Writer;
import com.bytedance.bitsail.base.connector.writer.v1.state.EmptyState;
import com.bytedance.bitsail.common.configuration.BitSailConfiguration;
import com.bytedance.bitsail.common.model.ColumnInfo;
import com.bytedance.bitsail.common.row.Row;
import com.bytedance.bitsail.connector.clickhouse.error.ClickhouseErrorCode;
import com.bytedance.bitsail.connector.clickhouse.option.ClickhouseWriterOptions;
import com.bytedance.bitsail.connector.clickhouse.source.reader.ClickhouseSourceReader;
import com.bytedance.bitsail.connector.clickhouse.util.ClickhouseConnectionHolder;

import com.clickhouse.jdbc.ClickHouseConnection;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import ru.yandex.clickhouse.domain.ClickHouseDataType;

import java.io.IOException;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicInteger;

import static com.bytedance.bitsail.connector.clickhouse.constant.ClickhouseConstants.CLICKHOUSE_DECIMAL_INPUT_TYPE;

public class ClickhouseWriter<CommitT> implements Writer<Row, CommitT, EmptyState> {
private static final Logger LOG = LoggerFactory.getLogger(ClickhouseSourceReader.class);

private final int subTaskId;

private final String jdbcUrl;
private final List<ColumnInfo> columnInfos;
private final String dbName;
private final String tableName;

private final String userName;
private final String password;
/*private final String writeMode;*/
/*private final String writerParallelismNum;*/
private final Integer batchSize;

private final String insertSql;

private final transient ClickhouseConnectionHolder connectionHolder;

private final List<Row> writeBuffer;
/*private final List<Row> commitBuffer;*/
private final AtomicInteger printCount;

/**
* Ensure there is only one connection activated.
*/
private transient ClickHouseConnection connection;
/**
* Ensure there is only one statement activated.
*/
private transient PreparedStatement statement;

public ClickhouseWriter(BitSailConfiguration jobConf, Writer.Context<EmptyState> context) {
this.subTaskId = context.getIndexOfSubTaskId();

this.jdbcUrl = jobConf.getNecessaryOption(ClickhouseWriterOptions.JDBC_URL,
ClickhouseErrorCode.REQUIRED_VALUE);
this.userName = jobConf.getNecessaryOption(ClickhouseWriterOptions.USER_NAME,
ClickhouseErrorCode.REQUIRED_VALUE);
this.password = jobConf.getNecessaryOption(ClickhouseWriterOptions.PASSWORD,
ClickhouseErrorCode.REQUIRED_VALUE);

this.dbName = jobConf.getNecessaryOption(ClickhouseWriterOptions.DB_NAME,
ClickhouseErrorCode.REQUIRED_VALUE);
this.tableName = jobConf.getNecessaryOption(ClickhouseWriterOptions.TABLE_NAME,
ClickhouseErrorCode.REQUIRED_VALUE);

/*this.writeMode = jobConf.get(ClickhouseWriterOptions.TABLE_NAME);*/
/*this.writerParallelismNum = jobConf.get(ClickhouseWriterOptions.WRITER_PARALLELISM_NUM);*/
this.batchSize = jobConf.get(ClickhouseWriterOptions.BATCH_SIZE);

this.columnInfos = jobConf.getNecessaryOption(ClickhouseWriterOptions.COLUMNS,
ClickhouseErrorCode.REQUIRED_VALUE);

connectionHolder = new ClickhouseConnectionHolder(jobConf, true);
this.connection = connectionHolder.connect();

insertSql = this.insertSql();
try {
this.statement = connection.prepareStatement(insertSql);
} catch (SQLException e) {
throw new RuntimeException("Failed to prepare statement.", e);
}

this.writeBuffer = new ArrayList<>(batchSize);
/*this.commitBuffer = new ArrayList<>(batchSize);*/
this.printCount = new AtomicInteger(0);

LOG.info("Clickhouse sink writer {} is initialized.", subTaskId);
}

@Override
public void write(Row element) throws IOException {
this.writeBuffer.add(element);

if (writeBuffer.size() == batchSize) {
this.flush(false);
}
printCount.incrementAndGet();
}

@Override
public void flush(boolean endOfInput) throws IOException {
/*this.commitBuffer.addAll(this.writeBuffer);*/

try {
for (Row ele : writeBuffer) {
for (int i = 0; i < this.columnInfos.size(); i++) {
statement.setObject(i + 1, ele.getField(i));
}
statement.addBatch();
}
statement.executeBatch();
} catch (SQLException e) {
throw new RuntimeException("Failed to write batch.", e);
}

writeBuffer.clear();
if (endOfInput) {
LOG.info("all records are sent to commit buffer.");
}
}

@Override
public List<CommitT> prepareCommit() throws IOException {
return Collections.emptyList();
}

@Override
public void close() throws IOException {
Writer.super.close();

try {
this.statement.close();
this.connection.close();
} catch (SQLException e) {
throw new RuntimeException("Failed to close clickhouse writer.", e);
}
}

private String insertSql() {
String inputFields = this.columnInfos.stream().map(col -> {
if (Objects.isNull(col.getNullable()) || !col.getNullable()) {
return String.format("%s %s", col.getName(), this.getClickhouseType(col.getType()));
} else {
return String.format("%s Nullable(%s)", col.getName(), this.getClickhouseType(col.getType()));
}
}).reduce((n1, n2) -> String.format("%s, %s", n1, n2)).get();

StringBuffer sql = new StringBuffer("INSERT INTO ");
sql.append(String.format("%s.%s", this.dbName, this.tableName))
.append(" SELECT ")
.append(this.columnInfos.stream().map(col -> col.getName()).reduce((n1, n2) -> String.format("%s, %s", n1, n2)).get())
.append(" FROM input('")
.append(inputFields)
.append("')");

return sql.toString();
}

private String getClickhouseType(String inputType) {
ClickHouseDataType ckType = ClickHouseDataType.fromTypeString(inputType);
String strType = ckType.toString();

if (strType.toLowerCase().contains("decimal")) {
return CLICKHOUSE_DECIMAL_INPUT_TYPE;
}
return strType;
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import com.bytedance.bitsail.common.configuration.BitSailConfiguration;
import com.bytedance.bitsail.connector.clickhouse.error.ClickhouseErrorCode;
import com.bytedance.bitsail.connector.clickhouse.option.ClickhouseReaderOptions;
import com.bytedance.bitsail.connector.clickhouse.option.ClickhouseWriterOptions;

import com.clickhouse.jdbc.ClickHouseConnection;
import com.clickhouse.jdbc.ClickHouseDataSource;
Expand Down Expand Up @@ -57,6 +58,22 @@ public ClickhouseConnectionHolder(BitSailConfiguration jobConf) {
}
}

public ClickhouseConnectionHolder(BitSailConfiguration jobConf, Boolean isWriter) {
this.jdbcUrl = jobConf.getNecessaryOption(ClickhouseWriterOptions.JDBC_URL,
ClickhouseErrorCode.REQUIRED_VALUE);
this.dbName = jobConf.getNecessaryOption(ClickhouseWriterOptions.DB_NAME,
ClickhouseErrorCode.REQUIRED_VALUE);

this.userName = jobConf.get(ClickhouseWriterOptions.USER_NAME);
this.password = jobConf.get(ClickhouseWriterOptions.PASSWORD);

this.connectionProperties = new Properties();
if (jobConf.fieldExists(ClickhouseWriterOptions.CUSTOMIZED_CONNECTION_PROPERTIES)) {
jobConf.get(ClickhouseWriterOptions.CUSTOMIZED_CONNECTION_PROPERTIES)
.forEach(connectionProperties::setProperty);
}
}

/**
* Refrence: <a href="https://github.com/ClickHouse/clickhouse-jdbc/tree/develop/clickhouse-jdbc#upgrade-to-032">Clickhouse-jdbc</a>
*/
Expand Down
Loading