Java 类com.facebook.presto.client.QueryResults 实例源码

@Override
public void addResults(QueryResults results)
{
    if (types.get() == null && results.getColumns() != null) {
        types.set(getTypes(results.getColumns()));
    }

    if (results.getData() != null) {
        checkState(types.get() != null, "Data without types received!");
        List<Column> columns = results.getColumns();
        for (List<Object> fields : results.getData()) {
            ImmutableMap.Builder<String, Object> builder = ImmutableMap.builder();
            for (int i = 0; i < fields.size(); i++) {
                Type type = types.get().get(i);
                Object value = convertValue(fields.get(i), type);
                if (value != null) {
                    builder.put(columns.get(i).getName(), value);
                }
            }

            producer.send(new KeyedMessage<>(topicName, count.getAndIncrement(), builder.build()));
        }
    }
}

@Override
public void addResults(QueryResults results)
{
    if (!loggedUri.getAndSet(true)) {
        log.info("Query %s: %s?pretty", results.getId(), results.getInfoUri());
    }

    if (types.get() == null && results.getColumns() != null) {
        types.set(getTypes(results.getColumns()));
    }

    if (results.getData() != null) {
        checkState(types.get() != null, "data received without types");
        rows.addAll(transform(results.getData(), dataToRow(timeZoneKey, types.get())));
    }
}

@Inject
public ExecuteResource(
        HttpServerInfo serverInfo,
        AccessControl accessControl,
        SessionPropertyManager sessionPropertyManager,
        @ForExecute HttpClient httpClient,
        QueryIdGenerator queryIdGenerator,
        JsonCodec<QueryResults> queryResultsCodec)
{
    this.serverInfo = requireNonNull(serverInfo, "serverInfo is null");
    this.accessControl = requireNonNull(accessControl, "accessControl is null");
    this.sessionPropertyManager = requireNonNull(sessionPropertyManager, "sessionPropertyManager is null");
    this.httpClient = requireNonNull(httpClient, "httpClient is null");
    this.queryIdGenerator = requireNonNull(queryIdGenerator, "queryIdGenerator is null");
    this.queryResultsCodec = requireNonNull(queryResultsCodec, "queryResultsCodec is null");
}

public synchronized QueryResults getResults(long token, UriInfo uriInfo, Duration maxWaitTime)
        throws InterruptedException
{
    // is the a repeated request for the last results?
    String requestedPath = uriInfo.getAbsolutePath().getPath();
    if (lastResultPath != null && requestedPath.equals(lastResultPath)) {
        // tell query manager we are still interested in the query
        queryManager.getQueryInfo(queryId);
        queryManager.recordHeartbeat(queryId);
        return lastResult;
    }

    if (token < resultId.get()) {
        throw new WebApplicationException(Status.GONE);
    }

    // if this is not a request for the next results, return not found
    if (lastResult.getNextUri() == null || !requestedPath.equals(lastResult.getNextUri().getPath())) {
        // unknown token
        throw new WebApplicationException(Status.NOT_FOUND);
    }

    return getNextResults(uriInfo, maxWaitTime);
}

public BenchmarkQueryRunner(int warm, int runs, boolean debug, int maxFailures, URI serverUri, Optional<HostAndPort> socksProxy)
{
    checkArgument(warm >= 0, "warm is negative");
    this.warm = warm;

    checkArgument(runs >= 1, "runs must be at least 1");
    this.runs = runs;

    checkArgument(maxFailures >= 0, "maxFailures must be at least 0");
    this.maxFailures = maxFailures;

    this.debug = debug;

    this.queryResultsCodec = jsonCodec(QueryResults.class);

    requireNonNull(socksProxy, "socksProxy is null");
    HttpClientConfig httpClientConfig = new HttpClientConfig();
    if (socksProxy.isPresent()) {
        httpClientConfig.setSocksProxy(socksProxy.get());
    }

    this.httpClient = new JettyHttpClient(httpClientConfig.setConnectTimeout(new Duration(10, TimeUnit.SECONDS)));

    nodes = getAllNodes(requireNonNull(serverUri, "serverUri is null"));
}

