使用Druid的sql parser做一個表數據血緣分析工具

前言

大數據場景下，每天可能都要在離線集群，運行大量的任務來支持業務、運營的分析查詢。任務越來越多的時候，就會有越來越多的依賴關系，每一個任務都需要等需要的input表生產出來后，再去生產自己的output表。最開始的時候，依賴關系自然是可以通過管理員來管理，隨着任務量的加大，就需要一個分析工具來解析任務的inputs、outs，並且自行依賴上生產inputs表的那些任務。本文就介紹一個使用druid parser，來解析SQL的input、output的血緣分析工具。

建議對druid比較陌生的同學可以先看下druid的官方文檔。

做一次sql的血緣分析的流程

• 解析sql，拿到抽象語法樹
• 遍歷抽象語法樹，得到from、to

使用druid解析sql到語法樹

druid提供了簡單、快速的SQL解析工具，可以很簡單拿到一段SQL的AST(抽象語法樹)。而druid對語法樹提供了多種的SQLStatement，使遍歷語法樹更加容易。

 SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, JdbcConstants.HIVE);
 SQLStatement stmt= parser.parseStatementList().get(0);

從語法樹中取出from和to

拿到語法樹之后，想辦法把from、to從語法樹中取出來就大功告成。

最初的寫法

最開始，就是簡單的遍歷一下語法樹的節點，取出from表和to表的表名。

    /** * 根據create或者insert的sql取出from、to * @param sql * @return * @throws ParserException */
    private static Map<String, Set<String>> getFromTo(String sql) throws ParserException {
        SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, JdbcConstants.HIVE);
        SQLStatement stmt= parser.parseStatementList().get(0);

        Set<String> from = new HashSet<>();
        Set<String> to = new HashSet<>();
        if (stmt instanceof SQLInsertStatement) {
            SQLInsertStatement istmt = (SQLInsertStatement) stmt;
            to.add(istmt.getTableSource().toString().toUpperCase());

            SQLTableSource sts = istmt.getQuery().getQueryBlock().getFrom();
            from = getFromTableFromTableSource(sts);
        } else if (stmt instanceof SQLCreateTableStatement) {
            SQLCreateTableStatement cstmt = (SQLCreateTableStatement) stmt;
            to.add(cstmt.getTableSource().toString().toUpperCase());

            SQLTableSource sts = cstmt.getSelect().getQueryBlock().getFrom();
            from = getFromTableFromTableSource(sts);
        }

        Map<String, Set<String>> fromTo = new HashMap<>(4);
        fromTo.put("from", from);
        fromTo.put("to", to);
        return fromTo;
    }

    private static Set<String> getFromTableFromTableSource (SQLTableSource sts) {
        Set<String> from = new HashSet<>();
        if (sts instanceof SQLJoinTableSource) {
            from = getFromTableFromJoinSource((SQLJoinTableSource)sts);
        } else {
            from.add(sts.toString().toUpperCase());
        }
        return from;
    }

    private static Set<String> getFromTableFromJoinSource (SQLJoinTableSource sjts) {
        Set<String> result = new HashSet<>();
        getFromTable(result, sjts);
        return result;
    }

    // 遞歸獲取join的表list
    private static void getFromTable (Set<String> fromList, SQLJoinTableSource sjts) {
        SQLTableSource left = sjts.getLeft();
        if (left instanceof SQLJoinTableSource) {
            getFromTable(fromList, (SQLJoinTableSource)left);
        } else {
            fromList.add(left.toString().toUpperCase());
        }
        SQLTableSource right = sjts.getRight();
        if (right instanceof SQLJoinTableSource) {
            getFromTable(fromList, (SQLJoinTableSource)right);
        } else {
            fromList.add(right.toString().toUpperCase());
        }
    }

用druid更好的實現

因為是為了快速完成，所以寫的取出from、to表的部分還是存在很大的問題的。只能支持一條sql，只能支持簡單的sql語句，比如union all或者子查詢就有些無力。於是又看了一下文檔，其實druid是提供了visitor方法來遍歷語法樹的，而且提供了一個簡單的SchemaStatVisitor，可以取出Sql中所有用到的表。於是就可以寫成這種格式。

public static Map<String, TreeSet<String>> getFromTo (String sql) throws ParserException {
        List<SQLStatement> stmts = SQLUtils.parseStatements(sql, JdbcConstants.HIVE);
        TreeSet<String> fromSet = new TreeSet<>();
        TreeSet<String> toSet = new TreeSet<>();
        if (stmts == null) {
            return null;
        }

        String database="DEFAULT";
        for (SQLStatement stmt : stmts) {
            SchemaStatVisitor statVisitor = SQLUtils.createSchemaStatVisitor(JdbcConstants.HIVE);
            if (stmt instanceof SQLUseStatement) {
                database = ((SQLUseStatement) stmt).getDatabase().getSimpleName().toUpperCase();
            }
            stmt.accept(statVisitor);
            Map<Name, TableStat> tables = statVisitor.getTables();
            if (tables != null) {
                final String db = database;
                tables.forEach((tableName, stat) -> {
                    if (stat.getCreateCount() > 0 || stat.getInsertCount() > 0) {
                        String to = tableName.getName().toUpperCase();
                        if (!to.contains("."))
                            to = db + "." + to;
                        toSet.add(to);
                    } else if (stat.getSelectCount() > 0) {
                        String from = tableName.getName().toUpperCase();
                        if (!from.contains("."))
                            from = db + "." + from;
                        fromSet.add(from);
                    }
                });
            }
        }