自定义Mybatis-plus插件(限制最大查询数量)

自定义Mybatis-plus插件(限制最大查询数量

需求背景

解决思路

1.经分析最后决定，应限制一次查询返回的最大结果数量不应该超出1万，对于一次返回结果大于限制的时候应该抛出异常，而不应该截取(limit 10000最大结果(结果需求不匹配。

步骤一

1.1 定义拦截器PreCheckBigQueryInnerInterceptor

public class PreCheckBigQueryInnerInterceptor implements InnerInterceptor {}

1.2 重写willDoQuery方法

 public boolean willDoQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql throws SQLException {
        // 解析sql
        Statement stmt = CCJSqlParserUtil.parse(boundSql.getSql(;
        if (stmt instanceof Select {
            PlainSelect selectStmt = (PlainSelect ((Select stmt.getSelectBody(;
            if (Objects.nonNull(selectStmt.getLimit( {
                //包含limit查询
                return true;
            }
            for (SelectItem selectItem : selectStmt.getSelectItems( {
                //计数查询 count(;
                SelectExpressionItem selectExpressionItem = (SelectExpressionItem selectItem;
                if (selectExpressionItem.getExpression( instanceof Function {
                    //包含function查询
                    return true;
                }
            }
            Long aLong = doQueryCount(executor, ms, parameter, rowBounds, resultHandler, boundSql;
            if (aLong == 0L {
                return false;
            }
            if (aLong > 20 {
                throw new RuntimeException("单个查询结果大于20条!!!";
            }
        }
        return true;
    }

1.3 代码解析

1.3.1 利用CCJSqlParserUtil解析sql，并判断sql类型，只对Select的SQL拦击.

1.3.2 对于已有limit的sql查询,直接放行.

1.3.3 对于包含function查询（例如count（1）计算，max（）...）,直接放行.

1.3.4 否则判断为大结果查询，执行（doQueryCount）与查询数量.

1.3.5 对于大于指定数量的结果，抛出异常.

1.4 定义doQueryCount方法

private Long doQueryCount(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql throws SQLException {
        MappedStatement countMs = buildAutoCountMappedStatement(ms;
        String countSqlStr = autoCountSql(true, boundSql.getSql(;
        PluginUtils.MPBoundSql mpBoundSql = PluginUtils.mpBoundSql(boundSql;
        BoundSql countSql = new BoundSql(countMs.getConfiguration(, countSqlStr, mpBoundSql.parameterMappings(, parameter;
        PluginUtils.setAdditionalParameter(countSql, mpBoundSql.additionalParameters(;
        CacheKey cacheKey = executor.createCacheKey(countMs, parameter, rowBounds, countSql;
        Object result = executor.query(countMs, parameter, rowBounds, resultHandler, cacheKey, countSql.get(0;
        System.out.println(result;
        return (result == null ? 0L : Long.parseLong(result.toString(;
    }

代码解读：参考PaginationInnerInterceptor（mybatis-plus）分页插件

1.4.1：构造MappedStatement对象buildAutoCountMappedStatement(ms，MappedStatement相当于一个存储 SQL 语句、输入参数和输出结果映射等信息的封装体，它对应一条 SQL 语句，并包含了该 SQL 语句执行所需的所有信息。如下代码

<mapper namespace="com.example.UserMapper">
   <select id="selectAllUsers" resultType="com.example.User">
       SELECT * FROM user
   </select>
</mapper>

注意：必须重新构造，不能直接使用入参中的ms

1.4.2：autoCountSql(true, boundSql.getSql( 定义并优化计数查询语句

String.format("SELECT COUNT(1 FROM (%s TOTAL", originalSql;

1.4.3: 执行查询executor.query

步骤二

1.1 注册拦截器PreCheckBigQueryInnerInterceptor

@Bean
public MybatisPlusInterceptor mybatisPlusInterceptor( {
    MybatisPlusInterceptor interceptor = new MybatisPlusInterceptor(;
    interceptor.addInnerInterceptor(new PaginationInnerInterceptor(DbType.MYSQL;//分页插件(Mybatis-plus
    interceptor.addInnerInterceptor(new BlockAttackInnerInterceptor(;//防止全表更新(Mybatis-plus
    interceptor.addInnerInterceptor(new PreCheckBigQueryInnerInterceptor(;//防止全表查询(自定义插件
    return interceptor;
}

知识小结:

MybatisPlusInterceptor

public class MybatisPlusInterceptor implements Interceptor {
    @Setter
    private List<InnerInterceptor> interceptors = new ArrayList<>(;
}

他是基于mybatis的Interceptor接口做的拦截器,上文中我们注册拦截器PreCheckBigQueryInnerInterceptor的拦截器其实添加到MybatisPlusInterceptor.interceptors集合中。

为啥重写willDoQuery见代码而不是beforeQuery

 public Object intercept(Invocation invocation throws Throwable {
       ......
                for (InnerInterceptor query : interceptors {
                    if (!query.willDoQuery(executor, ms, parameter, rowBounds, resultHandler, boundSql {
                        return Collections.emptyList(;
                    }
                    query.beforeQuery(executor, ms, parameter, rowBounds, resultHandler, boundSql;
                }
     ......
        return invocation.proceed(;
 }

2.1 willDoQuery先于beforeQuery方法，且一定会执行

编程笔记 » 自定义Mybatis-plus插件(限制最大查询数量)