com.tesora.dve.common.DBHelper.java Source code

Java tutorial

Introduction

Here is the source code for com.tesora.dve.common.DBHelper.java

Source

package com.tesora.dve.common;

/*
 * #%L
 * Tesora Inc.
 * Database Virtualization Engine
 * %%
 * Copyright (C) 2011 - 2014 Tesora Inc.
 * %%
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License, version 3,
 * as published by the Free Software Foundation.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Affero General Public License for more details.
 * 
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 * #L%
 */

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.Charset;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.SQLWarning;
import java.sql.Statement;
import java.util.List;
import java.util.Properties;

import org.apache.commons.lang.StringUtils;
import org.apache.log4j.Logger;

import com.tesora.dve.exceptions.PECodingException;
import com.tesora.dve.exceptions.PEException;

public class DBHelper {
    public static final String CONN_ROOT = PEConstants.PROP_JPA_JDBC_PREFIX;

    public static final String CONN_DRIVER_CLASS = PEConstants.PROP_FULL_JDBC_DRIVER;
    public static final String CONN_URL = PEConstants.PROP_FULL_JDBC_URL;
    public static final String CONN_USER = PEConstants.PROP_FULL_JDBC_USER;
    public static final String CONN_PASSWORD = PEConstants.PROP_FULL_JDBC_PASSWORD;
    public static final String CONN_DBNAME = PEConstants.PROP_DBNAME;

    private static Logger logger = Logger.getLogger(DBHelper.class);

    protected PEUrl connUrl;
    protected String connUserName;
    protected String connPwd;
    protected String connDBName = null;

    protected Connection connection = null;
    protected Statement stmt = null;
    protected PreparedStatement prepStmt = null;
    protected String prepQuery;

    protected int rowCount = 0;
    protected long lastInsertID = 0L;
    protected ResultSet resultSet = null;

    protected boolean enableResultSetStreaming = false;
    protected boolean useBufferedQuery = false;
    protected boolean disconnectAfterProcess = true;

    @SuppressWarnings("resource")
    public static Charset getConnectionCharset(final DBHelper dbHelper) throws SQLException {
        final Connection connection = dbHelper.getConnection();
        if ((connection != null) && !connection.isClosed()) {
            try (final ResultSet rs = connection.createStatement()
                    .executeQuery("SHOW VARIABLES WHERE Variable_name = 'character_set_connection'")) {
                rs.last();
                final String connectionCharsetName = rs.getString(2);
                return Charset.forName(PEStringUtils.dequote(connectionCharsetName));
            }
        }

        throw new PECodingException("Database connection is not open.");
    }

    /**
     * @param props
     * @throws PEException
     */
    public DBHelper(Properties props) throws PEException {
        this(props.getProperty(CONN_URL), props.getProperty(CONN_USER), props.getProperty(CONN_PASSWORD),
                props.getProperty(CONN_DBNAME));
    }

    /**
     * @param url
     * @param userName
     * @param password
     * @param database
     * @throws PEException
     */
    public DBHelper(String url, String userName, String password, String database) throws PEException {
        this(PEUrl.fromUrlString(url), userName, password, database);
    }

    /**
     * @param url
     * @param userName
     * @param password
     * @throws PEException
     */
    public DBHelper(String url, String userName, String password) throws PEException {
        this(PEUrl.fromUrlString(url), userName, password);
    }

    /**
     * @param url
     * @param userName
     * @param password
     * @param database
     * @throws PEException
     */
    public DBHelper(PEUrl url, String userName, String password, String database) throws PEException {
        this(url, userName, password);

        this.connDBName = database;
    }

    /**
     * @param url
     * @param userName
     * @param password
     * @throws PEException
     */
    public DBHelper(PEUrl url, String userName, String password) throws PEException {
        this.connUrl = url;
        this.connUserName = userName;
        this.connPwd = password;
        this.connDBName = connUrl.getPath();
    }

    public static DBType urlToDBType(String url) throws PEException {
        return DBType.fromDriverClass(urlToDriverClass(PEUrl.fromUrlString(url)));
    }

    public static String urlToDriverClass(String url) throws PEException {
        return urlToDriverClass(PEUrl.fromUrlString(url));
    }

