Mybatis源码之执行数据库脚本工具阅读

Mybatis里边有个ScriptRunner类，这个类是用来执行脚本的，我们可以直接用来调用去执行一些写好的脚本，就不用自己再写一套实现了，以下是源码，我已经加了相关的注释，相信大家一看就懂

/**
 *    Copyright 2009-2016 the original author or authors.
 *
 *    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.
 */
package org.apache.ibatis.jdbc;

import java.io.BufferedReader;
import java.io.PrintWriter;
import java.io.Reader;
import java.io.UnsupportedEncodingException;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.SQLWarning;
import java.sql.Statement;

/**
 * @author Clinton Begin
 */
public class ScriptRunner {

  /**
   * 行分隔符
   */
  private static final String LINE_SEPARATOR = System.getProperty("line.separator", "\n");

  /**
   * 默认分隔符
   */
  private static final String DEFAULT_DELIMITER = ";";

  /**
   * 数据库连接类
   */
  private Connection connection;

  /**
   * 发生错误是否停下
   */
  private boolean stopOnError;
  /**
   * 是否抛出警告信息
   */
  private boolean throwWarning;
  /**
   * 是否自动提交事务
   */
  private boolean autoCommit;
  /**
   * 是否批量发送脚本
   */
  private boolean sendFullScript;
  /**
   * 是否开启回车换行替换，相当于replaceAll("\r\n", "\n");
   */
  private boolean removeCRs;
  /**
   * 是否开启转意替换
   */
  private boolean escapeProcessing = true;

  /**
   * 日志记录
   */
  private PrintWriter logWriter = new PrintWriter(System.out);
  private PrintWriter errorLogWriter = new PrintWriter(System.err);

  private String delimiter = DEFAULT_DELIMITER;
  private boolean fullLineDelimiter = false;

  public ScriptRunner(Connection connection) {
    this.connection = connection;
  }

  public void setStopOnError(boolean stopOnError) {
    this.stopOnError = stopOnError;
  }

  public void setThrowWarning(boolean throwWarning) {
    this.throwWarning = throwWarning;
  }

  public void setAutoCommit(boolean autoCommit) {
    this.autoCommit = autoCommit;
  }

  public void setSendFullScript(boolean sendFullScript) {
    this.sendFullScript = sendFullScript;
  }

  public void setRemoveCRs(boolean removeCRs) {
    this.removeCRs = removeCRs;
  }

  /**
   * @since 3.1.1
   */
  public void setEscapeProcessing(boolean escapeProcessing) {
    this.escapeProcessing = escapeProcessing;
  }

  public void setLogWriter(PrintWriter logWriter) {
    this.logWriter = logWriter;
  }

  public void setErrorLogWriter(PrintWriter errorLogWriter) {
    this.errorLogWriter = errorLogWriter;
  }

  public void setDelimiter(String delimiter) {
    this.delimiter = delimiter;
  }

  public void setFullLineDelimiter(boolean fullLineDelimiter) {
    this.fullLineDelimiter = fullLineDelimiter;
  }

  /**
   * 脚本执行
   * @param reader
     */
  public void runScript(Reader reader) {
    //设置事务提交方式
    setAutoCommit();

    try {
      //是否批量执行脚本
      if (sendFullScript) {
        //批量执行脚本
        executeFullScript(reader);
      } else {
        //一句一句执行脚本
        executeLineByLine(reader);
      }
    } finally {
      //释放连接
      rollbackConnection();
    }
  }

  /**
   * 批量执行脚本
   * @param reader        源
     */
  private void executeFullScript(Reader reader) {
    StringBuilder script = new StringBuilder();
    try {
      BufferedReader lineReader = new BufferedReader(reader);
      String line;
      while ((line = lineReader.readLine()) != null) {
        script.append(line);
        script.append(LINE_SEPARATOR);
      }
      String command = script.toString();
      println(command);
      executeStatement(command);
      commitConnection();
    } catch (Exception e) {
      String message = "Error executing: " + script + ".  Cause: " + e;
      printlnError(message);
      throw new RuntimeSqlException(message, e);
    }
  }

  /**
   * 逐句执行脚本
   * @param reader        源
     */
  private void executeLineByLine(Reader reader) {
    StringBuilder command = new StringBuilder();
    try {
      BufferedReader lineReader = new BufferedReader(reader);
      String line;
      while ((line = lineReader.readLine()) != null) {
        /**
         * 脚本处理-具体逻辑看实现
         */
        command = handleLine(command, line);
      }
      commitConnection();
      //校验是否有结尾标识
      checkForMissingLineTerminator(command);
    } catch (Exception e) {
      String message = "Error executing: " + command + ".  Cause: " + e;
      printlnError(message);
      throw new RuntimeSqlException(message, e);
    }
  }

  public void closeConnection() {
    try {
      connection.close();
    } catch (Exception e) {
      // ignore
    }
  }

