Mybatis SQL数量限制插件

梦里梦外; 2022-09-04 02:48 256阅读 0赞

### Mybatis SQL数量限制插件 ###

*   *  插件背景
     *  插件功能
     *   *  支持灵活配置插件
         *  支持个别特殊方法插件效果不生效
     *  插件代码
     *  总结

## 插件背景 ##

起因是一次线上事故，现有项目对数据库的操作依赖编码者自己的行为规范，有可能出现考虑不当对全表进行查询，数据量极大的情况会引发频繁GC带来一系列的问题  
为了避免该问题，基于mybatis开发一个插件，默认限制查询的条数，默认开启该功能，默认1000条。

## 插件功能 ##

### 支持灵活配置插件 ###

可以通过配置`application.properties`文件中的配置，灵活控制插件，如果有接入动态配置中心，同样可以把一下配置添加进去，这样可以达到运行时灵活控制插件的效果，参数详情见上文：

`mybatis.limit.size.enable`:是否开启插件功能。  
`mybatis.limit.size`:插件限制查询数量，如果不配置使用默认：1000条。

### 支持个别特殊方法插件效果不生效 ###

该插件加载之后，默认会拦截所有查询语句，同时会过滤掉已经存在limit或者如count等统计类的sql，但是仍然有极个别业务场景需要绕开此拦截，因此提供了注解，支持在全局开启limit插件的情况下，特殊接口不走拦截。

**`@NotSupportDefaultLimit`** \-> 该注解在插件开启的情况下，可以针对某个不需要查询限制的接口单独设置屏蔽此功能

**注意**  
现有已知的不兼容的场景：  
代码中使用RowBounds进行逻辑分页，接口会报错，因为mybatis的RowBounds是一次性将数据查出来，在内存中进行分页的，而mybatis插件是无法区分该种形式，也就无法兼容。

## 插件代码 ##

1、添加一个配置类

@Configuration
    @PropertySource(value = { "classpath:application.properties"},encoding = "utf-8")
    public class LimitProperties { 
        /** * 默认限制数量 */
        private final static int defaultLimitSize = 1000;
        /** * 插件的开关 */
        @Value("${mybatis.limit.size.enable}")
        private Boolean enable;
        /** * 配置限制数量 */
        @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处理器，用来修改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、第一步定义一个插件

/** * mybatis插件：查询数量限制 * 使用方法详见： * 拦截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、下面定义一个注解，用来设置哪些接口不需要数量限制

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

5、使用AOP去拦截不需要数量控制限制的接口

/** * AOP：对@NotSupportDefaultLimit进行环绕通知 * 主要作用：为当前线程set一个标志，标记当前线程不需要数量限制 */
    @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;
            }
        }
    }

/** * 记录当前线程执行SQL，是否 不需要数量限制 */
    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;
        }
    }

## 总结 ##

大功告成，下面会分析一下mybatis插件的原理是什么？

[Mybatis插件原理][Mybatis]

[Mybatis]: https://editor.csdn.net/md?not_checkout=1&articleId=119572031