Examples with DruidPlanner org.apache.druid.sql.calcite.planner.DruidPlanner used on opensource projects

Example 6 with DruidPlanner

use of org.apache.druid.sql.calcite.planner.DruidPlanner in project druid by druid-io.

the class SqlBenchmark method querySql.

@Benchmark
@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
public void querySql(Blackhole blackhole) throws Exception {
    final Map<String, Object> context = ImmutableMap.of(QueryContexts.VECTORIZE_KEY, vectorize, QueryContexts.VECTORIZE_VIRTUAL_COLUMNS_KEY, vectorize);
    final String sql = QUERIES.get(Integer.parseInt(query));
    try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(context, sql)) {
        final PlannerResult plannerResult = planner.plan();
        final Sequence<Object[]> resultSequence = plannerResult.run();
        final Object[] lastRow = resultSequence.accumulate(null, (accumulated, in) -> in);
        blackhole.consume(lastRow);
    }
}

Example 7 with DruidPlanner

use of org.apache.druid.sql.calcite.planner.DruidPlanner in project druid by druid-io.

the class SqlLifecycle method plan.

/**
 * Plan the query to enable execution.
 *
 * If successful, the lifecycle will first transition from {@link State#AUTHORIZED} to {@link State#PLANNED}.
 */
public void plan() throws RelConversionException {
    transition(State.AUTHORIZED, State.PLANNED);
    Preconditions.checkNotNull(plannerContext, "Cannot plan, plannerContext is null");
    try (DruidPlanner planner = plannerFactory.createPlannerWithContext(plannerContext)) {
        this.plannerResult = planner.plan();
    }// we can't collapse catch clauses since SqlPlanningException has type-sensitive constructors.
     catch (SqlParseException e) {
        throw new SqlPlanningException(e);
    } catch (ValidationException e) {
        throw new SqlPlanningException(e);
    }
}

Example 8 with DruidPlanner

use of org.apache.druid.sql.calcite.planner.DruidPlanner in project druid by druid-io.

the class SqlLifecycleTest method testStateTransition.

