net.dongliu.dbutils.Database Maven / Gradle / Ivy

Go to download
package net.dongliu.dbutils;

import net.dongliu.dbutils.exception.UncheckedSQLException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Arrays;
import java.util.Objects;
import java.util.function.Consumer;
import java.util.function.Function;

/**
 * Executes SQL queries.
 *
 * @see ResultSetHandler
 */
public class Database {

    private static final Logger logger = LoggerFactory.getLogger(Database.class);

    // thread local used to save transaction connection
    private final ThreadLocal transactionInfoTL = new ThreadLocal<>();

    private final SqlExecutor sqlExecutor = new SqlExecutor();

    /**
     * The DataSource to retrieve connections from.
     */
    private final DataSource dataSource;

    /**
     * Constructor to provide a DataSource.
     */
    private Database(DataSource dataSource) {
        this.dataSource = Objects.requireNonNull(dataSource);
    }

    /**
     * Create a database from data source
     */
    public static Database create(DataSource dataSource) {
        return new Database(dataSource);
    }

    /**
     * Use internal non-pooled data source
     */
    public static Database create(String jdbcUrl, String user, String password, String driverClass) {
        return new Database(SimpleDataSource.create(jdbcUrl, user, password, driverClass));
    }

    /**
     * Fetch a Connection object.
     * For transaction, connection is save by thread-local, just retrieve and return
     *
     * @return An initialized Connection.
     */
    private ConnectionInfo retrieveConnection() {
        TransactionInfo transactionInfo = transactionInfoTL.get();
        if (transactionInfo != null) {
            // in a transaction now
            if (logger.isDebugEnabled()) {
                logger.debug("Get connection for transaction: " + System.identityHashCode(transactionInfo.connection));
            }
            return new ConnectionInfo(false, transactionInfo.connection);
        }
        try {
            Connection connection = this.dataSource.getConnection();
            if (logger.isDebugEnabled()) {
                logger.debug("Get connection: " + System.identityHashCode(connection));
            }
            return new ConnectionInfo(true, connection);
        } catch (SQLException e) {
            throw new UncheckedSQLException(e);
        }
    }

    /**
     * Start a transaction.
     * Only work for current thread.
     */
    public void startTransaction() {
        logger.debug("Start Transaction");
        Connection connection;
        boolean autoCommit;
        try {
            connection = this.dataSource.getConnection();
            autoCommit = connection.getAutoCommit();
            connection.setAutoCommit(false);
        } catch (SQLException e) {
            throw new UncheckedSQLException(e);
        }
        transactionInfoTL.set(new TransactionInfo(connection, autoCommit));
    }

    /**
     * Roll back transaction
     */
    public void rollback() {
        TransactionInfo transactionInfo = transactionInfoTL.get();
        if (transactionInfo == null) {
            throw new UncheckedSQLException("Not in a transaction");
        }
        logger.debug("Rollback Transaction");
        try {
            transactionInfo.connection.rollback();
        } catch (SQLException e) {
            throw new UncheckedSQLException(e);
        } finally {
            restoreConnection(transactionInfo);
        }
    }

    /**
     * Commit transaction
     */
    public void commit() {
        TransactionInfo transactionInfo = transactionInfoTL.get();
        if (transactionInfo == null) {
            throw new UncheckedSQLException("Not in a transaction");
        }
        logger.debug("Commit Transaction");
        try {
            transactionInfo.connection.commit();
        } catch (SQLException e) {
            throw new UncheckedSQLException(e);
        } finally {
            restoreConnection(transactionInfo);
        }
    }

    /**
     * called when transaction ended
     */
    private void restoreConnection(TransactionInfo transactionInfo) {
        transactionInfoTL.remove();
        try (Connection connection = transactionInfo.connection) {
            connection.setAutoCommit(transactionInfo.autoCommit);
        } catch (SQLException e) {
            throw new UncheckedSQLException(e);
        }
    }

    /**
     * Wrap function to run in a transaction
     *
     * @return the result returned by function
     */
    public  T applyWithTransaction(Function function) {
        this.startTransaction();
        try {
            T value = function.apply(this);
            this.commit();
            return value;
        } catch (Throwable e) {
            this.rollback();
            throw e;
        }
    }

