@GET public List<BasicQueryInfo> getAllQueryInfo(@QueryParam("state") String stateFilter) { QueryState expectedState = stateFilter == null ? null : QueryState.valueOf(stateFilter.toUpperCase(Locale.ENGLISH)); ImmutableList.Builder<BasicQueryInfo> builder = new ImmutableList.Builder<>(); for (BasicQueryInfo queryInfo : queryManager.getQueries()) { if (stateFilter == null || queryInfo.getState() == expectedState) { builder.add(queryInfo); } } return builder.build(); }
private void cancelAllQueries() { QueryManager queryManager = coordinator.getQueryManager(); for (BasicQueryInfo queryInfo : queryManager.getQueries()) { if (!queryInfo.getState().isDone()) { queryManager.cancelQuery(queryInfo.getQueryId()); } } }
@GET @Path("{queryId}") public Response getQueryInfo(@PathParam("queryId") QueryId queryId) { requireNonNull(queryId, "queryId is null"); try { QueryInfo queryInfo = queryManager.getFullQueryInfo(queryId); return Response.ok(queryInfo).build(); } catch (NoSuchElementException e) { return Response.status(Status.GONE).build(); } }
private Response failQuery(QueryId queryId, PrestoException queryException) { requireNonNull(queryId, "queryId is null"); try { QueryState state = queryManager.getQueryState(queryId); // check before killing to provide the proper error code (this is racy) if (state.isDone()) { return Response.status(Status.CONFLICT).build(); } queryManager.failQuery(queryId, queryException); // verify if the query was failed (if not, we lost the race) if (!queryException.getErrorCode().equals(queryManager.getQueryInfo(queryId).getErrorCode())) { return Response.status(Status.CONFLICT).build(); } return Response.status(Status.OK).build(); } catch (NoSuchElementException e) { return Response.status(Status.GONE).build(); } }
private synchronized ListenableFuture<?> getFutureStateChange() { // ensure the query has been submitted submissionFuture.submitQuery(); // if query query submission has not finished, wait for it to finish if (!submissionFuture.isDone()) { return submissionFuture; } // if the exchange client is open, wait for data if (!exchangeClient.isClosed()) { return exchangeClient.isBlocked(); } // otherwise, wait for the query to finish queryManager.recordHeartbeat(queryId); try { return queryDoneFuture(queryManager.getQueryState(queryId)); } catch (NoSuchElementException e) { return immediateFuture(null); } }
@Test(timeOut = 60_000L) public void testFailQuery() throws Exception { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); QueryId queryId = queryManager.createQueryId(); queryManager.createQuery( queryId, new TestingSessionContext(TEST_SESSION), "SELECT * FROM lineitem") .get(); // wait until query starts running while (true) { QueryState state = queryManager.getQueryState(queryId); if (state.isDone()) { fail("unexpected query state: " + state); } if (state == RUNNING) { break; } Thread.sleep(100); } // cancel query queryManager.failQuery(queryId, new PrestoException(GENERIC_INTERNAL_ERROR, "mock exception")); QueryInfo queryInfo = queryManager.getFullQueryInfo(queryId); assertEquals(queryInfo.getState(), FAILED); assertEquals(queryInfo.getErrorCode(), GENERIC_INTERNAL_ERROR.toErrorCode()); assertNotNull(queryInfo.getFailureInfo()); assertEquals(queryInfo.getFailureInfo().getMessage(), "mock exception"); }
@Test public void testMetadataIsClearedAfterQueryCanceled() throws Exception { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); QueryId queryId = queryManager.createQueryId(); queryManager.createQuery( queryId, new TestingSessionContext(TEST_SESSION), "SELECT * FROM lineitem") .get(); // wait until query starts running while (true) { QueryInfo queryInfo = queryManager.getFullQueryInfo(queryId); if (queryInfo.getState().isDone()) { assertEquals(queryInfo.getState(), FAILED); throw queryInfo.getFailureInfo().toException(); } if (queryInfo.getState() == RUNNING) { break; } Thread.sleep(100); } // cancel query queryManager.cancelQuery(queryId); assertEquals(metadataManager.getCatalogsByQueryId().size(), 0); }
@Test public void testQueryLoggingCount() { QueryManager queryManager = ((DistributedQueryRunner) getQueryRunner()).getCoordinator().getQueryManager(); executeExclusively(() -> { assertUntilTimeout( () -> assertEquals( queryManager.getQueries().stream() .map(BasicQueryInfo::getQueryId) .map(queryManager::getFullQueryInfo) .filter(info -> !info.isFinalQueryInfo()) .collect(toList()), ImmutableList.of()), new Duration(1, MINUTES)); // We cannot simply get the number of completed queries as soon as all the queries are completed, because this counter may not be up-to-date at that point. // The completed queries counter is updated in a final query info listener, which is called eventually. // Therefore, here we wait until the value of this counter gets stable. long beforeCompletedQueriesCount = waitUntilStable(() -> queryManager.getStats().getCompletedQueries().getTotalCount(), new Duration(5, SECONDS)); long beforeSubmittedQueriesCount = queryManager.getStats().getSubmittedQueries().getTotalCount(); assertUpdate("CREATE TABLE test_query_logging_count AS SELECT 1 foo_1, 2 foo_2_4", 1); assertQuery("SELECT foo_1, foo_2_4 FROM test_query_logging_count", "SELECT 1, 2"); assertUpdate("DROP TABLE test_query_logging_count"); assertQueryFails("SELECT * FROM test_query_logging_count", ".*Table .* does not exist"); // TODO: Figure out a better way of synchronization assertUntilTimeout( () -> assertEquals(queryManager.getStats().getCompletedQueries().getTotalCount() - beforeCompletedQueriesCount, 4), new Duration(1, MINUTES)); assertEquals(queryManager.getStats().getSubmittedQueries().getTotalCount() - beforeSubmittedQueriesCount, 4); }); }
public static void waitForQueryState(DistributedQueryRunner queryRunner, QueryId queryId, Set<QueryState> expectedQueryStates) throws InterruptedException { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); do { // Heartbeat all the running queries, so they don't die while we're waiting for (BasicQueryInfo queryInfo : queryManager.getQueries()) { if (queryInfo.getState() == RUNNING) { queryManager.recordHeartbeat(queryInfo.getQueryId()); } } MILLISECONDS.sleep(500); } while (!expectedQueryStates.contains(queryManager.getQueryState(queryId))); }
List<QueryInfo> queryInfos = queryManager.getQueries().stream() .map(BasicQueryInfo::getQueryId) .map(queryId -> { try { return queryManager.getFullQueryInfo(queryId);
session = queryManager.getFullQueryInfo(queryId).getSession().toSession(sessionPropertyManager); serde = new PagesSerdeFactory(blockEncodingSerde, isExchangeCompressionEnabled(session)).createPagesSerde(); queryManager.failQuery(queryId, cause); QueryInfo queryInfo = queryManager.getFullQueryInfo(queryId); queryManager.recordHeartbeat(queryId);
@GET @Path("{queryId}") @Produces(MediaType.APPLICATION_JSON) public QueryStateInfo getQueryStateInfo(@PathParam("queryId") String queryId) throws WebApplicationException { try { return getQueryStateInfo(queryManager.getQueryInfo(new QueryId(queryId))); } catch (NoSuchElementException e) { throw new WebApplicationException(NOT_FOUND); } } }
@DELETE @Path("{queryId}") public void cancelQuery(@PathParam("queryId") QueryId queryId) { requireNonNull(queryId, "queryId is null"); queryManager.cancelQuery(queryId); }
public static Query create( SessionContext sessionContext, String query, QueryManager queryManager, SessionPropertyManager sessionPropertyManager, ExchangeClient exchangeClient, Executor dataProcessorExecutor, ScheduledExecutorService timeoutExecutor, BlockEncodingSerde blockEncodingSerde) { Query result = new Query(sessionContext, query, queryManager, sessionPropertyManager, exchangeClient, dataProcessorExecutor, timeoutExecutor, blockEncodingSerde); // register listeners after submission finishes addSuccessCallback(result.submissionFuture, () -> { result.queryManager.addOutputInfoListener(result.getQueryId(), result::setQueryOutputInfo); result.queryManager.addStateChangeListener(result.getQueryId(), state -> { if (state.isDone()) { QueryInfo queryInfo = queryManager.getFullQueryInfo(result.getQueryId()); result.closeExchangeClientIfNecessary(queryInfo); } }); }); return result; }
@Test(timeOut = 60_000L) public void testLegacyQueryContext() throws Exception { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); QueryId queryId = queryManager.createQueryId(); queryManager.createQuery( queryId, new TestingSessionContext(TEST_SESSION), "SELECT * FROM lineitem") .get(); waitForQueryState(queryRunner, queryId, RUNNING); // cancel query queryManager.failQuery(queryId, new PrestoException(GENERIC_INTERNAL_ERROR, "mock exception")); // assert that LegacyQueryContext is used instead of the DefaultQueryContext SqlTaskManager taskManager = (SqlTaskManager) queryRunner.getServers().get(0).getTaskManager(); assertInstanceOf(taskManager.getQueryContext(queryId), LegacyQueryContext.class); } }
private synchronized void submitQuery() { if (querySubmissionFuture != null) { return; } querySubmissionFuture = queryManager.createQuery(queryId, sessionContext, this.query); Futures.addCallback(querySubmissionFuture, new FutureCallback<Object>() { @Override public void onSuccess(Object result) { set(null); } @Override public void onFailure(Throwable t) { setException(t); } }, directExecutor()); }
private synchronized Optional<QueryResults> getCachedResult(long token, UriInfo uriInfo) { // is the a repeated request for the last results? String requestedPath = uriInfo.getAbsolutePath().getPath(); if (requestedPath.equals(lastResultPath)) { if (submissionFuture.isDone()) { // tell query manager we are still interested in the query queryManager.recordHeartbeat(queryId); } return Optional.of(lastResult); } if (token < resultId.get()) { throw new WebApplicationException(Response.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(Response.Status.NOT_FOUND); } return Optional.empty(); }
private Query( SessionContext sessionContext, String query, QueryManager queryManager, SessionPropertyManager sessionPropertyManager, ExchangeClient exchangeClient, Executor resultsProcessorExecutor, ScheduledExecutorService timeoutExecutor, BlockEncodingSerde blockEncodingSerde) { requireNonNull(sessionContext, "sessionContext is null"); requireNonNull(query, "query is null"); requireNonNull(queryManager, "queryManager is null"); requireNonNull(sessionPropertyManager, "sessionPropertyManager is null"); requireNonNull(exchangeClient, "exchangeClient is null"); requireNonNull(resultsProcessorExecutor, "resultsProcessorExecutor is null"); requireNonNull(timeoutExecutor, "timeoutExecutor is null"); requireNonNull(blockEncodingSerde, "serde is null"); this.queryManager = queryManager; this.sessionPropertyManager = sessionPropertyManager; queryId = queryManager.createQueryId(); submissionFuture = new QuerySubmissionFuture(queryId, query, sessionContext, queryManager); this.exchangeClient = exchangeClient; this.resultsProcessorExecutor = resultsProcessorExecutor; this.timeoutExecutor = timeoutExecutor; this.blockEncodingSerde = blockEncodingSerde; }
if (queryManager.getQueryState(queryId) == QueryState.FAILED) { query.dispose();
private Response failQuery(QueryId queryId, PrestoException queryException) { requireNonNull(queryId, "queryId is null"); try { QueryState state = queryManager.getQueryState(queryId); // check before killing to provide the proper error code (this is racy) if (state.isDone()) { return Response.status(Status.CONFLICT).build(); } queryManager.failQuery(queryId, queryException); // verify if the query was failed (if not, we lost the race) if (!queryException.getErrorCode().equals(queryManager.getQueryInfo(queryId).getErrorCode())) { return Response.status(Status.CONFLICT).build(); } return Response.status(Status.OK).build(); } catch (NoSuchElementException e) { return Response.status(Status.GONE).build(); } }