c5d543cbaa6d26824f079be41773133e8838ed6a
[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         @Override
39         public int doUpdate(MappedStatement ms, Object parameterObject) throws SQLException {
40                 final Configuration configuration = ms.getConfiguration();
41                 final StatementHandler handler = configuration.newStatementHandler(this, ms, parameterObject, RowBounds.DEFAULT,
42                                 null, null);
43                 final Statement stmt = reuseUpdateStatement(handler, ms, parameterObject);
44                 handler.batch(stmt);
45                 return BATCH_UPDATE_RETURN_VALUE;
46         }
47
48         @Override
49         public <E> List<E> doQuery(MappedStatement ms, Object parameterObject, RowBounds rowBounds,
50                         ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
51                 Statement stmt = null;
52                 try {
53                         flushStatements();
54                         final Configuration configuration = ms.getConfiguration();
55                         final StatementHandler handler = configuration.newStatementHandler(wrapper, ms, parameterObject, rowBounds,
56                                         resultHandler, boundSql);
57                         stmt = reuseQueryStatement(handler, ms.getStatementLog());
58                         return handler.<E> query(stmt, resultHandler);
59                 } finally {
60                         closeStatement(stmt);
61                 }
62         }
63
64         @Override
65         public List<BatchResult> doFlushStatements(boolean isRollback) throws SQLException {
66
67                 try {
68                         final List<BatchResult> results = new ArrayList<>();
69                         if (isRollback) {
70                                 return Collections.emptyList();
71                         } else {
72
73                                 long count = 0;
74                                 for (Map.Entry<String, Statement> entry : statementMap.entrySet()) {
75                                         final Statement stmt = entry.getValue();
76                                         final String sql = entry.getKey();
77                                         final BatchResult batchResult = batchResultMap.get(sql);
78                                         if (batchResult != null) {
79
80                                                 try {
81                                                         batchResult.setUpdateCounts(stmt.executeBatch());
82                                                         MappedStatement ms = batchResult.getMappedStatement();
83                                                         List<Object> parameterObjects = batchResult.getParameterObjects();
84                                                         KeyGenerator keyGenerator = ms.getKeyGenerator();
85                                                         if (Jdbc3KeyGenerator.class.equals(keyGenerator.getClass())) {
86                                                                 Jdbc3KeyGenerator jdbc3KeyGenerator = (Jdbc3KeyGenerator) keyGenerator;
87                                                                 jdbc3KeyGenerator.processBatch(ms, stmt, parameterObjects);
88                                                         } else if (!NoKeyGenerator.class.equals(keyGenerator.getClass())) { // issue #141
89                                                                 for (Object parameter : parameterObjects) {
90                                                                         keyGenerator.processAfter(this, ms, stmt, parameter);
91                                                                 }
92                                                         }
93                                         } catch (BatchUpdateException e) {
94                                             StringBuilder message = new StringBuilder();
95                                             message.append(batchResult.getMappedStatement().getId())
96                                                 .append(" (batch index #")
97                                                 .append(count + 1)
98                                                 .append(")")
99                                                 .append(" failed.");
100                                             if (count > 0) {
101                                                 message.append(" ")
102                                                         .append(count)
103                                                         .append(" prior sub executor(s) completed successfully, but will be rolled back.");
104                                             }
105                                             
106                                             throw new BatchExecutorException(message.toString(), e, results, batchResult); 
107                                         }
108                                                 
109                                                 results.add(batchResult);
110                                         }
111
112                                         count = count + 1;
113                                 }
114
115                                 return results;
116                         }
117                 } finally {
118                         for (Statement stmt : statementMap.values()) {
119                                 closeStatement(stmt);
120                         }
121                         statementMap.clear();
122                         batchResultMap.clear();
123                 }
124         }
125
126         private Statement reuseQueryStatement(StatementHandler handler, Log statementLog) throws SQLException {
127                 final BoundSql boundSql = handler.getBoundSql();
128                 final String sql = boundSql.getSql();
129                 final Statement stmt;
130                 
131                 if (hasStatementFor(sql)) {
132                         stmt = getStatement(sql);
133                 } else {
134                         final Connection connection = getConnection(statementLog);
135                         stmt = handler.prepare(connection);
136                         putStatement(sql, stmt);
137                 }
138                 
139                 handler.parameterize(stmt);
140                 
141                 return stmt;
142         }
143         
144         private Statement reuseUpdateStatement(StatementHandler handler, MappedStatement ms, Object parameterObject) throws SQLException {
145                 final BoundSql boundSql = handler.getBoundSql();
146                 final String sql = boundSql.getSql();
147                 final Statement stmt;
148
149                 if (hasStatementFor(sql)) {
150                         stmt = getStatement(sql);
151
152                         final BatchResult batchResult = batchResultMap.get(sql);
153                         batchResult.addParameterObject(parameterObject);
154
155                 } else {
156                         final Connection connection = getConnection(ms.getStatementLog());
157                         stmt = handler.prepare(connection);
158
159                         batchResultMap.put(sql, new BatchResult(ms, sql, parameterObject));
160                         putStatement(sql, stmt);
161                 }
162
163                 handler.parameterize(stmt);
164                 
165                 return stmt;
166         }
167
168         private boolean hasStatementFor(String sql) {
169                 try {
170                         return statementMap.keySet().contains(sql) && !statementMap.get(sql).getConnection().isClosed();
171                 } catch (SQLException e) {
172                         return false;
173                 }
174         }
175
176         private Statement getStatement(String s) {
177                 return statementMap.get(s);
178         }
179
180         private void putStatement(String sql, Statement stmt) {
181                 statementMap.put(sql, stmt);
182         }
183
184 }