（四）PrestoDB源码解析(一)

presto-cli模块

概要

1. 该模块主要负责查询SQL的客户端功能，利用RESTful请求发送给Coordinator实现SQL语句的查询。
2. PrestoDB可以通过presto-cli客户端和JDBC连接这两种方法实现SQL的查询。
3. presto-cli模块查询流程：通过用户输入SQL将语句组装成一个RESTful请求，发送给Coordinator执行该SQL,并启动查询方法，分批查询结果和用于下一个RESTful请求的nextUrl地址，如果地址为空则停止查询，否则再次发送HTTP请求进行结果查询，不断返回给客户端并在前端显示。

源码

1. 进入客户端主方法：Presto类（该方法主要是客户端的启动方法）

//这是进入客户端的主方法
public final class Presto
{
    private Presto() {}

    public static void main(String[] args)
    {
        //根据传递的参数初始化一个Console对象，该对象保存了启动cli时的所有参数
        Console console = singleCommand(Console.class).parse(args);
        //如果启动的时候有--help等，返回帮助信息然后退出
        if (console.helpOption.showHelpIfRequested() ||
                console.versionOption.showVersionIfRequested()) {
            return;
        }

        System.exit(console.run() ? 0 : 1);
    }
}

2. 进入实际运行的线程run：Console类（该方法主要包装，获取参数并提交查询）

   public boolean run()
    {
        //获取所有客户端传入的参数，参数已经被组装到ClientSession
        ClientSession session = clientOptions.toClientSession();
        //如果Cli启动的时候指定了--execute，则hasQuery为true，说明需要直接执行shell传入的参数
        boolean hasQuery = !isNullOrEmpty(clientOptions.execute);
        //如果Cli启动的时候指定了-f或者--file，则isFromFile为true，说明需要直接执行文件中的SQL
        boolean isFromFile = !isNullOrEmpty(clientOptions.file);

        if (!hasQuery && !isFromFile) {
            AnsiConsole.systemInstall();
        }

        //初始化日志
        initializeLogging(clientOptions.logLevelsFile);

        //获取--execute对应的sql，这里指：如果通过--execute在运行客户端jar包时就指定了sql的话直接从这里拿sql
        String query = clientOptions.execute;
        if (hasQuery) {
            query += ";";
        }

        if (isFromFile) {
            if (hasQuery) {
                //在启动时不能同时指定execute和file
                throw new RuntimeException("both --execute and --file specified");
            }
            try {
                //读取文件中的sql
                query = Files.asCharSource(new File(clientOptions.file), UTF_8).read();
                hasQuery = true;
            }
            catch (IOException e) {
                throw new RuntimeException(format("Error reading from file %s: %s", clientOptions.file, e.getMessage()));
            }
        }

        // abort any running query if the CLI is terminated
        AtomicBoolean exiting = new AtomicBoolean();
        ThreadInterruptor interruptor = new ThreadInterruptor();
        CountDownLatch exited = new CountDownLatch(1);
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            exiting.set(true);
            interruptor.interrupt();
            awaitUninterruptibly(exited, EXIT_DELAY.toMillis(), MILLISECONDS);
        }));

        //生成一个查询的包装类，后面的查询会通过该类启动和执行
        try (QueryRunner queryRunner = new QueryRunner(
                session,
                clientOptions.debug,
                Optional.ofNullable(clientOptions.socksProxy),
                Optional.ofNullable(clientOptions.httpProxy),
                Optional.ofNullable(clientOptions.keystorePath),
                Optional.ofNullable(clientOptions.keystorePassword),
                Optional.ofNullable(clientOptions.truststorePath),
                Optional.ofNullable(clientOptions.truststorePassword),
                Optional.ofNullable(clientOptions.accessToken),
                Optional.ofNullable(clientOptions.user),
                clientOptions.password ? Optional.of(getPassword()) : Optional.empty(),
                Optional.ofNullable(clientOptions.krb5Principal),
                Optional.ofNullable(clientOptions.krb5RemoteServiceName),
                Optional.ofNullable(clientOptions.krb5ConfigPath),
                Optional.ofNullable(clientOptions.krb5KeytabPath),
                Optional.ofNullable(clientOptions.krb5CredentialCachePath),
                !clientOptions.krb5DisableRemoteServiceHostnameCanonicalization)) {
            if (hasQuery) {
                //启动Cli的时候指定--execute或--file参数，则执行该方法，否则就直接启动Cli，接受终端用户在Cli的输入，并提交查询
                return executeCommand(queryRunner, query, clientOptions.outputFormat, clientOptions.ignoreErrors);
            }
            //启动Cli窗口，接受用户输入，并分析SQL语句并提交
            runConsole(queryRunner, exiting);
            return true;
        }
        finally {
            exited.countDown();
            interruptor.close();
        }
    }

