close bytedance#109:first edition

bitsail-common/src/main/java/com/bytedance/bitsail/common/model/ColumnInfo.java

@@ -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) {
@@ -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;
   }

bitsail-connectors/connector-clickhouse/src/main/java/com/bytedance/bitsail/connector/clickhouse/constant/ClickhouseConstants.java

@@ -18,4 +18,6 @@
 
 public class ClickhouseConstants {
   public static String CLICKHOUSE_CONNECTOR_NAME = "clickhouse";
+
+  public static String CLICKHOUSE_DECIMAL_INPUT_TYPE = "Decimal(76, 76)";
 }

bitsail-connectors/connector-clickhouse/src/main/java/com/bytedance/bitsail/connector/clickhouse/option/ClickhouseWriterOptions.java

@@ -0,0 +1,57 @@
+/*
+ *     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 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);
+}

bitsail-connectors/connector-clickhouse/src/main/java/com/bytedance/bitsail/connector/clickhouse/sink/ClickhouseSink.java

@@ -0,0 +1,38 @@
+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());
+  }
+}
+

bitsail-connectors/connector-clickhouse/src/main/java/com/bytedance/bitsail/connector/clickhouse/sink/ClickhouseWriter.java

@@ -0,0 +1,180 @@
+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;
+  }
+
+}

bitsail-connectors/connector-clickhouse/src/main/java/com/bytedance/bitsail/connector/clickhouse/util/ClickhouseConnectionHolder.java

@@ -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;
@@ -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>
    */