快捷導(dǎo)航

Mybatis?SQL數(shù)量限制插件的實(shí)現(xiàn)

更新時(shí)間：2025年07月29日 10:41:25 作者：鏘鏘忒

本文介紹了一種應(yīng)對(duì)大數(shù)據(jù)查詢風(fēng)險(xiǎn)的Mybatis插件,通過(guò)配置限制查詢結(jié)果數(shù)量,并支持動(dòng)態(tài)調(diào)整和特殊接口繞過(guò),具有一定的參考價(jià)值,感興趣的可以了解一下

插件背景

起因是一次線上事故，現(xiàn)有項(xiàng)目對(duì)數(shù)據(jù)庫(kù)的操作依賴編碼者自己的行為規(guī)范，有可能出現(xiàn)考慮不當(dāng)對(duì)全表進(jìn)行查詢，數(shù)據(jù)量極大的情況會(huì)引發(fā)頻繁GC帶來(lái)一系列的問(wèn)題
為了避免該問(wèn)題，基于mybatis開發(fā)一個(gè)插件，默認(rèn)限制查詢的條數(shù)，默認(rèn)開啟該功能，默認(rèn)1000條。

插件功能

支持靈活配置插件

可以通過(guò)配置application.properties文件中的配置，靈活控制插件，如果有接入動(dòng)態(tài)配置中心，同樣可以把一下配置添加進(jìn)去，這樣可以達(dá)到運(yùn)行時(shí)靈活控制插件的效果，參數(shù)詳情見上文：

mybatis.limit.size.enable:是否開啟插件功能。
mybatis.limit.size:插件限制查詢數(shù)量，如果不配置使用默認(rèn)：1000條。

支持個(gè)別特殊方法插件效果不生效

該插件加載之后，默認(rèn)會(huì)攔截所有查詢語(yǔ)句，同時(shí)會(huì)過(guò)濾掉已經(jīng)存在limit或者如count等統(tǒng)計(jì)類的sql，但是仍然有極個(gè)別業(yè)務(wù)場(chǎng)景需要繞開此攔截，因此提供了注解，支持在全局開啟limit插件的情況下，特殊接口不走攔截。

@NotSupportDefaultLimit -> 該注解在插件開啟的情況下，可以針對(duì)某個(gè)不需要查詢限制的接口單獨(dú)設(shè)置屏蔽此功能

注意
現(xiàn)有已知的不兼容的場(chǎng)景：
代碼中使用RowBounds進(jìn)行邏輯分頁(yè)，接口會(huì)報(bào)錯(cuò)，因?yàn)閙ybatis的RowBounds是一次性將數(shù)據(jù)查出來(lái)，在內(nèi)存中進(jìn)行分頁(yè)的，而mybatis插件是無(wú)法區(qū)分該種形式，也就無(wú)法兼容。

插件代碼

1、添加一個(gè)配置類

@Configuration
@PropertySource(value = {"classpath:application.properties"},encoding = "utf-8")
public class LimitProperties {
    /**
     * 默認(rèn)限制數(shù)量
     */
    private final static int defaultLimitSize = 1000;
    /**
     * 插件的開關(guān)
     */
    @Value("${mybatis.limit.size.enable}")
    private Boolean enable;
    /**
     * 配置限制數(shù)量
     */
    @Value("${mybatis.limit.size}")
    private Integer size;

    public LimitProperties() {
    }

    public boolean isOffline() {
        return this.enable != null && !this.enable;
    }

    public int limit() {
        return this.size != null && this.size > 0 ? this.size : defaultLimitSize;
    }

    public void setSize(Integer size) {
        this.size = size;
    }
}

2、定義SQL處理器，用來(lái)修改SQL

public class SqlHandler {
    private static final String LIMIT_SQL_TEMPLATE = "%s limit %s;";
    private static final List<String> KEY_WORD = Arrays.asList("count", "limit", "sum", "avg", "min", "max");
    private BoundSql boundSql;
    private String originSql;
    private Boolean needOverride;
    private String newSql;

    public static SqlHandler build(BoundSql boundSql, int size) {
        String originSql = boundSql.getSql().toLowerCase();
        SqlHandler handler = new SqlHandler(boundSql, originSql);
        if (!containsKeyWord(handler.getOriginSql())) {
            handler.setNeedOverride(Boolean.TRUE);
            String newSql = String.format(LIMIT_SQL_TEMPLATE, originSql.replace(";", ""), size);
            handler.setNewSql(newSql);
        }

        return handler;
    }

    private SqlHandler(BoundSql boundSql, String originSql) {
        this.needOverride = Boolean.FALSE;
        this.boundSql = boundSql;
        this.originSql = originSql;
    }

    public boolean needOverride() {
        return this.needOverride;
    }

    public static boolean containsKeyWord(String sql) {
        Iterator var1 = KEY_WORD.iterator();

        String keyWord;
        do {
            if (!var1.hasNext()) {
                return Boolean.FALSE;
            }

            keyWord = (String)var1.next();
        } while(!sql.contains(keyWord));

        return Boolean.TRUE;
    }