private static List<Column> getColumns(StatementClient client)
        throws SQLException
{
    while (client.isValid()) {
        List<Column> columns = client.current().getColumns();
        if (columns != null) {
            return columns;
        }
        client.advance();
    }

    QueryResults results = client.finalResults();
    if (!client.isFailed()) {
        throw new SQLException(format("Query has no columns (#%s)", results.getId()));
    }
    throw resultsException(results);
}

public QueryRunner(
        ClientSession session,
        JsonCodec<QueryResults> queryResultsCodec,
        Optional<HostAndPort> socksProxy,
        Optional<String> keystorePath,
        Optional<String> keystorePassword,
        Optional<String> kerberosPrincipal,
        Optional<String> kerberosRemoteServiceName,
        boolean authenticationEnabled,
        KerberosConfig kerberosConfig)
{
    this.session = new AtomicReference<>(requireNonNull(session, "session is null"));
    this.queryResultsCodec = requireNonNull(queryResultsCodec, "queryResultsCodec is null");
    this.httpClient = new JettyHttpClient(
            getHttpClientConfig(socksProxy, keystorePath, keystorePassword, kerberosPrincipal, kerberosRemoteServiceName, authenticationEnabled),
            kerberosConfig,
            Optional.<JettyIoPool>empty(),
            ImmutableList.<HttpRequestFilter>of());
}

public static QueryRunner create(
        ClientSession session,
        Optional<HostAndPort> socksProxy,
        Optional<String> keystorePath,
        Optional<String> keystorePassword,
        Optional<String> kerberosPrincipal,
        Optional<String> kerberosRemoteServiceName,
        boolean authenticationEnabled,
        KerberosConfig kerberosConfig)
{
    return new QueryRunner(
            session,
            jsonCodec(QueryResults.class),
            socksProxy,
            keystorePath,
            keystorePassword,
            kerberosPrincipal,
            kerberosRemoteServiceName,
            authenticationEnabled,
            kerberosConfig);
}

private List<String> queryMetadata(String query)
{
    ImmutableList.Builder<String> cache = ImmutableList.builder();
    try (StatementClient client = queryRunner.startInternalQuery(query)) {
        while (client.isValid() && !Thread.currentThread().isInterrupted()) {
            QueryResults results = client.current();
            if (results.getData() != null) {
                for (List<Object> row : results.getData()) {
                    cache.add((String) row.get(0));
                }
            }
            client.advance();
        }
    }
    return cache.build();
}

private static Response getQueryResults(Query query, Optional<Long> token, UriInfo uriInfo, Duration wait)
        throws InterruptedException
{
    QueryResults queryResults;
    if (token.isPresent()) {
        queryResults = query.getResults(token.get(), uriInfo, wait);
    }
    else {
        queryResults = query.getNextResults(uriInfo, wait);
    }

    ResponseBuilder response = Response.ok(queryResults);

    // add set session properties
    query.getSetSessionProperties().entrySet().stream()
            .forEach(entry -> response.header(PRESTO_SET_SESSION, entry.getKey() + '=' + entry.getValue()));

    // add clear session properties
    query.getResetSessionProperties().stream()
            .forEach(name -> response.header(PRESTO_CLEAR_SESSION, name));

    // add new transaction ID
    query.getStartedTransactionId()
            .ifPresent(transactionId -> response.header(PRESTO_STARTED_TRANSACTION_ID, transactionId));

    // add clear transaction ID directive
    if (query.isClearTransactionId()) {
        response.header(PRESTO_CLEAR_TRANSACTION_ID, true);
    }

    return response.build();
}

private QueryExecutor(String userAgent, JsonCodec<QueryResults> queryResultsCodec, JsonCodec<ServerInfo> serverInfoCodec, HostAndPort socksProxy)
{
    requireNonNull(userAgent, "userAgent is null");
    requireNonNull(queryResultsCodec, "queryResultsCodec is null");
    requireNonNull(serverInfoCodec, "serverInfoCodec is null");

    this.queryInfoCodec = queryResultsCodec;
    this.serverInfoCodec = serverInfoCodec;
    this.httpClient = new JettyHttpClient(
            new HttpClientConfig()
                    .setConnectTimeout(new Duration(10, TimeUnit.SECONDS))
                    .setSocksProxy(socksProxy),
            new JettyIoPool("presto-jdbc", new JettyIoPoolConfig()),
            ImmutableSet.of(new UserAgentRequestFilter(userAgent)));
}

