zhanghongqing-knowledge-log/src/main/java/com/zdjizhi/utils/ck/ClickhouseSink.java

package com.zdjizhi.utils.ck;

import cn.hutool.core.io.IoUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.log.Log;
import cn.hutool.log.LogFactory;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import ru.yandex.clickhouse.BalancedClickhouseDataSource;
import ru.yandex.clickhouse.settings.ClickHouseProperties;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static com.zdjizhi.common.FlowWriteConfig.*;

public class ClickhouseSink extends RichSinkFunction<List<Map<String, Object>>> {

    private static final Log log = LogFactory.get();

    private Connection connection;
    private PreparedStatement preparedStatement;
    public String sink;


    public ClickhouseSink(String sink) {
        this.sink = sink;
    }

    public String getSink() {
        return sink;
    }

    public void setSink(String sink) {
        this.sink = sink;
    }

    @Override
    public void invoke(List<Map<String, Object>> logs, Context context) throws Exception {
        executeInsert(logs, getSink());
    }

    @Override
    public void open(Configuration parameters) throws Exception {

        try {
//            Class.forName("ru.yandex.clickhouse.ClickHouseDriver");
//            connection = DriverManager.getConnection("jdbc:clickhouse://" + CK_HOSTS + "/" + CK_DATABASE, CK_USERNAME, CK_PIN);
            ClickHouseProperties properties = new ClickHouseProperties();
            properties.setDatabase(CK_DATABASE);
            properties.setUser(CK_USERNAME);
            properties.setPassword(CK_PIN);
//            properties.setKeepAliveTimeout(5);
            properties.setConnectionTimeout(CK_CONNECTION_TIMEOUT);
            properties.setSocketTimeout(CK_SOCKET_TIMEOUT);

            BalancedClickhouseDataSource dataSource = new BalancedClickhouseDataSource("jdbc:clickhouse://" + CK_HOSTS, properties);
            dataSource.scheduleActualization(10, TimeUnit.SECONDS);//开启检测
            connection = dataSource.getConnection();

            log.info("get clickhouse connection success");
        } catch (SQLException e) {
            log.error("clickhouse connection error ,{}", e);
        }

    }

    @Override
    public void close() throws Exception {
        IoUtil.close(preparedStatement);
        IoUtil.close(connection);
    }

    public void executeInsert(List<Map<String, Object>> data, String tableName) {

        try {
            List<String> keys = new LinkedList<>(data.get(0).keySet());
            connection.setAutoCommit(false);
            preparedStatement = connection.prepareStatement(preparedSql(keys, tableName));
            int count = 0;
            for (Map<String, Object> map : data) {
                List<Object> values = new LinkedList<>(map.values());
                for (int i = 1; i <= values.size(); i++) {
                    Object val = values.get(i - 1);
                    if (val instanceof Long) {
                        preparedStatement.setLong((i), Long.valueOf(StrUtil.toString(val)));
                    } else if (val instanceof Integer) {
                        preparedStatement.setInt((i), Integer.valueOf(StrUtil.toString(val)));
                    } else if (val instanceof Boolean) {
                        preparedStatement.setBoolean((i), Boolean.valueOf(StrUtil.toString(val)));
                    } else {
                        preparedStatement.setString((i), StrUtil.toString(val));
                    }
                }
                preparedStatement.addBatch();
                count++;
                //1w提交一次
                if (count % SINK_BATCH == 0) {
                    preparedStatement.executeBatch();
                    connection.commit();
                    preparedStatement.clearBatch();
                    count = 0;
                }
            }
            if (count > 0) {
                preparedStatement.executeBatch();
                connection.commit();
            }

        } catch (Exception ex) {
            log.error("ClickhouseSink插入报错", ex);
        }
    }


    public static String preparedSql(List<String> fields, String tableName) {

        String placeholders = fields.stream()
                .filter(Objects::nonNull)
                .map(f -> "?")
                .collect(Collectors.joining(", "));
        String columns = fields.stream()
                .filter(Objects::nonNull)
                .collect(Collectors.joining(", "));
        String sql = StrUtil.concat(true, "INSERT INTO ", CK_DATABASE, ".", tableName,
                "(", columns, ") VALUES (", placeholders, ")");
        log.debug(sql);
        return sql;
    }


}