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

Example 1 with PrepareResult

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

the class DruidStatement method prepare.

public DruidStatement prepare(final String query, final long maxRowCount, final AuthenticationResult authenticationResult) {
    synchronized (lock) {
        try {
            ensure(State.NEW);
            sqlLifecycle.initialize(query, queryContext);
            sqlLifecycle.validateAndAuthorize(authenticationResult);
            this.authenticationResult = authenticationResult;
            PrepareResult prepareResult = sqlLifecycle.prepare();
            this.maxRowCount = maxRowCount;
            this.query = query;
            List<AvaticaParameter> params = new ArrayList<>();
            final RelDataType parameterRowType = prepareResult.getParameterRowType();
            for (RelDataTypeField field : parameterRowType.getFieldList()) {
                RelDataType type = field.getType();
                params.add(createParameter(field, type));
            }
            this.signature = Meta.Signature.create(createColumnMetaData(prepareResult.getRowType()), query, params, Meta.CursorFactory.ARRAY, // We only support SELECT
            Meta.StatementType.SELECT);
            this.state = State.PREPARED;
        } catch (Throwable t) {
            return closeAndPropagateThrowable(t);
        }
        return this;
    }
}

Example 2 with PrepareResult

use of org.apache.druid.sql.calcite.planner.PrepareResult 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 3 with PrepareResult

use of org.apache.druid.sql.calcite.planner.PrepareResult 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);
}