3. 进入提交查询方法executeCommand：Console类（该方法主要是将许多sql一条条完成执行、查询）

    //这个方法是带有--execute或--file时对SQL语句进行解析并执行
    private static boolean executeCommand(QueryRunner queryRunner, String query, OutputFormat outputFormat, boolean ignoreErrors)
    {
        boolean success = true;
        //StatementSplitter在构造函数里就对SQL以；进行切分
        StatementSplitter splitter = new StatementSplitter(query);
        for (Statement split : splitter.getCompleteStatements()) {
            if (!isEmptyStatement(split.statement())) {
                //对一条SQL进行处理
                if (!process(queryRunner, split.statement(), outputFormat, () -> {}, false)) {
                    if (!ignoreErrors) {
                        return false;
                    }
                    success = false;
                }
            }
        }
        if (!isEmptyStatement(splitter.getPartialStatement())) {
            System.err.println("Non-terminated statement: " + splitter.getPartialStatement());
            return false;
        }
        return success;
    }

4. 进入执行具体某条SQL的方法process：Console类（该方法主要是进行初始执行某一条SQL，并通过初始执行返回信息不断查询结果）

 private static boolean process(QueryRunner queryRunner, String sql, OutputFormat outputFormat, Runnable schemaChanged, boolean interactive)
    {
        String finalSql;
        try {
            finalSql = preprocessQuery(
                    Optional.ofNullable(queryRunner.getSession().getCatalog()),
                    Optional.ofNullable(queryRunner.getSession().getSchema()),
                    sql);
        }
        catch (QueryPreprocessorException e) {
            System.err.println(e.getMessage());
            if (queryRunner.isDebug()) {
                e.printStackTrace();
            }
            return false;
        }

        //执行查询
        //先进行初始执行
        try (Query query = queryRunner.startQuery(finalSql)) {
            //后续循环发送Restful请求，分批获取查询结果
            boolean success = query.renderOutput(System.out, outputFormat, interactive);

            ClientSession session = queryRunner.getSession();

            // update catalog and schema if present
            if (query.getSetCatalog().isPresent() || query.getSetSchema().isPresent()) {
                session = ClientSession.builder(session)
                        .withCatalog(query.getSetCatalog().orElse(session.getCatalog()))
                        .withSchema(query.getSetSchema().orElse(session.getSchema()))
                        .build();
                schemaChanged.run();
            }

            // update transaction ID if necessary
            if (query.isClearTransactionId()) {
                session = stripTransactionId(session);
            }

            ClientSession.Builder builder = ClientSession.builder(session);

            if (query.getStartedTransactionId() != null) {
                builder = builder.withTransactionId(query.getStartedTransactionId());
            }

            // update path if present
            if (query.getSetPath().isPresent()) {
                builder = builder.withPath(query.getSetPath().get());
            }

            // update session properties if present
            if (!query.getSetSessionProperties().isEmpty() || !query.getResetSessionProperties().isEmpty()) {
                Map sessionProperties = new HashMap<>(session.getProperties());
                sessionProperties.putAll(query.getSetSessionProperties());
                sessionProperties.keySet().removeAll(query.getResetSessionProperties());
                builder = builder.withProperties(sessionProperties);
            }

            // update session roles
            if (!query.getSetRoles().isEmpty()) {
                Map roles = new HashMap<>(session.getRoles());
                roles.putAll(query.getSetRoles());
                builder = builder.withRoles(roles);
            }

            // update prepared statements if present
            if (!query.getAddedPreparedStatements().isEmpty() || !query.getDeallocatedPreparedStatements().isEmpty()) {
                Map preparedStatements = new HashMap<>(session.getPreparedStatements());
                preparedStatements.putAll(query.getAddedPreparedStatements());
                preparedStatements.keySet().removeAll(query.getDeallocatedPreparedStatements());
                builder = builder.withPreparedStatements(preparedStatements);
            }

            session = builder.build();
            queryRunner.setSession(session);

            return success;
        }
        catch (RuntimeException e) {
            System.err.println("Error running command: " + e.getMessage());
            if (queryRunner.isDebug()) {
                e.printStackTrace();
            }
            return false;
        }
    }

