public QueryInfo getQueryInfo(QueryId queryId) { return coordinator.getQueryManager().getFullQueryInfo(queryId); }
private Set<String> captureQueries(Callable<?> action) throws Exception { Set<QueryId> queryIdsBefore = server.getQueryManager().getQueries().stream() .map(BasicQueryInfo::getQueryId) .collect(toImmutableSet()); action.call(); return server.getQueryManager().getQueries().stream() .filter(queryInfo -> !queryIdsBefore.contains(queryInfo.getQueryId())) .map(BasicQueryInfo::getQuery) .collect(toImmutableSet()); }
@Override public Plan createPlan(Session session, String sql, WarningCollector warningCollector) { QueryId queryId = executeWithQueryId(session, sql).getQueryId(); Plan queryPlan = getQueryPlan(queryId); coordinator.getQueryManager().cancelQuery(queryId); return queryPlan; }
public static void waitForQueryCount(DistributedQueryRunner queryRunner, Set<QueryState> countingStates, int expectedCount) throws InterruptedException { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); while (queryManager.getQueries().stream() .filter(q -> countingStates.contains(q.getState())).count() != expectedCount) { MILLISECONDS.sleep(500); } }
private void assertResourceGroup(DistributedQueryRunner queryRunner, Session session, String query, ResourceGroupId expectedResourceGroup) throws InterruptedException { QueryId queryId = createQuery(queryRunner, session, query); waitForQueryState(queryRunner, queryId, ImmutableSet.of(RUNNING, FINISHED)); Optional<ResourceGroupId> resourceGroupId = queryRunner.getCoordinator().getQueryManager().getFullQueryInfo(queryId).getResourceGroupId(); assertTrue(resourceGroupId.isPresent(), "Query should have a resource group"); assertEquals(resourceGroupId.get(), expectedResourceGroup, format("Expected: '%s' resource group, found: %s", expectedResourceGroup, resourceGroupId.get())); }
private void cancelAllQueries() { QueryManager queryManager = coordinator.getQueryManager(); for (BasicQueryInfo queryInfo : queryManager.getQueries()) { if (!queryInfo.getState().isDone()) { queryManager.cancelQuery(queryInfo.getQueryId()); } } }
public static void cancelQuery(DistributedQueryRunner queryRunner, QueryId queryId) { queryRunner.getCoordinator().getQueryManager().cancelQuery(queryId); }
@Test public void testQueryTypeBasedSelection() throws InterruptedException { Session session = testSessionBuilder() .setCatalog("tpch") .setSchema("sf100000") .build(); QueryId queryId = createQuery(queryRunner, session, "EXPLAIN " + LONG_LASTING_QUERY); waitForQueryState(queryRunner, queryId, ImmutableSet.of(RUNNING, FINISHED)); Optional<ResourceGroupId> resourceGroupId = queryRunner.getCoordinator().getQueryManager().getFullQueryInfo(queryId).getResourceGroupId(); assertTrue(resourceGroupId.isPresent(), "Query should have a resource group"); assertEquals(resourceGroupId.get(), createResourceGroupId("explain")); }
private void assertResourceGroupWithClientTags(Set<String> clientTags, ResourceGroupId expectedResourceGroup) throws InterruptedException { Session session = testSessionBuilder() .setCatalog("tpch") .setSchema("sf100000") .setSource("client_tags") .setClientTags(clientTags) .build(); QueryId queryId = createQuery(queryRunner, session, LONG_LASTING_QUERY); waitForQueryState(queryRunner, queryId, ImmutableSet.of(RUNNING, FINISHED)); Optional<ResourceGroupId> resourceGroupId = queryRunner.getCoordinator().getQueryManager().getFullQueryInfo(queryId).getResourceGroupId(); assertTrue(resourceGroupId.isPresent(), "Query should have a resource group"); assertEquals(resourceGroupId.get(), expectedResourceGroup, format("Expected: '%s' resource group, found: %s", expectedResourceGroup, resourceGroupId.get())); } }
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))); }
private void testNoLeak(@Language("SQL") String query) throws Exception { Map<String, String> properties = ImmutableMap.<String, String>builder() .put("task.verbose-stats", "true") .build(); try (DistributedQueryRunner queryRunner = createQueryRunner(TINY_SESSION, properties)) { executor.submit(() -> queryRunner.execute(query)).get(); for (BasicQueryInfo info : queryRunner.getCoordinator().getQueryManager().getQueries()) { assertEquals(info.getState(), FINISHED); } // Make sure we didn't leak any memory on the workers for (TestingPrestoServer worker : queryRunner.getServers()) { Optional<MemoryPool> reserved = worker.getLocalMemoryManager().getReservedPool(); assertTrue(reserved.isPresent()); assertEquals(reserved.get().getMaxBytes(), reserved.get().getFreeBytes()); MemoryPool general = worker.getLocalMemoryManager().getGeneralPool(); assertEquals(general.getMaxBytes(), general.getFreeBytes()); } } }
private void waitForQueryToBeKilled(DistributedQueryRunner queryRunner) throws InterruptedException { while (true) { for (BasicQueryInfo info : queryRunner.getCoordinator().getQueryManager().getQueries()) { if (info.getState().isDone()) { assertNotNull(info.getErrorCode()); assertEquals(info.getErrorCode().getCode(), CLUSTER_OUT_OF_MEMORY.toErrorCode().getCode()); return; } } MILLISECONDS.sleep(10); } }
@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); }); }
@Test(timeOut = 60_000L) public void testQueryCpuLimit() throws Exception { try (DistributedQueryRunner queryRunner = builder().setSingleExtraProperty("query.max-cpu-time", "1ms").build()) { QueryId queryId = createQuery(queryRunner, TEST_SESSION, "SELECT COUNT(*) FROM lineitem"); waitForQueryState(queryRunner, queryId, FAILED); QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); BasicQueryInfo queryInfo = queryManager.getQueryInfo(queryId); assertEquals(queryInfo.getState(), FAILED); assertEquals(queryInfo.getErrorCode(), EXCEEDED_CPU_LIMIT.toErrorCode()); } } }
@Test(timeOut = 60_000) public void testSelectorPriority() throws Exception { InternalResourceGroupManager manager = queryRunner.getCoordinator().getResourceGroupManager().get(); QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); DbResourceGroupConfigurationManager dbConfigurationManager = (DbResourceGroupConfigurationManager) manager.getConfigurationManager(); QueryId firstQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, firstQuery, RUNNING); Optional<ResourceGroupId> resourceGroup = queryManager.getFullQueryInfo(firstQuery).getResourceGroupId(); assertTrue(resourceGroup.isPresent()); assertEquals(resourceGroup.get().toString(), "global.user-user.dashboard-user"); // create a new resource group that rejects all queries submitted to it dao.insertResourceGroup(8, "reject-all-queries", "1MB", 0, 0, 0, null, null, null, null, null, 3L, TEST_ENVIRONMENT); // add a new selector that has a higher priority than the existing dashboard selector and that routes queries to the "reject-all-queries" resource group dao.insertSelector(8, 200, "user.*", "(?i).*dashboard.*", null, null, null); // reload the configuration dbConfigurationManager.load(); QueryId secondQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, secondQuery, FAILED); resourceGroup = queryManager.getFullQueryInfo(secondQuery).getResourceGroupId(); assertTrue(resourceGroup.isPresent()); assertEquals(resourceGroup.get(), createResourceGroupId("global", "user-user", "reject-all-queries")); }
public static QueryId createQuery(DistributedQueryRunner queryRunner, Session session, String sql) { QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); getFutureValue(queryManager.createQuery(session.getQueryId(), new TestingSessionContext(session), sql)); return session.getQueryId(); }
@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); }
private void testRejection() throws Exception { try (DistributedQueryRunner queryRunner = createQueryRunner()) { queryRunner.installPlugin(new ResourceGroupManagerPlugin()); queryRunner.getCoordinator().getResourceGroupManager().get().setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_config_dashboard.json"))); QueryId queryId = createQuery(queryRunner, newRejectionSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, queryId, FAILED); QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); assertEquals(queryManager.getQueryInfo(queryId).getErrorCode(), QUERY_REJECTED.toErrorCode()); } }
@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); } }
@Test public void testNonLeafGroup() throws Exception { Session session = testSessionBuilder() .setCatalog("tpch") .setSchema("sf100000") .setSource("non-leaf") .build(); QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); InternalResourceGroupManager manager = queryRunner.getCoordinator().getResourceGroupManager().get(); DbResourceGroupConfigurationManager dbConfigurationManager = (DbResourceGroupConfigurationManager) manager.getConfigurationManager(); int originalSize = getSelectors(queryRunner).size(); // Add a selector for a non leaf group dao.insertSelector(3, 100, "user.*", "(?i).*non-leaf.*", null, null, null); dbConfigurationManager.load(); while (getSelectors(queryRunner).size() != originalSize + 1) { MILLISECONDS.sleep(500); } // Submit query with side effect of creating resource groups QueryId firstDashboardQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY); waitForQueryState(queryRunner, firstDashboardQuery, RUNNING); cancelQuery(queryRunner, firstDashboardQuery); waitForQueryState(queryRunner, firstDashboardQuery, FAILED); // Submit a query to a non-leaf resource group QueryId invalidResourceGroupQuery = createQuery(queryRunner, session, LONG_LASTING_QUERY); waitForQueryState(queryRunner, invalidResourceGroupQuery, FAILED); assertEquals(queryRunner.getQueryInfo(invalidResourceGroupQuery).getErrorCode(), INVALID_RESOURCE_GROUP.toErrorCode()); }