io.github.shanqiang.sp.output.MysqlOutputTable Maven / Gradle / Ivy

Go to download
package io.github.shanqiang.sp.output;

import com.google.common.base.Strings;
import com.google.common.base.Throwables;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.mysql.cj.jdbc.MysqlDataSource;
import io.github.shanqiang.exception.UnknownTypeException;
import io.github.shanqiang.sp.StreamProcessing;
import io.github.shanqiang.table.Column;
import io.github.shanqiang.table.Table;
import io.github.shanqiang.table.Type;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Types;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

import static io.github.shanqiang.util.ScalarUtil.toDouble;
import static io.github.shanqiang.util.ScalarUtil.toInteger;
import static io.github.shanqiang.util.ScalarUtil.toLong;
import static io.github.shanqiang.util.ScalarUtil.toStr;
import static java.lang.String.format;
import static java.util.Objects.requireNonNull;

public class MysqlOutputTable extends AbstractOutputTable {
    private static final Logger logger = LoggerFactory.getLogger(MysqlOutputTable.class);

    private final String jdbcUrl;
    private final String tableName;
    private final String userName;
    private final String password;
    private final String createIndex;
    private final int maxRetryTimes;
    private final int batchSize;
    private final boolean autoDropTable;
    protected final Map columnTypeMap;
    private final String insertPrefix;

    private final ThreadPoolExecutor threadPoolExecutor;

    public MysqlOutputTable(String jdbcUrl,
                            String userName,
                            String password,
                            String tableName,
                            String createIndex,
                            boolean autoDropTable,
                            Map columnTypeMap) throws IOException {
        this(Runtime.getRuntime().availableProcessors(),
                1000,
                jdbcUrl,
                userName,
                password,
                tableName,
                createIndex,
                1,
                autoDropTable,
                columnTypeMap);
    }

    public MysqlOutputTable(int thread,
                            int batchSize,
                            String jdbcUrl,
                            String userName,
                            String password,
                            String tableName,
                            String createIndex,
                            int maxRetryTimes,
                            boolean autoDropTable,
                            Map columnTypeMap) throws IOException {
        super(thread, "|MysqlOutputTable|" + tableName);
        this.jdbcUrl = requireNonNull(jdbcUrl);
        this.userName = requireNonNull(userName);
        this.password = requireNonNull(password);
        this.tableName = requireNonNull(tableName);
        this.createIndex = createIndex;
        this.maxRetryTimes = requireNonNull(maxRetryTimes);
        this.batchSize = batchSize;
        this.autoDropTable = autoDropTable;
        this.columnTypeMap = requireNonNull(columnTypeMap);
        this.insertPrefix = "insert into " + tableName + " values ";

        threadPoolExecutor = new ThreadPoolExecutor(thread,
                thread,
                0,
                TimeUnit.SECONDS,
                new ArrayBlockingQueue<>(1),
                new ThreadFactoryBuilder().setNameFormat("mysql-output-%d").build());

        createTable();
    }

    private Connection connect() throws SQLException {
        MysqlDataSource dataSource = new MysqlDataSource();
        dataSource.setUrl(jdbcUrl);
        dataSource.setUser(userName);
        dataSource.setPassword(password);
        dataSource.setAutoReconnect(true);
        try {
            return dataSource.getConnection();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }

    private void createTable() throws IOException {
        StringBuilder fieldsBuilder = new StringBuilder();
        for (String columnName : columnTypeMap.keySet()) {
            Type type = columnTypeMap.get(columnName);
            if (type == Type.VARBYTE) {
                fieldsBuilder.append(columnName).append(" ").append("longtext").append(",");
            } else {
                fieldsBuilder.append(columnName).append(" ").append(type).append(",");
            }
        }

        String fieldsSchema = fieldsBuilder.toString();
        if (fieldsSchema.length() > 0) {
            fieldsSchema = fieldsSchema.substring(0, fieldsSchema.length() - 1);
        }

        /* 表不存在则创建表 */
        String createTableSql = format("CREATE TABLE IF NOT EXISTS %s (%s) ",
                tableName, fieldsSchema);

        logger.info(">>> create table sql: " + createTableSql);

        int retryCount = 0;
        while (retryCount < maxRetryTimes) {
            try {
                Connection connection = connect();
                if (autoDropTable) {
                    PreparedStatement preparedStatement = connection.prepareStatement("DROP TABLE IF EXISTS " + tableName);
                    preparedStatement.execute();
                }
                PreparedStatement preparedStatement = connection.prepareStatement(createTableSql);
                preparedStatement.execute();
                if (!Strings.isNullOrEmpty(createIndex)) {
                    preparedStatement = connection.prepareStatement(createIndex);
                    preparedStatement.execute();
                }

                return;
            } catch (Throwable t) {
                logger.error(">>> create table error: {}, has retried {} times", Throwables.getStackTraceAsString(t), retryCount);
                retryCount++;
                if (retryCount >= maxRetryTimes) {
                    throw new IOException(t);
                }
                try {
                    Thread.sleep(1 * 1000L);
                } catch (Throwable t2) {
                    logger.error("retry sleep error!", t2);
                }
            }
        }

        throw new IOException(">>> create mysql table error for " + tableName + ", we have tried " + maxRetryTimes + " times");
    }

    @Override
    public void produce(Table table) throws InterruptedException {
        putTable(table);
    }

    private void setValues(PreparedStatement preparedStatement, List