在前几天,我们使用了MyBatis plus
的SQL
注入器成功注入了我们想要的SQL
写法。
MyBatisPlus的SQL注入器 | 半月无霜 (banmoon.top)
现在我又新增了一个方法,来看看
其他代码就不贴了,去上一篇文章那看,这边只贴具体的方法实现
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);
}
}
还有一点需要注意,这边的主键要么都要有值,要么都没值(数据库自动生成)
我是半月,你我一同共勉!!!