    /**
     * Wrap no return value function to run in a transaction
     */
    public void runWithTransaction(Consumer consumer) {
        this.startTransaction();
        try {
            consumer.accept(this);
            this.commit();
        } catch (Throwable e) {
            this.rollback();
            throw e;
        }
    }

    /**
     * Execute a batch of SQL INSERT, UPDATE, or DELETE queries.
     *
     * @param sql    The SQL to execute.
     * @param params An array of query replacement parameters.  Each row in
     *               this array is one set of batch replacement values.
     * @return The number of rows updated per statement.
     */
    public int[] batchUpdate(String sql, Object[][] params) {
        if (logger.isDebugEnabled()) {
            logger.debug("Batch update, sql:" + sql + ", params:" + Arrays.deepToString(params));
        }
        ConnectionInfo ci = this.retrieveConnection();
        try {
            return sqlExecutor.batch(ci.connection, ci.autoClose, sql, params);
        } catch (SQLException e) {
            throw new UncheckedSQLException(e);
        }
    }


    /**
     * Executes the given SELECT SQL query and returns a result object.
     *
     * @param sql    The SQL statement to execute.
     * @param params Initialize the PreparedStatement's IN parameters with
     *               this array.
     */
    public ResultSetHolder query(String sql, Object... params) {
        if (logger.isDebugEnabled()) {
            logger.debug("Query, sql:" + sql + ", params:" + Arrays.deepToString(params));
        }
        ConnectionInfo ci = this.retrieveConnection();
        return rsh -> sqlExecutor.query(ci.connection, ci.autoClose, sql, rsh, params);
    }

    /**
     * Executes the given INSERT, UPDATE, or DELETE SQL statement.
     *
     * @param sql    The SQL statement to execute.
     * @param params Initializes the PreparedStatement's IN (i.e. '?')
     *               parameters.
     * @return The number of rows updated.
     */
    public int update(String sql, Object... params) {
        if (logger.isDebugEnabled()) {
            logger.debug("Update, sql:" + sql + ", params:" + Arrays.deepToString(params));
        }
        ConnectionInfo ci = this.retrieveConnection();
        try {
            return sqlExecutor.update(ci.connection, ci.autoClose, sql, params);
        } catch (SQLException e) {
            throw new UncheckedSQLException(e);
        }
    }

    /**
     * Executes the given INSERT SQL statement.
     *
     * @param sql    The SQL statement to execute.
     * @param params Initializes the PreparedStatement's IN (i.e. '?')
     */
    public ResultSetHolder insert(String sql, Object... params) {
        if (logger.isDebugEnabled()) {
            logger.debug("Insert, sql:" + sql + ", params:" + Arrays.deepToString(params));
        }
        ConnectionInfo ci = this.retrieveConnection();
        return rsh -> sqlExecutor.insert(ci.connection, ci.autoClose, sql, rsh, params);
    }

    /**
     * Executes the given batch of INSERT SQL statements.
     *
     * @param sql    The SQL statement to execute.
     * @param params Initializes the PreparedStatement's IN (i.e. '?')
     */
    public ResultSetHolder batchInsert(String sql, Object[][] params) {
        if (logger.isDebugEnabled()) {
            logger.debug("Batch insert, sql:" + sql + ", params:" + Arrays.deepToString(params));
        }
        ConnectionInfo ci = this.retrieveConnection();
        return rsh -> sqlExecutor.insertBatch(ci.connection, ci.autoClose, sql, rsh, params);
    }

    private static class ConnectionInfo {
        private final boolean autoClose;
        private final Connection connection;

        public ConnectionInfo(boolean autoClose, Connection connection) {
            this.autoClose = autoClose;
            this.connection = connection;
        }
    }

    private static class TransactionInfo {
        private final Connection connection;
        // backup auto commit setting
        private final boolean autoCommit;

        public TransactionInfo(Connection connection, boolean autoCommit) {
            this.connection = connection;
            this.autoCommit = autoCommit;
        }
    }
}