MyBatisPlus的SQL注入器批量插入更新方法

一、介绍

在前几天，我们使用了MyBatis plus的SQL注入器成功注入了我们想要的SQL写法。

现在我又新增了一个方法，来看看

二、代码

其他代码就不贴了，去上一篇文章那看，这边只贴具体的方法实现

package com.banmoon.business.mybatis.method;

import com.baomidou.mybatisplus.annotation.IdType;
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.sql.SqlInjectionUtils;
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;

public class InsertOnDuplicateKeyUpdateBatchMethod extends InsertOnDuplicateKeyUpdateMethod {

    private static final String METHOD_NAME = "insertOnDuplicateKeyUpdateBatch";

    private static final String METHOD_PARAM_NAME = "entityList";

    public InsertOnDuplicateKeyUpdateBatchMethod() {
        super(METHOD_NAME);
    }

    @Override
    public MappedStatement injectMappedStatement(Class<?> mapperClass, Class<?> modelClass, TableInfo tableInfo) {
        KeyGenerator keyGenerator = NoKeyGenerator.INSTANCE;
        List<TableFieldInfo> fieldList = tableInfo.getFieldList();
        String insertSqlColumn = getKeyInsertSqlColumn(tableInfo) + this.filterTableFieldInfo(fieldList, null, TableFieldInfo::getInsertSqlColumn, EMPTY);
        String columnScript = LEFT_BRACKET + insertSqlColumn.substring(0, insertSqlColumn.length() - 1) + RIGHT_BRACKET;
        String insertSqlProperty = tableInfo.getKeyInsertSqlProperty(false, ENTITY_DOT, false) +
                                           this.filterTableFieldInfo(fieldList, null, i -> i.getInsertSqlProperty(ENTITY_DOT), EMPTY);
        insertSqlProperty = LEFT_BRACKET + insertSqlProperty.substring(0, insertSqlProperty.length() - 1) + RIGHT_BRACKET;
        String valuesScript = SqlScriptUtils.convertForeach(insertSqlProperty, METHOD_PARAM_NAME, null, ENTITY, COMMA);
        String keyProperty = null;
        String keyColumn = null;
        // 表包含主键处理逻辑,如果不包含主键当普通字段处理
        if (tableInfo.havePK()) {
            if (tableInfo.getIdType() == IdType.AUTO) {
                /* 自增主键 */
                keyGenerator = Jdbc3KeyGenerator.INSTANCE;
                keyProperty = tableInfo.getKeyProperty();
                // 去除转义符
                keyColumn = SqlInjectionUtils.removeEscapeCharacter(tableInfo.getKeyColumn());
            } else {
                if (null != tableInfo.getKeySequence()) {
                    keyGenerator = TableInfoHelper.genKeyGenerator(this.methodName, tableInfo, builderAssistant);
                    keyProperty = tableInfo.getKeyProperty();
                    keyColumn = tableInfo.getKeyColumn();
                }
            }
        }
        String duplicateKeyUpdateScript = generateDuplicateKeyUpdateScript(tableInfo);
        String sql = String.format(SQL_TEMPLATE, tableInfo.getTableName(), columnScript, valuesScript, duplicateKeyUpdateScript);
        SqlSource sqlSource = super.createSqlSource(configuration, sql, modelClass);
        return this.addInsertMappedStatement(mapperClass, modelClass, methodName, sqlSource, keyGenerator, keyProperty, keyColumn);
    }

    private String getKeyInsertSqlColumn(TableInfo tableInfo) {
        String keyColumn = tableInfo.getKeyColumn();
        String keyProperty = tableInfo.getKeyProperty();
        if (tableInfo.havePK()) {
            if (tableInfo.getIdType() == IdType.AUTO) {
                return SqlScriptUtils.convertIf(keyColumn + COMMA, String.format("%s[0].%s != null", METHOD_PARAM_NAME, keyProperty), true);
            }
            return keyColumn + COMMA + NEWLINE;
        }
        return EMPTY;
    }
}

测试一下

package com.banmoon;

import cn.hutool.core.collection.CollUtil;
import com.banmoon.entity.UserEntity;
import com.banmoon.mapper.UserMapper;
import org.junit.Assert;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;

import javax.annotation.Resource;
import java.util.List;

@SpringBootTest
@RunWith(SpringRunner.class)
public class ServerTest {

    @Resource
    private UserMapper userMapper;

    @Test
    public void testBatch() {
        UserEntity userEntity = new UserEntity();
        userEntity.setId(6);
        userEntity.setUsername("测试1");
        userEntity.setPassword("1234");
        userEntity.setStatus(1);
        UserEntity userEntity1 = new UserEntity();
        userEntity1.setId(7);
        userEntity1.setUsername("测试2");
        userEntity1.setPassword("1234");
        userEntity1.setStatus(1);
        List<UserEntity> list = CollUtil.newArrayList(userEntity, userEntity1);

        int i = userMapper.insertOnDuplicateKeyUpdateBatch(list);
        Assert.assertEquals(list.size(), i);

        userEntity.setUsername("测试覆盖1");
        userEntity1.setUsername("测试覆盖2");
        i = userMapper.insertOnDuplicateKeyUpdateBatch(list);
        Assert.assertEquals(list.size() * 2, i);
    }

}