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

Example 1 with MaterializedResultWithPlan

use of io.prestosql.testing.QueryRunner.MaterializedResultWithPlan in project hetu-core by openlookeng.

the class QueryAssertions method assertQuery.

private static void assertQuery(QueryRunner actualQueryRunner, Session actualQuerySession, @Language("SQL") String actual, H2QueryRunner h2QueryRunner, Session expectedQuerySession, @Language("SQL") String expected, boolean ensureOrdering, boolean compareUpdate, Optional<Consumer<Plan>> planAssertion) {
    long start = System.nanoTime();
    Optional<MaterializedResult> actualResultsOp = Optional.empty();
    Plan queryPlan = null;
    if (planAssertion.isPresent()) {
        try {
            MaterializedResultWithPlan resultWithPlan = actualQueryRunner.executeWithPlan(actualQuerySession, actual, WarningCollector.NOOP);
            queryPlan = resultWithPlan.getQueryPlan();
            actualResultsOp = Optional.of(resultWithPlan.getMaterializedResult().toTestTypes());
        } catch (RuntimeException ex) {
            fail("Execution of 'actual' query failed: " + actual, ex);
        }
    } else {
        try {
            actualResultsOp = Optional.of(actualQueryRunner.execute(actualQuerySession, 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();
    Optional<MaterializedResult> expectedResultsOp = Optional.empty();
    try {
        expectedResultsOp = Optional.of(h2QueryRunner.execute(expectedQuerySession, expected, actualResultsOp.get().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 presto: %s, h2: %s, total: %s", actualTime, nanosSince(expectedStart), totalTime);
    }
    if (actualResultsOp.get().getUpdateType().isPresent() || actualResultsOp.get().getUpdateCount().isPresent()) {
        if (!actualResultsOp.get().getUpdateType().isPresent()) {
            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 = actualResultsOp.get().getMaterializedRows();
    List<MaterializedRow> expectedRows = expectedResultsOp.get().getMaterializedRows();
    if (compareUpdate) {
        if (!actualResultsOp.get().getUpdateType().isPresent()) {
            fail("update type not present for query: \n" + actual);
        }
        if (!actualResultsOp.get().getUpdateCount().isPresent()) {
            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), actualResultsOp.get().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.prestosql.testing.QueryRunner.MaterializedResultWithPlan in project hetu-core by openlookeng.

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 presto: %s", queryTime);
    }
    if (planAssertion.isPresent()) {
        planAssertion.get().accept(queryPlan);
    }
    if (!results.getUpdateType().isPresent()) {
        fail("update type is not set");
    }
    if (results.getUpdateCount().isPresent()) {
        if (!count.isPresent()) {
            fail("update count should not be present");
        }
        assertEquals(results.getUpdateCount().getAsLong(), count.getAsLong(), "update count");
    } else if (count.isPresent()) {
        fail("update count is not present");
    }
}