Examples with QueryResults io.trino.client.QueryResults used on opensource projects

Example 1 with QueryResults

use of io.trino.client.QueryResults in project trino by trinodb.

the class TestProgressMonitor method newQueryResults.

private String newQueryResults(Integer partialCancelId, Integer nextUriId, List<Column> responseColumns, List<List<Object>> data, String state) {
    String queryId = "20160128_214710_00012_rk68b";
    QueryResults queryResults = new QueryResults(queryId, server.url("/query.html?" + queryId).uri(), partialCancelId == null ? null : server.url(format("/v1/statement/partialCancel/%s.%s", queryId, partialCancelId)).uri(), nextUriId == null ? null : server.url(format("/v1/statement/%s/%s", queryId, nextUriId)).uri(), responseColumns, data, new StatementStats(state, state.equals("QUEUED"), true, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, null), null, ImmutableList.of(), null, null);
    return QUERY_RESULTS_CODEC.toJson(queryResults);
}

Example 2 with QueryResults

use of io.trino.client.QueryResults in project trino by trinodb.

the class Query method getNextResult.

private synchronized QueryResults getNextResult(long token, UriInfo uriInfo, DataSize targetResultSize) {
    // check if the result for the token have already been created
    Optional<QueryResults> cachedResult = getCachedResult(token);
    if (cachedResult.isPresent()) {
        return cachedResult.get();
    }
    verify(nextToken.isPresent(), "Cannot generate next result when next token is not present");
    verify(token == nextToken.getAsLong(), "Expected token to equal next token");
    // get the query info before returning
    // force update if query manager is closed
    QueryInfo queryInfo = queryManager.getFullQueryInfo(queryId);
    queryManager.recordHeartbeat(queryId);
    closeExchangeClientIfNecessary(queryInfo);
    // fetch result data from exchange
    QueryResultRows resultRows = removePagesFromExchange(queryInfo, targetResultSize.toBytes());
    if ((queryInfo.getUpdateType() != null) && (updateCount == null)) {
        // grab the update count for non-queries
        Optional<Long> updatedRowsCount = resultRows.getUpdateCount();
        updateCount = updatedRowsCount.orElse(null);
    }
    // (3) cached query result needs client acknowledgement to discard
    if (queryInfo.getState() != FAILED && (!queryInfo.isFinalQueryInfo() || !exchangeClient.isFinished() || (queryInfo.getOutputStage().isPresent() && !resultRows.isEmpty()))) {
        nextToken = OptionalLong.of(token + 1);
    } else {
        nextToken = OptionalLong.empty();
        // the client is not coming back, make sure the exchangeClient is closed
        exchangeClient.close();
    }
    URI nextResultsUri = null;
    URI partialCancelUri = null;
    if (nextToken.isPresent()) {
        long nextToken = this.nextToken.getAsLong();
        nextResultsUri = createNextResultsUri(uriInfo, nextToken);
        partialCancelUri = findCancelableLeafStage(queryInfo).map(stage -> createPartialCancelUri(stage, uriInfo, nextToken)).orElse(null);
    }
    // update catalog, schema, and path
    setCatalog = queryInfo.getSetCatalog();
    setSchema = queryInfo.getSetSchema();
    setPath = queryInfo.getSetPath();
    // update setSessionProperties
    setSessionProperties = queryInfo.getSetSessionProperties();
    resetSessionProperties = queryInfo.getResetSessionProperties();
    // update setRoles
    setRoles = queryInfo.getSetRoles();
    // update preparedStatements
    addedPreparedStatements = queryInfo.getAddedPreparedStatements();
    deallocatedPreparedStatements = queryInfo.getDeallocatedPreparedStatements();
    // update startedTransactionId
    startedTransactionId = queryInfo.getStartedTransactionId();
    clearTransactionId = queryInfo.isClearTransactionId();
    // first time through, self is null
    QueryResults queryResults = new QueryResults(queryId.toString(), getQueryInfoUri(queryInfoUrl, queryId, uriInfo), partialCancelUri, nextResultsUri, resultRows.getColumns().orElse(null), // client excepts null that indicates "no data"
    resultRows.isEmpty() ? null : resultRows, toStatementStats(queryInfo), toQueryError(queryInfo, typeSerializationException), mappedCopy(queryInfo.getWarnings(), Query::toClientWarning), queryInfo.getUpdateType(), updateCount);
    // cache the new result
    lastToken = token;
    lastResult = queryResults;
    return queryResults;
}

Example 3 with QueryResults

