基于语法树解析Spark SQL 获取访问的表/字段/UDF等信息

最近有需求，需要拆解Spark SQL中的表，字段等信息，然后配合Ranger实现一些权限校验。

其实难度不大，就是需要根据语法树做一些递归拆解，然后就能拆解出一段SQL中的相关信息，再创建一些数据结构bean对象用于配合校验。

下面是部分源码(全部本人原创)，不涉及业务信息。

对于Presto我也做了拆解，欢迎沟通交流。

package test.server.utils;

import com.fasterxml.jackson.databind.ObjectMapper;
import test.ranger.AccessType;
import test.ranger.ColumnAccess;
import test.ranger.ColumnsWithTable;
import test.ranger.PartitionData;
import test.server.service.impl.SparkSqlJob;
import lombok.extern.slf4j.Slf4j;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okio.Buffer;
import okio.BufferedSource;
import org.antlr.v4.runtime.*;
import org.antlr.v4.runtime.tree.ParseTree;
import org.antlr.v4.runtime.tree.RuleNode;
import org.antlr.v4.runtime.tree.TerminalNode;
import org.apache.spark.sql.catalyst.parser.SqlBaseParser;

import java.io.IOException;
import java.nio.charset.Charset;
import java.nio.charset.UnsupportedCharsetException;
import java.sql.SQLException;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static java.lang.String.format;

@Slf4j
public class SparkSqlParser {

    private final ObjectMapper mapper = new ObjectMapper();
    private Pattern pattern = Pattern.compile(".*(?i)set\\s+(.+)=(.+)");
    private String currentDB = "default";
    private Map<String, ColumnAccess> items = new HashMap<>();
    private Set<String> udfList = new HashSet<>();
    private Set<String> tempTable = new HashSet<>();
    private OkHttpClient client = new OkHttpClient().newBuilder().build();
    private String metadataUrl;
    private String metadataToken;

    public SparkSqlParser(String metadataUrl, String metadataToken) {
        this.metadataUrl = metadataUrl;
        this.metadataToken = metadataToken;
    }

    public Map<String, ColumnAccess> getItems() {
        return items;
    }

    public Set<String> getUdfList() {
        return udfList;
    }

    public Map<String, String> getConfigFromSql(String originalSql) throws IOException {
        List<String> commands = SparkSqlJob.prepare(originalSql);
        HashMap<String, String> configMap = new HashMap<>();
        for (String command : commands) {
            if (command.trim().toLowerCase().startsWith("set")) {
                try {
                    Matcher m = pattern.matcher(command);
                    if (m.find()) {
                        String key = m.group(1).trim();
                        String value = m.group(2).trim();
                        configMap.put(key, value);
                    }
                } catch (Exception e) {
                    log.warn(e.getMessage(), e);
                }
            }
        }
        return configMap;
    }

    public void reset(){
        items.clear();
        udfList.clear();
        tempTable.clear();
    }

    public void visit(ParseTree t) throws Exception {
        if (t instanceof TerminalNode){
        } else if (t instanceof