基于druid语法树解析SQL语句增加limit的一种方法

基于druid语法树解析SQL语句增加limit的一种方法

在系统服务中，通常会有面向业务人员查询数据库的页面服务，方便查询并分析数据。
但是如果不对查询数量做限制的话，返回的数据量过大的话会导致服务挂掉。
所以我们需要一个通用的对sql做长度限制的方法。

    public String parseSqlAndLimit(String dbName, String originSql, Integer defaultLimit) {
        logger.info("-- 原始sql:[{}],数据库[{}]", desensiSql(originSql), dbName);
        List<SQLStatement> list = SQLUtils.parseStatements(originSql, "mysql");
        // 0.获取sql解析结果
        SQLSelectStatement sqlStatement1 = (SQLSelectStatement) list.get(0);
        SQLSelect select = sqlStatement1.getSelect();
        // 设置limit
        setSelectLimit(select.getQuery(), defaultLimit);
        return sqlStatement1.toString().replace("###.", "");
    }
    
	    private void setSelectLimit(SQLSelectQuery query) {
        // 1.union类型的查询
        if (query instanceof SQLUnionQuery) {
            SQLUnionQuery unionQuery = (SQLUnionQuery) query;
            // 判断语句是否有limit限制
            unionQuery.setLimit(defaultLimit(unionQuery.getLimit()));
        } else if (query instanceof MySqlSelectQueryBlock) {
            MySqlSelectQueryBlock queryBlock = (MySqlSelectQueryBlock) query;
            queryBlock.setLimit(defaultLimit(queryBlock.getLimit()));
        } else {
            logger.info("未解析2[{}]:", query.getClass());
            throw new ParserException("暂不支持,请联系管理员2!");
        }
    }


    /**
     * 功能描述: 默认limit限制方法
     *
     * @date 2023/04/14 18:27
     */
    public SQLLimit defaultLimit(SQLLimit limit) {
        if (limit != null) {
            SQLExpr rowCount = limit.getRowCount();
            if (rowCount instanceof SQLIntegerExpr) {
                Number number = ((SQLIntegerExpr) rowCount).getNumber();
                // 前端输入的步长不允许大于最大长度,防止机器挂掉
                if (number.longValue() > DbMappingConst.MAX_LIMIT) {
                    limit.setRowCount(DbMappingConst.MAX_LIMIT);
                }
                return limit;
            } else {
                logger.info("未解析3[{}]:", rowCount.getClass());
                throw new ParserException("暂不支持,请联系管理员3!");
            }
        } else {
            SQLLimit sqlLimit = new SQLLimit();
            sqlLimit.setRowCount(DbMappingConst.MAX_LIMIT);
            return sqlLimit;
        }
    }

druid版本，必须是 1.2.6 以上，一下版本对limit支持不友好。

 <dependency>
     <groupId>com.alibaba</groupId>
     <artifactId>druid-spring-boot-starter</artifactId>
     <version>1.2.6</version>
 </dependency>