實現MyBatisPlus自定義sql注入器

目標:新增mysql下的 插入更新的語法

INSERT INTO %s %s VALUES %s ON DUPLICATE KEY UPDATE %s

新增方法類，新增的方法名稱為insertOrUpdate和insertOrUpdateBatch方法，但其mapper層的方法名為insertOrUpdate方法

import com.baomidou.mybatisplus.annotation.IdType;
import com.baomidou.mybatisplus.core.injector.AbstractMethod;
import com.baomidou.mybatisplus.core.metadata.TableFieldInfo;
import com.baomidou.mybatisplus.core.metadata.TableInfo;
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.baomidou.mybatisplus.core.toolkit.sql.SqlScriptUtils;
import org.apache.ibatis.executor.keygen.Jdbc3KeyGenerator;
import org.apache.ibatis.executor.keygen.KeyGenerator;
import org.apache.ibatis.executor.keygen.NoKeyGenerator;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlSource;

import java.util.List;
import java.util.Objects;

import static java.util.stream.Collectors.joining;

public class InsertOrUpdate extends AbstractMethod {

    public InsertOrUpdate() {
        super(MyBatisPlusMethod.INSERT_OR_UPDATE.getMethod());
    }

    public InsertOrUpdate(String method) {
        super(method);
    }

    @Override
    public MappedStatement injectMappedStatement(Class<?> mapperClass, Class<?> modelClass, TableInfo tableInfo) {
        KeyGenerator keyGenerator = NoKeyGenerator.INSTANCE;
        MyBatisPlusMethod sqlMethod = MyBatisPlusMethod.INSERT_OR_UPDATE;
        String columnScript = SqlScriptUtils.convertTrim(tableInfo.getAllInsertSqlColumnMaybeIf(ENTITY_DOT),
                LEFT_BRACKET, RIGHT_BRACKET, null, COMMA);
        String valuesScript = SqlScriptUtils.convertTrim(tableInfo.getAllInsertSqlPropertyMaybeIf(ENTITY_DOT),
                LEFT_BRACKET, RIGHT_BRACKET, null, COMMA);
        String setScript = SqlScriptUtils.convertTrim(this.getUpdateValuePart(tableInfo, ENTITY_DOT),
                null, null, null, COMMA);
        String keyProperty = null;
        String keyColumn = null;
        // 表包含主鍵處理邏輯,如果不包含主鍵當普通欄位處理
        if (StringUtils.isNotBlank(tableInfo.getKeyProperty())) {
            if (tableInfo.getIdType() == IdType.AUTO) {
                /* 自增主鍵 */
                keyGenerator = Jdbc3KeyGenerator.INSTANCE;
                keyProperty = tableInfo.getKeyProperty();
                keyColumn = tableInfo.getKeyColumn();
            } else if (null != tableInfo.getKeySequence()) {
                keyGenerator = TableInfoHelper.genKeyGenerator(this.methodName, tableInfo, builderAssistant);
                keyProperty = tableInfo.getKeyProperty();
                keyColumn = tableInfo.getKeyColumn();
            }
        }
        String sql = String.format(sqlMethod.getSql(), tableInfo.getTableName(), columnScript, valuesScript, setScript);
        SqlSource sqlSource = languageDriver.createSqlSource(configuration, sql, modelClass);
        return this.addInsertMappedStatement(mapperClass, modelClass, getMethod(sqlMethod), sqlSource, keyGenerator, keyProperty, keyColumn);
    }

    protected String getMethod(MyBatisPlusMethod sqlMethod) {
        return StringUtils.isBlank(methodName) ? sqlMethod.getMethod() : this.methodName;
    }

    protected String getUpdateValuePart(TableInfo tableInfo, String prefix) {
        List<TableFieldInfo> fieldList = tableInfo.getFieldList();
        return fieldList.stream().map(i -> i.getSqlSet(false, prefix))
                .filter(Objects::nonNull).collect(joining(NEWLINE));
    }
}

新增該方法的列舉類

public enum MyBatisPlusMethod {
    /**
     * 插入
     */
    INSERT_OR_UPDATE("insertOrUpdate", "插入更新一條資料（選擇欄位插入）", "<script>\nINSERT INTO %s %s VALUES %s ON DUPLICATE KEY UPDATE %s\n</script>");

    private final String method;
    private final String desc;
    private final String sql;