5. 进入开始初始执行SQL方法startQuery：Console类（该方法是开始包装HTTP请求并发送执行SQL）

   //开始包装请求，并执行
    public Query startQuery(String query)
    {
        return new Query(startInternalQuery(session.get(), query), debug);
    }

    public StatementClient startInternalQuery(String query)
    {
        return startInternalQuery(stripTransactionId(session.get()), query);
    }

    //建立HTTP请求，包装请求，并发送请求
    private StatementClient startInternalQuery(ClientSession session, String query)
    {
        OkHttpClient.Builder builder = httpClient.newBuilder();
        sslSetup.accept(builder);
        OkHttpClient client = builder.build();

        return newStatementClient(client, session, query);
    }

6. 进入向Coordinator包装并发送请求的工厂方法newStatementClient：StatementClientFactory类

public final class StatementClientFactory
{
    private StatementClientFactory() {}

    public static StatementClient newStatementClient(OkHttpClient httpClient, ClientSession session, String query)
    {
        return new StatementClientV1(httpClient, session, query);
    }
}

7. 进入具体包装并发送请求的方法StatementClientV1：StatementClientV1类

 public StatementClientV1(OkHttpClient httpClient, ClientSession session, String query)
    {
        requireNonNull(httpClient, "httpClient is null");
        requireNonNull(session, "session is null");
        requireNonNull(query, "query is null");

        this.httpClient = httpClient;
        this.timeZone = session.getTimeZone();
        this.query = query;
        this.requestTimeoutNanos = session.getClientRequestTimeout();
        this.user = session.getUser();
        this.clientCapabilities = Joiner.on(",").join(ClientCapabilities.values());

        //在presto中是用向Coordinator提供的RESTful服务结构发送HTTP请求来执行查询的，下面的方法时用来生成请求内容，发起SQL查询的HTTP请求
        Request request = buildQueryRequest(session, query);
        //执行
        JsonResponse response = JsonResponse.execute(QUERY_RESULTS_CODEC, httpClient, request);
        if ((response.getStatusCode() != HTTP_OK) || !response.hasValue()) {
            state.compareAndSet(State.RUNNING, State.CLIENT_ERROR);
            throw requestFailedException("starting query", request, response);
        }

        //将返回结果设置到客户端的currentResults
        processResponse(response.getHeaders(), response.getValue());
    }

