View Javadoc
1   /**
2    *    Copyright 2009-2020 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  
18  import java.io.BufferedReader;
19  import java.io.PrintWriter;
20  import java.io.Reader;
21  import java.sql.Connection;
22  import java.sql.ResultSet;
23  import java.sql.ResultSetMetaData;
24  import java.sql.SQLException;
25  import java.sql.SQLWarning;
26  import java.sql.Statement;
27  import java.util.regex.Matcher;
28  import java.util.regex.Pattern;
29  
30  /**
31   * @author Clinton Begin
32   */
33  public class ScriptRunner {
34  
35    private static final String LINE_SEPARATOR = System.lineSeparator();
36  
37    private static final String DEFAULT_DELIMITER = ";";
38  
39    private static final Pattern DELIMITER_PATTERN = Pattern.compile("^\\s*((--)|(//))?\\s*(//)?\\s*@DELIMITER\\s+([^\\s]+)", Pattern.CASE_INSENSITIVE);
40  
41    private final Connection connection;
42  
43    private boolean stopOnError;
44    private boolean throwWarning;
45    private boolean autoCommit;
46    private boolean sendFullScript;
47    private boolean removeCRs;
48    private boolean escapeProcessing = true;
49  
50    private PrintWriter logWriter = new PrintWriter(System.out);
51    private PrintWriter errorLogWriter = new PrintWriter(System.err);
52  
53    private String delimiter = DEFAULT_DELIMITER;
54    private boolean fullLineDelimiter;
55  
56    public ScriptRunner(Connection connection) {
57      this.connection = connection;
58    }
59  
60    public void setStopOnError(boolean stopOnError) {
61      this.stopOnError = stopOnError;
62    }
63  
64    public void setThrowWarning(boolean throwWarning) {
65      this.throwWarning = throwWarning;
66    }
67  
68    public void setAutoCommit(boolean autoCommit) {
69      this.autoCommit = autoCommit;
70    }
71  
72    public void setSendFullScript(boolean sendFullScript) {
73      this.sendFullScript = sendFullScript;
74    }
75  
76    public void setRemoveCRs(boolean removeCRs) {
77      this.removeCRs = removeCRs;
78    }
79  
80    /**
81     * @since 3.1.1
82     */
83    public void setEscapeProcessing(boolean escapeProcessing) {
84      this.escapeProcessing = escapeProcessing;
85    }
86  
87    public void setLogWriter(PrintWriter logWriter) {
88      this.logWriter = logWriter;
89    }
90  
91    public void setErrorLogWriter(PrintWriter errorLogWriter) {
92      this.errorLogWriter = errorLogWriter;
93    }
94  
95    public void setDelimiter(String delimiter) {
96      this.delimiter = delimiter;
97    }
98  
99    public void setFullLineDelimiter(boolean fullLineDelimiter) {
100     this.fullLineDelimiter = fullLineDelimiter;
101   }
102 
103   public void runScript(Reader reader) {
104     setAutoCommit();
105 
106     try {
107       if (sendFullScript) {
108         executeFullScript(reader);
109       } else {
110         executeLineByLine(reader);
111       }
112     } finally {
113       rollbackConnection();
114     }
115   }
116 
117   private void executeFullScript(Reader reader) {
118     StringBuilder script = new StringBuilder();
119     try {
120       BufferedReader lineReader = new BufferedReader(reader);
121       String line;
122       while ((line = lineReader.readLine()) != null) {
123         script.append(line);
124         script.append(LINE_SEPARATOR);
125       }
126       String command = script.toString();
127       println(command);
128       executeStatement(command);
129       commitConnection();
130     } catch (Exception e) {
131       String message = "Error executing: " + script + ".  Cause: " + e;
132       printlnError(message);
133       throw new RuntimeSqlException(message, e);
134     }
135   }
136 
137   private void executeLineByLine(Reader reader) {
138     StringBuilder command = new StringBuilder();
139     try {
140       BufferedReader lineReader = new BufferedReader(reader);
141       String line;
142       while ((line = lineReader.readLine()) != null) {
143         handleLine(command, line);
144       }
145       commitConnection();
146       checkForMissingLineTerminator(command);
147     } catch (Exception e) {
148       String message = "Error executing: " + command + ".  Cause: " + e;
149       printlnError(message);
150       throw new RuntimeSqlException(message, e);
151     }
152   }
153 
154   /**
155    * @deprecated Since 3.5.4, this method is deprecated. Please close the {@link Connection} outside of this class.
156    */
157   @Deprecated
158   public void closeConnection() {
159     try {
160       connection.close();
161     } catch (Exception e) {
162       // ignore
163     }
164   }
165 
166   private void setAutoCommit() {
167     try {
168       if (autoCommit != connection.getAutoCommit()) {
169         connection.setAutoCommit(autoCommit);
170       }
171     } catch (Throwable t) {
172       throw new RuntimeSqlException("Could not set AutoCommit to " + autoCommit + ". Cause: " + t, t);
173     }
174   }
175 
176   private void commitConnection() {
177     try {
178       if (!connection.getAutoCommit()) {
179         connection.commit();
180       }
181     } catch (Throwable t) {
182       throw new RuntimeSqlException("Could not commit transaction. Cause: " + t, t);
183     }
184   }
185 
186   private void rollbackConnection() {
187     try {
188       if (!connection.getAutoCommit()) {
189         connection.rollback();
190       }
191     } catch (Throwable t) {
192       // ignore
193     }
194   }
195 
196   private void checkForMissingLineTerminator(StringBuilder command) {
197     if (command != null && command.toString().trim().length() > 0) {
198       throw new RuntimeSqlException("Line missing end-of-line terminator (" + delimiter + ") => " + command);
199     }
200   }
201 
202   private void handleLine(StringBuilder command, String line) throws SQLException {
203     String trimmedLine = line.trim();
204     if (lineIsComment(trimmedLine)) {
205       Matcher matcher = DELIMITER_PATTERN.matcher(trimmedLine);
206       if (matcher.find()) {
207         delimiter = matcher.group(5);
208       }
209       println(trimmedLine);
210     } else if (commandReadyToExecute(trimmedLine)) {
211       command.append(line.substring(0, line.lastIndexOf(delimiter)));
212       command.append(LINE_SEPARATOR);
213       println(command);
214       executeStatement(command.toString());
215       command.setLength(0);
216     } else if (trimmedLine.length() > 0) {
217       command.append(line);
218       command.append(LINE_SEPARATOR);
219     }
220   }
221 
222   private boolean lineIsComment(String trimmedLine) {
223     return trimmedLine.startsWith("//") || trimmedLine.startsWith("--");
224   }
225 
226   private boolean commandReadyToExecute(String trimmedLine) {
227     // issue #561 remove anything after the delimiter
228     return !fullLineDelimiter && trimmedLine.contains(delimiter) || fullLineDelimiter && trimmedLine.equals(delimiter);
229   }
230 
231   private void executeStatement(String command) throws SQLException {
232     Statement statement = connection.createStatement();
233     try {
234       statement.setEscapeProcessing(escapeProcessing);
235       String sql = command;
236       if (removeCRs) {
237         sql = sql.replace("\r\n", "\n");
238       }
239       try {
240         boolean hasResults = statement.execute(sql);
241         while (!(!hasResults && statement.getUpdateCount() == -1)) {
242           checkWarnings(statement);
243           printResults(statement, hasResults);
244           hasResults = statement.getMoreResults();
245         }
246       } catch (SQLWarning e) {
247         throw e;
248       } catch (SQLException e) {
249         if (stopOnError) {
250           throw e;
251         } else {
252           String message = "Error executing: " + command + ".  Cause: " + e;
253           printlnError(message);
254         }
255       }
256     } finally {
257       try {
258         statement.close();
259       } catch (Exception ignored) {
260         // Ignore to workaround a bug in some connection pools
261         // (Does anyone know the details of the bug?)
262       }
263     }
264   }
265 
266   private void checkWarnings(Statement statement) throws SQLException {
267     if (!throwWarning) {
268       return;
269     }
270     // In Oracle, CREATE PROCEDURE, FUNCTION, etc. returns warning
271     // instead of throwing exception if there is compilation error.
272     SQLWarning warning = statement.getWarnings();
273     if (warning != null) {
274       throw warning;
275     }
276   }
277 
278   private void printResults(Statement statement, boolean hasResults) {
279     if (!hasResults) {
280       return;
281     }
282     try (ResultSet rs = statement.getResultSet()) {
283       ResultSetMetaData md = rs.getMetaData();
284       int cols = md.getColumnCount();
285       for (int i = 0; i < cols; i++) {
286         String name = md.getColumnLabel(i + 1);
287         print(name + "\t");
288       }
289       println("");
290       while (rs.next()) {
291         for (int i = 0; i < cols; i++) {
292           String value = rs.getString(i + 1);
293           print(value + "\t");
294         }
295         println("");
296       }
297     } catch (SQLException e) {
298       printlnError("Error printing results: " + e.getMessage());
299     }
300   }
301 
302   private void print(Object o) {
303     if (logWriter != null) {
304       logWriter.print(o);
305       logWriter.flush();
306     }
307   }
308 
309   private void println(Object o) {
310     if (logWriter != null) {
311       logWriter.println(o);
312       logWriter.flush();
313     }
314   }
315 
316   private void printlnError(Object o) {
317     if (errorLogWriter != null) {
318       errorLogWriter.println(o);
319       errorLogWriter.flush();
320     }
321   }
322 
323 }