ScriptRunner.java

  1. /*
  2.  *    Copyright 2009-2021 the original author or authors.
  3.  *
  4.  *    Licensed under the Apache License, Version 2.0 (the "License");
  5.  *    you may not use this file except in compliance with the License.
  6.  *    You may obtain a copy of the License at
  7.  *
  8.  *       http://www.apache.org/licenses/LICENSE-2.0
  9.  *
  10.  *    Unless required by applicable law or agreed to in writing, software
  11.  *    distributed under the License is distributed on an "AS IS" BASIS,
  12.  *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13.  *    See the License for the specific language governing permissions and
  14.  *    limitations under the License.
  15.  */
  16. package org.apache.ibatis.jdbc;

  17. import java.io.BufferedReader;
  18. import java.io.PrintWriter;
  19. import java.io.Reader;
  20. import java.sql.Connection;
  21. import java.sql.ResultSet;
  22. import java.sql.ResultSetMetaData;
  23. import java.sql.SQLException;
  24. import java.sql.SQLWarning;
  25. import java.sql.Statement;
  26. import java.util.regex.Matcher;
  27. import java.util.regex.Pattern;

  28. /**
  29.  * This is an internal testing utility.<br>
  30.  * You are welcome to use this class for your own purposes,<br>
  31.  * but if there is some feature/enhancement you need for your own usage,<br>
  32.  * please make and modify your own copy instead of sending us an enhancement request.<br>
  33.  *
  34.  * @author Clinton Begin
  35.  */
  36. public class ScriptRunner {

  37.   private static final String LINE_SEPARATOR = System.lineSeparator();

  38.   private static final String DEFAULT_DELIMITER = ";";

  39.   private static final Pattern DELIMITER_PATTERN = Pattern.compile("^\\s*((--)|(//))?\\s*(//)?\\s*@DELIMITER\\s+([^\\s]+)", Pattern.CASE_INSENSITIVE);

  40.   private final Connection connection;

  41.   private boolean stopOnError;
  42.   private boolean throwWarning;
  43.   private boolean autoCommit;
  44.   private boolean sendFullScript;
  45.   private boolean removeCRs;
  46.   private boolean escapeProcessing = true;

  47.   private PrintWriter logWriter = new PrintWriter(System.out);
  48.   private PrintWriter errorLogWriter = new PrintWriter(System.err);

  49.   private String delimiter = DEFAULT_DELIMITER;
  50.   private boolean fullLineDelimiter;

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

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

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

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

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

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

  69.   /**
  70.    * Sets the escape processing.
  71.    *
  72.    * @param escapeProcessing
  73.    *          the new escape processing
  74.    * @since 3.1.1
  75.    */
  76.   public void setEscapeProcessing(boolean escapeProcessing) {
  77.     this.escapeProcessing = escapeProcessing;
  78.   }

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

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

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

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

  91.   public void runScript(Reader reader) {
  92.     setAutoCommit();

  93.     try {
  94.       if (sendFullScript) {
  95.         executeFullScript(reader);
  96.       } else {
  97.         executeLineByLine(reader);
  98.       }
  99.     } finally {
  100.       rollbackConnection();
  101.     }
  102.   }

  103.   private void executeFullScript(Reader reader) {
  104.     StringBuilder script = new StringBuilder();
  105.     try {
  106.       BufferedReader lineReader = new BufferedReader(reader);
  107.       String line;
  108.       while ((line = lineReader.readLine()) != null) {
  109.         script.append(line);
  110.         script.append(LINE_SEPARATOR);
  111.       }
  112.       String command = script.toString();
  113.       println(command);
  114.       executeStatement(command);
  115.       commitConnection();
  116.     } catch (Exception e) {
  117.       String message = "Error executing: " + script + ".  Cause: " + e;
  118.       printlnError(message);
  119.       throw new RuntimeSqlException(message, e);
  120.     }
  121.   }

  122.   private void executeLineByLine(Reader reader) {
  123.     StringBuilder command = new StringBuilder();
  124.     try {
  125.       BufferedReader lineReader = new BufferedReader(reader);
  126.       String line;
  127.       while ((line = lineReader.readLine()) != null) {
  128.         handleLine(command, line);
  129.       }
  130.       commitConnection();
  131.       checkForMissingLineTerminator(command);
  132.     } catch (Exception e) {
  133.       String message = "Error executing: " + command + ".  Cause: " + e;
  134.       printlnError(message);
  135.       throw new RuntimeSqlException(message, e);
  136.     }
  137.   }

  138.   /**
  139.    * @deprecated Since 3.5.4, this method is deprecated. Please close the {@link Connection} outside of this class.
  140.    */
  141.   @Deprecated
  142.   public void closeConnection() {
  143.     try {
  144.       connection.close();
  145.     } catch (Exception e) {
  146.       // ignore
  147.     }
  148.   }

  149.   private void setAutoCommit() {
  150.     try {
  151.       if (autoCommit != connection.getAutoCommit()) {
  152.         connection.setAutoCommit(autoCommit);
  153.       }
  154.     } catch (Throwable t) {
  155.       throw new RuntimeSqlException("Could not set AutoCommit to " + autoCommit + ". Cause: " + t, t);
  156.     }
  157.   }

  158.   private void commitConnection() {
  159.     try {
  160.       if (!connection.getAutoCommit()) {
  161.         connection.commit();
  162.       }
  163.     } catch (Throwable t) {
  164.       throw new RuntimeSqlException("Could not commit transaction. Cause: " + t, t);
  165.     }
  166.   }

  167.   private void rollbackConnection() {
  168.     try {
  169.       if (!connection.getAutoCommit()) {
  170.         connection.rollback();
  171.       }
  172.     } catch (Throwable t) {
  173.       // ignore
  174.     }
  175.   }

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

  181.   private void handleLine(StringBuilder command, String line) throws SQLException {
  182.     String trimmedLine = line.trim();
  183.     if (lineIsComment(trimmedLine)) {
  184.       Matcher matcher = DELIMITER_PATTERN.matcher(trimmedLine);
  185.       if (matcher.find()) {
  186.         delimiter = matcher.group(5);
  187.       }
  188.       println(trimmedLine);
  189.     } else if (commandReadyToExecute(trimmedLine)) {
  190.       command.append(line, 0, line.lastIndexOf(delimiter));
  191.       command.append(LINE_SEPARATOR);
  192.       println(command);
  193.       executeStatement(command.toString());
  194.       command.setLength(0);
  195.     } else if (trimmedLine.length() > 0) {
  196.       command.append(line);
  197.       command.append(LINE_SEPARATOR);
  198.     }
  199.   }

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

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

  207.   private void executeStatement(String command) throws SQLException {
  208.     try (Statement statement = connection.createStatement()) {
  209.       statement.setEscapeProcessing(escapeProcessing);
  210.       String sql = command;
  211.       if (removeCRs) {
  212.         sql = sql.replace("\r\n", "\n");
  213.       }
  214.       try {
  215.         boolean hasResults = statement.execute(sql);
  216.         while (!(!hasResults && statement.getUpdateCount() == -1)) {
  217.           checkWarnings(statement);
  218.           printResults(statement, hasResults);
  219.           hasResults = statement.getMoreResults();
  220.         }
  221.       } catch (SQLWarning e) {
  222.         throw e;
  223.       } catch (SQLException e) {
  224.         if (stopOnError) {
  225.           throw e;
  226.         } else {
  227.           String message = "Error executing: " + command + ".  Cause: " + e;
  228.           printlnError(message);
  229.         }
  230.       }
  231.     }
  232.   }

  233.   private void checkWarnings(Statement statement) throws SQLException {
  234.     if (!throwWarning) {
  235.       return;
  236.     }
  237.     // In Oracle, CREATE PROCEDURE, FUNCTION, etc. returns warning
  238.     // instead of throwing exception if there is compilation error.
  239.     SQLWarning warning = statement.getWarnings();
  240.     if (warning != null) {
  241.       throw warning;
  242.     }
  243.   }

  244.   private void printResults(Statement statement, boolean hasResults) {
  245.     if (!hasResults) {
  246.       return;
  247.     }
  248.     try (ResultSet rs = statement.getResultSet()) {
  249.       ResultSetMetaData md = rs.getMetaData();
  250.       int cols = md.getColumnCount();
  251.       for (int i = 0; i < cols; i++) {
  252.         String name = md.getColumnLabel(i + 1);
  253.         print(name + "\t");
  254.       }
  255.       println("");
  256.       while (rs.next()) {
  257.         for (int i = 0; i < cols; i++) {
  258.           String value = rs.getString(i + 1);
  259.           print(value + "\t");
  260.         }
  261.         println("");
  262.       }
  263.     } catch (SQLException e) {
  264.       printlnError("Error printing results: " + e.getMessage());
  265.     }
  266.   }

  267.   private void print(Object o) {
  268.     if (logWriter != null) {
  269.       logWriter.print(o);
  270.       logWriter.flush();
  271.     }
  272.   }

  273.   private void println(Object o) {
  274.     if (logWriter != null) {
  275.       logWriter.println(o);
  276.       logWriter.flush();
  277.     }
  278.   }

  279.   private void printlnError(Object o) {
  280.     if (errorLogWriter != null) {
  281.       errorLogWriter.println(o);
  282.       errorLogWriter.flush();
  283.     }
  284.   }

  285. }