MybatisPlus Sql Inject魔法🪄

本篇只讲应用，喜欢原理的可以不用看

只保证Mysql可用，没有兼容其余数据库的方言

代码笔者线上环境已用，非纸上谈兵

背景介绍

笔者日常的工作有些业务会遇到唯一索引约束。注意到：

业务功能相似，要求当唯一键存在时根据唯一键更新，否则直接插入
Mysql DUPLICATE KEY UPDATE天然支持上述功能

由此引发笔者做出我司公共组件，用来提高开发效率。

常见做法

保证业务安全，还要保证平稳更新与写入，遇到上述需求一般方案是使用分布式锁。步骤如下:

根据唯一索引查询，如果数据存在直接更新，完成写入
数据不存在则加分布锁
根据唯一索引查询，如果数据存在直接更新，反之插入数据，完成写入
解锁

简化代码如下：

java 复制代码

public void duplicateKeyUpdate(NeedInsertModel model) {
    NeedInsertModel dbModel = findByUnique(model);
    /* 这里没考虑此刻恰好数据被别的线程删除的场景 */
    if (Objects.nonNull(dbModel)) {
        updateByUnique(model);
        return;
    }
    
    String lockKey = lockKey();
    RLock lock = redisClient.getLock(lockKey);
    lock.lock();

    try {
        dbModel = findByUnique(model);
        
        if (Objects.nonNull(dbModel)) {
            updateByUnique(model);
        } else  {
            insert(model);
        }
    } finally {
        lock.unlock();
    }
}

可以看到为了达成当唯一键存在时根据唯一键更新，否则直接插入这一简单目的，在多线程，多进程场景下会产生很多套版代码(没什么不好，只是笔者懒惰成性)。

Sql Inject新思路

我想到MybatisPlus的源码中这么写道：

只要继承该接口，就自动具备基础的CRUD功能，这是因为框架帮你生成了代理对象。可是问题是Mybatis Plus怎么知道要生成哪些代理方法，其中的代理逻辑又是哪里定义的呢？后来查资料发现，Mybatis Plus定义了一系列的com.baomidou.mybatisplus.core.injector.AbstractMethod 对象来定制具体的逻辑，也就是生成SQL的逻辑。每一个实现类对应BaseMapper的一个方法。选取最简单SelectById分析，其余的原理相同，其实就是拼接SQL语句。看到这里，笔者当时就想，我直接按照官方的规范定制一个AbstractMethod 的实现不就可以一劳永逸嘛。

经过研究代码如下:

java 复制代码

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.Objects;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;

/**
 * @author Raphael
 * @since 2025/7/15 19:49
 */
public class DuplicateInserter extends AbstractMethod {

    /**
     * 创建时间应该不要被更新
     */
    private static final String CREATE_TIME = "create_time";

    /**
     * 新注入的方法名
     */
    private static final String METHOD_NAME = "duplicateUpdate";

    /**
     * 更新字段集的sql片段
     */
    private static final String SEGMENT = " = VALUES(";

    /**
     * SQL模板：INSERT INTO 表名 字段集合 VALUES 值集合 ON DUPLICATE KEY UPDATE 更新字段集
     */
    private static final String FORMAT = "<script>" +
            "\nINSERT INTO %s %s VALUES %s ON DUPLICATE KEY UPDATE %s\n" +
            "</script>";

    public DuplicateInserter() {
        super(METHOD_NAME);
    }

