public <T> T inTransaction(Function<Session, T> transactionSessionConsumer) { return inTransaction(defaultSession, transactionSessionConsumer); }
@Override public MaterializedResultWithPlan executeWithPlan(Session session, String sql, WarningCollector warningCollector) { return inTransaction(session, transactionSession -> executeInternal(transactionSession, sql)); }
protected Plan plan(String sql, LogicalPlanner.Stage stage, boolean forceSingleNode) { try { return queryRunner.inTransaction(transactionSession -> queryRunner.createPlan(transactionSession, sql, stage, forceSingleNode, WarningCollector.NOOP)); } catch (RuntimeException e) { throw new AssertionError("Planning failed for SQL: " + sql, e); } }
public ConnectorId getCurrentConnectorId() { return queryRunner.inTransaction(transactionSession -> metadata.getCatalogHandle(transactionSession, session.getCatalog().get())).get(); } }
@Benchmark public List<Plan> planQueries(BenchmarkData benchmarkData) { return benchmarkData.queryRunner.inTransaction(transactionSession -> { LogicalPlanner.Stage stage = LogicalPlanner.Stage.valueOf(benchmarkData.stage.toUpperCase()); return benchmarkData.queries.stream() .map(query -> benchmarkData.queryRunner.createPlan(transactionSession, query, stage, false, WarningCollector.NOOP)) .collect(toImmutableList()); }); }
public ConnectorId getCurrentConnectorId() { return queryRunner.inTransaction(transactionSession -> queryRunner.getMetadata().getCatalogHandle(transactionSession, transactionSession.getCatalog().get())).get(); }
private void assertInvalidSpatialPartitioning(Session session, String sql, String expectedMessageRegExp) { LocalQueryRunner queryRunner = getQueryRunner(); try { queryRunner.inTransaction(session, transactionSession -> { queryRunner.createPlan(transactionSession, sql, LogicalPlanner.Stage.OPTIMIZED_AND_VALIDATED, false, WarningCollector.NOOP); return null; }); fail(format("Expected query to fail: %s", sql)); } catch (PrestoException ex) { assertEquals(ex.getErrorCode(), INVALID_SPATIAL_PARTITIONING.toErrorCode()); if (!nullToEmpty(ex.getMessage()).matches(expectedMessageRegExp)) { fail(format("Expected exception message '%s' to match '%s' for query: %s", ex.getMessage(), expectedMessageRegExp, sql), ex); } } }
@TearDown(Level.Invocation) public void dropPointsTable() { queryRunner.inTransaction(queryRunner.getDefaultSession(), transactionSession -> { Metadata metadata = queryRunner.getMetadata(); Optional<TableHandle> tableHandle = metadata.getTableHandle(transactionSession, QualifiedObjectName.valueOf("memory.default.points")); assertTrue(tableHandle.isPresent(), "Table memory.default.points does not exist"); metadata.dropTable(transactionSession, tableHandle.get()); return null; }); }
private String getPlanText(Session session, String sql) { return localQueryRunner.inTransaction(session, transactionSession -> { Plan plan = localQueryRunner.createPlan(transactionSession, sql, LogicalPlanner.Stage.OPTIMIZED_AND_VALIDATED, WarningCollector.NOOP); return PlanPrinter.textLogicalPlan( plan.getRoot(), plan.getTypes(), localQueryRunner.getMetadata().getFunctionRegistry(), plan.getStatsAndCosts(), transactionSession, 0, false); }); } }
protected void assertPlanWithSession(@Language("SQL") String sql, Session session, boolean forceSingleNode, PlanMatchPattern pattern, Consumer<Plan> planValidator) { queryRunner.inTransaction(session, transactionSession -> { Plan actualPlan = queryRunner.createPlan(transactionSession, sql, LogicalPlanner.Stage.OPTIMIZED_AND_VALIDATED, forceSingleNode, WarningCollector.NOOP); PlanAssert.assertPlan(transactionSession, queryRunner.getMetadata(), queryRunner.getStatsCalculator(), actualPlan, pattern); planValidator.accept(actualPlan); return null; }); }
protected void assertPlan(String sql, LogicalPlanner.Stage stage, PlanMatchPattern pattern, List<PlanOptimizer> optimizers) { queryRunner.inTransaction(transactionSession -> { Plan actualPlan = queryRunner.createPlan(transactionSession, sql, optimizers, stage, WarningCollector.NOOP); PlanAssert.assertPlan(transactionSession, queryRunner.getMetadata(), queryRunner.getStatsCalculator(), actualPlan, pattern); return null; }); }
private void assertPlan(String sql, LogicalPlanner.Stage stage, PlanMatchPattern pattern) { queryRunner.inTransaction(transactionSession -> { Plan actualPlan = queryRunner.createPlan(transactionSession, sql, stage, WarningCollector.NOOP); PlanAssert.assertPlan(transactionSession, queryRunner.getMetadata(), queryRunner.getStatsCalculator(), actualPlan, pattern); return null; }); } }
private void validatePlan(Function<PlanBuilder, PlanNode> planProvider) { PlanBuilder builder = new PlanBuilder(idAllocator, metadata); PlanNode planNode = planProvider.apply(builder); TypeProvider types = builder.getTypes(); getQueryRunner().inTransaction(session -> { // metadata.getCatalogHandle() registers the catalog for the transaction session.getCatalog().ifPresent(catalog -> metadata.getCatalogHandle(session, catalog)); new ValidateStreamingAggregations().validate(planNode, session, metadata, sqlParser, types, WarningCollector.NOOP); return null; }); } }
private void validatePlan(PlanNode root, boolean forceSingleNode) { getQueryRunner().inTransaction(session -> { // metadata.getCatalogHandle() registers the catalog for the transaction session.getCatalog().ifPresent(catalog -> metadata.getCatalogHandle(session, catalog)); new ValidateAggregationsWithDefaultValues(forceSingleNode).validate(root, session, metadata, SQL_PARSER, TypeProvider.empty(), WarningCollector.NOOP); return null; }); } }
@Test(timeOut = 1000) public void optimizerTimeoutsOnNonConvergingPlan() { PlanOptimizer optimizer = new IterativeOptimizer( new RuleStatsRecorder(), queryRunner.getStatsCalculator(), queryRunner.getCostCalculator(), ImmutableSet.of(new NonConvergingRule())); try { queryRunner.inTransaction(transactionSession -> { queryRunner.createPlan(transactionSession, "SELECT * FROM nation", ImmutableList.of(optimizer), WarningCollector.NOOP); fail("The optimizer should not converge"); return null; }); } catch (PrestoException ex) { assertEquals(ex.getErrorCode(), OPTIMIZER_TIMEOUT.toErrorCode()); } }
protected void assertPlanWithSession(@Language("SQL") String sql, Session session, boolean forceSingleNode, PlanMatchPattern pattern) { queryRunner.inTransaction(session, transactionSession -> { Plan actualPlan = queryRunner.createPlan(transactionSession, sql, LogicalPlanner.Stage.OPTIMIZED_AND_VALIDATED, forceSingleNode, WarningCollector.NOOP); PlanAssert.assertPlan(transactionSession, queryRunner.getMetadata(), queryRunner.getStatsCalculator(), actualPlan, pattern); return null; }); }
public static void assertPlannerWarnings(LocalQueryRunner queryRunner, @Language("SQL") String sql, Map<String, String> sessionProperties, List<WarningCode> expectedWarnings, Optional<List<Rule<?>>> rules) { Session.SessionBuilder sessionBuilder = testSessionBuilder() .setCatalog(queryRunner.getDefaultSession().getCatalog().get()) .setSchema(queryRunner.getDefaultSession().getSchema().get()); sessionProperties.forEach(sessionBuilder::setSystemProperty); WarningCollector warningCollector = new DefaultWarningCollector(new WarningCollectorConfig()); try { queryRunner.inTransaction(sessionBuilder.build(), transactionSession -> { if (rules.isPresent()) { createPlan(queryRunner, transactionSession, sql, warningCollector, rules.get()); } else { queryRunner.createPlan(transactionSession, sql, LogicalPlanner.Stage.CREATED, false, warningCollector); } return null; }); } catch (SemanticException e) { // ignore } Set<WarningCode> warnings = warningCollector.getWarnings().stream() .map(PrestoWarning::getWarningCode) .collect(toImmutableSet()); for (WarningCode expectedWarning : expectedWarnings) { if (!warnings.contains(expectedWarning)) { fail("Expected warning: " + expectedWarning); } } }
private void assertInvalidSpatialPartitioning(Session session, String sql, String expectedMessageRegExp) { LocalQueryRunner queryRunner = getQueryRunner(); try { queryRunner.inTransaction(session, transactionSession -> { queryRunner.createPlan(transactionSession, sql, LogicalPlanner.Stage.OPTIMIZED_AND_VALIDATED, false, WarningCollector.NOOP); return null; }); fail(format("Expected query to fail: %s", sql)); } catch (PrestoException ex) { assertEquals(ex.getErrorCode(), INVALID_SPATIAL_PARTITIONING.toErrorCode()); if (!nullToEmpty(ex.getMessage()).matches(expectedMessageRegExp)) { fail(format("Expected exception message '%s' to match '%s' for query: %s", ex.getMessage(), expectedMessageRegExp, sql), ex); } } }
@TearDown(Level.Invocation) public void dropPointsTable() { queryRunner.inTransaction(queryRunner.getDefaultSession(), transactionSession -> { Metadata metadata = queryRunner.getMetadata(); Optional<TableHandle> tableHandle = metadata.getTableHandle(transactionSession, QualifiedObjectName.valueOf("memory.default.points")); assertTrue(tableHandle.isPresent(), "Table memory.default.points does not exist"); metadata.dropTable(transactionSession, tableHandle.get()); return null; }); }
private String getPlanText(Session session, String sql) { return localQueryRunner.inTransaction(session, transactionSession -> { Plan plan = localQueryRunner.createPlan(transactionSession, sql, LogicalPlanner.Stage.OPTIMIZED_AND_VALIDATED, WarningCollector.NOOP); return PlanPrinter.textLogicalPlan( plan.getRoot(), plan.getTypes(), localQueryRunner.getMetadata().getFunctionRegistry(), plan.getStatsAndCosts(), transactionSession, 0, false); }); } }