1
2
3
4
5
6
7
8
9
10
11
12
13
14
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
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
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
156
157 @Deprecated
158 public void closeConnection() {
159 try {
160 connection.close();
161 } catch (Exception e) {
162
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
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
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
261
262 }
263 }
264 }
265
266 private void checkWarnings(Statement statement) throws SQLException {
267 if (!throwWarning) {
268 return;
269 }
270
271
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 }