Examples with QueryInfo io.trino.execution.QueryInfo used on opensource projects

Example 1 with QueryInfo

use of io.trino.execution.QueryInfo in project trino by trinodb.

the class WorkerResource method getThreads.

@ResourceSecurity(WEB_UI)
@GET
@Path("{nodeId}/task/{taskId}")
public Response getThreads(@PathParam("taskId") TaskId task, @PathParam("nodeId") String nodeId, @Context HttpServletRequest servletRequest, @Context HttpHeaders httpHeaders) {
    QueryId queryId = task.getQueryId();
    Optional<QueryInfo> queryInfo = dispatchManager.getFullQueryInfo(queryId);
    if (queryInfo.isPresent()) {
        try {
            checkCanViewQueryOwnedBy(sessionContextFactory.extractAuthorizedIdentity(servletRequest, httpHeaders, alternateHeaderName), queryInfo.get().getSession().toIdentity(), accessControl);
            return proxyJsonResponse(nodeId, "v1/task/" + task);
        } catch (AccessDeniedException e) {
            throw new ForbiddenException();
        }
    }
    return Response.status(Status.GONE).build();
}

Example 2 with QueryInfo

use of io.trino.execution.QueryInfo 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 QueryInfo

use of io.trino.execution.QueryInfo in project trino by trinodb.

the class QueryResource method getQueryInfo.

@ResourceSecurity(AUTHENTICATED_USER)
@GET
@Path("{queryId}")
public Response getQueryInfo(@PathParam("queryId") QueryId queryId, @Context HttpServletRequest servletRequest, @Context HttpHeaders httpHeaders) {
    requireNonNull(queryId, "queryId is null");
    Optional<QueryInfo> queryInfo = dispatchManager.getFullQueryInfo(queryId);
    if (queryInfo.isEmpty()) {
        return Response.status(Status.GONE).build();
    }
    try {
        checkCanViewQueryOwnedBy(sessionContextFactory.extractAuthorizedIdentity(servletRequest, httpHeaders, alternateHeaderName), queryInfo.get().getSession().toIdentity(), accessControl);
        return Response.ok(queryInfo.get()).build();
    } catch (AccessDeniedException e) {
        throw new ForbiddenException();
    }
}

Example 4 with QueryInfo

use of io.trino.execution.QueryInfo in project trino by trinodb.

the class TestQueryResource method testGetQueryInfoDispatchFailure.

@Test
public void testGetQueryInfoDispatchFailure() {
    String queryId = runToCompletion("SELECT");
    QueryInfo info = getQueryInfo(queryId);
    assertFalse(info.isScheduled());
    assertNotNull(info.getFailureInfo());
    assertEquals(info.getFailureInfo().getErrorCode(), SYNTAX_ERROR.toErrorCode());
    server.getAccessControl().deny(privilege("query", VIEW_QUERY));
    try {
        assertThatThrownBy(() -> getQueryInfo(queryId)).isInstanceOf(UnexpectedResponseException.class).matches(throwable -> ((UnexpectedResponseException) throwable).getStatusCode() == 403);
    } finally {
        server.getAccessControl().reset();
    }
}

Example 5 with QueryInfo

use of io.trino.execution.QueryInfo in project trino by trinodb.

the class TestQueryResource method testGetQueryInfoExecutionFailure.

@Test
public void testGetQueryInfoExecutionFailure() {
    String queryId = runToCompletion("SELECT cast(rand() AS integer) / 0");
    QueryInfo info = getQueryInfo(queryId);
    assertTrue(info.isScheduled());
    assertNotNull(info.getFailureInfo());
    assertEquals(info.getFailureInfo().getErrorCode(), DIVISION_BY_ZERO.toErrorCode());
}