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