  private void setAutoCommit() {
    try {
      if (autoCommit != connection.getAutoCommit()) {
        connection.setAutoCommit(autoCommit);
      }
    } catch (Throwable t) {
      throw new RuntimeSqlException("Could not set AutoCommit to " + autoCommit + ". Cause: " + t, t);
    }
  }

  private void commitConnection() {
    try {
      if (!connection.getAutoCommit()) {
        connection.commit();
      }
    } catch (Throwable t) {
      throw new RuntimeSqlException("Could not commit transaction. Cause: " + t, t);
    }
  }

  private void rollbackConnection() {
    try {
      if (!connection.getAutoCommit()) {
        connection.rollback();
      }
    } catch (Throwable t) {
      // ignore
    }
  }

  private void checkForMissingLineTerminator(StringBuilder command) {
    if (command != null && command.toString().trim().length() > 0) {
      throw new RuntimeSqlException("Line missing end-of-line terminator (" + delimiter + ") => " + command);
    }
  }

  /**
   * 处理命令
   * @param command
   * @param line
   * @return
   * @throws SQLException
   * @throws UnsupportedEncodingException
     */
  private StringBuilder handleLine(StringBuilder command, String line) throws SQLException, UnsupportedEncodingException {
    String trimmedLine = line.trim();
    if (lineIsComment(trimmedLine)) {
        //注释处理
        final String cleanedString = trimmedLine.substring(2).trim().replaceFirst("//", "");
        if(cleanedString.toUpperCase().startsWith("@DELIMITER")) {
            delimiter = cleanedString.substring(11,12);
            return command;
        }
      println(trimmedLine);
    } else if (commandReadyToExecute(trimmedLine)) {
      //是否是语句结尾处理，是则将脚本执行，否则继续拼接直到遇到结束符
      command.append(line.substring(0, line.lastIndexOf(delimiter)));
      command.append(LINE_SEPARATOR);
      println(command);
      executeStatement(command.toString());
      command.setLength(0);
    } else if (trimmedLine.length() > 0) {
      //不是语句结尾，继续拼装
      command.append(line);
      command.append(LINE_SEPARATOR);
    }
    return command;
  }

  private boolean lineIsComment(String trimmedLine) {
    return trimmedLine.startsWith("//") || trimmedLine.startsWith("--");
  }

  private boolean commandReadyToExecute(String trimmedLine) {
    // issue #561 remove anything after the delimiter
    return !fullLineDelimiter && trimmedLine.contains(delimiter) || fullLineDelimiter && trimmedLine.equals(delimiter);
  }

  /**
   * 执行数据库脚本
   * @param command           脚本
   * @throws SQLException
     */
  private void executeStatement(String command) throws SQLException {
    boolean hasResults = false;
    Statement statement = connection.createStatement();
    //是否开启转义
    statement.setEscapeProcessing(escapeProcessing);
    String sql = command;
    //换行回车转换
    if (removeCRs) {
      sql = sql.replaceAll("\r\n", "\n");
    }
    if (stopOnError) {
      //如果执行期间遇到了错误则停止
      hasResults = statement.execute(sql);
      if (throwWarning) {
        // In Oracle, CRATE PROCEDURE, FUNCTION, etc. returns warning
        // instead of throwing exception if there is compilation error.
        SQLWarning warning = statement.getWarnings();
        if (warning != null) {
          throw warning;
        }
      }
    } else {
      //如果执行期间遇到了错误则继续执行并抛出异常
      try {
        hasResults = statement.execute(sql);
      } catch (SQLException e) {
        String message = "Error executing: " + command + ".  Cause: " + e;
        printlnError(message);
      }
    }
    //打印执行结果
    printResults(statement, hasResults);
    try {
      statement.close();
    } catch (Exception e) {
      // Ignore to workaround a bug in some connection pools
    }
  }

  private void printResults(Statement statement, boolean hasResults) {
    try {
      if (hasResults) {
        ResultSet rs = statement.getResultSet();
        if (rs != null) {
          ResultSetMetaData md = rs.getMetaData();
          int cols = md.getColumnCount();
          for (int i = 0; i < cols; i++) {
            String name = md.getColumnLabel(i + 1);
            print(name + "\t");
          }
          println("");
          while (rs.next()) {
            for (int i = 0; i < cols; i++) {
              String value = rs.getString(i + 1);
              print(value + "\t");
            }
            println("");
          }
        }
      }
    } catch (SQLException e) {
      printlnError("Error printing results: " + e.getMessage());
    }
  }

  private void print(Object o) {
    if (logWriter != null) {
      logWriter.print(o);
      logWriter.flush();
    }
  }

  private void println(Object o) {
    if (logWriter != null) {
      logWriter.println(o);
      logWriter.flush();
    }
  }

