Examples with MaterializedResultWithPlan io.trino.testing.QueryRunner.MaterializedResultWithPlan used on opensource projects

Example 1 with MaterializedResultWithPlan

use of io.trino.testing.QueryRunner.MaterializedResultWithPlan in project trino by trinodb.

the class QueryAssertions method assertQuery.

private static void assertQuery(QueryRunner actualQueryRunner, Session session, @Language("SQL") String actual, H2QueryRunner h2QueryRunner, @Language("SQL") String expected, boolean ensureOrdering, boolean compareUpdate, Optional<Consumer<Plan>> planAssertion) {
    long start = System.nanoTime();
    MaterializedResult actualResults = null;
    Plan queryPlan = null;
    if (planAssertion.isPresent()) {
        try {
            MaterializedResultWithPlan resultWithPlan = actualQueryRunner.executeWithPlan(session, actual, WarningCollector.NOOP);
            queryPlan = resultWithPlan.getQueryPlan();
            actualResults = resultWithPlan.getMaterializedResult().toTestTypes();
        } catch (RuntimeException ex) {
            fail("Execution of 'actual' query failed: " + actual, ex);
        }
    } else {
        try {
            actualResults = actualQueryRunner.execute(session, actual).toTestTypes();
        } catch (RuntimeException ex) {
            fail("Execution of 'actual' query failed: " + actual, ex);
        }
    }
    if (planAssertion.isPresent()) {
        planAssertion.get().accept(queryPlan);
    }
    Duration actualTime = nanosSince(start);
    long expectedStart = System.nanoTime();
    MaterializedResult expectedResults = null;
    try {
        expectedResults = h2QueryRunner.execute(session, expected, actualResults.getTypes());
    } catch (RuntimeException ex) {
        fail("Execution of 'expected' query failed: " + expected, ex);
    }
    Duration totalTime = nanosSince(start);
    if (totalTime.compareTo(Duration.succinctDuration(1, SECONDS)) > 0) {
        log.info("FINISHED in Trino: %s, H2: %s, total: %s", actualTime, nanosSince(expectedStart), totalTime);
    }
    if (actualResults.getUpdateType().isPresent() || actualResults.getUpdateCount().isPresent()) {
        if (actualResults.getUpdateType().isEmpty()) {
            fail("update count present without update type for query: \n" + actual);
        }
        if (!compareUpdate) {
            fail("update type should not be present (use assertUpdate) for query: \n" + actual);
        }
    }
    List<MaterializedRow> actualRows = actualResults.getMaterializedRows();
    List<MaterializedRow> expectedRows = expectedResults.getMaterializedRows();
    if (compareUpdate) {
        if (actualResults.getUpdateType().isEmpty()) {
            fail("update type not present for query: \n" + actual);
        }
        if (actualResults.getUpdateCount().isEmpty()) {
            fail("update count not present for query: \n" + actual);
        }
        assertEquals(actualRows.size(), 1, "For query: \n " + actual + "\n:");
        assertEquals(expectedRows.size(), 1, "For query: \n " + actual + "\n:");
        MaterializedRow row = expectedRows.get(0);
        assertEquals(row.getFieldCount(), 1, "For query: \n " + actual + "\n:");
        assertEquals(row.getField(0), actualResults.getUpdateCount().getAsLong(), "For query: \n " + actual + "\n:");
    }
    if (ensureOrdering) {
        if (!actualRows.equals(expectedRows)) {
            assertEquals(actualRows, expectedRows, "For query: \n " + actual + "\n:");
        }
    } else {
        assertEqualsIgnoreOrder(actualRows, expectedRows, "For query: \n " + actual);
    }
}

Example 2 with MaterializedResultWithPlan

use of io.trino.testing.QueryRunner.MaterializedResultWithPlan in project trino by trinodb.

the class QueryAssertions method assertUpdate.

public static void assertUpdate(QueryRunner queryRunner, Session session, @Language("SQL") String sql, OptionalLong count, Optional<Consumer<Plan>> planAssertion) {
    long start = System.nanoTime();
    MaterializedResult results;
    Plan queryPlan;
    if (planAssertion.isPresent()) {
        MaterializedResultWithPlan resultWithPlan = queryRunner.executeWithPlan(session, sql, WarningCollector.NOOP);
        queryPlan = resultWithPlan.getQueryPlan();
        results = resultWithPlan.getMaterializedResult().toTestTypes();
    } else {
        queryPlan = null;
        results = queryRunner.execute(session, sql);
    }
    Duration queryTime = nanosSince(start);
    if (queryTime.compareTo(Duration.succinctDuration(1, SECONDS)) > 0) {
        log.info("FINISHED in Trino: %s", queryTime);
    }
    if (planAssertion.isPresent()) {
        planAssertion.get().accept(queryPlan);
    }
    if (results.getUpdateType().isEmpty()) {
        fail("update type is not set");
    }
    if (results.getUpdateCount().isPresent()) {
        if (count.isEmpty()) {
            fail("expected no update count, but got " + results.getUpdateCount().getAsLong());
        }
        assertEquals(results.getUpdateCount().getAsLong(), count.getAsLong(), "update count");
    } else if (count.isPresent()) {
        fail("update count is not present");
    }
}