    public static String urlToDriverClass(PEUrl url) throws PEException {
        if (url.getSubProtocol().equalsIgnoreCase(PEConstants.MYSQL_SUBPROTOCOL))
            return PEConstants.MYSQL_DRIVER_CLASS;

        throw new PEException("Cannot determine JDBC driver class for '" + url.getSubProtocol() + "'");
    }

    public static String loadDriverForURL(String url) throws PEException {
        return loadDriver(urlToDriverClass(url));
    }

    public static String loadDriverForURL(PEUrl url) throws PEException {
        return loadDriver(urlToDriverClass(url));
    }

    public static String loadDriver(String driverClass) throws PEException {
        try {
            Class.forName(driverClass);
        } catch (Exception e) {
            logger.error("JDBC driver class '" + driverClass + "' not found");

            throw new PEException("JDBC driver class '" + driverClass + "' not found");
        }
        return driverClass;
    }

    public Connection getConnection() {
        return this.connection;
    }

    public PEUrl getUrl() {
        return connUrl;
    }

    public String getUserName() {
        return connUserName;
    }

    public String getPassword() {
        return connPwd;
    }

    public String getDBName() {
        return connDBName;
    }

    public void setDBName(String connDBName) {
        this.connDBName = connDBName;
    }

    public void disconnect() {
        logger.info("Disconnecting from " + this.connUrl + " as " + this.connUserName);

        try {
            closeStatement();

            if (connection != null)
                connection.close();
        } catch (SQLException e) {
            logger.error(e, e);
        }
    }

    public DBHelper connect() throws PEException {
        return connect(new Properties());
    }

    public DBHelper connect(Properties props) throws PEException {
        props.put("user", connUserName);
        props.put("password", connPwd);

        return doConnect(props);
    }

    private DBHelper doConnect(Properties props) throws PEException {

        DBHelper.loadDriverForURL(connUrl);

        try {
            logger.info("Connecting to " + this.connUrl + " as " + this.connUserName);

            connection = DriverManager.getConnection(connUrl.getURL(), props);
        } catch (Exception se) {
            throw new PEException("Error connecting to database '" + connUrl + "' - " + se.getMessage(), se);
        }

        try {
            if (this.connDBName != null) {
                connection.setCatalog(connDBName);
                logger.info("Database set to " + connDBName);
            }
        } catch (Exception se) {
            throw new PEException(
                    "Error using database '" + connDBName + "' on '" + connUrl + "' - " + se.getMessage(), se);
        }
        return this;
    }

    public void closeStatement() throws SQLException {
        if (resultSet != null) {
            resultSet.close();
            resultSet = null;
        }
        if (stmt != null) {
            stmt.close();
            stmt = null;
        }
        if (prepStmt != null) {
            prepStmt.close();
            prepStmt = null;
        }
    }

    /**
     * @param query
     *            The query to run
     * @return <b>true</b> if the first result is a ResultSet object;
     *         <b>false</b> if it is an update count or there are no results
     * @throws SQLException
     */
    public boolean executeQuery(String query) throws SQLException {
        closeStatement();

        checkConnected();

        if (logger.isDebugEnabled())
            logger.debug("ExecuteQuery '" + query + "' on " + getUrl());

        stmt = connection.createStatement();
        lastInsertID = 0L;

        if (isEnableResultSetStreaming()) {
            stmt.setFetchSize(Integer.MIN_VALUE);
        }
        boolean ret = stmt.execute(query, Statement.RETURN_GENERATED_KEYS);

        if (!ret) {
            // when stmt.execute returns false it means no result set is
            // expected get the number of rows affected
            rowCount = stmt.getUpdateCount();

            printLine(rowCount + " rows affected");

            ResultSet rs = stmt.getGeneratedKeys();
            if (rs.next()) {
                lastInsertID = rs.getLong(1);
            }
        } else {
            // a stmt returning a result set was run
            resultSet = stmt.getResultSet();
            if (useBufferedQuery)
                resultSet.setFetchSize(Integer.MAX_VALUE);
        }
        return ret;
    }

    /**
     * @param line
     */
    protected void printLine(String line) {
        // TODO - implement this?
    }