@Test
public void testStateTransition() throws ValidationException, SqlParseException, RelConversionException, IOException {
    SqlLifecycle lifecycle = sqlLifecycleFactory.factorize();
    final String sql = "select 1 + ?";
    final Map<String, Object> queryContext = Collections.emptyMap();
    Assert.assertEquals(SqlLifecycle.State.NEW, lifecycle.getState());
    // test initialize
    lifecycle.initialize(sql, queryContext);
    Assert.assertEquals(SqlLifecycle.State.INITIALIZED, lifecycle.getState());
    List<TypedValue> parameters = ImmutableList.of(new SqlParameter(SqlType.BIGINT, 1L).getTypedValue());
    lifecycle.setParameters(parameters);
    // setting parameters should not change the state
    Assert.assertEquals(SqlLifecycle.State.INITIALIZED, lifecycle.getState());
    // test authorization
    DruidPlanner mockPlanner = EasyMock.createMock(DruidPlanner.class);
    PlannerContext mockPlannerContext = EasyMock.createMock(PlannerContext.class);
    ValidationResult validationResult = new ValidationResult(Collections.emptySet());
    EasyMock.expect(plannerFactory.createPlanner(EasyMock.eq(sql), EasyMock.anyObject())).andReturn(mockPlanner).once();
    EasyMock.expect(mockPlanner.getPlannerContext()).andReturn(mockPlannerContext).once();
    mockPlannerContext.setAuthenticationResult(CalciteTests.REGULAR_USER_AUTH_RESULT);
    EasyMock.expectLastCall();
    mockPlannerContext.setParameters(parameters);
    EasyMock.expectLastCall();
    EasyMock.expect(plannerFactory.getAuthorizerMapper()).andReturn(CalciteTests.TEST_AUTHORIZER_MAPPER).once();
    mockPlannerContext.setAuthorizationResult(Access.OK);
    EasyMock.expectLastCall();
    EasyMock.expect(mockPlanner.validate()).andReturn(validationResult).once();
    mockPlanner.close();
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext);
    lifecycle.validateAndAuthorize(CalciteTests.REGULAR_USER_AUTH_RESULT);
    Assert.assertEquals(SqlLifecycle.State.AUTHORIZED, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext);
    // test prepare
    PrepareResult mockPrepareResult = EasyMock.createMock(PrepareResult.class);
    EasyMock.expect(plannerFactory.createPlannerWithContext(EasyMock.eq(mockPlannerContext))).andReturn(mockPlanner).once();
    EasyMock.expect(mockPlanner.prepare()).andReturn(mockPrepareResult).once();
    mockPlanner.close();
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult);
    lifecycle.prepare();
    // prepare doens't change lifecycle state
    Assert.assertEquals(SqlLifecycle.State.AUTHORIZED, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult);
    // test plan
    PlannerResult mockPlanResult = EasyMock.createMock(PlannerResult.class);
    EasyMock.expect(plannerFactory.createPlannerWithContext(EasyMock.eq(mockPlannerContext))).andReturn(mockPlanner).once();
    EasyMock.expect(mockPlanner.plan()).andReturn(mockPlanResult).once();
    mockPlanner.close();
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    lifecycle.plan();
    Assert.assertEquals(mockPlannerContext, lifecycle.getPlannerContext());
    Assert.assertEquals(SqlLifecycle.State.PLANNED, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    // test execute
    EasyMock.expect(mockPlanResult.run()).andReturn(Sequences.simple(ImmutableList.of(new Object[] { 2L }))).once();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    lifecycle.execute();
    Assert.assertEquals(SqlLifecycle.State.EXECUTING, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    // test emit
    EasyMock.expect(mockPlannerContext.getSqlQueryId()).andReturn("id").once();
    CopyOnWriteArrayList<String> nativeQueryIds = new CopyOnWriteArrayList<>(ImmutableList.of("id"));
    EasyMock.expect(mockPlannerContext.getNativeQueryIds()).andReturn(nativeQueryIds).times(2);
    EasyMock.expect(mockPlannerContext.getAuthenticationResult()).andReturn(CalciteTests.REGULAR_USER_AUTH_RESULT).once();
    serviceEmitter.emit(EasyMock.anyObject(ServiceEventBuilder.class));
    EasyMock.expectLastCall();
    serviceEmitter.emit(EasyMock.anyObject(ServiceEventBuilder.class));
    EasyMock.expectLastCall();
    requestLogger.logSqlQuery(EasyMock.anyObject());
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    lifecycle.finalizeStateAndEmitLogsAndMetrics(null, null, 10);
    Assert.assertEquals(SqlLifecycle.State.DONE, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
}

Example 9 with DruidPlanner

use of org.apache.druid.sql.calcite.planner.DruidPlanner in project druid by druid-io.

the class SqlLifecycleTest method testStateTransitionHttpRequest.

@Test
public void testStateTransitionHttpRequest() throws ValidationException, SqlParseException, RelConversionException, IOException {
    // this test is a duplicate of testStateTransition except with a slight variation of how validate and authorize
    // is run
    SqlLifecycle lifecycle = sqlLifecycleFactory.factorize();
    final String sql = "select 1 + ?";
    final Map<String, Object> queryContext = Collections.emptyMap();
    Assert.assertEquals(SqlLifecycle.State.NEW, lifecycle.getState());
    // test initialize
    lifecycle.initialize(sql, queryContext);
    Assert.assertEquals(SqlLifecycle.State.INITIALIZED, lifecycle.getState());
    List<TypedValue> parameters = ImmutableList.of(new SqlParameter(SqlType.BIGINT, 1L).getTypedValue());
    lifecycle.setParameters(parameters);
    // setting parameters should not change the state
    Assert.assertEquals(SqlLifecycle.State.INITIALIZED, lifecycle.getState());
    // test authorization
    DruidPlanner mockPlanner = EasyMock.createMock(DruidPlanner.class);
    PlannerContext mockPlannerContext = EasyMock.createMock(PlannerContext.class);
    ValidationResult validationResult = new ValidationResult(Collections.emptySet());
    EasyMock.expect(plannerFactory.createPlanner(EasyMock.eq(sql), EasyMock.anyObject())).andReturn(mockPlanner).once();
    EasyMock.expect(mockPlanner.getPlannerContext()).andReturn(mockPlannerContext).once();
    mockPlannerContext.setAuthenticationResult(CalciteTests.REGULAR_USER_AUTH_RESULT);
    EasyMock.expectLastCall();
    mockPlannerContext.setParameters(parameters);
    EasyMock.expectLastCall();
    EasyMock.expect(plannerFactory.getAuthorizerMapper()).andReturn(CalciteTests.TEST_AUTHORIZER_MAPPER).once();
    mockPlannerContext.setAuthorizationResult(Access.OK);
    EasyMock.expectLastCall();
    EasyMock.expect(mockPlanner.validate()).andReturn(validationResult).once();
    mockPlanner.close();
    EasyMock.expectLastCall();
    HttpServletRequest request = EasyMock.createMock(HttpServletRequest.class);
    EasyMock.expect(request.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)).andReturn(CalciteTests.REGULAR_USER_AUTH_RESULT).times(2);
    EasyMock.expect(request.getAttribute(AuthConfig.DRUID_ALLOW_UNSECURED_PATH)).andReturn(null).once();
    EasyMock.expect(request.getAttribute(AuthConfig.DRUID_AUTHORIZATION_CHECKED)).andReturn(null).once();
    request.setAttribute(AuthConfig.DRUID_AUTHORIZATION_CHECKED, true);
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, request);
    lifecycle.validateAndAuthorize(request);
    Assert.assertEquals(SqlLifecycle.State.AUTHORIZED, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, request);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, request);
    // test prepare
    PrepareResult mockPrepareResult = EasyMock.createMock(PrepareResult.class);
    EasyMock.expect(plannerFactory.createPlannerWithContext(EasyMock.eq(mockPlannerContext))).andReturn(mockPlanner).once();
    EasyMock.expect(mockPlanner.prepare()).andReturn(mockPrepareResult).once();
    mockPlanner.close();
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult);
    lifecycle.prepare();
    // prepare doens't change lifecycle state
    Assert.assertEquals(SqlLifecycle.State.AUTHORIZED, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult);
    // test plan
    PlannerResult mockPlanResult = EasyMock.createMock(PlannerResult.class);
    EasyMock.expect(plannerFactory.createPlannerWithContext(EasyMock.eq(mockPlannerContext))).andReturn(mockPlanner).once();
    EasyMock.expect(mockPlanner.plan()).andReturn(mockPlanResult).once();
    mockPlanner.close();
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    lifecycle.plan();
    Assert.assertEquals(mockPlannerContext, lifecycle.getPlannerContext());
    Assert.assertEquals(SqlLifecycle.State.PLANNED, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    // test execute
    EasyMock.expect(mockPlanResult.run()).andReturn(Sequences.simple(ImmutableList.of(new Object[] { 2L }))).once();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    lifecycle.execute();
    Assert.assertEquals(SqlLifecycle.State.EXECUTING, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    // test emit
    EasyMock.expect(mockPlannerContext.getSqlQueryId()).andReturn("id").once();
    CopyOnWriteArrayList<String> nativeQueryIds = new CopyOnWriteArrayList<>(ImmutableList.of("id"));
    EasyMock.expect(mockPlannerContext.getNativeQueryIds()).andReturn(nativeQueryIds).times(2);
    EasyMock.expect(mockPlannerContext.getAuthenticationResult()).andReturn(CalciteTests.REGULAR_USER_AUTH_RESULT).once();
    serviceEmitter.emit(EasyMock.anyObject(ServiceEventBuilder.class));
    EasyMock.expectLastCall();
    serviceEmitter.emit(EasyMock.anyObject(ServiceEventBuilder.class));
    EasyMock.expectLastCall();
    requestLogger.logSqlQuery(EasyMock.anyObject());
    EasyMock.expectLastCall();
    EasyMock.replay(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    lifecycle.finalizeStateAndEmitLogsAndMetrics(null, null, 10);
    Assert.assertEquals(SqlLifecycle.State.DONE, lifecycle.getState());
    EasyMock.verify(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
    EasyMock.reset(plannerFactory, serviceEmitter, requestLogger, mockPlanner, mockPlannerContext, mockPrepareResult, mockPlanResult);
}