org.codehaus.mojo.dbupgrade.sqlexec.DefaultSQLExec.java Source code

Java tutorial

Introduction

Here is the source code for org.codehaus.mojo.dbupgrade.sqlexec.DefaultSQLExec.java

Source

package org.codehaus.mojo.dbupgrade.sqlexec;

import java.io.BufferedOutputStream;
import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.PrintStream;
import java.io.Reader;
import java.io.StringReader;
import java.sql.Connection;
import java.sql.Driver;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.SQLWarning;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Properties;
import java.util.StringTokenizer;

import org.apache.commons.dbutils.DbUtils;
import org.apache.commons.io.IOUtils;
import org.codehaus.plexus.util.IOUtil;
import org.codehaus.plexus.util.StringUtils;

/*
 * Copyright 2000-2010 The Apache Software Foundation
 * 
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 * 
 * http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software distributed under the License
 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 * or implied. See the License for the specific language governing permissions and limitations under
 * the License.
 */

/**
 * Executes SQL against a database. Extracted from sql-maven-plugin 1.3
 */
public class DefaultSQLExec implements SQLExec {
    private SQLExecConfig config;

    private PrintStream outLog = System.out;

    public DefaultSQLExec(SQLExecConfig config) {
        this.config = config;
    }

    // //////////////////////////////// Internal properties//////////////////////

    private int successfulStatements = 0;

    private int totalStatements = 0;

    /**
     * Database connection
     */
    private Connection conn = null;

    /**
     * SQL statement
     */
    private Statement statement = null;

    /**
     * Add a SQL transaction to execute
     */
    private Transaction createTransaction() {
        Transaction t = new Transaction();
        // transactions.add( t );
        return t;
    }

    /**
     * Add sql command to transactions list.
     *
     */
    private List<Transaction> addCommandToTransactions(List<Transaction> transactions, String sqlCommand) {
        Transaction t = createTransaction();

        t.addText(sqlCommand);

        transactions.add(t);

        return transactions;

    }

    /**
     * Add user sql fileset to transation list
     *
     */
    private List<Transaction> addFileSetToTransactions(List<Transaction> transactions, FileSet fileset) {
        String[] includedFiles;
        if (fileset != null) {
            fileset.scan();
            includedFiles = fileset.getIncludedFiles();
        } else {
            includedFiles = new String[0];
        }

        for (int j = 0; j < includedFiles.length; j++) {
            Transaction t = createTransaction();
            t.setSrc(new File(fileset.getBasedir(), includedFiles[j]));
            transactions.add(t);
        }

        return transactions;
    }

    /**
     * Add user input of srcFiles to transaction list.
     *
     * @throws SQLException
     */
    private List<Transaction> addFilesToTransactions(List<Transaction> transactions, File[] files)
            throws SQLException {
        for (int i = 0; files != null && i < files.length; ++i) {
            if (files[i] != null && !files[i].exists()) {
                throw new SQLException(files[i].getPath() + " not found.");
            }

            Transaction t = createTransaction();
            t.setSrc(files[i]);
            transactions.add(t);
        }

        return transactions;
    }

    /**
     * Sort the transaction list.
     */
    private void sortTransactions(List<Transaction> transactions) {
        if (SQLExecConfig.FILE_SORTING_ASC.equalsIgnoreCase(this.config.getOrderFile())) {
            Collections.sort(transactions);
        } else if (SQLExecConfig.FILE_SORTING_DSC.equalsIgnoreCase(this.config.getOrderFile())) {
            Collections.sort(transactions, Collections.reverseOrder());
        }
    }

    private void handleWindowsDomainUser(Properties driverProperties) {
        if ("net.sourceforge.jtds.jdbc.Driver".equals(config.getDriver())) {
            String[] tokens = StringUtils.split(config.getUsername(), "\\");

            if (tokens != null && tokens.length == 2) {
                driverProperties.put("user", tokens[1]);
                driverProperties.put("domain", tokens[0]);
            }
        }
    }

    private Driver createJDBCDriver() {
        Driver driverInstance = null;

        try {
            Class<?> dc = Class.forName(config.getDriver());
            driverInstance = (Driver) dc.newInstance();
        } catch (ClassNotFoundException e) {
            throw new RuntimeException("Driver class not found: " + config.getDriver(), e);
        } catch (Exception e) {
            throw new RuntimeException("Failure loading driver: " + config.getDriver(), e);
        }

        return driverInstance;
    }