8. 进入生成并包装HTTP请求内容buildQueryRequest：StatementClientV1类

 //创建HTTP请求的内容
    private Request buildQueryRequest(ClientSession session, String query)
    {
        //如果启动客户端时通过--server指定了请求Coordinator的IP，则直接拿过来用
        HttpUrl url = HttpUrl.get(session.getServer());
        if (url == null) {
            throw new ClientException("Invalid server URL: " + session.getServer());
        }
        url = url.newBuilder().encodedPath("/v1/statement").build();

        //请求地址是：Coordinator的IP+端口+/v1/statement
        Request.Builder builder = prepareRequest(url)
                .post(RequestBody.create(MEDIA_TYPE_TEXT, query));

        if (session.getSource() != null) {
            builder.addHeader(PRESTO_SOURCE, session.getSource());
        }

        session.getTraceToken().ifPresent(token -> builder.addHeader(PRESTO_TRACE_TOKEN, token));

        if (session.getClientTags() != null && !session.getClientTags().isEmpty()) {
            builder.addHeader(PRESTO_CLIENT_TAGS, Joiner.on(",").join(session.getClientTags()));
        }
        if (session.getClientInfo() != null) {
            builder.addHeader(PRESTO_CLIENT_INFO, session.getClientInfo());
        }
        if (session.getCatalog() != null) {
            builder.addHeader(PRESTO_CATALOG, session.getCatalog());
        }
        if (session.getSchema() != null) {
            builder.addHeader(PRESTO_SCHEMA, session.getSchema());
        }
        if (session.getPath() != null) {
            builder.addHeader(PRESTO_PATH, session.getPath());
        }
        builder.addHeader(PRESTO_TIME_ZONE, session.getTimeZone().getId());
        if (session.getLocale() != null) {
            builder.addHeader(PRESTO_LANGUAGE, session.getLocale().toLanguageTag());
        }

        Map property = session.getProperties();
        for (Entry entry : property.entrySet()) {
            builder.addHeader(PRESTO_SESSION, entry.getKey() + "=" + entry.getValue());
        }

        Map resourceEstimates = session.getResourceEstimates();
        for (Entry entry : resourceEstimates.entrySet()) {
            builder.addHeader(PRESTO_RESOURCE_ESTIMATE, entry.getKey() + "=" + entry.getValue());
        }

        Map roles = session.getRoles();
        for (Entry entry : roles.entrySet()) {
            builder.addHeader(PrestoHeaders.PRESTO_ROLE, entry.getKey() + '=' + urlEncode(entry.getValue().toString()));
        }

        Map statements = session.getPreparedStatements();
        for (Entry entry : statements.entrySet()) {
            builder.addHeader(PRESTO_PREPARED_STATEMENT, urlEncode(entry.getKey()) + "=" + urlEncode(entry.getValue()));
        }

        builder.addHeader(PRESTO_TRANSACTION_ID, session.getTransactionId() == null ? "NONE" : session.getTransactionId());

        builder.addHeader(PRESTO_CLIENT_CAPABILITIES, clientCapabilities);

        return builder.build();
    }

9. 进入执行完SQL后获取到返回信息后不断循环查询的方法renderOutput：Console类

    public boolean renderOutput(PrintStream out, OutputFormat outputFormat, boolean interactive)
    {
        Thread clientThread = Thread.currentThread();
        SignalHandler oldHandler = Signal.handle(SIGINT, signal -> {
            if (ignoreUserInterrupt.get() || client.isClientAborted()) {
                return;
            }
            client.close();
            clientThread.interrupt();
        });
        try {
            //执行完后，通过这个方法不断循环查询结果，并动态显示在终端
            return renderQueryOutput(out, outputFormat, interactive);
        }
        finally {
            Signal.handle(SIGINT, oldHandler);
            Thread.interrupted(); // clear interrupt status
        }
    }

    //具体循环的方法
    private boolean renderQueryOutput(PrintStream out, OutputFormat outputFormat, boolean interactive)
    {
        StatusPrinter statusPrinter = null;
        @SuppressWarnings("resource")
        PrintStream errorChannel = interactive ? out : System.err;
        WarningsPrinter warningsPrinter = new PrintStreamWarningsPrinter(System.err);

        //不断循环查询结果
        if (interactive) {
            statusPrinter = new StatusPrinter(client, out, debug);
            statusPrinter.printInitialStatusUpdates();
        }
        else {
            processInitialStatusUpdates(warningsPrinter);
        }

        // if running or finished
        if (client.isRunning() || (client.isFinished() && client.finalStatusInfo().getError() == null)) {
            QueryStatusInfo results = client.isRunning() ? client.currentStatusInfo() : client.finalStatusInfo();
            if (results.getUpdateType() != null) {
                renderUpdate(errorChannel, results);
            }
            else if (results.getColumns() == null) {
                errorChannel.printf("Query %s has no columns\n", results.getId());
                return false;
            }
            else {
                renderResults(out, outputFormat, interactive, results.getColumns());
            }
        }

        checkState(!client.isRunning());

        if (statusPrinter != null) {
            // Print all warnings at the end of the query
            new PrintStreamWarningsPrinter(System.err).print(client.finalStatusInfo().getWarnings(), true, true);
            statusPrinter.printFinalInfo();
        }
        else {
            // Print remaining warnings separated
            warningsPrinter.print(client.finalStatusInfo().getWarnings(), true, true);
        }

        if (client.isClientAborted()) {
            errorChannel.println("Query aborted by user");
            return false;
        }
        if (client.isClientError()) {
            errorChannel.println("Query is gone (server restarted?)");
            return false;
        }

        verify(client.isFinished());
        if (client.finalStatusInfo().getError() != null) {
            renderFailure(errorChannel);
            return false;
        }

        return true;
    }