    public void executeFromStream(InputStream is) throws SQLException {
        final String cmdDelimeter = ";";
        final String comment = "--";
        String command = "";
        String separator = "";

        BufferedReader reader = null;
        try {
            reader = new BufferedReader(new InputStreamReader(is));

            String line = null;
            while ((line = reader.readLine()) != null) {

                if (line.equalsIgnoreCase("quit") || line.equalsIgnoreCase("exit")) {
                    break;
                }

                String trimmed = line.trim();
                if (trimmed.startsWith(comment))
                    continue;
                if (trimmed.endsWith(cmdDelimeter)) {
                    command = command + separator + line;
                    try {
                        executeQuery(command);
                    } catch (SQLException se) {
                        System.out.println(command);
                        throw se;
                    }
                    command = "";
                } else {
                    command = command + separator + line;
                    separator = " ";
                }
            }
        } catch (IOException e) {
            logger.error(e, e);
            e.printStackTrace();
        } finally {
            if (disconnectAfterProcess)
                this.disconnect();

            if (reader != null)
                try {
                    reader.close();
                } catch (Exception e) {
                    // do nothing
                }
        }
    }

    public int getRowCount() {
        return rowCount;
    }

    public long getLastInsertID() {
        return lastInsertID;
    }

    public ResultSet getResultSet() {
        return resultSet;
    }

    public void setResultSet(ResultSet resultSet) {
        this.resultSet = resultSet;
    }

    public void prepare(String query) throws SQLException {
        closeStatement();

        checkConnected();

        prepQuery = query;
        prepStmt = connection.prepareStatement(query);
    }

    public PreparedStatement prepareIndependent(String query) throws SQLException {
        if (resultSet != null) {
            resultSet.close();
            resultSet = null;
        }

        checkConnected();

        return connection.prepareStatement(query);
    }

    private boolean executePrepared(PreparedStatement ps, String query, List<Object> params)
            throws SQLException, PEException {

        if (ps == null) {
            throw new PEException("A prepared statement is not available to be executed - call prepare first");
        }

        if (logger.isDebugEnabled() && query != null)
            logger.debug("Command: " + query);

        for (int i = 0; i < params.size(); i++) {
            ps.setObject(i + 1, params.get(i));
        }

        boolean ret = ps.execute();
        if (!ret) {
            // when stmt.execute returns false it means no result set is
            // expected get the number of rows affected
            rowCount = ps.getUpdateCount();

            printLine(rowCount + " rows affected");

        } else {
            // a prepStmt returning a result set was run
            resultSet = ps.getResultSet();
            if (useBufferedQuery)
                resultSet.setFetchSize(Integer.MAX_VALUE);
        }
        return ret;

    }

    public boolean executePrepared(List<Object> params) throws SQLException, PEException {
        return executePrepared(prepStmt, prepQuery, params);
    }

    public boolean executePrepared(PreparedStatement ps, List<Object> params) throws SQLException, PEException {
        return executePrepared(ps, null, params);
    }

    public boolean isEnableResultSetStreaming() {
        return enableResultSetStreaming;
    }

    public void setEnableResultSetStreaming(boolean enableResultSetStreaming) {
        this.enableResultSetStreaming = enableResultSetStreaming;
    }

    public boolean useBufferedQuery() {
        return useBufferedQuery;
    }

    public void setBufferedQuery(boolean useBufferedQuery) {
        this.useBufferedQuery = useBufferedQuery;
    }

    public boolean disconnectAfterProcess() {
        return disconnectAfterProcess;
    }

    public void setDisconnectAfterProcess(boolean disconnectAfterProcess) {
        this.disconnectAfterProcess = disconnectAfterProcess;
    }

    public int getNumWarnings() throws SQLException {
        int count = 0;
        if (stmt != null) {
            SQLWarning warning = stmt.getWarnings();
            while (warning != null) {
                // TODO: filter for our test code until SHOW is completely fixed
                if (!StringUtils.equalsIgnoreCase(warning.getMessage(), "SHOW WARNINGS is not supported")) {
                    count++;
                }
                warning = warning.getNextWarning();
            }
            return count;
        }
        return 0;
    }

    private void checkConnected() throws SQLException {
        if (connection == null)
            throw new SQLException("Not connected");
    }
}