use of io.trino.client.QueryResults in project trino by trinodb.

the class TestQueryStateInfoResource method setUp.

@BeforeClass
public void setUp() {
    server = TestingTrinoServer.create();
    server.installPlugin(new TpchPlugin());
    server.createCatalog("tpch", "tpch");
    client = new JettyHttpClient();
    Request request1 = preparePost().setUri(uriBuilderFrom(server.getBaseUrl()).replacePath("/v1/statement").build()).setBodyGenerator(createStaticBodyGenerator(LONG_LASTING_QUERY, UTF_8)).setHeader(TRINO_HEADERS.requestUser(), "user1").build();
    queryResults = client.execute(request1, createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC));
    client.execute(prepareGet().setUri(queryResults.getNextUri()).build(), createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC));
    Request request2 = preparePost().setUri(uriBuilderFrom(server.getBaseUrl()).replacePath("/v1/statement").build()).setBodyGenerator(createStaticBodyGenerator(LONG_LASTING_QUERY, UTF_8)).setHeader(TRINO_HEADERS.requestUser(), "user2").build();
    QueryResults queryResults2 = client.execute(request2, createJsonResponseHandler(jsonCodec(QueryResults.class)));
    client.execute(prepareGet().setUri(queryResults2.getNextUri()).build(), createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC));
    // queries are started in the background, so they may not all be immediately visible
    long start = System.nanoTime();
    while (Duration.nanosSince(start).compareTo(new Duration(5, MINUTES)) < 0) {
        List<BasicQueryInfo> queryInfos = client.execute(prepareGet().setUri(uriBuilderFrom(server.getBaseUrl()).replacePath("/v1/query").build()).setHeader(TRINO_HEADERS.requestUser(), "unknown").build(), createJsonResponseHandler(listJsonCodec(BasicQueryInfo.class)));
        if (queryInfos.size() == 2) {
            if (queryInfos.stream().allMatch(info -> info.getState() == RUNNING)) {
                break;
            }
            List<ErrorCode> errorCodes = queryInfos.stream().filter(info -> info.getState() == FAILED).map(BasicQueryInfo::getErrorCode).collect(toImmutableList());
            if (!errorCodes.isEmpty()) {
                fail("setup queries failed with: " + errorCodes);
            }
        }
    }
}

Example 4 with QueryResults

use of io.trino.client.QueryResults in project trino by trinodb.

the class TestQueryResource method testIdempotentResults.

@Test
public void testIdempotentResults() {
    String sql = "SELECT * FROM tpch.tiny.lineitem";
    Request request = preparePost().setHeader(TRINO_HEADERS.requestUser(), "user").setUri(uriBuilderFrom(server.getBaseUrl().resolve("/v1/statement")).build()).setBodyGenerator(createStaticBodyGenerator(sql, UTF_8)).build();
    QueryResults queryResults = client.execute(request, createJsonResponseHandler(jsonCodec(QueryResults.class)));
    URI uri = queryResults.getNextUri();
    while (uri != null) {
        QueryResults attempt1 = client.execute(prepareGet().setHeader(TRINO_HEADERS.requestUser(), "user").setUri(uri).build(), createJsonResponseHandler(jsonCodec(QueryResults.class)));
        QueryResults attempt2 = client.execute(prepareGet().setHeader(TRINO_HEADERS.requestUser(), "user").setUri(uri).build(), createJsonResponseHandler(jsonCodec(QueryResults.class)));
        assertEquals(attempt2.getData(), attempt1.getData());
        uri = attempt1.getNextUri();
    }
}

Example 5 with QueryResults

use of io.trino.client.QueryResults in project trino by trinodb.

the class TestQueryResource method startQuery.

private String startQuery(String sql) {
    URI uri = uriBuilderFrom(server.getBaseUrl()).replacePath("/v1/statement").build();
    Request request = preparePost().setUri(uri).setBodyGenerator(createStaticBodyGenerator(sql, UTF_8)).setHeader(TRINO_HEADERS.requestUser(), "user").build();
    QueryResults queryResults = client.execute(request, createJsonResponseHandler(jsonCodec(QueryResults.class)));
    while (queryResults.getNextUri() != null && !queryResults.getStats().getState().equals(RUNNING.toString())) {
        request = prepareGet().setHeader(TRINO_HEADERS.requestUser(), "user").setUri(queryResults.getNextUri()).build();
        queryResults = client.execute(request, createJsonResponseHandler(jsonCodec(QueryResults.class)));
    }
    return queryResults.getId();
}