    @Override
    public MappedStatement injectMappedStatement(Class<?> mapperClass, Class<?> modelClass, TableInfo tableInfo) {
        String insertColumns = tableInfo.getAllInsertSqlColumnMaybeIf(EMPTY);
        String insertValues = tableInfo.getAllInsertSqlPropertyMaybeIf(EMPTY);
        String insertColumnsTrim = SqlScriptUtils.convertTrim(
            insertColumns, LEFT_BRACKET, RIGHT_BRACKET,
            null, COMMA
        );
        String insertValuesTrim = SqlScriptUtils.convertTrim(
            insertValues, LEFT_BRACKET, RIGHT_BRACKET,
            null, COMMA
        );

        String keyProperty = tableInfo.getKeyProperty(), keyColumn = tableInfo.getKeyColumn();

        /* 过滤掉主键和create_time字段 */
        Predicate<TableFieldInfo> needConcat = field
            -> (
                !Objects.equals(field.getColumn(), keyColumn)
                && !Objects.equals(field.getColumn(), CREATE_TIME)
        );

        /* 构建 "column = VALUES(column)" 形式的字符串 */
        Function<TableFieldInfo, String> stringFunc = field
            -> field.getColumn() + SEGMENT + field.getColumn() + RIGHT_BRACKET;

        /* 构建 ON DUPLICATE KEY UPDATE 后面的 SET 子句 */
        String updateSet = tableInfo
            .getFieldList()
            .stream()
            .filter(needConcat)
            .map(stringFunc)
            .collect(Collectors.joining(COMMA));

        KeyGenerator keyGenerator = NoKeyGenerator.INSTANCE;
        /* 表包含主键处理逻辑,如果不包含主键当普通字段处理 */
        if (StringUtils.isNotBlank(tableInfo.getKeyProperty())) {
            if (tableInfo.getIdType() == IdType.AUTO) {
                /* 自增主键 */
                keyGenerator = Jdbc3KeyGenerator.INSTANCE;
            } else if (null != tableInfo.getKeySequence()) {
                keyGenerator = TableInfoHelper.genKeyGenerator(this.methodName, tableInfo, builderAssistant);
            }
        }

        String sql = String.format(FORMAT, tableInfo.getTableName(), insertColumnsTrim, insertValuesTrim, updateSet);
        SqlSource sqlSource = languageDriver.createSqlSource(configuration, sql, modelClass);
        return this.addInsertMappedStatement(
            mapperClass, modelClass, METHOD_NAME,
            sqlSource, keyGenerator, keyProperty, keyColumn
        );
    }

}

然而并没有什么作用，因为你没有将自己的AbstractMethod定制实现注册到Mybatis Plus框架，我们需要找到一个切口。

java 复制代码

public class ‌StrengthenSqlInjector extends DefaultSqlInjector {

    @Override
    public List<AbstractMethod> getMethodList(Class<?> mapperClass, TableInfo tableInfo) {
        List<AbstractMethod> methodList = super.getMethodList(mapperClass, tableInfo);
        /* ⚠️ ⚠️ ⚠️ ：注册自己的定制实现 */
        methodList.add(new DuplicateInserter());
        return methodList;
    }
    
}

还需要替换MybatisPlus自带的DefaultSqlInjector

java 复制代码

@Configuration
@EnableTransactionManagement
public class MybatisPlusAutoConfiguration implements MybatisPlusPropertiesCustomizer {

    @Override
    public void customize(MybatisPlusProperties properties) {
        properties.getGlobalConfig()
            .setSqlInjector(new StrengthenSqlInjector());
    }

}

如此便可完成了。你问我怎么使用？仅仅只需要在业务Mapper中声明一下duplicateUpdate即可，框架会自动帮你生成支持DUPLICATE KEY UPDATE语法的SQL。

java 复制代码

@Mapper
public interface BusinessMapper 
    extends BaseMapper<BusinessModel> {
    
    void duplicateUpdate(BusinessModel model);

}

温馨提示

⚠️ ⚠️ ⚠️接下来这段话非常重要

因为我们是替换MybatisPlus自带的DefaultSqlInjector，注意这里是替换逻辑，因此假设你的系统中有多个MybatisPlusPropertiesCustomizer的Bean那么只会有最后一个生效，因此如果你有多个自定义实现最好全部都放置在一起，只有一个StrengthenSqlInjector最好。

假如没听懂那么等你踩坑了就知道了，祝你好运......