com.huaweicloud.dws.client.worker.ConnectionProvider Maven / Gradle / Ivy

Go to download
package com.huaweicloud.dws.client.worker;

import com.huaweicloud.dws.client.DwsConfig;
import com.huaweicloud.dws.client.exception.InvalidException;
import com.huaweicloud.dws.client.model.CreateTempTableMode;
import com.huaweicloud.dws.client.model.TableName;
import com.huaweicloud.dws.client.model.TableSchema;
import com.huaweicloud.dws.client.util.JdbcUtil;
import lombok.extern.slf4j.Slf4j;

import java.io.Closeable;
import java.io.Serializable;
import java.sql.Connection;
import java.sql.Driver;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.util.Enumeration;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @ProjectName: dws-connector
 * @ClassName: JdbcConnectionProvider
 * @Description: JDBC连接管理
 * @Date: 2023/1/16 14:23
 * @Version: 1.0
 */
@Slf4j
public class ConnectionProvider implements Closeable, Serializable {
    private transient Driver loadedDriver;
    private transient Connection connection;
    private final DwsConfig config;
    private long connectionStartTime = System.currentTimeMillis();
    private long lastActive = System.currentTimeMillis();

    private Map TEMP_TABLES_CACHE = new ConcurrentHashMap<>();

    static {
        DriverManager.getDrivers();
    }

    public void refresh() {
        this.lastActive = System.currentTimeMillis();
    }

    public long getLastActive() {
        return lastActive;
    }

    private static Driver loadDriver(String driverName)
            throws SQLException, ClassNotFoundException {
        Enumeration drivers = DriverManager.getDrivers();
        while (drivers.hasMoreElements()) {
            Driver driver = drivers.nextElement();
            if (driver.getClass().getName().equals(driverName)) {
                return driver;
            }
        }
        Class clazz =
                Class.forName(driverName, true, Thread.currentThread().getContextClassLoader());
        try {
            return (Driver) clazz.newInstance();
        } catch (Exception ex) {
            throw new SQLException("Fail to create driver of class " + driverName, ex);
        }
    }

    private Driver getLoadedDriver() throws SQLException, ClassNotFoundException {
        if (loadedDriver == null) {
            loadedDriver = loadDriver(config.getDriverName());
        }
        return loadedDriver;
    }

    public ConnectionProvider(DwsConfig config) {
        this.config = config;
    }

    /**
     * 检查连接是否可用
     */
    public boolean isConnectionValid() throws SQLException {
        return connection != null && connection.isValid(config.getConnectionCheckTimeoutSeconds());
    }


    /**
     * 获取已有连接 或者初始化一个新连接，如果使用时间大于配置时间 将强制刷新
     */
    public Connection getOrInitConnection() throws SQLException, ClassNotFoundException {
        Connection tmp = getConnection();
        if (!tmp.getAutoCommit()) {
            tmp.setAutoCommit(true);
        }
        return tmp;
    }

    private Connection getConnection() throws SQLException, ClassNotFoundException {
        if (connection != null) {
            return connectionUseTime() > config.getConnectionMaxUseTimeSeconds() || !isConnectionValid() ? restConnection() :
                    connection;
        }
        initConnection();
        return connection;
    }

    private void initConnection() throws SQLException, ClassNotFoundException {
        TEMP_TABLES_CACHE.clear();
        Driver driver = getLoadedDriver();
        Properties info = new Properties();
        info.setProperty("user", config.getUsername());
        info.setProperty("password", config.getPassword());
        connection = driver.connect(config.getUrl(), info);
        if (connection == null) {
            // Throw same exception as DriverManager.getConnection when no driver found to match
            // caller expectation.
            throw new SQLException(
                    "No suitable driver found for " + config.getUrl(), "08001");
        }
        JdbcUtil.executeSql(connection, String.format("SET statement_timeout = %s", config.getTimeOutMs()));
        connectionStartTime = System.currentTimeMillis();
    }

    /** Close possible existing connection. */
    @Override
    public void close() {
        if (connection != null) {
            log.warn("connection close. url = {}, use time = {}, idle time = {}",
                    config.getUrl(), connectionUseTime(), (System.currentTimeMillis() - lastActive));
            try {
                connection.close();
            } catch (SQLException e) {
                log.warn("JDBC connection close failed.", e);
            } finally {
                connection = null;
                TEMP_TABLES_CACHE.clear();
            }
        }
    }


    /**
     * 重置连接，先关闭再获取
     */
    public Connection restConnection() throws SQLException, ClassNotFoundException {
        close();
        return getOrInitConnection();
    }

    private long connectionUseTime() {
        return (System.currentTimeMillis() - connectionStartTime) / 1000;
    }

    public String getTempTable(TableSchema tableSchema) {
        return TEMP_TABLES_CACHE.computeIfAbsent(tableSchema, table -> {
            TableName tableName = table.getTableName();
            String name = String.format("tmp_%s_%s", tableName.getTableName(), System.currentTimeMillis());
            CreateTempTableMode mode = config.getTableConfig(tableName).getCreateTempTableMode();
            String sql = mode == CreateTempTableMode.LIKE ? JdbcUtil.getCreateTempTableSql(tableName.getFullName(), name)
                    : JdbcUtil.getCreateTempTableAsSql(tableName.getFullName(), name);
            log.info("init temp table {}, sql = {}", name, sql);
            try {
                JdbcUtil.executeSql(getConnection(), sql);
            } catch (Exception e) {
                log.error("init temp table fail. {}", name, e);
                throw new InvalidException(e.getMessage());
            }
            return name;
        });
    }

    public void cleanTempTable(TableSchema tableSchema) {
        TEMP_TABLES_CACHE.remove(tableSchema);
    }
}