    public BoundSql getBoundSql() {
        return this.boundSql;
    }

    public void setBoundSql(BoundSql boundSql) {
        this.boundSql = boundSql;
    }

    public String getOriginSql() {
        return this.originSql;
    }

    public void setOriginSql(String originSql) {
        this.originSql = originSql;
    }

    public Boolean getNeedOverride() {
        return this.needOverride;
    }

    public void setNeedOverride(Boolean needOverride) {
        this.needOverride = needOverride;
    }

    public String getNewSql() {
        return this.newSql;
    }

    public void setNewSql(String newSql) {
        this.newSql = newSql;
    }
}

3、第一步定義一個(gè)插件

/**
 * mybatis插件：查詢數(shù)量限制
 * 使用方法詳見：
 * 攔截mybatis的：Executor.query()
 * @see Executor#query(org.apache.ibatis.mapping.MappedStatement, java.lang.Object, org.apache.ibatis.session.RowBounds, org.apache.ibatis.session.ResultHandler)
 * @see Executor#query(org.apache.ibatis.mapping.MappedStatement, java.lang.Object, org.apache.ibatis.session.RowBounds, org.apache.ibatis.session.ResultHandler, org.apache.ibatis.cache.CacheKey, org.apache.ibatis.mapping.BoundSql)
 */
@Component
@Intercepts({@Signature(
        type = Executor.class,
        method = "query",
        args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}
), @Signature(
        type = Executor.class,
        method = "query",
        args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}
)})
public class LimitInterceptor implements Interceptor {
    @Autowired
    LimitProperties limitProperties;

    public LimitInterceptor() {
    }

    public Object intercept(Invocation invocation) throws Throwable {
        if (!this.limitProperties.isOffline() && !LimitThreadLocal.isNotSupport()) {
            Object[] args = invocation.getArgs();
            MappedStatement ms = (MappedStatement) args[0];
            Object parameter = args[1];
            BoundSql boundSql = args.length == 4 ? ms.getBoundSql(parameter) : (BoundSql) args[5];
            SqlHandler sqlHandler = SqlHandler.build(boundSql, this.limitProperties.limit());
            if (!sqlHandler.needOverride()) {
                return invocation.proceed();
            } else {
                // 需要覆蓋
                Executor executor = (Executor) invocation.getTarget();
                RowBounds rowBounds = (RowBounds) args[2];
                ResultHandler resultHandler = (ResultHandler) args[3];
                CacheKey cacheKey = args.length == 4 ? executor.createCacheKey(ms, parameter, rowBounds, boundSql) : (CacheKey) args[4];
                MetaObject metaObject = SystemMetaObject.forObject(boundSql);
                metaObject.setValue("sql", sqlHandler.getNewSql());
                return executor.query(ms, parameter, rowBounds, resultHandler, cacheKey, boundSql);
            }
        } else {
            return invocation.proceed();
        }
    }

    public Object plugin(Object o) {
        return Plugin.wrap(o, this);
    }

    public void setProperties(Properties properties) {
    }
}

4、下面定義一個(gè)注解，用來(lái)設(shè)置哪些接口不需要數(shù)量限制

@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD})
@Documented
public @interface NotSupportDefaultLimit {
}

5、使用AOP去攔截不需要數(shù)量控制限制的接口

/**
 * AOP：對(duì)@NotSupportDefaultLimit進(jìn)行環(huán)繞通知
 * 主要作用：為當(dāng)前線程set一個(gè)標(biāo)志，標(biāo)記當(dāng)前線程不需要數(shù)量限制
 */
@Aspect
@Component
public class LimitSupportAop {
    @Autowired
    LimitProperties limitProperties;

    public LimitSupportAop() {
    }

    @Around("@annotation(com.jkys.vitamin.rpc.mybatis.NotSupportDefaultLimit)")
    public Object around(ProceedingJoinPoint proceedingJoinPoint) throws Throwable {
        if (this.limitProperties.isOffline()) {
            return proceedingJoinPoint.proceed();
        } else {
            Object var2;
            try {
                LimitThreadLocal.tryAcquire();
                var2 = proceedingJoinPoint.proceed();
            } finally {
                LimitThreadLocal.tryRelease();
            }

            return var2;
        }
    }
}

/**
 * 記錄當(dāng)前線程執(zhí)行SQL，是否 不需要數(shù)量限制
 */
public class LimitThreadLocal {
    private static final ThreadLocal<Integer> times = new ThreadLocal();

    public LimitThreadLocal() {
    }

    public static void tryAcquire() {
        Integer time = (Integer)times.get();
        if (time == null || time < 0) {
            time = 0;
        }

        times.set(time + 1);
    }

    public static void tryRelease() {
        Integer time = (Integer)times.get();
        if (time != null && time > 0) {
            times.set(time - 1);
        }

        if ((Integer)times.get() <= 0) {
            times.remove();
        }

    }

    public static boolean isSupport() {
        return times.get() == null || (Integer)times.get() <= 0;
    }

    public static boolean isNotSupport() {
        return times.get() != null && (Integer)times.get() > 0;
    }
}