    private Connection createConnection(Driver driverInstance, Properties driverProperties) throws SQLException {
        Connection connection = null;

        for (int i = 0; i < config.getConnectionRetries(); ++i) {
            try {
                connection = driverInstance.connect(config.getUrl(), driverProperties);
                if (connection == null) {
                    // Driver doesn't understand the URL
                    throw new RuntimeException("No suitable Driver for " + config.getUrl());
                }

                break;

            } catch (SQLException e) {
                if (i < config.getConnectionRetries()) {
                    try {
                        Thread.sleep(config.getConnectionRetryDelay());
                    } catch (Exception iex) {
                        throw new SQLException("Unable to connect to " + config.getUrl(), iex);
                    }
                    continue;
                }

                throw new SQLException("Unable to connect to " + config.getUrl(), e);
            }

        }

        return connection;

    }

    /**
     * parse driverProperties into Properties set
     *
     * @return
     * @throws SQLException
     */
    protected Properties getDriverProperties() {
        //set as protected scopy for unit test purpose
        Properties properties = new Properties();

        if (!StringUtils.isEmpty(this.config.getDriverProperties())) {
            String[] tokens = StringUtils.split(this.config.getDriverProperties(), ",");
            for (int i = 0; i < tokens.length; ++i) {
                String[] keyValueTokens = StringUtils.split(tokens[i].trim(), "=");
                if (keyValueTokens.length != 2) {
                    throw new RuntimeException(
                            "Invalid JDBC Driver properties: " + this.config.getDriverProperties());
                }

                properties.setProperty(keyValueTokens[0], keyValueTokens[1]);

            }
        }

        return properties;
    }

    /**
     * read in lines and execute them
     */
    private void runStatements(Reader reader, PrintStream out) throws SQLException, IOException {
        String line;

        StringBuffer sql = new StringBuffer();

        BufferedReader in = new BufferedReader(reader);

        while ((line = in.readLine()) != null) {
            if (!config.isKeepFormat()) {
                line = line.trim();
            }

            if (!config.isKeepFormat()) {
                if (line.startsWith("#")) {
                    continue;
                }
                if (line.startsWith("//")) {
                    continue;
                }
                if (line.startsWith("--")) {
                    continue;
                }
                StringTokenizer st = new StringTokenizer(line);
                if (st.hasMoreTokens()) {
                    String token = st.nextToken();
                    if ("REM".equalsIgnoreCase(token)) {
                        continue;
                    }
                }
            }

            if (!config.isKeepFormat()) {
                sql.append(" ").append(line);
            } else {
                sql.append("\n").append(line);
            }

            // SQL defines "--" as a comment to EOL
            // and in Oracle it may contain a hint
            // so we cannot just remove it, instead we must end it
            if (!config.isKeepFormat()) {
                if (SqlSplitter.containsSqlEnd(line, config.getDelimiter()) == SqlSplitter.NO_END) {
                    sql.append("\n");
                }
            }

            DelimiterType delimiterType = this.config.getDelimiterType();
            String delimiter = this.config.getDelimiter();

            if ((delimiterType.equals(DelimiterType.NORMAL) && SqlSplitter.containsSqlEnd(line, delimiter) > 0)
                    || (delimiterType.equals(DelimiterType.ROW) && line.trim().equals(delimiter))) {
                execSQL(sql.substring(0, sql.length() - delimiter.length()), out);
                sql.setLength(0); // clean buffer
            }
        }

        // Catch any statements not followed by ;
        if (!sql.toString().equals("")) {
            execSQL(sql.toString(), out);
        }
    }

    /**
     * Exec the sql statement.
     */
    private void execSQL(String sql, PrintStream out) throws SQLException {
        // Check and ignore empty statements
        if ("".equals(sql.trim())) {
            return;
        }

        if (config.isVerbose()) {
            out.append(sql).append("\n");
        }

        ResultSet resultSet = null;
        try {
            totalStatements++;

            boolean ret;
            int updateCount, updateCountTotal = 0;

            ret = statement.execute(sql);
            updateCount = statement.getUpdateCount();
            resultSet = statement.getResultSet();
            do {
                if (!ret) {
                    if (updateCount != -1) {
                        updateCountTotal += updateCount;
                    }
                } else {
                    if (config.isPrintResultSet()) {
                        printResultSet(resultSet, out);
                    }
                }
                ret = statement.getMoreResults();
                if (ret) {
                    updateCount = statement.getUpdateCount();
                    resultSet = statement.getResultSet();
                }
            } while (ret);

            if (config.isPrintResultSet()) {
                StringBuffer line = new StringBuffer();
                line.append(updateCountTotal).append(" rows affected");
                out.println(line);
            }

            SQLWarning warning = conn.getWarnings();
            while (warning != null) {
                warning = warning.getNextWarning();
            }
            conn.clearWarnings();
            successfulStatements++;
        } catch (SQLException e) {
            if (SQLExecConfig.ON_ERROR_ABORT.equalsIgnoreCase(config.getOnError())) {
                throw new SQLException("Unable to execute: " + sql, e);
            }
        } finally {
            if (resultSet != null) {
                resultSet.close();
            }
        }
    }

