cryptoe commented on code in PR #14416: URL: https://github.com/apache/druid/pull/14416#discussion_r1243039033
########## extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java: ########## @@ -0,0 +1,876 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.msq.sql.resources; + + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableMap; +import com.google.common.io.CountingOutputStream; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.inject.Inject; +import org.apache.calcite.plan.RelOptPlanner; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.druid.client.indexing.TaskPayloadResponse; +import org.apache.druid.client.indexing.TaskStatusResponse; +import org.apache.druid.common.exception.SanitizableException; +import org.apache.druid.common.guava.FutureUtils; +import org.apache.druid.guice.annotations.MSQ; +import org.apache.druid.indexer.TaskLocation; +import org.apache.druid.indexer.TaskState; +import org.apache.druid.indexer.TaskStatusPlus; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.Pair; +import org.apache.druid.java.util.common.RE; +import org.apache.druid.java.util.common.StringUtils; +import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.java.util.common.guava.Yielder; +import org.apache.druid.java.util.common.guava.Yielders; +import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.msq.indexing.MSQControllerTask; +import org.apache.druid.msq.indexing.TaskReportMSQDestination; +import org.apache.druid.msq.sql.MSQTaskQueryMaker; +import org.apache.druid.msq.sql.MSQTaskSqlEngine; +import org.apache.druid.msq.sql.SqlStatementState; +import org.apache.druid.msq.sql.entity.ColNameAndType; +import org.apache.druid.msq.sql.entity.ResultSetInformation; +import org.apache.druid.msq.sql.entity.SqlStatementResult; +import org.apache.druid.query.BadQueryException; +import org.apache.druid.query.ExecutionMode; +import org.apache.druid.query.QueryCapacityExceededException; +import org.apache.druid.query.QueryContexts; +import org.apache.druid.query.QueryException; +import org.apache.druid.query.QueryInterruptedException; +import org.apache.druid.query.QueryTimeoutException; +import org.apache.druid.query.QueryUnsupportedException; +import org.apache.druid.query.ResourceLimitExceededException; +import org.apache.druid.rpc.indexing.OverlordClient; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.server.QueryResponse; +import org.apache.druid.server.initialization.ServerConfig; +import org.apache.druid.server.security.Access; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.security.AuthorizationUtils; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.server.security.ForbiddenException; +import org.apache.druid.sql.DirectStatement; +import org.apache.druid.sql.HttpStatement; +import org.apache.druid.sql.SqlPlanningException; +import org.apache.druid.sql.SqlRowTransformer; +import org.apache.druid.sql.SqlStatementFactory; +import org.apache.druid.sql.calcite.planner.ColumnMappings; +import org.apache.druid.sql.http.ResultFormat; +import org.apache.druid.sql.http.SqlQuery; +import org.apache.druid.sql.http.SqlResource; +import org.apache.http.HttpStatus; + +import javax.servlet.http.HttpServletRequest; +import javax.ws.rs.Consumes; +import javax.ws.rs.DELETE; +import javax.ws.rs.GET; +import javax.ws.rs.POST; +import javax.ws.rs.Path; +import javax.ws.rs.PathParam; +import javax.ws.rs.Produces; +import javax.ws.rs.QueryParam; +import javax.ws.rs.core.Context; +import javax.ws.rs.core.MediaType; +import javax.ws.rs.core.Response; +import javax.ws.rs.core.StreamingOutput; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + + +@Path("/druid/v2/sql/statements/") +public class SqlStatementResource +{ + + private static final Logger log = new Logger(SqlStatementResource.class); + private final SqlStatementFactory msqSqlStatementFactory; + private final ServerConfig serverConfig; + private final AuthorizerMapper authorizerMapper; + private final ObjectMapper jsonMapper; + private final OverlordClient overlordClient; + + + @Inject + public SqlStatementResource( + final @MSQ SqlStatementFactory msqSqlStatementFactory, + final ServerConfig serverConfig, + final AuthorizerMapper authorizerMapper, + final ObjectMapper jsonMapper, + final OverlordClient overlordClient + ) + { + this.msqSqlStatementFactory = msqSqlStatementFactory; + this.serverConfig = serverConfig; + this.authorizerMapper = authorizerMapper; + this.jsonMapper = jsonMapper; + this.overlordClient = overlordClient; + } + + + @GET + @Path("/enabled") + @Produces(MediaType.APPLICATION_JSON) + public Response doGetEnabled(@Context final HttpServletRequest request) + { + // All authenticated users are authorized for this API: check an empty resource list. + final Access authResult = AuthorizationUtils.authorizeAllResourceActions( + request, + Collections.emptyList(), + authorizerMapper + ); + + if (!authResult.isAllowed()) { + throw new ForbiddenException(authResult.toString()); + } + + return Response.ok(ImmutableMap.of("enabled", true)).build(); + } + + @POST + @Produces(MediaType.APPLICATION_JSON) + @Consumes(MediaType.APPLICATION_JSON) + public Response doPost(final SqlQuery sqlQuery, @Context final HttpServletRequest req) + { + final HttpStatement stmt = msqSqlStatementFactory.httpStatement(sqlQuery, req); + final String sqlQueryId = stmt.sqlQueryId(); + final String currThreadName = Thread.currentThread().getName(); + try { + ExecutionMode executionMode = QueryContexts.getAsEnum( + QueryContexts.CTX_EXECUTION_MODE, + sqlQuery.getContext().get(QueryContexts.CTX_EXECUTION_MODE), + ExecutionMode.class + ); + if (ExecutionMode.ASYNC != executionMode) { + return buildNonOkResponse( + HttpStatus.SC_UNPROCESSABLE_ENTITY, + new QueryException( + QueryException.UNSUPPORTED_OPERATION_ERROR_CODE, + StringUtils.format( + "The statement sql api only supports sync mode[%s]. Please set context parameter [%s=%s] in the payload", + ExecutionMode.ASYNC, + QueryContexts.CTX_EXECUTION_MODE, + ExecutionMode.ASYNC + ), + null, + null, + null + ), + stmt.sqlQueryId() + ); + } + + + Thread.currentThread().setName(StringUtils.format("statement_sql[%s]", sqlQueryId)); + + final DirectStatement.ResultSet plan = stmt.plan(); + final QueryResponse<Object[]> response = plan.run(); + final Sequence<Object[]> sequence = response.getResults(); + final SqlRowTransformer rowTransformer = plan.createRowTransformer(); + + final boolean isTaskStruct = MSQTaskSqlEngine.TASK_STRUCT_FIELD_NAMES.equals(rowTransformer.getFieldList()); + + if (isTaskStruct) { + return buildTaskResponse(sequence, stmt.query().authResult().getIdentity()); + } else { + // Used for EXPLAIN + return buildStandardResponse(sequence, sqlQuery, sqlQueryId, rowTransformer); + } + } + // Kitchen-sinking the errors since they are all unchecked. + // Just copied from SqlResource. + catch (QueryCapacityExceededException cap) { + stmt.reporter().failed(cap); + return buildNonOkResponse(QueryCapacityExceededException.STATUS_CODE, cap, sqlQueryId); + } + catch (QueryUnsupportedException unsupported) { + stmt.reporter().failed(unsupported); + return buildNonOkResponse(QueryUnsupportedException.STATUS_CODE, unsupported, sqlQueryId); + } + catch (QueryTimeoutException timeout) { + stmt.reporter().failed(timeout); + return buildNonOkResponse(QueryTimeoutException.STATUS_CODE, timeout, sqlQueryId); + } + catch (SqlPlanningException | ResourceLimitExceededException e) { + stmt.reporter().failed(e); + return buildNonOkResponse(BadQueryException.STATUS_CODE, e, sqlQueryId); + } + catch (ForbiddenException e) { + // No request logs for forbidden queries; same as SqlResource + throw (ForbiddenException) serverConfig.getErrorResponseTransformStrategy() + .transformIfNeeded(e); // let ForbiddenExceptionMapper handle this + } + catch (RelOptPlanner.CannotPlanException e) { + stmt.reporter().failed(e); + SqlPlanningException spe = new SqlPlanningException( + SqlPlanningException.PlanningError.UNSUPPORTED_SQL_ERROR, + e.getMessage() + ); + return buildNonOkResponse(BadQueryException.STATUS_CODE, spe, sqlQueryId); + } + // Calcite throws a java.lang.AssertionError which is type Error not Exception. Using Throwable catches both. + catch (Throwable e) { + stmt.reporter().failed(e); + log.noStackTrace().warn(e, "Failed to handle query: %s", sqlQueryId); + + return buildNonOkResponse( + Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(), + QueryInterruptedException.wrapIfNeeded(e), + sqlQueryId + ); + } + finally { + stmt.close(); + Thread.currentThread().setName(currThreadName); + } + } + + + @GET + @Path("/{id}") + @Produces(MediaType.APPLICATION_JSON) + public Response doGetStatus( + @PathParam("id") final String queryId, @Context final HttpServletRequest req + ) + { + try { + AuthorizationUtils.authorizeAllResourceActions(req, Collections.emptyList(), authorizerMapper); + final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req); + + Optional<SqlStatementResult> sqlStatementResult = getStatementStatus( + queryId, + authenticationResult.getIdentity(), + true + ); + if (sqlStatementResult.isPresent()) { + return Response.ok().entity(sqlStatementResult.get()).build(); + } else { + return Response.status(Response.Status.NOT_FOUND).build(); + } + } + catch (ForbiddenException e) { + throw (ForbiddenException) serverConfig.getErrorResponseTransformStrategy().transformIfNeeded(e); + } + catch (QueryException e) { + return buildNonOkResponse(Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(), e, queryId); + } + catch (Exception e) { + log.noStackTrace().warn(e, "Failed to handle query: %s", queryId); + throw e; + } + } + + @GET + @Path("/{id}/results") + @Produces(MediaType.APPLICATION_JSON) + public Response doGetResults( + @PathParam("id") final String queryId, + @QueryParam("offset") Long offset, + @QueryParam("numRows") Long numberOfRows, + @QueryParam("sizeInBytes") Long size, + @QueryParam("timeout") Integer timeout, + @Context final HttpServletRequest req + ) + { + if (offset != null && offset < 0) { + return buildNonOkResponse( + Response.Status.PRECONDITION_FAILED.getStatusCode(), + new QueryException( + null, + "offset cannot be negative. Please pass a positive number.", + null, + null, + null + ), + queryId + ); + } + if (numberOfRows != null && numberOfRows < 0) { + return buildNonOkResponse( + Response.Status.PRECONDITION_FAILED.getStatusCode(), + new QueryException( + null, + "numRows cannot be negative. Please pass a positive number.", + null, + null, + null + ), + queryId + ); + } + + final long start = offset == null ? 0 : offset; + final long last = getLastIndex(numberOfRows, start); + + try { + AuthorizationUtils.authorizeAllResourceActions(req, Collections.emptyList(), authorizerMapper); + final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req); + + TaskStatusResponse taskResponse = overlordWork(overlordClient.taskStatus(queryId)); + if (taskResponse == null) { + return Response.status(Response.Status.NOT_FOUND).build(); + } + + TaskStatusPlus statusPlus = taskResponse.getStatus(); + if (statusPlus == null || !MSQControllerTask.TYPE.equals(statusPlus.getType())) { + return Response.status(Response.Status.NOT_FOUND).build(); + } + SqlStatementState sqlStatementState = getSqlStatementState(statusPlus); + + if (sqlStatementState == SqlStatementState.RUNNING || sqlStatementState == SqlStatementState.ACCEPTED) { + return buildNonOkResponse(Response.Status.NOT_FOUND.getStatusCode(), new QueryException( + null, + StringUtils.format( + "Query is [%s]. Please wait for it to complete.", + sqlStatementState + ), + null, + null, + null + ), queryId); + } else if (sqlStatementState == SqlStatementState.FAILED) { + return buildNonOkResponse( + Response.Status.NOT_FOUND.getStatusCode(), + new QueryException(null, statusPlus.getErrorMsg(), null, null, null), + queryId + ); + } else { + MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, authenticationResult.getIdentity()); + Optional<List<ColNameAndType>> signature = getSignature(msqControllerTask); + if (!signature.isPresent()) { + return Response.ok().build(); + } + Optional<List<Object>> results = getResults(getPayload(overlordWork(overlordClient.taskReportAsMap(queryId)))); Review Comment: Yes. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
