Antlr词法分析之技巧——修改某个token

在上一次的博客Antlr词法分析之技巧——保留空白符中，

我们演示了如何通过词法分析、语法分析解析出SQL中的所有表名，然后给没有带库名前缀的表名添加库名前缀。

这一次我们要做一个更有意思的功能。

很多人都知道Hue这个工具，它支持各种数据库的网页查询。

数据库可能很大，如果用户没有指定limit，那么可能会把网页卡死。

那如果我们想要实现一个类似hue的产品，当用户在查询一张大表的时候，如果他没有写limit，我们希望能给他补上limit。

可能的查询SQL如下：

--单表查询SQL
select id from a.b.c 
--多表查询SQL
select b.id,b.job_id ,b.name,a.job_name from 
  (select id,job_name from zt_mysql.dev_center.process_info )a 
   right join 
  (select id,job_id,name from  zt_mysql.dev_center.task_info)b
  on a.id=b.job_id

那么如何才能实现这个功能呢？本篇博客的主角Antlr隆重登场了。

研究发现，limit涉及的语法大概有这些

statement
    : query                                                            #statementDefault


query
    :  with? queryNoWith
    ;

queryNoWith
    : queryTerm
      (ORDER BY sortItem (',' sortItem)*)?
      (OFFSET offset=rowCount (ROW | ROWS)?)?
      ( (LIMIT limit=limitRowCount)
      | (FETCH (FIRST | NEXT) (fetchFirst=rowCount)? (ROW | ROWS) (ONLY | WITH TIES))
      )?
    ;

limitRowCount
    : ALL
    | rowCount
    ;

rowCount
    : INTEGER_VALUE
    | QUESTION_MARK
    ;


queryPrimary
    : querySpecification                   #queryPrimaryDefault
    | TABLE qualifiedName                  #table
    | VALUES expression (',' expression)*  #inlineTable
    | '(' queryNoWith ')'                  #subquery
    ;

我们实现一个简单的listener方法，把它找到

@Slf4j
public class TrinoListener extends SqlBaseBaseListener {

    public TokenStreamRewriter rewriter;

    public TrinoListener(TokenStream tokens) {
        this.rewriter = new TokenStreamRewriter(tokens);
    }

    

    @Override
    public void exitQueryNoWith(SqlBaseParser.QueryNoWithContext ctx) {
        TerminalNode limit = ctx.LIMIT();
        SqlBaseParser.LimitRowCountContext limitRowCountContext = ctx.limitRowCount();
        RuleContext parent = ctx.getParent();
        RuleContext grandParent = parent.getParent();
        //没有limit，添加
            if (limit == null && limitRowCountContext == null) {
                rewriter.insertAfter(ctx.stop, " LIMIT 100");
            } 
        

    }
}

把刚才的单表SQL填进去试试看

 public static void main(String[] args) {
        String sql="select id from a.b.c";
        System.out.println(sql);

        CodePointCharStream charStream = CharStreams.fromString(sql.toUpperCase());
        SqlBaseLexer sqlBaseLexer = new SqlBaseLexer(charStream);
        CommonTokenStream tokenStream = new CommonTokenStream(sqlBaseLexer);
        SqlBaseParser sqlBaseParser = new SqlBaseParser(tokenStream);
        SqlBaseParser.SingleStatementContext tree = sqlBaseParser.singleStatement();
        TrinoListener trinoListener = new TrinoListener(tokenStream);
        ParseTreeWalker walker = new ParseTreeWalker();
        walker.walk(trinoListener,tree);
        String text = trinoListener.rewriter.getText();
        System.out.println(text);
    }

打印下结果

select id from a.b.c
16:12:34.225 [main] INFO com.example.sql.trino.TrinoListener - A.B.C
SELECT ID FROM A.B.C LIMIT 100

成功了！

再试试多表SQL

SELECT B.ID,B.JOB_ID ,B.NAME,A.JOB_NAME FROM 
(SELECT ID,JOB_NAME FROM ZT_MYSQL.DEV_CENTER.PROCESS_INFO LIMIT 100 )A 
RIGHT JOIN 
(SELECT ID,JOB_ID,NAME FROM  ZT_MYSQL.DEV_CENTER.TASK_INFO LIMIT 100)B
ON A.ID=B.JOB_ID LIMIT 100

发现有点不对劲，不仅最外层的查询添加了limit，连里面的子查询也添加了。这不符合我们的需求，

所以我们需要把这两者区分开来。

修改我们的listener实现

public void exitQueryNoWith(SqlBaseParser.QueryNoWithContext ctx) {
        TerminalNode limit = ctx.LIMIT();
        SqlBaseParser.LimitRowCountContext limitRowCountContext = ctx.limitRowCount();
        RuleContext parent = ctx.getParent();
        RuleContext grandParent = parent.getParent();
        //排除掉子查询subQuery中的limit
        if (grandParent instanceof SqlBaseParser.StatementDefaultContext) {
            //没有limit，添加
            if (limit == null && limitRowCountContext == null) {
                rewriter.insertAfter(ctx.stop, " LIMIT 100");
            } 
        }

    }

在执行一遍，打印一下

SELECT B.ID,B.JOB_ID ,B.NAME,A.JOB_NAME FROM 
(SELECT ID,JOB_NAME FROM ZT_MYSQL.DEV_CENTER.PROCESS_INFO )A 
RIGHT JOIN 
(SELECT ID,JOB_ID,NAME FROM  ZT_MYSQL.DEV_CENTER.TASK_INFO)B
ON A.ID=B.JOB_ID LIMIT 100

发现子查询中的limit已经不会添加了，离成功又进了一步！

现在还有一个问题，如果客户填了limit，但是limit 10000怎么办？这种情况也需要考虑

我们需要捕捉到用户的limit rowCount,然后改成我们要限制的值

继续修改我们的listener

public void exitQueryNoWith(SqlBaseParser.QueryNoWithContext ctx) {
        TerminalNode limit = ctx.LIMIT();
        SqlBaseParser.LimitRowCountContext limitRowCountContext = ctx.limitRowCount();
        RuleContext parent = ctx.getParent();
        RuleContext grandParent = parent.getParent();
        //排除掉子查询subQuery中的limit
        if (grandParent instanceof SqlBaseParser.StatementDefaultContext) {
            //没有limit，添加
            if (limit == null && limitRowCountContext == null) {
                rewriter.insertAfter(ctx.stop, " LIMIT 100");
            } else if (limit != null && limitRowCountContext != null) {
                //有limit，修改
                SqlBaseParser.RowCountContext rowCountContext = limitRowCountContext.rowCount();
                TerminalNode terminalNode = rowCountContext.INTEGER_VALUE();
                terminalNode.getSymbol();
                String text = terminalNode.getText();
                int integer = Integer.parseInt(text);
                if (integer > 100) {
                    rewriter.replace(terminalNode.getSymbol(), 100);
                }
            }
        }

    }

对客户提供的rowCount进行判断，如果超出了我们的限定值，则进行替换操作。

我们提供了新的SQL：

select id from a.b.c limit 10000
经过替换后变成了

SELECT ID FROM A.B.C LIMIT 100

是不是很酷！

rewriter不仅可以insert、replace，还支持delete，基本我们满足我们大部分需求了！