10. 进入到具体实现循环的方法processInitialStatusUpdates：Query类

    private void processInitialStatusUpdates(WarningsPrinter warningsPrinter)
    {
      //如果客户端有效且客户端的上一批数据处理完毕就循环获取下一批数据结果
        while (client.isRunning() && (client.currentData().getData() == null)) {
            warningsPrinter.print(client.currentStatusInfo().getWarnings(), true, false);
            client.advance();
        }
        List warnings;
        if (client.isRunning()) {
            warnings = client.currentStatusInfo().getWarnings();
        }
        else {
            warnings = client.finalStatusInfo().getWarnings();
        }
        warningsPrinter.print(warnings, false, true);
    }

11. 进入循环查询结果的具体方法advance：Query类（客户端是通过不断调用advance方法获取nextUri的，并通过nextUri判断查询是否结束，如果nextUri为空则查询结束，不为空则拿着nextUri去继续接着请求）

//该方法是通过coordinator返回的nextUri向coordinator发送请求，分批获取查询结果
    @Override
    public boolean advance()
    {
        if (!isRunning()) {
            return false;
        }

        //获取请求下一批结果的url，这个url是由coordinator生成的，并返回给客户端，他请求的是StatementResource类中的{queryId}/{token}接口获取数据
        URI nextUri = currentStatusInfo().getNextUri();
        if (nextUri == null) {
            //如果查询结果为空，则停止循环查询结果
            state.compareAndSet(State.RUNNING, State.FINISHED);
            return false;
        }

        Request request = prepareRequest(HttpUrl.get(nextUri)).build();

        Exception cause = null;
        long start = System.nanoTime();
        long attempts = 0;

        while (true) {
            if (isClientAborted()) {
                return false;
            }

            Duration sinceStart = Duration.nanosSince(start);
            if (attempts > 0 && sinceStart.compareTo(requestTimeoutNanos) > 0) {
                state.compareAndSet(State.RUNNING, State.CLIENT_ERROR);
                throw new RuntimeException(format("Error fetching next (attempts: %s, duration: %s)", attempts, sinceStart), cause);
            }

            if (attempts > 0) {
                // back-off on retry
                try {
                    MILLISECONDS.sleep(attempts * 100);
                }
                catch (InterruptedException e) {
                    try {
                        close();
                    }
                    finally {
                        Thread.currentThread().interrupt();
                    }
                    state.compareAndSet(State.RUNNING, State.CLIENT_ERROR);
                    throw new RuntimeException("StatementClient thread was interrupted");
                }
            }
            attempts++;

            JsonResponse response;
            try {
                response = JsonResponse.execute(QUERY_RESULTS_CODEC, httpClient, request);
            }
            catch (RuntimeException e) {
                cause = e;
                continue;
            }

            if ((response.getStatusCode() == HTTP_OK) && response.hasValue()) {
                processResponse(response.getHeaders(), response.getValue());
                return true;
            }

            if (response.getStatusCode() != HTTP_UNAVAILABLE) {
                state.compareAndSet(State.RUNNING, State.CLIENT_ERROR);
                throw requestFailedException("fetching next", request, response);
            }
        }
    }