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

Example 1 with QueryManager

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

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

the class TestLateMaterializationQueries method assertLazyQuery.

private void assertLazyQuery(@Language("SQL") String sql) {
    QueryManager queryManager = getDistributedQueryRunner().getCoordinator().getQueryManager();
    ResultWithQueryId<MaterializedResult> workProcessorResultResultWithQueryId = getDistributedQueryRunner().executeWithQueryId(lateMaterialization(), sql);
    QueryInfo workProcessorQueryInfo = queryManager.getFullQueryInfo(workProcessorResultResultWithQueryId.getQueryId());
    ResultWithQueryId<MaterializedResult> noWorkProcessorResultResultWithQueryId = getDistributedQueryRunner().executeWithQueryId(noLateMaterialization(), sql);
    QueryInfo noWorkProcessorQueryInfo = queryManager.getFullQueryInfo(noWorkProcessorResultResultWithQueryId.getQueryId());
    // ensure results are correct
    MaterializedResult expected = computeExpected(sql, workProcessorResultResultWithQueryId.getResult().getTypes());
    assertEqualsIgnoreOrder(workProcessorResultResultWithQueryId.getResult(), expected, "For query: \n " + sql);
    assertEqualsIgnoreOrder(noWorkProcessorResultResultWithQueryId.getResult(), expected, "For query: \n " + sql);
    // ensure work processor query processed less input data
    long workProcessorProcessedInputBytes = workProcessorQueryInfo.getQueryStats().getProcessedInputDataSize().toBytes();
    long noWorkProcessorProcessedInputBytes = noWorkProcessorQueryInfo.getQueryStats().getProcessedInputDataSize().toBytes();
    assertTrue(workProcessorProcessedInputBytes < noWorkProcessorProcessedInputBytes, "Expected work processor query to process less input data");
}

Example 3 with QueryManager

use of io.trino.execution.QueryManager 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());
}

Example 4 with QueryManager

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

the class ExecutingStatementResource method getQuery.

protected Query getQuery(QueryId queryId, String slug, long token) {
    Query query = queries.get(queryId);
    if (query != null) {
        if (!query.isSlugValid(slug, token)) {
            throw queryNotFound();
        }
        return query;
    }
    // this is the first time the query has been accessed on this coordinator
    Session session;
    Slug querySlug;
    try {
        session = queryManager.getQuerySession(queryId);
        querySlug = queryManager.getQuerySlug(queryId);
        if (!querySlug.isValid(EXECUTING_QUERY, slug, token)) {
            throw queryNotFound();
        }
    } catch (NoSuchElementException e) {
        throw queryNotFound();
    }
    query = queries.computeIfAbsent(queryId, id -> Query.create(session, querySlug, queryManager, queryInfoUrlFactory.getQueryInfoUrl(queryId), directExchangeClientSupplier, responseExecutor, timeoutExecutor, blockEncodingSerde));
    return query;
}

Example 5 with QueryManager

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

the class BaseDeltaLakeConnectorSmokeTest method testInputDataSize.

@Test
public void testInputDataSize() {
    DistributedQueryRunner queryRunner = (DistributedQueryRunner) getQueryRunner();
    queryRunner.installPlugin(new TestingHivePlugin());
    queryRunner.createCatalog("hive", "hive", ImmutableMap.of("hive.metastore.uri", dockerizedDataLake.getTestingHadoop().getMetastoreAddress(), "hive.allow-drop-table", "true"));
    String hiveTableName = "foo_hive";
    queryRunner.execute(format("CREATE TABLE hive.%s.%s (foo_id bigint, bar_id bigint, data varchar) WITH (format = 'PARQUET', external_location = '%s')", SCHEMA, hiveTableName, getLocationForTable(bucketName, "foo")));
    ResultWithQueryId<MaterializedResult> deltaResult = queryRunner.executeWithQueryId(broadcastJoinDistribution(true), "SELECT * FROM foo");
    assertEquals(deltaResult.getResult().getRowCount(), 2);
    ResultWithQueryId<MaterializedResult> hiveResult = queryRunner.executeWithQueryId(broadcastJoinDistribution(true), format("SELECT * FROM %s.%s.%s", "hive", SCHEMA, hiveTableName));
    assertEquals(hiveResult.getResult().getRowCount(), 2);
    QueryManager queryManager = queryRunner.getCoordinator().getQueryManager();
    assertThat(queryManager.getFullQueryInfo(deltaResult.getQueryId()).getQueryStats().getProcessedInputDataSize()).as("delta processed input data size").isGreaterThan(DataSize.ofBytes(0)).isEqualTo(queryManager.getFullQueryInfo(hiveResult.getQueryId()).getQueryStats().getProcessedInputDataSize());
    queryRunner.execute(format("DROP TABLE hive.%s.%s", SCHEMA, hiveTableName));
}