static SQLException resultsException(QueryResults results)
{
    QueryError error = requireNonNull(results.getError());
    String message = format("Query failed (#%s): %s", results.getId(), error.getMessage());
    Throwable cause = (error.getFailureInfo() == null) ? null : error.getFailureInfo().toException();
    return new SQLException(message, error.getSqlState(), error.getErrorCode(), cause);
}

private void renderUpdate(PrintStream out, QueryResults results)
{
    String status = results.getUpdateType();
    if (results.getUpdateCount() != null) {
        long count = results.getUpdateCount();
        status += format(": %s row%s", count, (count != 1) ? "s" : "");
    }
    out.println(status);
    discardResults();
}

public void renderFailure(PrintStream out)
{
    QueryResults results = client.finalResults();
    QueryError error = results.getError();
    checkState(error != null);

    out.printf("Query %s failed: %s%n", results.getId(), error.getMessage());
    if (client.isDebug() && (error.getFailureInfo() != null)) {
        error.getFailureInfo().toException().printStackTrace(out);
    }
    if (error.getErrorLocation() != null) {
        renderErrorLocation(client.getQuery(), error.getErrorLocation(), out);
    }
    out.println();
}

private List<HiveColumn> queryColumns(String query)
{
    final ImmutableList.Builder<HiveColumn> cache = ImmutableList.builder();
    QueryRunner queryRunner = queryRunnerFactory.create();
    QueryClient queryClient = new QueryClient(queryRunner, io.dropwizard.util.Duration.seconds(60), query);

    try {
        queryClient.executeWith(new Function<StatementClient, Void>() {
            @Nullable
            @Override
            public Void apply(StatementClient client)
            {
                QueryResults results = client.current();
                if (results.getData() != null) {
                    for (List<Object> row : results.getData()) {
                        Column column = new Column((String) row.get(0), (String) row.get(1), new ClientTypeSignature(TypeSignature.parseTypeSignature((String)row.get(1))));
                        boolean isNullable = (Boolean) row.get(2);
                        boolean isPartition = (Boolean) row.get(3);

                        cache.add(HiveColumn.fromColumn(column, isNullable, isPartition));
                    }
                }

                return null;
            }
        });
    }
    catch (QueryClient.QueryTimeOutException e) {
        log.error("Caught timeout loading columns", e);
    }

    return cache.build();
}

private List<List<Object>> queryRows(String query)
{
    final ImmutableList.Builder<List<Object>> cache = ImmutableList.builder();
    QueryRunner queryRunner = queryRunnerFactory.create();
    QueryClient queryClient = new QueryClient(queryRunner, io.dropwizard.util.Duration.seconds(60), query);

    try {
        queryClient.executeWith(new Function<StatementClient, Void>() {
            @Nullable
            @Override
            public Void apply(StatementClient client)
            {
                QueryResults results = client.current();
                if (results.getData() != null) {
                    cache.addAll(results.getData());
                }

                return null;
            }
        });
    }
    catch (QueryClient.QueryTimeOutException e) {
        log.error("Caught timeout loading columns", e);
    }

    return cache.build();
}

@Override
public void addResults(QueryResults results)
{
    if (types.get() == null && results.getColumns() != null) {
        types.set(getTypes(results.getColumns()));
    }

    if (results.getData() != null) {
        checkState(types.get() != null, "Data without types received!");
        List<Column> columns = results.getColumns();
        for (List<Object> fields : results.getData()) {
            String redisKey = tableName + ":" + count.getAndIncrement();

            try (Jedis jedis = jedisPool.getResource()) {
                switch (dataFormat) {
                    case "string":
                        ImmutableMap.Builder<String, Object> builder = ImmutableMap.builder();
                        for (int i = 0; i < fields.size(); i++) {
                            Type type = types.get().get(i);
                            Object value = convertValue(fields.get(i), type);
                            if (value != null) {
                                builder.put(columns.get(i).getName(), value);
                            }
                        }
                        jedis.set(redisKey, jsonEncoder.toString(builder.build()));
                        break;
                    case "hash":
                        // add keys to zset
                        String redisZset = "keyset:" + tableName;
                        jedis.zadd(redisZset, count.get(), redisKey);
                        // add values to Hash
                        for (int i = 0; i < fields.size(); i++) {
                            jedis.hset(redisKey, columns.get(i).getName(), fields.get(i).toString());
                        }
                        break;
                    default:
                        throw new AssertionError("unhandled value type: " + dataFormat);
                }
            }
        }
    }
}

