com.github.wz2cool.dynamic.mybatis.mapper.batch.MapperBatchAction Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of mybatis-dynamic-query Show documentation

There is a newer version: 3.2.32

Show newest version

package com.github.wz2cool.dynamic.mybatis.mapper.batch;

import com.github.wz2cool.dynamic.mybatis.mapper.DynamicQueryMapper;
import org.apache.ibatis.executor.BatchResult;
import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.function.Consumer;

/**
 * @author Frank
 **/
public class MapperBatchAction> {

    private final SqlSessionFactory sqlSessionFactory;
    private final Class mapperClass;
    private int batchSize = 100;

    private List> actions = new ArrayList<>();

    public MapperBatchAction(Class mapperClass, SqlSessionFactory sqlSessionFactory) {
        this.sqlSessionFactory = sqlSessionFactory;
        this.mapperClass = mapperClass;
    }

    public MapperBatchAction(Class mapperClass, SqlSessionFactory sqlSessionFactory, int batchSize) {
        this.sqlSessionFactory = sqlSessionFactory;
        this.mapperClass = mapperClass;
        this.batchSize = batchSize;
    }

    public static > MapperBatchAction create(
            Class mapperClass, SqlSessionFactory sqlSessionFactory) {
        return new MapperBatchAction<>(mapperClass, sqlSessionFactory);
    }

    public static > MapperBatchAction create(
            Class mapperClass, SqlSessionFactory sqlSessionFactory, int batchSize) {
        return new MapperBatchAction<>(mapperClass, sqlSessionFactory, batchSize);
    }

    public MapperBatchAction addAction(Consumer action) {
        actions.add(action);
        return this;
    }

    /**
     * do batch action and return batch result.
     *
     * @return batch result.
     */
    public List doBatchActionWithResults() {
        List result = new ArrayList<>();
        try (SqlSession sqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH, false)) {
            final M mapper = sqlSession.getMapper(mapperClass);
            for (int i = 0; i < actions.size(); i++) {
                actions.get(i).accept(mapper);

                if (i != 0 && (i + 1) % batchSize == 0) {
                    final List batchResults = sqlSession.flushStatements();
                    result.addAll(batchResults);
                }
            }
            final List batchResults = sqlSession.flushStatements();
            result.addAll(batchResults);
        }
        return result;
    }

    /**
     * do batch action
     *
     * @return effect rows (update counts)
     */
    public int doBatchActions() {
        final List batchResults = doBatchActionWithResults();
        return batchResults.stream().flatMapToInt(x -> Arrays.stream(x.getUpdateCounts())).sum();
    }
}