  private void printlnError(Object o) {
    if (errorLogWriter != null) {
      errorLogWriter.println(o);
      errorLogWriter.flush();
    }
  }

}

以下代码是我写的测试类，可以完美的执行

package zxltest;

import org.apache.ibatis.datasource.unpooled.UnpooledDataSource;
import org.apache.ibatis.io.Resources;
import org.apache.ibatis.jdbc.ScriptRunner;

import javax.sql.DataSource;
import java.io.IOException;
import java.io.Reader;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Properties;

/**
 * Created by zxl on 2016/11/25.
 */
public class ScriptRunnerTest {

    public static final String MYSQL_PROPERTIES = "zxltest/databases/mysql/mysql-derby.properties";
    public static final String MYSQL_DDL = "zxltest/databases/mysql/mysql-derby-schema.sql";
    public static final String MYSQL_DATA_LOAD = "zxltest/databases/mysql/mysql-derby-dataload.sql";
    public static final String MYSQL_DATA_QUERY = "zxltest/databases/mysql/mysql-derby-dataquery.sql";

    /**
     * 创建无连接池的数据源
     * @param resource
     * @return
     * @throws IOException
     */
    public static UnpooledDataSource createUnpooledDataSource(String resource) throws IOException {
        Properties props = Resources.getResourceAsProperties(resource);
        UnpooledDataSource ds = new UnpooledDataSource();
        ds.setDriver(props.getProperty("driver"));
        ds.setUrl(props.getProperty("url"));
        ds.setUsername(props.getProperty("username"));
        ds.setPassword(props.getProperty("password"));
        return ds;
    }

    /**
     * 运行脚本
     * @param ds
     * @param resource
     * @throws IOException
     * @throws SQLException
     */
    public static void runScript(DataSource ds, String resource) throws IOException, SQLException {
        Connection connection = ds.getConnection();
        try {
            ScriptRunner runner = new ScriptRunner(connection);
            runner.setAutoCommit(true);
            runner.setStopOnError(true);
            runner.setRemoveCRs(true);
//            runner.setLogWriter(null);
//            runner.setSendFullScript(true);
//            runner.setErrorLogWriter(null);
            runScript(runner, resource);
        } finally {
            connection.close();
        }
    }

    /**
     * 批量执行sql 操作数据
     * @param ds
     * @param resource
     * @throws IOException
     * @throws SQLException
     */
    public static void runFullScript(DataSource ds, String resource) throws IOException, SQLException {
        Connection connection = ds.getConnection();
        try {
            ScriptRunner runner = new ScriptRunner(connection);
            runner.setAutoCommit(true);
            runner.setStopOnError(true);
            runner.setRemoveCRs(true);
            runner.setLogWriter(null);
            runner.setSendFullScript(true);
//            runner.setErrorLogWriter(null);
            runScript(runner, resource);
        } finally {
            connection.close();
        }
    }

    /**
     * 运行脚本
     * @param runner
     * @param resource
     * @throws IOException
     * @throws SQLException
     */
    public static void runScript(ScriptRunner runner, String resource) throws IOException, SQLException {
        Reader reader = Resources.getResourceAsReader(resource);
        try {
            runner.runScript(reader);
        } finally {
            reader.close();
        }
    }

    public static void main(String[] args) {
        DataSource ds = null;
        try {
            ds = createUnpooledDataSource(MYSQL_PROPERTIES);
            runScript(ds, MYSQL_DDL);
            runScript(ds, MYSQL_DATA_LOAD);
            runScript(ds, MYSQL_DATA_QUERY);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}




mysql-derby.properties-->数据库连接属性文件

driver=com.mysql.jdbc.Driver
url=jdbc:mysql://127.0.0.1:3306/testdb?useUnicode=true&characterEncoding=utf-8&allowMultiQueries=true
username=root
password=root


mysql-derby-schema.sql-->数据库DLL脚本文件

--
-- 测试sql

DROP TABLE IF EXISTS tmp_user;

CREATE TABLE tmp_user (
id                INT(255) NOT NULL AUTO_INCREMENT COMMENT '主键',
username          VARCHAR(255) NOT NULL,
password          VARCHAR(255) NOT NULL,
email             VARCHAR(255) NOT NULL,
PRIMARY KEY (id)
);



mysql-derby-dataload.sql-->数据库插入数据脚本


INSERT INTO tmp_user (id,username, password, email) VALUES (101,'jim','********','jim@ibatis.apache.org');
INSERT INTO tmp_user (id,username, password, email) VALUES (102,'sally','********','sally@ibatis.apache.org');



mysql-derby-dataquery.sql-->数据库查询脚本文件

SELECT * FROM tmp_user;

主要逻辑：创建一个无池子的数据库链接，new 一个ScriptRunner实体类，将一些基本属性进行设置，读取脚本属性文件得到Reader实体，然后调用

runner.runScript(reader);即可执行脚本