2d6702a158e6513ccc2096692b32270de84ab7f1
[JavaForFun] /
1 package de.example.mybatis.executor;
2
3 import java.sql.BatchUpdateException;
4 import java.sql.Connection;
5 import java.sql.SQLException;
6 import java.sql.Statement;
7 import java.util.ArrayList;
8 import java.util.Collections;
9 import java.util.HashMap;
10 import java.util.List;
11 import java.util.Map;
12
13 import org.apache.ibatis.executor.BaseExecutor;
14 import org.apache.ibatis.executor.BatchExecutorException;
15 import org.apache.ibatis.executor.BatchResult;
16 import org.apache.ibatis.executor.keygen.Jdbc3KeyGenerator;
17 import org.apache.ibatis.executor.keygen.KeyGenerator;
18 import org.apache.ibatis.executor.keygen.NoKeyGenerator;
19 import org.apache.ibatis.executor.statement.StatementHandler;
20 import org.apache.ibatis.logging.Log;
21 import org.apache.ibatis.mapping.BoundSql;
22 import org.apache.ibatis.mapping.MappedStatement;
23 import org.apache.ibatis.session.Configuration;
24 import org.apache.ibatis.session.ResultHandler;
25 import org.apache.ibatis.session.RowBounds;
26 import org.apache.ibatis.transaction.Transaction;
27
28 public class ReuseBatchExecutor extends BaseExecutor {
29         public static final int BATCH_UPDATE_RETURN_VALUE = Integer.MIN_VALUE + 1002;
30
31         private final Map<String, Statement> statementMap = new HashMap<>();
32         private final Map<String, BatchResult> batchResultMap = new HashMap<>();
33
34         public ReuseBatchExecutor(Configuration configuration, Transaction transaction) {
35                 super(configuration, transaction);
36         }
37
38         public int doUpdate(MappedStatement ms, Object parameterObject) throws SQLException {
39                 final Configuration configuration = ms.getConfiguration();
40                 final StatementHandler handler = configuration.newStatementHandler(this, ms, parameterObject, RowBounds.DEFAULT,
41                                 null, null);
42                 final Statement stmt = reuseUpdateStatement(handler, ms, parameterObject);
43                 handler.batch(stmt);
44                 return BATCH_UPDATE_RETURN_VALUE;
45         }
46
47         public <E> List<E> doQuery(MappedStatement ms, Object parameterObject, RowBounds rowBounds,
48                         ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
49                 Statement stmt = null;
50                 try {
51                         flushStatements();
52                         final Configuration configuration = ms.getConfiguration();
53                         final StatementHandler handler = configuration.newStatementHandler(wrapper, ms, parameterObject, rowBounds,
54                                         resultHandler, boundSql);
55                         stmt = reuseQueryStatement(handler, ms.getStatementLog());
56                         return handler.<E> query(stmt, resultHandler);
57                 } finally {
58                         closeStatement(stmt);
59                 }
60         }
61
62         public List<BatchResult> doFlushStatements(boolean isRollback) throws SQLException {
63
64                 try {
65                         final List<BatchResult> results = new ArrayList<>();
66                         if (isRollback) {
67                                 return Collections.emptyList();
68                         } else {
69
70                                 long count = 0;
71                                 for (Map.Entry<String, Statement> entry : statementMap.entrySet()) {
72                                         final Statement stmt = entry.getValue();
73                                         final String sql = entry.getKey();
74                                         final BatchResult batchResult = batchResultMap.get(sql);
75                                         if (batchResult != null) {
76
77                                                 try {
78                                                         batchResult.setUpdateCounts(stmt.executeBatch());
79                                                         MappedStatement ms = batchResult.getMappedStatement();
80                                                         List<Object> parameterObjects = batchResult.getParameterObjects();
81                                                         KeyGenerator keyGenerator = ms.getKeyGenerator();
82                                                         if (Jdbc3KeyGenerator.class.equals(keyGenerator.getClass())) {
83                                                                 Jdbc3KeyGenerator jdbc3KeyGenerator = (Jdbc3KeyGenerator) keyGenerator;
84                                                                 jdbc3KeyGenerator.processBatch(ms, stmt, parameterObjects);
85                                                         } else if (!NoKeyGenerator.class.equals(keyGenerator.getClass())) { // issue #141
86                                                                 for (Object parameter : parameterObjects) {
87                                                                         keyGenerator.processAfter(this, ms, stmt, parameter);
88                                                                 }
89                                                         }
90                                         } catch (BatchUpdateException e) {
91                                             StringBuilder message = new StringBuilder();
92                                             message.append(batchResult.getMappedStatement().getId())
93                                                 .append(" (batch index #")
94                                                 .append(count + 1)
95                                                 .append(")")
96                                                 .append(" failed.");
97                                             if (count > 0) {
98                                                 message.append(" ")
99                                                         .append(count)
100                                                         .append(" prior sub executor(s) completed successfully, but will be rolled back.");
101                                             }
102                                             
103                                             throw new BatchExecutorException(message.toString(), e, results, batchResult); 
104                                         }
105                                                 
106                                                 results.add(batchResult);
107                                         }
108
109                                         count = count + 1;
110                                 }
111
112                                 return results;
113                         }
114                 } finally {
115                         for (Statement stmt : statementMap.values()) {
116                                 closeStatement(stmt);
117                         }
118                         statementMap.clear();
119                         batchResultMap.clear();
120                 }
121         }
122
123         private Statement reuseQueryStatement(StatementHandler handler, Log statementLog) throws SQLException {
124                 final BoundSql boundSql = handler.getBoundSql();
125                 final String sql = boundSql.getSql();
126                 final Statement stmt;
127                 
128                 if (hasStatementFor(sql)) {
129                         stmt = getStatement(sql);
130                 } else {
131                         final Connection connection = getConnection(statementLog);
132                         stmt = handler.prepare(connection);
133                         putStatement(sql, stmt);
134                 }
135                 
136                 handler.parameterize(stmt);
137                 
138                 return stmt;
139         }
140         
141         private Statement reuseUpdateStatement(StatementHandler handler, MappedStatement ms, Object parameterObject) throws SQLException {
142                 final BoundSql boundSql = handler.getBoundSql();
143                 final String sql = boundSql.getSql();
144                 final Statement stmt;
145
146                 if (hasStatementFor(sql)) {
147                         stmt = getStatement(sql);
148
149                         final BatchResult batchResult = batchResultMap.get(sql);
150                         batchResult.addParameterObject(parameterObject);
151
152                 } else {
153                         final Connection connection = getConnection(ms.getStatementLog());
154                         stmt = handler.prepare(connection);
155
156                         batchResultMap.put(sql, new BatchResult(ms, sql, parameterObject));
157                         putStatement(sql, stmt);
158                 }
159
160                 handler.parameterize(stmt);
161                 
162                 return stmt;
163         }
164
165         private boolean hasStatementFor(String sql) {
166                 try {
167                         return statementMap.keySet().contains(sql) && !statementMap.get(sql).getConnection().isClosed();
168                 } catch (SQLException e) {
169                         return false;
170                 }
171         }
172
173         private Statement getStatement(String s) {
174                 return statementMap.get(s);
175         }
176
177         private void putStatement(String sql, Statement stmt) {
178                 statementMap.put(sql, stmt);
179         }
180
181 }