    /**
     * print any results in the result set.
     *
     * @param rs the resultset to print information about
     * @param out the place to print results
     * @throws SQLException on SQL problems.
     */
    private void printResultSet(ResultSet rs, PrintStream out) throws SQLException {
        if (rs != null) {
            ResultSetMetaData md = rs.getMetaData();
            int columnCount = md.getColumnCount();
            StringBuffer line = new StringBuffer();
            if (config.isShowheaders()) {
                for (int col = 1; col < columnCount; col++) {
                    line.append(md.getColumnName(col));
                    line.append(",");
                }
                line.append(md.getColumnName(columnCount));
                out.println(line);
                line = new StringBuffer();
            }
            while (rs.next()) {
                boolean first = true;
                for (int col = 1; col <= columnCount; col++) {
                    String columnValue = rs.getString(col);
                    if (columnValue != null) {
                        columnValue = columnValue.trim();
                    }

                    if (first) {
                        first = false;
                    } else {
                        line.append(",");
                    }
                    line.append(columnValue);
                }
                out.println(line);
                line = new StringBuffer();
            }
        }
        out.println();
    }

    /**
     * Contains the definition of a new transaction element. Transactions allow several files or
     * blocks of statements to be executed using the same JDBC connection and commit operation in
     * between.
     */
    private class Transaction implements Comparable<Object> {
        private File tSrcFile = null;

        private String tSqlCommand = "";

        /**
         *
         */
        public void setSrc(File src) {
            this.tSrcFile = src;
        }

        /**
         *
         */
        public void addText(String sql) {
            if (!StringUtils.isBlank(sql)) {
                this.tSqlCommand += sql;
            }
        }

        /**
         *
         */
        private void runTransaction(PrintStream out) throws IOException, SQLException {
            if (tSqlCommand.length() != 0) {
                runStatements(new StringReader(tSqlCommand), out);
            }

            if (tSrcFile != null) {
                Reader reader = null;

                if (StringUtils.isEmpty(config.getEncoding())) {
                    reader = new FileReader(tSrcFile);
                } else {
                    reader = new InputStreamReader(new FileInputStream(tSrcFile), config.getEncoding());
                }

                try {
                    runStatements(reader, out);
                } finally {
                    reader.close();
                }
            }
        }

        public int compareTo(Object object) {
            Transaction transaction = (Transaction) object;

            if (transaction.tSrcFile == null) {
                if (this.tSrcFile == null) {
                    return 0;
                } else {
                    return Integer.MAX_VALUE;
                }
            } else {
                if (this.tSrcFile == null) {
                    return Integer.MIN_VALUE;
                } else {
                    return this.tSrcFile.compareTo(transaction.tSrcFile);
                }
            }
        }
    }

    // ////////////////////////////////////////////////////////////////////////////////////////////

    /**
     * Load the sql file and then it
     */
    private void execute(Connection conn, List<Transaction> transactions) throws SQLException {

        try {
            statement = conn.createStatement();
            statement.setEscapeProcessing(config.isEscapeProcessing());

            PrintStream out = System.out;
            try {
                if (config.getOutputFile() != null) {
                    out = new PrintStream(new BufferedOutputStream(
                            new FileOutputStream(config.getOutputFile().getAbsolutePath(), config.isAppend())));
                }

                // Process all transactions
                for (Transaction t : transactions) {
                    t.runTransaction(out);

                    if (!config.isAutocommit()) {
                        conn.commit();
                    }
                }

            } finally {
                if (out != null && out != System.out) {
                    out.close();
                }
            }
        } catch (IOException e) {
            throw new SQLException(e.getMessage(), e);
        } catch (SQLException e) {
            if (!config.isAutocommit() && conn != null
                    && SQLExecConfig.ON_ERROR_ABORT.equalsIgnoreCase(config.getOnError())) {
                this.rollbackQuietly();
            }
            throw new SQLException(e.getMessage(), e);
        } finally {
            DbUtils.closeQuietly(statement);
        }

        if (SQLExecConfig.ON_ERROR_ABORT_AFTER.equalsIgnoreCase(config.getOnError())
                && totalStatements != successfulStatements) {
            throw new SQLException("Some SQL statements failed to execute");
        }

    }

    ///////////////////////////////////////////////////////////////////////////////////////
    ///////////////////////////////////////////////////////////////////////////////////////

