Examples with ResultWithQueryId com.facebook.presto.tests.ResultWithQueryId used on opensource projects

Example 1 with ResultWithQueryId

use of com.facebook.presto.tests.ResultWithQueryId in project presto by prestodb.

the class TestTaskInfoResource method testGetTaskInfo.

@Test
public void testGetTaskInfo() {
    String sql = "SELECT * FROM tpch.sf1.customer WHERE tpch.sf1.customer.nationkey = 1";
    ResultWithQueryId<MaterializedResult> result = queryRunner.executeWithQueryId(queryRunner.getDefaultSession(), sql);
    QueryId queryId = result.getQueryId();
    Optional<StageInfo> stageInfo = server.getQueryManager().getFullQueryInfo(queryId).getOutputStage();
    if (stageInfo.isPresent()) {
        Stream<TaskInfo> latestTaskInfo = stageInfo.get().getAllStages().stream().flatMap(stage -> stage.getLatestAttemptExecutionInfo().getTasks().stream());
        Iterable<TaskInfo> iterableTaskInfo = latestTaskInfo::iterator;
        for (TaskInfo taskInfo : iterableTaskInfo) {
            URI taskURI = taskUri("v1/taskInfo/", taskInfo.getTaskId().toString());
            Request taskInfoRequest = prepareGet().setUri(taskURI).build();
            TaskInfo responseTaskInfo = client.execute(taskInfoRequest, createJsonResponseHandler(jsonCodec(TaskInfo.class)));
            compareTasks(taskInfo, responseTaskInfo);
        }
    } else {
        fail("StageInfo not present");
    }
}

Example 2 with ResultWithQueryId

use of com.facebook.presto.tests.ResultWithQueryId in project presto by prestodb.

the class TestStageInfo method testGetStageWithStageId.

@Test
public void testGetStageWithStageId() {
    String sql = "SELECT nationkey AS nationkey, count(*) AS count FROM tpch.sf1.supplier GROUP BY nationkey";
    ResultWithQueryId<MaterializedResult> result = queryRunner.executeWithQueryId(queryRunner.getDefaultSession(), sql);
    QueryId queryId = result.getQueryId();
    Optional<StageInfo> stageInfo = server.getQueryManager().getFullQueryInfo(queryId).getOutputStage();
    if (stageInfo.isPresent()) {
        List<StageInfo> allStages = stageInfo.get().getAllStages();
        for (StageInfo expectedStage : allStages) {
            Optional<StageInfo> actualStage = stageInfo.get().getStageWithStageId(expectedStage.getStageId());
            if (!actualStage.isPresent()) {
                fail("StageInfo with id " + expectedStage.getStageId() + "not found");
            }
            assertEquals(expectedStage.getStageId(), actualStage.get().getStageId());
        }
    }
}