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

2023-05-24,,

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

需求背景

​ 一次查询如果结果返回太多(1万或更多),往往会导致系统性能下降,有时更会内存不足,影响系统稳定性,故需要做限制。

解决思路

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

2.利用mybatis拦截器技术,统一拦截sql,并真对大结果的查询先做一次count查询。

步骤一

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插件(限制最大查询数量)的相关教程结束。

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

下载本文的Word格式文档,以方便收藏与打印。