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

Go to download
Show more of this group Show more artifacts with this name
Show all versions of dws-client Show documentation
There is a newer version: 2.0.0-r0
package com.huaweicloud.dws.client.worker;

import com.huaweicloud.dws.client.DwsConfig;
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.Properties;

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

    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 {
        if (config.getDriverName() == null) {
            connection =
                    DriverManager.getConnection(
                            config.getUrl(),
                            config.getUsername(),
                            config.getPassword());
        } else {
            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");
            }
        }
        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;
            }
        }
    }


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

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