    MyBatisPlusMethod(String method, String desc, String sql) {
        this.method = method;
        this.desc = desc;
        this.sql = sql;
    }

    public String getMethod() {
        return method;
    }

    public String getDesc() {
        return desc;
    }

    public String getSql() {
        return sql;
    }
}

繼承並實現MyBatisPlus的mapper、service層的方法

mapper

import com.baomidou.mybatisplus.core.toolkit.Constants;
import com.chinacreator.c2.dao.mapper.BaseMapper;
import org.apache.ibatis.annotations.Param;

public interface MyBatisPlusMapper<T> extends BaseMapper<T> {

    /**
     * 插入更新一條記錄
     *
     * @param entity 實體物件
     * @return
     */
    int insertOrUpdate(@Param(Constants.ENTITY) T entity);
}

service

package com.chinacreator.c2.chs.mp;

import com.baomidou.mybatisplus.extension.service.IService;
import org.springframework.transaction.annotation.Transactional;

import java.util.Collection;

public interface MyBatisPlusService<T> extends IService<T> {

    /**
     * 一句sql執行插入更新語句
     *
     * @param entityList 需要插入更新的實體
     * @return 是否成功
     */
    @Transactional(rollbackFor = Exception.class)
    default boolean insertOrUpdateBatch(Collection<T> entityList) {
        return insertOrUpdateBatch(entityList, DEFAULT_BATCH_SIZE);
    }

    boolean insertOrUpdateBatch(Collection<T> entityList, int batchSize);

    /**
     * 一句sql執行插入更新語句
     *
     * @param entity 需要插入更新的實體
     * @return 是否成功
     */
    boolean insertOrUpdate(T entity);

}

serviceImpl

import com.baomidou.mybatisplus.core.toolkit.Constants;
import com.baomidou.mybatisplus.core.toolkit.StringPool;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.baomidou.mybatisplus.extension.toolkit.SqlHelper;
import org.apache.ibatis.binding.MapperMethod;
import org.springframework.transaction.annotation.Transactional;

import java.util.Collection;

public class MyBatisPlusServiceImpl<M extends MyBatisPlusMapper<T>, T> extends ServiceImpl<M, T> implements MyBatisPlusService<T> {
    @Override
    @Transactional(rollbackFor = Exception.class)
    public boolean insertOrUpdateBatch(Collection<T> entityList, int batchSize) {
        String mapperStatementId = mapperClass.getName() + StringPool.DOT + MyBatisPlusMethod.INSERT_OR_UPDATE.getMethod();
        return SqlHelper.executeBatch(entityClass, log, entityList, batchSize, (sqlSession, entity) -> {
            MapperMethod.ParamMap<Object> param = new MapperMethod.ParamMap<>();
            param.put(Constants.ENTITY, entity);
            sqlSession.insert(mapperStatementId, param);
        });
    }

    @Override
    public boolean insertOrUpdate(T entity) {
        return SqlHelper.retBool(getBaseMapper().insertOrUpdate(entity));
    }
}

新增自定義的MyBatisPlus自定義注入器

import com.baomidou.mybatisplus.core.injector.AbstractMethod;
import com.baomidou.mybatisplus.core.injector.DefaultSqlInjector;
import com.baomidou.mybatisplus.core.metadata.TableInfo;
import com.chinacreator.c2.dao.mapper.InsertBatchMethod;

import java.util.List;

public class MyBatisPlusSqlInjector extends DefaultSqlInjector {

    /**
     * 如果只需增加方法，保留mybatis plus自帶方法， 可以先獲取super.getMethodList()，再新增add
     */
    @Override
    public List<AbstractMethod> getMethodList(Class<?> mapperClass, TableInfo tableInfo) {
        List<AbstractMethod> methodList = super.getMethodList(mapperClass, tableInfo);
        methodList.add(new InsertOrUpdate());
        return methodList;
    }
}

在MyBatisPlus的配置檔案類上 建立該自定義注入器的bean物件

import com.chinacreator.c2.chs.mp.MyBatisPlusSqlInjector;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Configuration;

@Slf4j
@Configuration
public class MyBatisPlusConfig {
    @Bean
    public MyBatisPlusSqlInjector customizedSqlInjector() {
        return new MyBatisPlusSqlInjector();
    }
}

對使用官方生成的mapper、service、serviceImpl檔案的繼承類上改為我這邊新增的

最後再呼叫service或mapper下的insertOrUpdate方法或insertOrUpdateBatch方法