private static String failureMessage(QueryResults results)
{
    return format("Query failed (#%s): %s", results.getId(), results.getError().getMessage());
}

@Test
public void testQuery()
        throws Exception
{
    // start query
    Request request = preparePost()
            .setUri(uriFor("/v1/statement"))
            .setBodyGenerator(createStaticBodyGenerator("show catalogs", UTF_8))
            .setHeader(PRESTO_USER, "user")
            .setHeader(PRESTO_SOURCE, "source")
            .setHeader(PRESTO_CATALOG, "catalog")
            .setHeader(PRESTO_SCHEMA, "schema")
            .addHeader(PRESTO_SESSION, QUERY_MAX_MEMORY + "=1GB")
            .addHeader(PRESTO_SESSION, DISTRIBUTED_JOIN + "=true," + HASH_PARTITION_COUNT + " = 43")
            .build();

    QueryResults queryResults = client.execute(request, createJsonResponseHandler(jsonCodec(QueryResults.class)));

    // get the query info
    QueryInfo queryInfo = server.getQueryManager().getQueryInfo(new QueryId(queryResults.getId()));

    // verify session properties
    assertEquals(queryInfo.getSession().getSystemProperties(), ImmutableMap.builder()
            .put(QUERY_MAX_MEMORY, "1GB")
            .put(DISTRIBUTED_JOIN, "true")
            .put(HASH_PARTITION_COUNT, "43")
            .build());

    ImmutableList.Builder<List<Object>> data = ImmutableList.builder();
    if (queryResults.getData() != null) {
        data.addAll(queryResults.getData());
    }

    while (queryResults.getNextUri() != null) {
        queryResults = client.execute(prepareGet().setUri(queryResults.getNextUri()).build(), createJsonResponseHandler(jsonCodec(QueryResults.class)));

        if (queryResults.getData() != null) {
            data.addAll(queryResults.getData());
        }
    }

    // only the system catalog exists by default
    List<List<Object>> rows = data.build();
    assertEquals(rows, ImmutableList.of(ImmutableList.of("system")));
}

static QueryExecutor create(String userAgent)
{
    return new QueryExecutor(userAgent, jsonCodec(QueryResults.class), jsonCodec(ServerInfo.class), getSystemSocksProxy());
}

private void renderQueryOutput(PrintStream out, OutputFormat outputFormat, boolean interactive)
{
    StatusPrinter statusPrinter = null;
    @SuppressWarnings("resource")
    PrintStream errorChannel = interactive ? out : System.err;

    if (interactive) {
        statusPrinter = new StatusPrinter(client, out);
        statusPrinter.printInitialStatusUpdates();
    }
    else {
        waitForData();
    }

    if ((!client.isFailed()) && (!client.isGone()) && (!client.isClosed())) {
        QueryResults results = client.isValid() ? client.current() : client.finalResults();
        if (results.getUpdateType() != null) {
            renderUpdate(out, results);
        }
        else if (results.getColumns() == null) {
            errorChannel.printf("Query %s has no columns\n", results.getId());
            return;
        }
        else {
            renderResults(out, outputFormat, interactive, results.getColumns());
        }
    }

    if (statusPrinter != null) {
        statusPrinter.printFinalInfo();
    }

    if (client.isClosed()) {
        errorChannel.println("Query aborted by user");
    }
    else if (client.isGone()) {
        errorChannel.println("Query is gone (server restarted?)");
    }
    else if (client.isFailed()) {
        renderFailure(errorChannel);
    }
}

