Mybatis-Plus的BaseMapper扩展

package cn.cjx913.mybatisplus.mapper;

import com.baomidou.mybatisplus.core.enums.SqlMethod;
import com.baomidou.mybatisplus.core.mapper.BaseMapper;
import com.baomidou.mybatisplus.core.metadata.TableInfo;
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
import com.baomidou.mybatisplus.core.toolkit.*;
import com.baomidou.mybatisplus.extension.conditions.query.LambdaQueryChainWrapper;
import com.baomidou.mybatisplus.extension.conditions.query.QueryChainWrapper;
import com.baomidou.mybatisplus.extension.conditions.update.LambdaUpdateChainWrapper;
import com.baomidou.mybatisplus.extension.conditions.update.UpdateChainWrapper;
import com.baomidou.mybatisplus.extension.toolkit.ChainWrappers;
import com.baomidou.mybatisplus.extension.toolkit.SqlHelper;
import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.reflection.ExceptionUtil;
import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.mybatis.spring.MyBatisExceptionTranslator;

import java.io.Serializable;
import java.lang.reflect.ParameterizedType;
import java.util.Collection;
import java.util.Objects;
import java.util.function.Function;

public interface CBaseMapper<T> extends BaseMapper <T> {
    default LambdaQueryChainWrapper <T> lambdaQuery() {
        return ChainWrappers.lambdaQueryChain(this);
    }

    default QueryChainWrapper <T> query() {
        return ChainWrappers.queryChain(this);
    }

    default LambdaUpdateChainWrapper <T> lambdaUpdate() {
        return ChainWrappers.lambdaUpdateChain(this);
    }

    default UpdateChainWrapper <T> update() {
        return ChainWrappers.updateChain(this);
    }

    default int insertBatch(Collection <T> entityList) {
        return insertBatch(entityList, 1000);
    }

    default int insertBatch(Collection <T> entityList, int batchSize) {
        String sqlStatement = sqlStatement(SqlMethod.INSERT_ONE);
        int size = entityList.size();
        return executeBatch(sqlSession -> {
            int insertCount = 0;
            for (T entity : entityList) {
                sqlSession.insert(sqlStatement, entity);
                insertCount++;
                if ((insertCount % batchSize == 0) || insertCount == size) {
                    sqlSession.flushStatements();
                }
            }
            return insertCount;
        });
    }

    default int insertOrUpdate(T entity) {
        if (null != entity) {
            Class <?> cls = entity.getClass();
            TableInfo tableInfo = TableInfoHelper.getTableInfo(cls);
            Assert.notNull(tableInfo, "error: can not execute. because can not find cache of TableInfo for entity!");
            String keyProperty = tableInfo.getKeyProperty();
            Assert.notEmpty(keyProperty, "error: can not execute. because can not find column for id from entity!");
            Object idVal = ReflectionKit.getMethodValue(cls, entity, tableInfo.getKeyProperty());
            return StringUtils.checkValNull(idVal) || Objects.isNull(selectById((Serializable) idVal)) ?
                    insert(entity) : updateById(entity);
        }
        return 0;
    }

    default int insertOrUpdateBatch(Collection <T> entityList, int batchSize) {
        Assert.notEmpty(entityList, "error: entityList must not be empty");
        Class <?> cls = currentModelClass();
        TableInfo tableInfo = TableInfoHelper.getTableInfo(cls);
        Assert.notNull(tableInfo, "error: can not execute. because can not find cache of TableInfo for entity!");
        String keyProperty = tableInfo.getKeyProperty();
        Assert.notEmpty(keyProperty, "error: can not execute. because can not find column for id from entity!");
        int size = entityList.size();
        return executeBatch(sqlSession -> {
            int insertOrUpdateCount = 0;
            for (T entity : entityList) {
                Object idVal = ReflectionKit.getMethodValue(cls, entity, keyProperty);
                if (StringUtils.checkValNull(idVal) || Objects.isNull(this.selectById((Serializable) idVal))) {
                    int insert = sqlSession.insert(sqlStatement(SqlMethod.INSERT_ONE), entity);
                    insertOrUpdateCount ++;
                } else {
                    MapperMethod.ParamMap <T> param = new MapperMethod.ParamMap <>();
                    param.put(Constants.ENTITY, entity);
                    int update = sqlSession.update(sqlStatement(SqlMethod.UPDATE_BY_ID), param);
                    insertOrUpdateCount ++;
                }
                if ((insertOrUpdateCount % batchSize == 0) || insertOrUpdateCount == size) {
                    sqlSession.flushStatements();
                }
            }
            return insertOrUpdateCount;
        });
    }

    default int updateBatchById(Collection <T> entityList, int batchSize) {
        Assert.notEmpty(entityList, "error: entityList must not be empty");
        String sqlStatement = sqlStatement(SqlMethod.UPDATE_BY_ID);
        int size = entityList.size();
        return executeBatch(sqlSession -> {
            int updateCount = 0;
            for (T anEntityList : entityList) {
                MapperMethod.ParamMap <T> param = new MapperMethod.ParamMap <>();
                param.put(Constants.ENTITY, anEntityList);
                int update = sqlSession.update(sqlStatement, param);
                updateCount ++;
                if ((updateCount % batchSize == 0) || updateCount == size) {
                    sqlSession.flushStatements();
                }
            }
            return updateCount;
        });
    }

    default String sqlStatement(SqlMethod sqlMethod) {
        return SqlHelper.table(currentModelClass()).getSqlStatement(sqlMethod.getMethod());
    }

    /**
     * 建议重写该方法，返回T——数据库表实体类的Class类型
     *
     * @return T.class
     */
    default Class <T> currentModelClass() {
        try {
            String baseMapperClassTypeName = this.getClass().getGenericInterfaces()[0].getTypeName();
            Class <? extends BaseMapper> baseMapperClass = (Class <? extends BaseMapper>) Class.forName(baseMapperClassTypeName);
            ParameterizedType parameterizedType = (ParameterizedType) baseMapperClass.getGenericInterfaces()[0];
            String entityClassName = parameterizedType.getActualTypeArguments()[0].getTypeName();
            return (Class <T>) Class.forName(entityClassName);
        } catch (ClassNotFoundException e) {
            throw new RuntimeException("请重写该方法，返回T的Class类型", e);
        }
//        return (Class <T>) ReflectionKit.getSuperClassGenericType(getClass(), 1);
    }

    default <R> R executeBatch(Function <SqlSession, R> fun) {
        Class <T> tClass = currentModelClass();
        SqlHelper.clearCache(tClass);
        SqlSessionFactory sqlSessionFactory = SqlHelper.sqlSessionFactory(tClass);
        SqlSession sqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH);
        try {
            R apply = fun.apply(sqlSession);
            sqlSession.commit();
            return apply;
        } catch (Throwable t) {
            sqlSession.rollback();
            Throwable unwrapped = ExceptionUtil.unwrapThrowable(t);
            if (unwrapped instanceof RuntimeException) {
                MyBatisExceptionTranslator myBatisExceptionTranslator
                        = new MyBatisExceptionTranslator(sqlSessionFactory.getConfiguration().getEnvironment().getDataSource(), true);
                throw Objects.requireNonNull(myBatisExceptionTranslator.translateExceptionIfPossible((RuntimeException) unwrapped));
            }
            throw ExceptionUtils.mpe(unwrapped);
        } finally {
            sqlSession.close();
        }
    }


}