MybatisPlus Sql Inject魔法🪄

image

  • 本篇只講應用,喜歡原理的可以不用看
  • 只保證Mysql可用,沒有兼容其餘資料庫的方言
  • 代碼筆者線上環境已用,非紙上談兵

背景介紹

筆者日常的工作有些業務會遇到唯一索引約束。注意到:

  • 業務功能相似,要求當唯一鍵存在時根據唯一鍵更新,否則直接插入
  • Mysql DUPLICATE KEY UPDATE 天然支持上述功能

由此引發筆者做出我司公共組件,用來提高開發效率。

常見做法

保證業務安全,還要保證平穩更新與寫入,遇到上述需求一般方案是使用分佈式鎖。步驟如下:

  • 根據唯一索引查詢,如果資料存在直接更新,完成寫入
  • 資料不存在則加分佈鎖
  • 根據唯一索引查詢,如果資料存在直接更新,反之插入資料,完成寫入
  • 解鎖

image

簡化代碼如下:

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的原碼中這麼寫道:

image

只要繼承該介面,就自動具備基礎的CRUD功能,這是因為框架幫你生成了代理物件。可問題是Mybatis Plus怎麼知道要生成哪些代理方法,其中的代理邏輯又是在哪裡定義的呢?後來查資料發現,Mybatis Plus定義了一系列的 com.baomidou.mybatisplus.core.injector.AbstractMethod 物件來定制具體的邏輯,也就是生成SQL的邏輯。每一個實現類對應BaseMapper的一個方法。

image

選取最簡單SelectById分析,其餘的原理相同,其實就是拼接SQL語句。

image

看到這裡,筆者當時就想,我直接按照官方的規範定制一個 AbstractMethod 的實現不就可以一勞永逸嘛。
經過研究代碼如下:

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框架,我們需要找到一個切口。

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

@Configuration
@EnableTransactionManagement
public class MybatisPlusAutoConfiguration implements MybatisPlusPropertiesCustomizer {

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

如此便可完成了。你問我怎麼使用?僅僅只需要在業務Mapper中聲明一下duplicateUpdate即可,框架會自動幫你生成支持 DUPLICATE KEY UPDATE 語法的SQL。

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

溫馨提示

⚠️ ⚠️ ⚠️ 接下來這段話非常重要
因為我們是替換MybatisPlus自帶的DefaultSqlInjector,注意這裡是替換邏輯,因此假設你的系統中有多個MybatisPlusPropertiesCustomizer的Bean那麼只會有最後一個生效,因此如果你有多個自定義實現最好全部都放置在一起,只有一個StrengthenSqlInjector最好。
假如沒聽懂那麼等你踩坑了就知道了,祝你好运......


原文出處:https://juejin.cn/post/7546051941325094921

關於作者
精選技術文章翻譯,幫助開發者持續吸收新知。

共有 0 則留言


精選技術文章翻譯,幫助開發者持續吸收新知。
🏆 本月排行榜
🥇
站長阿川
📝9   💬8   ❤️13
429
🥈
我愛JS
📝1   💬6   ❤️4
88
🥉
酷豪
📝1   ❤️1
51
#4
AppleLily
📝1   💬4   ❤️1
39
#5
💬3  
10
評分標準:發文×10 + 留言×3 + 獲讚×5 + 點讚×1 + 瀏覽數÷10
本數據每小時更新一次