public void printFinalInfo()
{
    Duration wallTime = nanosSince(start);

    QueryResults results = client.finalResults();
    StatementStats stats = results.getStats();

    int nodes = stats.getNodes();
    if ((nodes == 0) || (stats.getTotalSplits() == 0)) {
        return;
    }

    // blank line
    out.println();

    // Query 12, FINISHED, 1 node
    String querySummary = String.format("Query %s, %s, %,d %s",
            results.getId(),
            stats.getState(),
            nodes,
            pluralize("node", nodes));
    out.println(querySummary);

    if (debug) {
        out.println(results.getInfoUri() + "?pretty");
    }

    // Splits: 1000 total, 842 done (84.20%)
    String splitsSummary = String.format("Splits: %,d total, %,d done (%.2f%%)",
            stats.getTotalSplits(),
            stats.getCompletedSplits(),
            percentage(stats.getCompletedSplits(), stats.getTotalSplits()));
    out.println(splitsSummary);

    if (debug) {
        // CPU Time: 565.2s total,   26K rows/s, 3.85MB/s
        Duration cpuTime = millis(stats.getCpuTimeMillis());
        String cpuTimeSummary = String.format("CPU Time: %.1fs total, %5s rows/s, %8s, %d%% active",
                cpuTime.getValue(SECONDS),
                formatCountRate(stats.getProcessedRows(), cpuTime, false),
                formatDataRate(bytes(stats.getProcessedBytes()), cpuTime, true),
                (int) percentage(stats.getCpuTimeMillis(), stats.getWallTimeMillis()));
        out.println(cpuTimeSummary);

        double parallelism = cpuTime.getValue(MILLISECONDS) / wallTime.getValue(MILLISECONDS);

        // Per Node: 3.5 parallelism, 83.3K rows/s, 0.7 MB/s
        String perNodeSummary = String.format("Per Node: %.1f parallelism, %5s rows/s, %8s",
                parallelism / nodes,
                formatCountRate((double) stats.getProcessedRows() / nodes, wallTime, false),
                formatDataRate(bytes(stats.getProcessedBytes() / nodes), wallTime, true));
        reprintLine(perNodeSummary);

        out.println(String.format("Parallelism: %.1f", parallelism));
    }

    // 0:32 [2.12GB, 15M rows] [67MB/s, 463K rows/s]
    String statsLine = String.format("%s [%s rows, %s] [%s rows/s, %s]",
            formatTime(wallTime),
            formatCount(stats.getProcessedRows()),
            formatDataSize(bytes(stats.getProcessedBytes()), true),
            formatCountRate(stats.getProcessedRows(), wallTime, false),
            formatDataRate(bytes(stats.getProcessedBytes()), wallTime, true));

    out.println(statsLine);

    // blank line
    out.println();
}

protected QueryRunner(ClientSession session, JsonCodec<QueryResults> queryResultsCodec, HttpClient httpClient)
{
    this.session = checkNotNull(session, "session is null");
    this.queryResultsCodec = checkNotNull(queryResultsCodec, "queryResultsCodec is null");
    this.httpClient = httpClient;
}

public QueryRunner create(String user, String schema)
{
    return new QueryRunner(sessionFactory.create(user, schema), jsonCodec(QueryResults.class), httpClient);
}

public QueryRunner create()
{
    return new QueryRunner(sessionFactory.create(), jsonCodec(QueryResults.class), httpClient);
}

private Map<String, List<String>> queryMetadata(String query)
{
    final Map<String, List<String>> cache = Maps.newHashMap();
    QueryRunner queryRunner = queryRunnerFactory.create();
    QueryClient queryClient = new QueryClient(queryRunner, io.dropwizard.util.Duration.seconds(60), query);

    try {
        queryClient.executeWith(new Function<StatementClient, Void>() {
            @Nullable
            @Override
            public Void apply(StatementClient client)
            {
                QueryResults results = client.current();
                if (results.getData() != null) {
                    for (List<Object> row : results.getData()) {
                        String schema = (String) row.get(1);
                        String table = (String) row.get(2);

                        if (EXCLUDED_SCHEMAS.contains(schema)) {
                            continue;
                        }

                        List<String> tables = cache.get(schema);

                        if (tables == null) {
                            tables = Lists.newArrayList();
                            cache.put(schema, tables);
                        }

                        tables.add(table);
                    }
                }

                return null;
            }
        });
    }
    catch (QueryClient.QueryTimeOutException e) {
        log.error("Caught timeout loading columns", e);
    }

    return ImmutableMap.copyOf(cache);
}

public QueryResults finalResults()
{
    return finalResults.get();
}

void addResults(QueryResults result);