    /**
     * Creates a new Connection as using the driver, url, userid and password specified.
     *
     * The calling method is responsible for closing the connection.
     *
     * @return Connection the newly created connection.
     * @throws SQLException if the UserId/Password/Url is not set or there is no suitable driver
     *             or the driver fails to load.
     * @throws SQLException if there is problem getting connection with valid url
     *
     */
    public Connection getConnection() throws SQLException {
        if (conn != null) {
            return conn;
        }

        Properties driverProperties = new Properties();

        driverProperties.put("user", config.getUsername());

        handleWindowsDomainUser(driverProperties);

        if (!config.isEnableAnonymousPassword()) {
            if (!StringUtils.isBlank(this.config.getPassword())) {
                driverProperties.put("password", this.config.getPassword());
            }
        }

        driverProperties.putAll(this.getDriverProperties());

        Driver driverInstance = this.createJDBCDriver();

        conn = this.createConnection(driverInstance, driverProperties);

        conn.setAutoCommit(config.isAutocommit());

        return conn;
    }

    public void execute(String sqlCommand) throws SQLException {
        execute(sqlCommand, new File[0], null);
    }

    public void execute(File[] srcFiles) throws SQLException {
        execute(null, srcFiles, null);
    }

    public void execute(FileSet fileset) throws SQLException {
        execute(null, null, fileset);
    }

    public void execute(String sqlCommand, File[] srcFiles, FileSet fileset) throws SQLException {

        List<Transaction> transactions = new ArrayList<Transaction>();

        successfulStatements = 0;

        totalStatements = 0;

        if (!StringUtils.isBlank(sqlCommand)) {
            addCommandToTransactions(transactions, sqlCommand);
        }

        if (srcFiles != null) {
            addFilesToTransactions(transactions, srcFiles);
        }

        if (fileset != null) {
            addFileSetToTransactions(transactions, fileset);
        }

        sortTransactions(transactions);

        for (int i = 0; i < config.getRepeats() / config.getTransactionsPerConnection(); ++i) {
            Connection connection = null;
            try {
                connection = getConnection();
            } catch (SQLException e) {
                if (!config.isSkipOnConnectionError()) {
                    throw new SQLException(e.getMessage(), e);
                } else {
                    // error on get connection and user asked to skip the rest
                    break;
                }
            }

            for (int j = 0; j < config.getTransactionsPerConnection(); ++j) {
                try {
                    this.execute(connection, transactions);
                    Thread.sleep(this.config.getSleepTimeBetweenRepeats());
                } catch (InterruptedException e) {
                }
            }
        }

    }

    public void execute(Reader reader) throws SQLException {
        try {
            statement = this.getConnection().createStatement();
            statement.setEscapeProcessing(config.isEscapeProcessing());

            this.runStatements(reader, outLog);
        } catch (IOException e) {
            throw new SQLException("Error reading SQL stream: " + e.getMessage(), e);
        } finally {
            DbUtils.closeQuietly(statement);
            statement = null;
        }

    }

    public void execute(File sqlFile) throws SQLException {
        Reader reader = null;
        try {
            reader = new FileReader(sqlFile);
            execute(reader);
        } catch (IOException e) {
            throw new SQLException("Error reading SQL stream: " + e.getMessage(), e);
        } finally {
            IOUtil.close(reader);
        }

    }

    public void execute(File sqlFile, boolean disableSQLParser) throws SQLException, IOException {
        if (!disableSQLParser) {
            this.execute(sqlFile);
        } else {
            InputStream is = null;
            Statement statement = null;
            String sql = null;

            try {
                is = new FileInputStream(sqlFile);
                sql = IOUtils.toString(is);

                statement = getConnection().createStatement();
                statement.setEscapeProcessing(false);

                if (statement.execute(sql)) {
                    //we expect a false return since the execution has no result set
                    throw new SQLException("Unable execute SQL Statement:" + sql);
                }

            } finally {
                DbUtils.closeQuietly(statement);
                IOUtils.closeQuietly(is);
            }
        }

    }

    public void commit() throws SQLException {
        if (!this.config.isAutocommit()) {
            this.getConnection().commit();
        }
    }

    public void rollback() {
        try {
            if (!this.config.isAutocommit()) {
                this.getConnection().rollback();
            }
        } catch (SQLException e) {
            //unexpected exception, throw runtime to get more attention
            throw new RuntimeException("Unable to rollback.");
        }
    }

    public void rollbackQuietly() {
        try {
            this.getConnection().rollback();
        } catch (SQLException e) {

        }
    }

    public void close() {
        DbUtils.closeQuietly(this.conn);
        this.conn = null;
    }

    public void execute(InputStream istream) throws SQLException {
        Reader reader = new InputStreamReader(istream);
        this.execute(reader);
    }

    /**
     * Number of SQL statements executed so far that caused errors.
     *
     * @return the number
     */
    public int getSuccessfulStatements() {
        return successfulStatements;
    }

    /**
     * Number of SQL statements executed so far, including the ones that caused errors.
     *
     * @return the number
     */
    public int getTotalStatements() {
        return totalStatements;
    }

}