Examples with ResultWithQueryId io.trino.testing.ResultWithQueryId used on opensource projects

Example 1 with ResultWithQueryId

use of io.trino.testing.ResultWithQueryId in project trino by trinodb.

the class TestMinWorkerRequirement method testMultipleRequiredWorkerNodesSessionOverride.

@Test
public void testMultipleRequiredWorkerNodesSessionOverride() throws Exception {
    ListeningExecutorService service = MoreExecutors.listeningDecorator(newFixedThreadPool(3));
    try (DistributedQueryRunner queryRunner = TpchQueryRunnerBuilder.builder().setNodeCount(1).build()) {
        Session session1 = testSessionBuilder().setSystemProperty(REQUIRED_WORKERS_COUNT, "2").setCatalog("tpch").setSchema("tiny").build();
        ListenableFuture<ResultWithQueryId<MaterializedResult>> queryFuture1 = service.submit(() -> queryRunner.executeWithQueryId(session1, "SELECT COUNT(*) from lineitem"));
        Session session2 = Session.builder(session1).setSystemProperty(REQUIRED_WORKERS_COUNT, "3").build();
        ListenableFuture<ResultWithQueryId<MaterializedResult>> queryFuture2 = service.submit(() -> queryRunner.executeWithQueryId(session2, "SELECT COUNT(*) from lineitem"));
        Session session3 = Session.builder(session1).setSystemProperty(REQUIRED_WORKERS_COUNT, "4").build();
        ListenableFuture<ResultWithQueryId<MaterializedResult>> queryFuture3 = service.submit(() -> queryRunner.executeWithQueryId(session3, "SELECT COUNT(*) from lineitem"));
        MILLISECONDS.sleep(1000);
        // None of the queries should run
        assertFalse(queryFuture1.isDone());
        assertFalse(queryFuture2.isDone());
        assertFalse(queryFuture3.isDone());
        queryRunner.addServers(1);
        assertEquals(queryRunner.getCoordinator().refreshNodes().getActiveNodes().size(), 2);
        // After adding 1 node, only 1st query should run
        MILLISECONDS.sleep(1000);
        assertTrue(queryFuture1.get().getResult().getRowCount() > 0);
        QueryManager queryManager = queryRunner.getCoordinator().getQueryManager();
        QueryInfo completedQueryInfo = queryManager.getFullQueryInfo(queryFuture1.get().getQueryId());
        assertTrue(completedQueryInfo.getQueryStats().getResourceWaitingTime().roundTo(SECONDS) >= 1);
        assertFalse(queryFuture2.isDone());
        assertFalse(queryFuture3.isDone());
        // After adding 2 nodes, 2nd and 3rd query should also run
        queryRunner.addServers(2);
        assertEquals(queryRunner.getCoordinator().refreshNodes().getActiveNodes().size(), 4);
        assertTrue(queryFuture2.get().getResult().getRowCount() > 0);
        completedQueryInfo = queryManager.getFullQueryInfo(queryFuture2.get().getQueryId());
        assertTrue(completedQueryInfo.getQueryStats().getResourceWaitingTime().roundTo(SECONDS) >= 2);
        assertTrue(queryFuture3.get().getResult().getRowCount() > 0);
        completedQueryInfo = queryManager.getFullQueryInfo(queryFuture3.get().getQueryId());
        assertTrue(completedQueryInfo.getQueryStats().getResourceWaitingTime().roundTo(SECONDS) >= 2);
    } finally {
        service.shutdown();
    }
}