Examples with DispatchManager io.trino.dispatcher.DispatchManager used on opensource projects

Example 1 with DispatchManager

use of io.trino.dispatcher.DispatchManager in project trino by trinodb.

the class BaseConnectorTest method testQueryLoggingCount.

@Test
public void testQueryLoggingCount() {
    skipTestUnless(hasBehavior(SUPPORTS_CREATE_TABLE));
    QueryManager queryManager = getDistributedQueryRunner().getCoordinator().getQueryManager();
    executeExclusively(() -> {
        assertEventually(new Duration(1, MINUTES), () -> assertEquals(queryManager.getQueries().stream().map(BasicQueryInfo::getQueryId).map(queryManager::getFullQueryInfo).filter(info -> !info.isFinalQueryInfo()).collect(toList()), ImmutableList.of()));
        // 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.
        DispatchManager dispatchManager = ((DistributedQueryRunner) getQueryRunner()).getCoordinator().getDispatchManager();
        long beforeCompletedQueriesCount = waitUntilStable(() -> dispatchManager.getStats().getCompletedQueries().getTotalCount(), new Duration(5, SECONDS));
        long beforeSubmittedQueriesCount = dispatchManager.getStats().getSubmittedQueries().getTotalCount();
        String tableName = "test_logging_count" + randomTableSuffix();
        assertUpdate("CREATE TABLE " + tableName + tableDefinitionForQueryLoggingCount());
        assertQueryReturnsEmptyResult("SELECT foo_1, foo_2_4 FROM " + tableName);
        assertUpdate("DROP TABLE " + tableName);
        assertQueryFails("SELECT * FROM " + tableName, ".*Table .* does not exist");
        // TODO: Figure out a better way of synchronization
        assertEventually(new Duration(1, MINUTES), () -> assertEquals(dispatchManager.getStats().getCompletedQueries().getTotalCount() - beforeCompletedQueriesCount, 4));
        assertEquals(dispatchManager.getStats().getSubmittedQueries().getTotalCount() - beforeSubmittedQueriesCount, 4);
    });
}

Example 2 with DispatchManager

use of io.trino.dispatcher.DispatchManager in project trino by trinodb.

the class TestQueues method testRejection.

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);
        DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager();
        assertEquals(dispatchManager.getQueryInfo(queryId).getErrorCode(), QUERY_REJECTED.toErrorCode());
    }
}

Example 3 with DispatchManager

use of io.trino.dispatcher.DispatchManager in project trino by trinodb.

the class TestMetadataManager method testMetadataIsClearedAfterQueryCanceled.

@Test
public void testMetadataIsClearedAfterQueryCanceled() throws Exception {
    DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager();
    QueryId queryId = dispatchManager.createQueryId();
    dispatchManager.createQuery(queryId, Slug.createNew(), TestingSessionContext.fromSession(TEST_SESSION), "SELECT * FROM lineitem").get();
    // wait until query starts running
    while (true) {
        BasicQueryInfo queryInfo = dispatchManager.getQueryInfo(queryId);
        if (queryInfo.getState().isDone()) {
            assertEquals(queryInfo.getState(), FAILED);
            throw dispatchManager.getDispatchInfo(queryId).get().getFailureInfo().get().toException();
        }
        if (queryInfo.getState() == RUNNING) {
            break;
        }
        Thread.sleep(100);
    }
    // cancel query
    dispatchManager.cancelQuery(queryId);
    assertEquals(metadataManager.getActiveQueryIds().size(), 0);
}

Example 4 with DispatchManager

use of io.trino.dispatcher.DispatchManager in project trino by trinodb.

the class TestQueuesDb method testRejection.

@Test(timeOut = 60_000)
public void testRejection() throws Exception {
    InternalResourceGroupManager<?> manager = queryRunner.getCoordinator().getResourceGroupManager().get();
    DbResourceGroupConfigurationManager dbConfigurationManager = (DbResourceGroupConfigurationManager) manager.getConfigurationManager();
    // Verify the query cannot be submitted
    QueryId queryId = createQuery(queryRunner, rejectingSession(), LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, queryId, FAILED);
    DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager();
    assertEquals(dispatchManager.getQueryInfo(queryId).getErrorCode(), QUERY_REJECTED.toErrorCode());
    int selectorCount = getSelectors(queryRunner).size();
    dao.insertSelector(4, 100_000, "user.*", null, "(?i).*reject.*", null, null, null);
    dbConfigurationManager.load();
    assertEquals(getSelectors(queryRunner).size(), selectorCount + 1);
    // Verify the query can be submitted
    queryId = createQuery(queryRunner, rejectingSession(), LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, queryId, RUNNING);
    dao.deleteSelector(4, "user.*", "(?i).*reject.*", null);
    dbConfigurationManager.load();
    // Verify the query cannot be submitted
    queryId = createQuery(queryRunner, rejectingSession(), LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, queryId, FAILED);
}

Example 5 with DispatchManager

use of io.trino.dispatcher.DispatchManager in project trino by trinodb.

the class TestQueuesDb method testSelectorPriority.

@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.*", null, "(?i).*dashboard.*", null, null, null);
    // reload the configuration
    dbConfigurationManager.load();
    QueryId secondQuery = createQuery(queryRunner, dashboardSession(), LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, secondQuery, FAILED);
    DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager();
    BasicQueryInfo basicQueryInfo = dispatchManager.getQueryInfo(secondQuery);
    assertEquals(basicQueryInfo.getErrorCode(), QUERY_QUEUE_FULL.toErrorCode());
}