001/**
002 * Copyright (c) 2022-2023, Mybatis-Flex (fuhai999@gmail.com).
003 * <p>
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 * <p>
008 * http://www.apache.org/licenses/LICENSE-2.0
009 * <p>
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016package com.mybatisflex.core.mybatis.executor;
017
018import com.mybatisflex.core.keygen.RowKeyGenerator;
019import org.apache.ibatis.cache.CacheKey;
020import org.apache.ibatis.executor.BatchExecutor;
021import org.apache.ibatis.executor.BatchExecutorException;
022import org.apache.ibatis.executor.BatchResult;
023import org.apache.ibatis.executor.keygen.Jdbc3KeyGenerator;
024import org.apache.ibatis.executor.keygen.KeyGenerator;
025import org.apache.ibatis.executor.keygen.NoKeyGenerator;
026import org.apache.ibatis.executor.statement.StatementHandler;
027import org.apache.ibatis.mapping.BoundSql;
028import org.apache.ibatis.mapping.MappedStatement;
029import org.apache.ibatis.session.Configuration;
030import org.apache.ibatis.session.RowBounds;
031import org.apache.ibatis.transaction.Transaction;
032
033import java.sql.BatchUpdateException;
034import java.sql.Connection;
035import java.sql.SQLException;
036import java.sql.Statement;
037import java.util.ArrayList;
038import java.util.Collections;
039import java.util.List;
040
041public class FlexBatchExecutor extends BatchExecutor implements CacheKeyBuilder {
042
043    private final List<Statement> statementList = new ArrayList<>();
044    private final List<BatchResult> batchResultList = new ArrayList<>();
045    private String currentSql;
046    private MappedStatement currentStatement;
047
048
049    public FlexBatchExecutor(Configuration configuration, Transaction transaction) {
050        super(configuration, transaction);
051    }
052
053    @Override
054    public CacheKey createCacheKey(MappedStatement ms, Object parameterObject, RowBounds rowBounds, BoundSql boundSql) {
055        return buildCacheKey(super.createCacheKey(ms, parameterObject, rowBounds, boundSql), parameterObject);
056    }
057
058    @Override
059    public int doUpdate(MappedStatement ms, Object parameterObject) throws SQLException {
060        final Configuration configuration = ms.getConfiguration();
061        final StatementHandler handler = configuration.newStatementHandler(this, ms, parameterObject, RowBounds.DEFAULT, null, null);
062        final BoundSql boundSql = handler.getBoundSql();
063        final String sql = boundSql.getSql();
064        final Statement stmt;
065        if (sql.equals(currentSql) && ms.equals(currentStatement)) {
066            int last = statementList.size() - 1;
067            stmt = statementList.get(last);
068            applyTransactionTimeout(stmt);
069            handler.parameterize(stmt);// fix Issues 322
070            BatchResult batchResult = batchResultList.get(last);
071            batchResult.addParameterObject(parameterObject);
072        } else {
073            Connection connection = getConnection(ms.getStatementLog());
074            stmt = handler.prepare(connection, transaction.getTimeout());
075            handler.parameterize(stmt);    // fix Issues 322
076            currentSql = sql;
077            currentStatement = ms;
078            statementList.add(stmt);
079            batchResultList.add(new BatchResult(ms, sql, parameterObject));
080        }
081        handler.batch(stmt);
082        return BATCH_UPDATE_RETURN_VALUE;
083    }
084
085
086    @Override
087    public List<BatchResult> doFlushStatements(boolean isRollback) throws SQLException {
088        try {
089            List<BatchResult> results = new ArrayList<>();
090            if (isRollback) {
091                return Collections.emptyList();
092            }
093            for (int i = 0, n = statementList.size(); i < n; i++) {
094                Statement stmt = statementList.get(i);
095                applyTransactionTimeout(stmt);
096                BatchResult batchResult = batchResultList.get(i);
097                try {
098                    batchResult.setUpdateCounts(stmt.executeBatch());
099                    MappedStatement ms = batchResult.getMappedStatement();
100                    List<Object> parameterObjects = batchResult.getParameterObjects();
101                    KeyGenerator keyGenerator = ms.getKeyGenerator();
102                    if (Jdbc3KeyGenerator.class.equals(keyGenerator.getClass())) {
103                        Jdbc3KeyGenerator jdbc3KeyGenerator = (Jdbc3KeyGenerator) keyGenerator;
104                        jdbc3KeyGenerator.processBatch(ms, stmt, parameterObjects);
105                    }
106                    // 修复批量插入并设置主键时出错
107                    // fixed https://gitee.com/mybatis-flex/mybatis-flex/issues/I6Y8ZU
108                    else if (RowKeyGenerator.class.equals(keyGenerator.getClass())
109                            && ((RowKeyGenerator) keyGenerator).hasGeneratedKeys()) {
110                        keyGenerator.processAfter(this, ms, stmt, parameterObjects);
111                    }
112                    // issue #141
113                    else if (!NoKeyGenerator.class.equals(keyGenerator.getClass())) {
114                        for (Object parameter : parameterObjects) {
115                            keyGenerator.processAfter(this, ms, stmt, parameter);
116                        }
117                    }
118                    // Close statement to close cursor #1109
119                    closeStatement(stmt);
120                } catch (BatchUpdateException e) {
121                    StringBuilder message = new StringBuilder();
122                    message.append(batchResult.getMappedStatement().getId())
123                            .append(" (batch index #")
124                            .append(i + 1)
125                            .append(")")
126                            .append(" failed.");
127                    if (i > 0) {
128                        message.append(" ")
129                                .append(i)
130                                .append(" prior sub executor(s) completed successfully, but will be rolled back.");
131                    }
132                    throw new BatchExecutorException(message.toString(), e, results, batchResult);
133                }
134                results.add(batchResult);
135            }
136            return results;
137        } finally {
138            for (Statement stmt : statementList) {
139                closeStatement(stmt);
140            }
141            currentSql = null;
142            statementList.clear();
143            batchResultList.clear();
144        }
145    }
146}