[GitHub] [druid] cryptoe opened a new pull request, #14416: [Draft] Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe opened a new pull request, #14416:
URL: https://github.com/apache/druid/pull/14416

   <!-- Thanks for trying to help us make Apache Druid be the best it can be! Please fill out as much of the following information as is possible (where relevant, and remove it when irrelevant) to help make the intention and scope of this PR clear in order to ease review. -->
   
   <!-- Please read the doc for contribution (https://github.com/apache/druid/blob/master/CONTRIBUTING.md) before making this PR. Also, once you open a PR, please _avoid using force pushes and rebasing_ since these make it difficult for reviewers to see what you've changed in response to their reviews. See [the 'If your pull request shows conflicts with master' section](https://github.com/apache/druid/blob/master/CONTRIBUTING.md#if-your-pull-request-shows-conflicts-with-master) for more details. -->
   
   The current SQL endpoint is more suited toward hot queries where the client blocks for the results. The current task endpoint is more of a fire-and-forget API, primarily meant for ingestion. 
   
   This PR aims to expose a new API called 
   "@Path("/druid/v2/sql/statements/")" which takes the same payload as the current "/druid/v2/sql" endpoint and allows users to fetch results in an async manner.
   
   The statement execution API runs a single SQL statement which can be:
   
       SELECT
   
       DDL
   
       CREATE [TEMPORARY|EXTERNAL] TABLE AS 
   
       INSERT 
   
   
   ### POST SQL statement
   
   `POST /sql/statements`
   
   Executes a SQL statement.
   
   
   * Request: The post payload is in the same format as https://druid.apache.org/docs/latest/querying/sql.html#http-post
    so that users can easily switch to the new async apis.
    Two additional context field:
    ```
   
       mode: Determines how the query results are fetched. It will have 3 values [SYNC, ASYNC, AUTO] .  For the MVP mode=ASYNC is only supported. 
   
       execTimeout: The maximum amount of time that the query can execute. Handy to prevent “runaway queries.” Must be lower than any system-configured value. This is distinct from the results expiration timeout. Optional.(not supported for initial version)
       
   ```
   
   * Response 
   
   ```
   #http 200
   {
     "state":"RUNNING",
     "queryID":"queryIDXX",
     "createdAt":"2023-03-07T12:58:17Z",
     "schema": [
       {
         "name": "col1",
         "type": "varchar"
       },
       {
         "name": "col2",
         "type": "varchar"
       }
     ]
   }
   ```
   ```
   #http 200
   {
     "state": "FINISHED",
     "createdAt": "2023-03-07T12:58:17Z",
     "queryID": "queryIDXX",
     "durationInMS": 20,
     "schema": [
       {
         "name": "col1",
         "type": "varchar"
       },
       {
         "name": "col2",
         "type": "varchar"
       }
     ],
     "result": {
       "totalRows": 12412,
       "totalSize": 12000,
       "format":"json",
       "header":true,
       "sampleRecords": []
     }
   }
   ```
   
   ### Query execution failures
   
   
   This sections is exactly similar as https://druid.apache.org/docs/latest/querying/querying.html#query-execution-failures
   If a query fails, Druid returns a response with an HTTP response code and a JSON object with the following structure:
   ```
   {
     "error" : "Query timeout",
     "errorMessage" : "Timeout waiting for task.",
     "errorClass" : "java.util.concurrent.TimeoutException",
     "host" : "druid1.example.com:8083"
   }
   ```
   
   ### Status 
   
   `GET /sql/statements/{id}`
   
   
   Returns the same response as the post API. 
   
   
   
   ### Results
   
   
   
   `GET /sql/statements/{id}/results?offset=x&numRows=y&size=s&timeout=z`
   
    
   
   Returns a page of results if they are read  Parameter names are preliminary: the goal here is to list them, not finalize names:
   ```
       id Query Id.
       offset. Row offset of the first row to fetch.
       numRows: number of rows to fetch
       size: Size limit of the results, in bytes. (Optional: size limit in rows. The smallest value wins.)
       timeout: The timeout, in MS, to wait for results. The request returns when the timeout is hit, regardless of size. The results can be zero rows.
   
   ```
   
   
   
   
   
   
   ### Close Query
   
   `DELETE /sql/statements/{id}`
   
   Cancels a running/accepted query. 
   
   
   <!-- Replace XXXX with the id of the issue fixed in this PR. Remove this section if there is no corresponding issue. Don't reference the issue in the title of this pull-request. -->
   
   <!-- If you are a committer, follow the PR action item checklist for committers:
   https://github.com/apache/druid/blob/master/dev/committer-instructions.md#pr-and-issue-action-item-checklist-for-committers. -->
   
   ### Description
   
   <!-- Describe the goal of this PR, what problem are you fixing. If there is a corresponding issue (referenced above), it's not necessary to repeat the description here, however, you may choose to keep one summary sentence. -->
   
   <!-- Describe your patch: what did you change in code? How did you fix the problem? -->
   
   <!-- If there are several relatively logically separate changes in this PR, create a mini-section for each of them. For example: -->
   
   
   ### TODO
   * Add more java docs
   * UT's
   
   <!-- Check the items by putting "x" in the brackets for the done things. Not all of these items apply to every PR. Remove the items which are not done or not relevant to the PR. None of the items from the checklist below are strictly necessary, but it would be very helpful if you at least self-review the PR. -->
   
   This PR has:
   
   - [ ] been self-reviewed.
      - [ ] using the [concurrency checklist](https://github.com/apache/druid/blob/master/dev/code-review/concurrency.md) (Remove this item if the PR doesn't have any relation to concurrency.)
   - [ ] added documentation for new or modified features or behaviors.
   - [ ] a release note entry in the PR description.
   - [ ] added Javadocs for most classes and all non-trivial methods. Linked related entities via Javadoc links.
   - [ ] added or updated version, license, or notice information in [licenses.yaml](https://github.com/apache/druid/blob/master/dev/license.md)
   - [ ] added comments explaining the "why" and the intent of the code wherever would not be obvious for an unfamiliar reader.
   - [ ] added unit tests or modified existing tests to cover new code paths, ensuring the threshold for [code coverage](https://github.com/apache/druid/blob/master/dev/code-review/code-coverage.md) is met.
   - [ ] added integration tests.
   - [ ] been tested in a test Druid cluster.
   


-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] abhishekagarwal87 commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["abhishekagarwal87 (via GitHub)" <gi...@apache.org>]

abhishekagarwal87 commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1233940733


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/SqlStatementState.java:
##########
@@ -0,0 +1,28 @@
+/*
+ * 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;
+
+public enum SqlStatementState

Review Comment:
   Please add description for different states. 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/ColNameAndType.java:
##########
@@ -0,0 +1,83 @@
+/*
+ * 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.entity;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Objects;
+
+
+public class ColNameAndType

Review Comment:
   Please add java docs. 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),

Review Comment:
   We are now using DruidException so you may have to switch to that one here. 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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();

Review Comment:
   can you add a comment that we haven't actually run the query? These results are not query results. 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          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 = overlordClient.getTaskStatus(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
+            ),
+            queryId
+        );
+      } else if (sqlStatementState == SqlStatementState.FAILED) {
+        return buildNonOkResponse(
+            Response.Status.NOT_FOUND.getStatusCode(),

Review Comment:
   it shouldn't be a 404 IMO. the status code should come from statusPlus itself. 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          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 = overlordClient.getTaskStatus(queryId);
+      if (taskResponse == null) {
+        return Response.status(Response.Status.NOT_FOUND).build();

Review Comment:
   It will be useful to add a message that includes queryId in it. 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          queryId
+      );
+    }
+
+    final long start = offset == null ? 0 : offset;
+    final long last = getLastIndex(numberOfRows, start);
+
+    try {
+      AuthorizationUtils.authorizeAllResourceActions(req, Collections.emptyList(), authorizerMapper);

Review Comment:
   it's also called with an empty list so it seems like it has some interesting purpose that I can't actually tell :) 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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)

Review Comment:
   whats the purpose of this API? 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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)

Review Comment:
   ```suggestion
     public Response isEnabled(@Context final HttpServletRequest request)
   ```



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/ResultSetInformation.java:
##########
@@ -0,0 +1,141 @@
+/*
+ * 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.entity;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import org.apache.druid.sql.http.ResultFormat;
+
+import javax.annotation.Nullable;
+import java.util.List;
+import java.util.Objects;
+
+public class ResultSetInformation
+{
+
+  @Nullable
+  private final Long totalRows;
+  @Nullable
+  private final Long totalSize;
+
+  private final ResultFormat resultFormat;
+
+  private final Boolean header;
+
+  @Nullable
+  private final List<Object> records;
+
+  @JsonCreator
+  public ResultSetInformation(
+      @Nullable
+      @JsonProperty ResultFormat resultFormat,
+      @Nullable
+      @JsonProperty Boolean header,
+      @JsonProperty @Nullable Long totalRows,
+      @JsonProperty @Nullable Long totalSize,
+      @JsonProperty("sampleRecords") @Nullable
+      List<Object> records
+  )
+  {
+    this.totalRows = totalRows;
+    this.totalSize = totalSize;
+    this.resultFormat = resultFormat;
+    this.header = header;
+    this.records = records;
+  }
+
+  @Nullable
+  @JsonProperty
+  @JsonInclude(JsonInclude.Include.NON_NULL)
+  public Long getTotalRows()
+  {
+    return totalRows;
+  }
+
+  @Nullable
+  @JsonProperty
+  @JsonInclude(JsonInclude.Include.NON_NULL)
+  public Long getTotalSize()
+  {
+    return totalSize;
+  }
+
+  @JsonProperty
+  @Nullable
+  @JsonInclude(JsonInclude.Include.NON_NULL)
+  public ResultFormat getResultFormat()
+  {
+    return resultFormat;
+  }
+
+  @JsonProperty
+  @Nullable
+  @JsonInclude(JsonInclude.Include.NON_NULL)
+  public Boolean getHeader()

Review Comment:
   rename -> hasHeader
   Also does it need to return a  `Boolean`? is null same as false? 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/ResultSetInformation.java:
##########
@@ -0,0 +1,141 @@
+/*
+ * 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.entity;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import org.apache.druid.sql.http.ResultFormat;
+
+import javax.annotation.Nullable;
+import java.util.List;
+import java.util.Objects;
+
+public class ResultSetInformation
+{
+
+  @Nullable
+  private final Long totalRows;
+  @Nullable
+  private final Long totalSize;
+
+  private final ResultFormat resultFormat;
+
+  private final Boolean header;

Review Comment:
   These are nullable too it seems. 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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);

Review Comment:
   do we need to check the result of this call? 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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);

Review Comment:
   are you missing auth checks here? 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          queryId
+      );
+    }
+
+    final long start = offset == null ? 0 : offset;
+    final long last = getLastIndex(numberOfRows, start);
+
+    try {
+      AuthorizationUtils.authorizeAllResourceActions(req, Collections.emptyList(), authorizerMapper);

Review Comment:
   this can happen before validation and call is not doing anything unless we validate the result. 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          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 = overlordClient.getTaskStatus(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();

Review Comment:
   it will be useful to add an error message along with response code. 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          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 = overlordClient.getTaskStatus(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
+            ),
+            queryId
+        );
+      } else if (sqlStatementState == SqlStatementState.FAILED) {
+        return buildNonOkResponse(
+            Response.Status.NOT_FOUND.getStatusCode(),
+            new QueryException(null, statusPlus.getErrorMsg(), null, null),
+            queryId
+        );
+      } else {
+        MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, authenticationResult.getIdentity());
+        Optional<List<ColNameAndType>> signature = getSignature(msqControllerTask);
+        Optional<List<Object>> results = getResults(overlordClient.getTaskReport(queryId));
+
+        return Response.ok(
+            (StreamingOutput) outputStream -> {
+              CountingOutputStream os = new CountingOutputStream(outputStream);
+
+              try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+                List<ColNameAndType> rowSignature = signature.get();
+                writer.writeResponseStart();
+
+                for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {
+                  writer.writeRowStart();
+                  for (int i = 0; i < rowSignature.size(); i++) {
+                    writer.writeRowField(rowSignature.get(i).getColName(), ((List) results.get().get(k)).get(i));
+                  }
+                  writer.writeRowEnd();
+                }
+
+                writer.writeResponseEnd();
+              }
+              catch (Exception e) {
+                log.error(e, "Unable to stream results back for query[%s]", queryId);
+                throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+              }
+            }
+        ).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;
+    }
+  }
+
+  /**
+   * Queries can be canceled while in anystate. Canceling a query that has already completed will remove its results.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@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(),
+          false
+      );
+      if (sqlStatementResult.isPresent()) {
+        switch (sqlStatementResult.get().getState()) {
+          case ACCEPTED:
+          case RUNNING:
+            overlordClient.cancelTask(queryId);
+            return Response.status(Response.Status.ACCEPTED).build();
+          case SUCCESS:
+          case FAILED:
+            // we would also want to clean up the
+            return Response.ok().build();
+          default:
+            throw new ISE("Illegal State[%s] encountered", sqlStatementResult.get().getState());
+        }
+
+      } 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;
+    }
+  }
+
+  private Response buildStandardResponse(
+      Sequence<Object[]> sequence,
+      SqlQuery sqlQuery,
+      String sqlQueryId,
+      SqlRowTransformer rowTransformer
+  ) throws IOException
+  {
+    final Yielder<Object[]> yielder0 = Yielders.each(sequence);
+
+    try {
+      final Response.ResponseBuilder responseBuilder = Response
+          .ok(
+              (StreamingOutput) outputStream -> {
+                CountingOutputStream os = new CountingOutputStream(outputStream);
+                Yielder<Object[]> yielder = yielder0;
+
+                try (final ResultFormat.Writer writer = sqlQuery.getResultFormat()
+                                                                .createFormatter(os, jsonMapper)) {
+                  writer.writeResponseStart();
+
+                  if (sqlQuery.includeHeader()) {
+                    writer.writeHeader(
+                        rowTransformer.getRowType(),
+                        sqlQuery.includeTypesHeader(),
+                        sqlQuery.includeSqlTypesHeader()
+                    );
+                  }
+
+                  while (!yielder.isDone()) {
+                    final Object[] row = yielder.get();
+                    writer.writeRowStart();
+                    for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
+                      final Object value = rowTransformer.transform(row, i);
+                      writer.writeRowField(rowTransformer.getFieldList().get(i), value);
+                    }
+                    writer.writeRowEnd();
+                    yielder = yielder.next(null);
+                  }
+
+                  writer.writeResponseEnd();
+                }
+                catch (Exception e) {
+                  log.error(e, "Unable to send SQL response [%s]", sqlQueryId);
+                  throw new RuntimeException(e);
+                }
+                finally {
+                  yielder.close();
+                }
+              }
+          );
+
+      if (sqlQuery.includeHeader()) {
+        responseBuilder.header(SqlResource.SQL_HEADER_RESPONSE_HEADER, SqlResource.SQL_HEADER_VALUE);
+      }
+
+      return responseBuilder.build();
+    }
+    catch (Throwable e) {
+      // make sure to close yielder if anything happened before starting to serialize the response.
+      yielder0.close();
+      throw e;
+    }
+  }
+
+  private Response buildTaskResponse(Sequence<Object[]> sequence, String user)
+  {
+    List<Object[]> rows = sequence.toList();
+    int numRows = rows.size();
+    if (numRows != 1) {
+      throw new RE("Expected a single row but got [%d] rows. Please check broker logs for more information.", numRows);
+    }
+    String taskId = (String) rows.get(0)[0];
+    try {
+      Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, user, true);
+
+      if (statementResult.isPresent()) {
+        return Response.status(Response.Status.OK).entity(statementResult.get()).build();
+      } else {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(
+                null,
+                StringUtils.format(
+                    "Unable to find associated task for query id %s",
+                    taskId
+                ),
+                null,
+                null
+            ),
+            taskId
+        );
+      }
+    }
+    catch (QueryException e) {
+      return buildNonOkResponse(Response.Status.OK.getStatusCode(), e, taskId);
+    }
+  }
+
+  private Response buildNonOkResponse(int status, SanitizableException e, String sqlQueryId)
+  {
+    // Though transformIfNeeded returns an exception, its purpose is to return
+    // a QueryException
+    Exception cleaned = serverConfig
+        .getErrorResponseTransformStrategy()
+        .transformIfNeeded(e);
+    return Response
+        .status(status)
+        .entity(cleaned)
+        .build();
+  }
+
+
+  private static Optional<List<ColNameAndType>> getSignature(
+      MSQControllerTask msqControllerTask
+  )
+  {
+    // only populate signature for select q's
+    if (msqControllerTask.getQuerySpec().getDestination().getClass() == TaskReportMSQDestination.class) {
+      ColumnMappings columnMappings = msqControllerTask.getQuerySpec().getColumnMappings();
+      List<SqlTypeName> sqlTypeNames = msqControllerTask.getSqlTypeNames();
+      if (sqlTypeNames == null) {
+        return Optional.empty();
+      }
+      List<ColNameAndType> signature = new ArrayList<>(columnMappings.size());
+      int index = 0;
+      for (String colName : columnMappings.getOutputColumnNames()) {
+        signature.add(new ColNameAndType(colName, sqlTypeNames.get(index).getName()));
+        index++;
+      }
+      return Optional.of(signature);
+    }
+    return Optional.empty();
+  }
+
+
+  private void checkTaskPayloadOrThrow(TaskPayloadResponse taskPayloadResponse, String queryId) throws QueryException
+  {
+    if (taskPayloadResponse == null || taskPayloadResponse.getPayload() == null) {
+      throw new QueryException(QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+                               StringUtils.format(
+                                   "Could not get payload details of query[%s] from the overlord",
+                                   queryId
+                               ), null, null
+      );
+    }
+
+    if (MSQControllerTask.class != taskPayloadResponse.getPayload().getClass()) {
+      throw new QueryException(QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+                               StringUtils.format(
+                                   "Fetched an unexpected payload details of query[%s] from the overlord.",
+                                   queryId
+                               ), null, null
+      );
+    }
+  }
+
+  private Optional<ResultSetInformation> getSampleResults(
+      String asyncResultId,
+      boolean isSelectQuery,
+      SqlStatementState sqlStatementState
+  )
+  {
+    // only populate sample results in case a select query is successful
+    if (isSelectQuery && sqlStatementState == SqlStatementState.SUCCESS) {
+      Map<String, Object> report = overlordClient.getTaskReport(asyncResultId);
+      Optional<List<Object>> rows = getResults(report);
+
+      if (rows.isPresent()) {
+        return Optional.of(new ResultSetInformation(
+            null,
+            false,
+            (long) rows.get().size(),
+            null,
+            rows.get()
+        ));
+      } else {
+        return Optional.empty();
+      }
+
+    } else {
+      return Optional.empty();
+    }
+  }
+
+
+  private Optional<SqlStatementResult> getStatementStatus(String queryId, String currentUser, boolean withResults)
+      throws QueryException, ForbiddenException
+  {
+    TaskStatusResponse taskResponse = overlordClient.getTaskStatus(queryId);
+    if (taskResponse == null) {
+      return Optional.empty();
+    }
+
+    TaskStatusPlus statusPlus = taskResponse.getStatus();
+    if (statusPlus == null || !MSQControllerTask.TYPE.equals(statusPlus.getType())) {
+      return Optional.empty();
+    }

Review Comment:
   is this expected? If not, can we log some warning? 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          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 = overlordClient.getTaskStatus(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
+            ),
+            queryId
+        );
+      } else if (sqlStatementState == SqlStatementState.FAILED) {
+        return buildNonOkResponse(
+            Response.Status.NOT_FOUND.getStatusCode(),
+            new QueryException(null, statusPlus.getErrorMsg(), null, null),
+            queryId
+        );
+      } else {

Review Comment:
   you can avoid this nested code by using following structure 
   if (abc)
     return xyz
   if (abcd)
     return something_else
   code for default



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] adarshsanjeev commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["adarshsanjeev (via GitHub)" <gi...@apache.org>]

adarshsanjeev commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1243654468


##########
sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java:
##########
@@ -2056,7 +2056,7 @@ private <T> T deserializeResponse(Response resp, Class<T> clazz) throws IOExcept
     return JSON_MAPPER.readValue(responseToByteArray(resp), clazz);
   }
 
-  private byte[] responseToByteArray(Response resp) throws IOException
+  public static byte[] responseToByteArray(Response resp) throws IOException

Review Comment:
   Minor nit: Could be pushed up to `CalciteTestBase` so that other classes can use it in the future



##########
server/src/main/java/org/apache/druid/rpc/indexing/OverlordClient.java:
##########
@@ -52,5 +53,7 @@ public interface OverlordClient
 
   ListenableFuture<Map<String, Object>> taskReportAsMap(String taskId);
 
+  ListenableFuture<TaskPayloadResponse> taskPayload(String taskId);

Review Comment:
   nit: This function might need a better name to avoid confusing it with the task payload containing results, since we have a payload() function in SqlStatementResource already. Maybe `taskDefinition`?



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] github-code-scanning[bot] commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["github-code-scanning[bot] (via GitHub)" <gi...@apache.org>]

github-code-scanning[bot] commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1243975717


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/SqlStatementResult.java:
##########
@@ -0,0 +1,185 @@
+/*
+ * 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.entity;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import org.apache.druid.error.ErrorResponse;
+import org.apache.druid.msq.sql.SqlStatementState;
+import org.joda.time.DateTime;
+
+import javax.annotation.Nullable;
+import java.util.List;
+import java.util.Objects;
+
+public class SqlStatementResult
+{
+
+  private final String queryId;
+
+  private final SqlStatementState state;
+
+  private final DateTime createdAt;
+
+  @Nullable
+  private final List<ColumnNameAndTypes> sqlRowSignature;
+
+  @Nullable
+  private final Long durationMs;
+
+  @Nullable
+  private final ResultSetInformation resultSetInformation;
+
+  @Nullable
+  private final ErrorResponse errorResponse;
+
+
+  @JsonCreator
+  public SqlStatementResult(
+      @JsonProperty("queryId")
+      String queryId,
+      @JsonProperty("state")
+      SqlStatementState state,
+      @JsonProperty("createdAt")
+      DateTime createdAt,
+      @Nullable @JsonProperty("schema")
+      List<ColumnNameAndTypes> sqlRowSignature,
+      @Nullable @JsonProperty("durationMs")
+      Long durationMs,
+      @Nullable @JsonProperty("result")
+      ResultSetInformation resultSetInformation,
+      @Nullable @JsonProperty("errorDetails")
+      ErrorResponse errorResponse
+
+  )
+  {
+    this.queryId = queryId;
+    this.state = state;
+    this.createdAt = createdAt;
+    this.sqlRowSignature = sqlRowSignature;
+    this.durationMs = durationMs;
+    this.resultSetInformation = resultSetInformation;
+    this.errorResponse = errorResponse;
+  }
+
+  @JsonProperty
+  public String getQueryId()
+  {
+    return queryId;
+  }
+
+  @JsonProperty
+  public SqlStatementState getState()
+  {
+    return state;
+  }
+
+  @JsonProperty
+  public DateTime getCreatedAt()
+  {
+    return createdAt;
+  }
+
+  @JsonProperty("schema")
+  @Nullable
+  @JsonInclude(JsonInclude.Include.NON_NULL)
+  public List<ColumnNameAndTypes> getSqlRowSignature()
+  {
+    return sqlRowSignature;
+  }
+
+  @JsonProperty
+  @Nullable
+  @JsonInclude(JsonInclude.Include.NON_NULL)
+  public Long getDurationMs()
+  {
+    return durationMs;
+  }
+
+  @JsonProperty("result")
+  @Nullable
+  @JsonInclude(JsonInclude.Include.NON_NULL)
+  public ResultSetInformation getResultSetInformation()
+  {
+    return resultSetInformation;
+  }
+
+  @JsonProperty("errorDetails")
+  @Nullable
+  @JsonInclude(JsonInclude.Include.NON_NULL)
+  public ErrorResponse getErrorResponse()
+  {
+    return errorResponse;
+  }
+
+
+  @Override
+  public boolean equals(Object o)
+  {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    SqlStatementResult that = (SqlStatementResult) o;
+    return Objects.equals(queryId, that.queryId) && state == that.state && Objects.equals(
+        createdAt,
+        that.createdAt
+    ) && Objects.equals(sqlRowSignature, that.sqlRowSignature) && Objects.equals(
+        durationMs,
+        that.durationMs
+    ) && Objects.equals(resultSetInformation, that.resultSetInformation) && Objects.equals(
+        errorResponse == null ? null : errorResponse.getAsMap(),
+        that.errorResponse == null ? null : errorResponse.getAsMap()

Review Comment:
   ## Dereferenced variable may be null
   
   Variable [errorResponse](1) may be null at this access as suggested by [this](2) null guard.
   
   [Show more details](https://github.com/apache/druid/security/code-scanning/5196)



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1244664179


##########
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))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColNameAndType> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {
+              writer.writeRowStart();
+              for (int i = 0; i < rowSignature.size(); i++) {
+                writer.writeRowField(rowSignature.get(i).getColName(), ((List) results.get().get(k)).get(i));
+              }
+              writer.writeRowEnd();
+            }
+
+            writer.writeResponseEnd();
+          }
+          catch (Exception e) {
+            log.error(e, "Unable to stream results back for query[%s]", queryId);
+            throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+          }
+        }).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;
+    }
+  }
+
+  /**
+   * Queries can be canceled while in anystate. Canceling a query that has already completed will remove its results.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@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(),
+          false
+      );
+      if (sqlStatementResult.isPresent()) {
+        switch (sqlStatementResult.get().getState()) {
+          case ACCEPTED:
+          case RUNNING:
+            overlordClient.cancelTask(queryId);
+            return Response.status(Response.Status.ACCEPTED).build();
+          case SUCCESS:
+          case FAILED:
+            // we would also want to clean up the
+            return Response.ok().build();
+          default:
+            throw new ISE("Illegal State[%s] encountered", sqlStatementResult.get().getState());
+        }
+
+      } 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;
+    }
+  }
+
+  private Response buildStandardResponse(
+      Sequence<Object[]> sequence,
+      SqlQuery sqlQuery,
+      String sqlQueryId,
+      SqlRowTransformer rowTransformer
+  ) throws IOException
+  {
+    final Yielder<Object[]> yielder0 = Yielders.each(sequence);
+
+    try {
+      final Response.ResponseBuilder responseBuilder = Response.ok((StreamingOutput) outputStream -> {
+        CountingOutputStream os = new CountingOutputStream(outputStream);
+        Yielder<Object[]> yielder = yielder0;
+
+        try (final ResultFormat.Writer writer = sqlQuery.getResultFormat().createFormatter(os, jsonMapper)) {
+          writer.writeResponseStart();
+
+          if (sqlQuery.includeHeader()) {
+            writer.writeHeader(
+                rowTransformer.getRowType(),
+                sqlQuery.includeTypesHeader(),
+                sqlQuery.includeSqlTypesHeader()
+            );
+          }
+
+          while (!yielder.isDone()) {
+            final Object[] row = yielder.get();
+            writer.writeRowStart();
+            for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
+              final Object value = rowTransformer.transform(row, i);
+              writer.writeRowField(rowTransformer.getFieldList().get(i), value);
+            }
+            writer.writeRowEnd();
+            yielder = yielder.next(null);
+          }
+
+          writer.writeResponseEnd();
+        }
+        catch (Exception e) {
+          log.error(e, "Unable to send SQL response [%s]", sqlQueryId);
+          throw new RuntimeException(e);
+        }
+        finally {
+          yielder.close();
+        }
+      });
+
+      if (sqlQuery.includeHeader()) {
+        responseBuilder.header(SqlResource.SQL_HEADER_RESPONSE_HEADER, SqlResource.SQL_HEADER_VALUE);
+      }
+
+      return responseBuilder.build();
+    }
+    catch (Throwable e) {
+      // make sure to close yielder if anything happened before starting to serialize the response.
+      yielder0.close();
+      throw e;
+    }
+  }
+
+  private Response buildTaskResponse(Sequence<Object[]> sequence, String user)
+  {
+    List<Object[]> rows = sequence.toList();
+    int numRows = rows.size();
+    if (numRows != 1) {
+      throw new RE("Expected a single row but got [%d] rows. Please check broker logs for more information.", numRows);
+    }
+    String taskId = (String) rows.get(0)[0];
+    try {
+      Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, user, true);
+
+      if (statementResult.isPresent()) {
+        return Response.status(Response.Status.OK).entity(statementResult.get()).build();
+      } else {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(
+                null,
+                StringUtils.format(
+                    "Unable to find associated task for query id %s",
+                    taskId
+                ),
+                null,
+                null,
+                null
+            ),
+            taskId
+        );
+      }
+    }
+    catch (QueryException e) {
+      return buildNonOkResponse(Response.Status.OK.getStatusCode(), e, taskId);
+    }
+  }
+
+  private Response buildNonOkResponse(int status, SanitizableException e, String sqlQueryId)
+  {
+    // Though transformIfNeeded returns an exception, its purpose is to return
+    // a QueryException
+    Exception cleaned = serverConfig.getErrorResponseTransformStrategy().transformIfNeeded(e);
+    return Response.status(status).entity(cleaned).build();
+  }
+
+
+  private static Optional<List<ColNameAndType>> getSignature(
+      MSQControllerTask msqControllerTask
+  )
+  {
+    // only populate signature for select q's
+    if (msqControllerTask.getQuerySpec().getDestination().getClass() == TaskReportMSQDestination.class) {
+      ColumnMappings columnMappings = msqControllerTask.getQuerySpec().getColumnMappings();
+      List<SqlTypeName> sqlTypeNames = msqControllerTask.getSqlTypeNames();
+      if (sqlTypeNames == null || sqlTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColumnType> nativeTypeNames = msqControllerTask.getNativeTypeNames();
+      if (nativeTypeNames == null || nativeTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColNameAndType> signature = new ArrayList<>(columnMappings.size());
+      int index = 0;
+      for (String colName : columnMappings.getOutputColumnNames()) {
+        signature.add(new ColNameAndType(
+            colName,
+            sqlTypeNames.get(index).getName(),
+            nativeTypeNames.get(index).asTypeString()
+        ));
+        index++;
+      }
+      return Optional.of(signature);
+    }
+    return Optional.empty();
+  }
+
+
+  private void checkTaskPayloadOrThrow(TaskPayloadResponse taskPayloadResponse, String queryId) throws QueryException
+  {
+    if (taskPayloadResponse == null || taskPayloadResponse.getPayload() == null) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Could not get payload details for query[%s] from the overlord",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+
+    if (MSQControllerTask.class != taskPayloadResponse.getPayload().getClass()) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Fetched an unexpected payload for query[%s] from the overlord.",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+  }
+
+  private Optional<ResultSetInformation> getSampleResults(
+      String asyncResultId,
+      boolean isSelectQuery,
+      String dataSource,
+      SqlStatementState sqlStatementState
+  )
+  {
+    if (sqlStatementState == SqlStatementState.SUCCESS) {
+      Map<String, Object> payload = getPayload(overlordWork(overlordClient.taskReportAsMap(asyncResultId)));
+      Optional<Pair<Long, Long>> rowsAndSize = getRowsAndSizeFromPayload(payload, isSelectQuery);
+      return Optional.of(new ResultSetInformation(
+          null,
+          // since the rows can be sampled, get the number of rows from counters
+          rowsAndSize.orElse(new Pair<>(null, null)).lhs,
+          rowsAndSize.orElse(new Pair<>(null, null)).rhs,
+          dataSource,
+          // only populate sample results in case a select query is successful
+          isSelectQuery ? getResults(payload).orElse(null) : null
+      ));
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  private Optional<Pair<Long, Long>> getRowsAndSizeFromPayload(Map<String, Object> payload, boolean isSelectQuery)
+  {
+    List stages = getList(payload, "stages");
+    if (stages == null || stages.isEmpty()) {
+      return Optional.empty();
+    } else {
+      int maxStage = stages.size() - 1; // Last stage output is the total number of rows returned to the end user.
+      Map<String, Object> counterMap = getMap(getMap(payload, "counters"), String.valueOf(maxStage));
+      Long rows = -1L;
+      Long sizeInBytes = -1L;
+      if (counterMap == null) {
+        return Optional.empty();
+      }
+      for (Map.Entry<String, Object> worker : counterMap.entrySet()) {
+        Object workerChannels = worker.getValue();
+        if (workerChannels == null || !(workerChannels instanceof Map)) {
+          return Optional.empty();
+        }
+        if (isSelectQuery) {
+          Object output = ((Map<?, ?>) workerChannels).get("output");
+          if (output != null && output instanceof Map) {
+            List<Integer> rowsPerChannel = (List<Integer>) ((Map<String, Object>) output).get("rows");
+            List<Integer> bytesPerChannel = (List<Integer>) ((Map<String, Object>) output).get("bytes");
+            for (Integer row : rowsPerChannel) {
+              rows = rows + row;
+            }
+            for (Integer bytes : bytesPerChannel) {
+              sizeInBytes = sizeInBytes + bytes;
+            }
+          }
+        } else {
+          Object output = ((Map<?, ?>) workerChannels).get("segmentGenerationProgress");
+          if (output != null && output instanceof Map) {
+            rows += (Integer) ((Map<String, Object>) output).get("rowsPushed");
+          }
+        }
+      }
+
+      return Optional.of(new Pair<>(rows == -1L ? null : rows + 1, sizeInBytes == -1L ? null : sizeInBytes + 1));
+    }
+  }
+
+
+  private Optional<SqlStatementResult> getStatementStatus(String queryId, String currentUser, boolean withResults)
+      throws QueryException, ForbiddenException
+  {
+    TaskStatusResponse taskResponse = overlordWork(overlordClient.taskStatus(queryId));
+    if (taskResponse == null) {
+      return Optional.empty();
+    }
+
+    TaskStatusPlus statusPlus = taskResponse.getStatus();
+    if (statusPlus == null || !MSQControllerTask.TYPE.equals(statusPlus.getType())) {
+      return Optional.empty();
+    }
+
+    // since we need the controller payload for auth checks.
+    MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, currentUser);
+    SqlStatementState sqlStatementState = getSqlStatementState(statusPlus);
+
+    if (SqlStatementState.FAILED == sqlStatementState) {
+      Map<String, Object> exceptionDetails = getQueryExceptionDetails(getPayload(overlordWork(overlordClient.taskReportAsMap(
+          queryId))));
+      Map<String, Object> exception = getMap(exceptionDetails, "error");
+      if (exceptionDetails == null || exception == null) {
+        return Optional.of(new SqlStatementResult(
+            queryId,
+            sqlStatementState,
+            taskResponse.getStatus().getCreatedTime(),
+            null,
+            taskResponse.getStatus().getDuration(),
+            null,
+            new QueryException(null, statusPlus.getErrorMsg(), null, null, null)
+        ));
+      }
+
+      final String errorMessage = String.valueOf(exception.getOrDefault("errorMessage", statusPlus.getErrorMsg()));
+      exception.remove("errorMessage");
+      String host = (String) exceptionDetails.getOrDefault("host", null);
+      Map<String, String> stringException = new HashMap<>();
+      for (Map.Entry<String, Object> exceptionKeys : exception.entrySet()) {
+        stringException.put(exceptionKeys.getKey(), String.valueOf(exceptionKeys.getValue()));
+      }
+      return Optional.of(new SqlStatementResult(
+          queryId,
+          sqlStatementState,
+          taskResponse.getStatus().getCreatedTime(),
+          null,
+          taskResponse.getStatus().getDuration(),
+          null,
+          new QueryException(null, errorMessage, null, host, stringException)
+      ));
+    } else {
+      Optional<List<ColNameAndType>> signature = getSignature(msqControllerTask);
+
+      return Optional.of(new SqlStatementResult(
+          queryId,
+          sqlStatementState,
+          taskResponse.getStatus().getCreatedTime(),
+          signature.orElse(null),
+          taskResponse.getStatus().getDuration(),
+          withResults ? getSampleResults(
+              queryId,
+              signature.isPresent(),
+              msqControllerTask.getDataSource(),
+              sqlStatementState
+          ).orElse(null) : null,
+          null
+      ));
+    }
+  }
+
+  private Map<String, Object> getQueryExceptionDetails(Map<String, Object> payload)
+  {
+    return getMap(getMap(payload, "status"), "errorReport");
+  }
+
+  private MSQControllerTask getMSQControllerTaskOrThrow(String queryId, String currentUser) throws ForbiddenException
+  {
+    TaskPayloadResponse taskPayloadResponse = overlordWork(overlordClient.taskPayload(queryId));
+    checkTaskPayloadOrThrow(taskPayloadResponse, queryId);
+
+    MSQControllerTask msqControllerTask = (MSQControllerTask) taskPayloadResponse.getPayload();
+    if (currentUser == null || !currentUser.equals(msqControllerTask.getQuerySpec()
+                                                                    .getQuery()
+                                                                    .getContext()
+                                                                    .get(MSQTaskQueryMaker.USER_KEY))) {
+      throw new ForbiddenException();
+    }
+    return msqControllerTask;
+  }
+
+  private <T> T overlordWork(final ListenableFuture<T> future)
+  {
+    try {
+      return FutureUtils.getUnchecked(future, true);
+    }
+    catch (RuntimeException e) {
+      throw new QueryException(null, "Unable to contact overlord " + e.getMessage(), null, null, null);

Review Comment:
   This can happen for n reasons. Changed it to a developer exception. 
   ``` 
   DruidException.forPersona(DruidException.Persona.DEVELOPER)
                             .ofCategory(DruidException.Category.UNCATEGORIZED)
                             .build("Unable to contact overlord " + e.getMessage());
       }
       ```



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1244665651


##########
sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java:
##########
@@ -2056,7 +2056,7 @@ private <T> T deserializeResponse(Response resp, Class<T> clazz) throws IOExcept
     return JSON_MAPPER.readValue(responseToByteArray(resp), clazz);
   }
 
-  private byte[] responseToByteArray(Response resp) throws IOException
+  public static byte[] responseToByteArray(Response resp) throws IOException

Review Comment:
    I generally avoid pushing stuff to calcite test base since its already a blowed up class.



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1245109564


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          queryId
+      );
+    }
+
+    final long start = offset == null ? 0 : offset;
+    final long last = getLastIndex(numberOfRows, start);
+
+    try {
+      AuthorizationUtils.authorizeAllResourceActions(req, Collections.emptyList(), authorizerMapper);

Review Comment:
   @abhishekagarwal87 I am following the pattern from
   
   [druid/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/SqlTaskResource.java](https://github.com/apache/druid/blob/baa64e6d8a71b1d3816dea6144ca4322d0432e53/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/SqlTaskResource.java#L110)
   
   We might need to add a new method to AuthorizationUtilsto set the req header but that can be done as part of a follow up PR.



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] github-code-scanning[bot] commented on a diff in pull request #14416: [Draft] Adding Interactive API's for MSQ engine

["github-code-scanning[bot] (via GitHub)" <gi...@apache.org>]

github-code-scanning[bot] commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1228335168


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,753 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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,

Review Comment:
   ## Useless parameter
   
   The parameter 'size' is never used.
   
   [Show more details](https://github.com/apache/druid/security/code-scanning/5054)



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,753 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          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 = overlordClient.getTaskStatus(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();
+      }
+
+      MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, authenticationResult.getIdentity());
+      SqlStatementState sqlStatementState = getSqlStatementState(statusPlus);
+
+      if (sqlStatementState == SqlStatementState.RUNNING || sqlStatementState == SqlStatementState.ACCEPTED) {
+        return buildNonOkResponse(
+            Response.Status.PRECONDITION_FAILED.getStatusCode(),
+            new QueryException(null, "Query not ready to retrieve results", null, null),
+            queryId
+        );
+      } else if (sqlStatementState == SqlStatementState.FAILED) {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(null, statusPlus.getErrorMsg(), null, null),
+            queryId
+        );
+      } else {
+        Optional<List<ColNameAndType>> signature = getSignature(msqControllerTask);
+        Optional<List<Object>> results = getResults(overlordClient.getTaskReport(queryId));
+
+        return Response.ok(
+            (StreamingOutput) outputStream -> {
+              CountingOutputStream os = new CountingOutputStream(outputStream);
+
+              try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+                List<ColNameAndType> rowSignature = signature.get();
+                writer.writeResponseStart();
+
+                for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {
+                  writer.writeRowStart();
+                  for (int i = 0; i < rowSignature.size(); i++) {
+                    writer.writeRowField(rowSignature.get(i).getColName(), ((List) results.get().get(k)).get(i));
+                  }
+                  writer.writeRowEnd();
+                }
+
+                writer.writeResponseEnd();
+              }
+              catch (Exception e) {
+                log.error(e, "Unable to stream results back for query[%s]", queryId);
+                throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+              }
+            }
+        ).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;
+    }
+  }
+
+  /**
+   * Queries can be canceled while in anystate. Canceling a query that has already completed will remove its results.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@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(),
+          false
+      );
+      if (sqlStatementResult.isPresent()) {
+        switch (sqlStatementResult.get().getState()) {
+          case ACCEPTED:
+          case RUNNING:
+            overlordClient.cancelTask(queryId);
+            return Response.status(Response.Status.ACCEPTED).build();
+          case SUCCESS:
+          case FAILED:
+            // we would also want to clean up the
+            return Response.ok().build();
+          default:
+            throw new ISE("Illegal State[%s] encountered", sqlStatementResult.get().getState());
+        }
+
+      } 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;
+    }
+  }
+
+  private Response buildStandardResponse(
+      Sequence<Object[]> sequence,
+      SqlQuery sqlQuery,
+      String sqlQueryId,
+      SqlRowTransformer rowTransformer
+  ) throws IOException
+  {
+    final Yielder<Object[]> yielder0 = Yielders.each(sequence);
+
+    try {
+      final Response.ResponseBuilder responseBuilder = Response
+          .ok(
+              (StreamingOutput) outputStream -> {
+                CountingOutputStream os = new CountingOutputStream(outputStream);
+                Yielder<Object[]> yielder = yielder0;
+
+                try (final ResultFormat.Writer writer = sqlQuery.getResultFormat()
+                                                                .createFormatter(os, jsonMapper)) {
+                  writer.writeResponseStart();
+
+                  if (sqlQuery.includeHeader()) {
+                    writer.writeHeader(
+                        rowTransformer.getRowType(),
+                        sqlQuery.includeTypesHeader(),
+                        sqlQuery.includeSqlTypesHeader()
+                    );
+                  }
+
+                  while (!yielder.isDone()) {
+                    final Object[] row = yielder.get();
+                    writer.writeRowStart();
+                    for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
+                      final Object value = rowTransformer.transform(row, i);
+                      writer.writeRowField(rowTransformer.getFieldList().get(i), value);
+                    }
+                    writer.writeRowEnd();
+                    yielder = yielder.next(null);
+                  }
+
+                  writer.writeResponseEnd();
+                }
+                catch (Exception e) {
+                  log.error(e, "Unable to send SQL response [%s]", sqlQueryId);
+                  throw new RuntimeException(e);
+                }
+                finally {
+                  yielder.close();
+                }
+              }
+          );
+
+      if (sqlQuery.includeHeader()) {
+        responseBuilder.header(SqlResource.SQL_HEADER_RESPONSE_HEADER, SqlResource.SQL_HEADER_VALUE);
+      }
+
+      return responseBuilder.build();
+    }
+    catch (Throwable e) {
+      // make sure to close yielder if anything happened before starting to serialize the response.
+      yielder0.close();
+      throw e;
+    }
+  }
+
+  private Response buildTaskResponse(Sequence<Object[]> sequence, String user)
+  {
+    List<Object[]> rows = sequence.toList();
+    int numRows = rows.size();
+    if (numRows != 1) {
+      throw new RE("Expected a single row but got [%d] rows. Please check broker logs for more information.", numRows);
+    }
+    String taskId = (String) rows.get(0)[0];
+    try {
+      Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, user, true);
+
+      if (statementResult.isPresent()) {
+        return Response.status(Response.Status.OK).entity(statementResult.get()).build();
+      } else {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(
+                null,
+                StringUtils.format(
+                    "Unable to find associated task for query id %s",
+                    taskId
+                ),
+                null,
+                null
+            ),
+            taskId
+        );
+      }
+    }
+    catch (QueryException e) {
+      return buildNonOkResponse(Response.Status.OK.getStatusCode(), e, taskId);
+    }
+  }
+
+  private Response buildNonOkResponse(int status, SanitizableException e, String sqlQueryId)

Review Comment:
   ## Useless parameter
   
   The parameter 'sqlQueryId' is never used.
   
   [Show more details](https://github.com/apache/druid/security/code-scanning/5056)



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,753 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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,

Review Comment:
   ## Useless parameter
   
   The parameter 'timeout' is never used.
   
   [Show more details](https://github.com/apache/druid/security/code-scanning/5055)



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,753 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          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 = overlordClient.getTaskStatus(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();
+      }
+
+      MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, authenticationResult.getIdentity());
+      SqlStatementState sqlStatementState = getSqlStatementState(statusPlus);
+
+      if (sqlStatementState == SqlStatementState.RUNNING || sqlStatementState == SqlStatementState.ACCEPTED) {
+        return buildNonOkResponse(
+            Response.Status.PRECONDITION_FAILED.getStatusCode(),
+            new QueryException(null, "Query not ready to retrieve results", null, null),
+            queryId
+        );
+      } else if (sqlStatementState == SqlStatementState.FAILED) {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(null, statusPlus.getErrorMsg(), null, null),
+            queryId
+        );
+      } else {
+        Optional<List<ColNameAndType>> signature = getSignature(msqControllerTask);
+        Optional<List<Object>> results = getResults(overlordClient.getTaskReport(queryId));
+
+        return Response.ok(
+            (StreamingOutput) outputStream -> {
+              CountingOutputStream os = new CountingOutputStream(outputStream);
+
+              try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+                List<ColNameAndType> rowSignature = signature.get();
+                writer.writeResponseStart();
+
+                for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {

Review Comment:
   ## User-controlled data in numeric cast
   
   This cast to a narrower type depends on a [user-provided value](1), potentially causing truncation.
   
   [Show more details](https://github.com/apache/druid/security/code-scanning/5053)



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,753 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          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 = overlordClient.getTaskStatus(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();
+      }
+
+      MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, authenticationResult.getIdentity());
+      SqlStatementState sqlStatementState = getSqlStatementState(statusPlus);
+
+      if (sqlStatementState == SqlStatementState.RUNNING || sqlStatementState == SqlStatementState.ACCEPTED) {
+        return buildNonOkResponse(
+            Response.Status.PRECONDITION_FAILED.getStatusCode(),
+            new QueryException(null, "Query not ready to retrieve results", null, null),
+            queryId
+        );
+      } else if (sqlStatementState == SqlStatementState.FAILED) {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(null, statusPlus.getErrorMsg(), null, null),
+            queryId
+        );
+      } else {
+        Optional<List<ColNameAndType>> signature = getSignature(msqControllerTask);
+        Optional<List<Object>> results = getResults(overlordClient.getTaskReport(queryId));
+
+        return Response.ok(
+            (StreamingOutput) outputStream -> {
+              CountingOutputStream os = new CountingOutputStream(outputStream);
+
+              try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+                List<ColNameAndType> rowSignature = signature.get();
+                writer.writeResponseStart();
+
+                for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {

Review Comment:
   ## Comparison of narrow type with wide type in loop condition
   
   Comparison between [expression](1) of type int and [expression](2) of wider type long.
   
   [Show more details](https://github.com/apache/druid/security/code-scanning/5058)



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1243040857


##########
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();

Review Comment:
   Signature should not be empty but that's for the MSQ engine to decide. 
   We can't return results if the signature is empty. So if the results call get's an empty signature, we just return an empty result. 
   



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1244665651


##########
sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java:
##########
@@ -2056,7 +2056,7 @@ private <T> T deserializeResponse(Response resp, Class<T> clazz) throws IOExcept
     return JSON_MAPPER.readValue(responseToByteArray(resp), clazz);
   }
 
-  private byte[] responseToByteArray(Response resp) throws IOException
+  public static byte[] responseToByteArray(Response resp) throws IOException

Review Comment:
    I generally avoid pushing stuff to calcite test base since its already a big class until necessary. We can probably do that if this method is being used in other places in the future. 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

Re: [PR] Adding Interactive API's for MSQ engine (druid)

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1299395739


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,670 @@
+/*
+ * 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.druid.client.indexing.TaskPayloadResponse;
+import org.apache.druid.client.indexing.TaskStatusResponse;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.error.ErrorResponse;
+import org.apache.druid.error.QueryExceptionCompat;
+import org.apache.druid.guice.annotations.MSQ;
+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.sql.MSQTaskQueryMaker;
+import org.apache.druid.msq.sql.MSQTaskSqlEngine;
+import org.apache.druid.msq.sql.SqlStatementState;
+import org.apache.druid.msq.sql.entity.ColumnNameAndTypes;
+import org.apache.druid.msq.sql.entity.ResultSetInformation;
+import org.apache.druid.msq.sql.entity.SqlStatementResult;
+import org.apache.druid.msq.util.SqlStatementResourceHelper;
+import org.apache.druid.query.ExecutionMode;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.query.QueryException;
+import org.apache.druid.rpc.indexing.OverlordClient;
+import org.apache.druid.server.QueryResponse;
+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.SqlRowTransformer;
+import org.apache.druid.sql.SqlStatementFactory;
+import org.apache.druid.sql.http.ResultFormat;
+import org.apache.druid.sql.http.SqlQuery;
+import org.apache.druid.sql.http.SqlResource;
+
+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.Collections;
+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 AuthorizerMapper authorizerMapper;
+  private final ObjectMapper jsonMapper;
+  private final OverlordClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final OverlordClient overlordClient
+  )
+  {
+    this.msqSqlStatementFactory = msqSqlStatementFactory;
+    this.authorizerMapper = authorizerMapper;
+    this.jsonMapper = jsonMapper;
+    this.overlordClient = overlordClient;
+  }
+
+
+  @GET
+  @Path("/enabled")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response isEnabled(@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(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              StringUtils.format(
+                                  "The statement sql api only supports sync mode[%s]. Please set context parameter [%s=%s] in the context payload",
+                                  ExecutionMode.ASYNC,
+                                  QueryContexts.CTX_EXECUTION_MODE,
+                                  ExecutionMode.ASYNC
+                              )
+                          )
+        );
+      }
+
+
+      Thread.currentThread().setName(StringUtils.format("statement_sql[%s]", sqlQueryId));
+
+      final DirectStatement.ResultSet plan = stmt.plan();
+      // in case the engine is async, the query is not run yet. We just return the taskID in case of non explain queries.
+      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);
+      }
+    }
+    catch (DruidException e) {
+      stmt.reporter().failed(e);
+      return buildNonOkResponse(e);
+    }
+    catch (QueryException queryException) {
+      stmt.reporter().failed(queryException);
+      final DruidException underlyingException = DruidException.fromFailure(new QueryExceptionCompat(queryException));
+      return buildNonOkResponse(underlyingException);
+    }
+    catch (ForbiddenException e) {
+      log.debug("Got forbidden request for reason [%s]", e.getErrorMessage());
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.USER)
+                        .ofCategory(DruidException.Category.FORBIDDEN)
+                        .build(Access.DEFAULT_ERROR_MESSAGE)
+      );
+    }
+    // Calcite throws java.lang.AssertionError at various points in planning/validation.
+    catch (AssertionError | Exception e) {
+      stmt.reporter().failed(e);
+      log.noStackTrace().warn(e, "Failed to handle query: %s", sqlQueryId);
+
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.DEVELOPER)
+                        .ofCategory(DruidException.Category.UNCATEGORIZED)
+                        .build(e.getMessage())
+      );
+    }
+    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 {
+      Access authResult = AuthorizationUtils.authorizeAllResourceActions(
+          req,
+          Collections.emptyList(),
+          authorizerMapper
+      );
+      if (!authResult.isAllowed()) {
+        throw new ForbiddenException(authResult.toString());
+      }
+      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 (DruidException e) {
+      return buildNonOkResponse(e);
+    }
+    catch (ForbiddenException e) {
+      log.debug("Got forbidden request for reason [%s]", e.getErrorMessage());
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.USER)
+                        .ofCategory(DruidException.Category.FORBIDDEN)
+                        .build(Access.DEFAULT_ERROR_MESSAGE)
+      );
+    }
+    catch (Exception e) {
+      log.noStackTrace().warn(e, "Failed to handle query: %s", queryId);
+      return buildNonOkResponse(DruidException.forPersona(DruidException.Persona.DEVELOPER)
+                                              .ofCategory(DruidException.Category.UNCATEGORIZED)
+                                              .build(e, "Failed to handle query: [%s]", queryId));
+    }
+  }
+
+  @GET
+  @Path("/{id}/results")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response doGetResults(
+      @PathParam("id") final String queryId,
+      @QueryParam("offset") Long offset,
+      @QueryParam("numRows") Long numberOfRows,
+      @Context final HttpServletRequest req
+  )
+  {
+    try {
+      Access authResult = AuthorizationUtils.authorizeAllResourceActions(
+          req,
+          Collections.emptyList(),
+          authorizerMapper
+      );
+      if (!authResult.isAllowed()) {
+        throw new ForbiddenException(authResult.toString());
+      }
+      final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req);
+
+      if (offset != null && offset < 0) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "offset cannot be negative. Please pass a positive number."
+                          )
+        );
+      }
+      if (numberOfRows != null && numberOfRows < 0) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "numRows cannot be negative. Please pass a positive number."
+                          )
+        );
+      }
+
+      final long start = offset == null ? 0 : offset;
+      final long last = SqlStatementResourceHelper.getLastIndex(numberOfRows, start);
+
+      TaskStatusResponse taskResponse = contactOverlord(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();
+      }
+
+      MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, authenticationResult.getIdentity());
+      SqlStatementState sqlStatementState = SqlStatementResourceHelper.getSqlStatementState(statusPlus);
+
+      if (sqlStatementState == SqlStatementState.RUNNING || sqlStatementState == SqlStatementState.ACCEPTED) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "Query[%s] is currently in [%s] state. Please wait for it to complete.",
+                              queryId,
+                              sqlStatementState
+                          )
+        );
+      } else if (sqlStatementState == SqlStatementState.FAILED) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "Query[%s] failed. Hit status api for more details.",
+                              queryId
+                          )
+        );
+      } else {
+        Optional<List<ColumnNameAndTypes>> signature = SqlStatementResourceHelper.getSignature(msqControllerTask);
+        if (!signature.isPresent()) {
+          return Response.ok().build();
+        }
+        Optional<List<Object>> results = SqlStatementResourceHelper.getResults(SqlStatementResourceHelper.getPayload(
+            contactOverlord(overlordClient.taskReportAsMap(queryId))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColumnNameAndTypes> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (long k = start; k < Math.min(last, results.get().size()); k++) {
+              writer.writeRowStart();
+              for (int i = 0; i < rowSignature.size(); i++) {
+                writer.writeRowField(
+                    rowSignature.get(i).getColName(),
+                    ((List) results.get().get(Math.toIntExact(k))).get(i)
+                );
+              }
+              writer.writeRowEnd();
+            }
+
+            writer.writeResponseEnd();
+          }
+          catch (Exception e) {
+            log.error(e, "Unable to stream results back for query[%s]", queryId);
+            throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+          }
+        }).build();
+
+      }
+    }
+    catch (DruidException e) {
+      return buildNonOkResponse(e);
+    }
+    catch (ForbiddenException e) {
+      log.debug("Got forbidden request for reason [%s]", e.getErrorMessage());
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.USER)
+                        .ofCategory(DruidException.Category.FORBIDDEN)
+                        .build(Access.DEFAULT_ERROR_MESSAGE)
+      );
+    }
+    catch (Exception e) {
+      log.noStackTrace().warn(e, "Failed to handle query: %s", queryId);
+      return buildNonOkResponse(DruidException.forPersona(DruidException.Persona.DEVELOPER)
+                                              .ofCategory(DruidException.Category.UNCATEGORIZED)
+                                              .build(e, "Failed to handle query: [%s]", queryId));
+    }
+  }
+
+  /**
+   * Queries can be canceled while in any {@link SqlStatementState}. Canceling a query that has already completed will be a no-op.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@PathParam("id") final String queryId, @Context final HttpServletRequest req)
+  {
+
+    try {
+      Access authResult = AuthorizationUtils.authorizeAllResourceActions(

Review Comment:
   @abhishekagarwal87 Raised a PR : https://github.com/apache/druid/pull/14878



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1244662494


##########
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))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColNameAndType> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {
+              writer.writeRowStart();
+              for (int i = 0; i < rowSignature.size(); i++) {
+                writer.writeRowField(rowSignature.get(i).getColName(), ((List) results.get().get(k)).get(i));
+              }
+              writer.writeRowEnd();
+            }
+
+            writer.writeResponseEnd();
+          }
+          catch (Exception e) {
+            log.error(e, "Unable to stream results back for query[%s]", queryId);
+            throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+          }
+        }).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;
+    }
+  }
+
+  /**
+   * Queries can be canceled while in anystate. Canceling a query that has already completed will remove its results.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@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(),
+          false
+      );
+      if (sqlStatementResult.isPresent()) {
+        switch (sqlStatementResult.get().getState()) {
+          case ACCEPTED:
+          case RUNNING:
+            overlordClient.cancelTask(queryId);
+            return Response.status(Response.Status.ACCEPTED).build();
+          case SUCCESS:
+          case FAILED:
+            // we would also want to clean up the
+            return Response.ok().build();
+          default:
+            throw new ISE("Illegal State[%s] encountered", sqlStatementResult.get().getState());
+        }
+
+      } 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;
+    }
+  }
+
+  private Response buildStandardResponse(
+      Sequence<Object[]> sequence,
+      SqlQuery sqlQuery,
+      String sqlQueryId,
+      SqlRowTransformer rowTransformer
+  ) throws IOException
+  {
+    final Yielder<Object[]> yielder0 = Yielders.each(sequence);
+
+    try {
+      final Response.ResponseBuilder responseBuilder = Response.ok((StreamingOutput) outputStream -> {
+        CountingOutputStream os = new CountingOutputStream(outputStream);
+        Yielder<Object[]> yielder = yielder0;
+
+        try (final ResultFormat.Writer writer = sqlQuery.getResultFormat().createFormatter(os, jsonMapper)) {
+          writer.writeResponseStart();
+
+          if (sqlQuery.includeHeader()) {
+            writer.writeHeader(
+                rowTransformer.getRowType(),
+                sqlQuery.includeTypesHeader(),
+                sqlQuery.includeSqlTypesHeader()
+            );
+          }
+
+          while (!yielder.isDone()) {
+            final Object[] row = yielder.get();
+            writer.writeRowStart();
+            for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
+              final Object value = rowTransformer.transform(row, i);
+              writer.writeRowField(rowTransformer.getFieldList().get(i), value);
+            }
+            writer.writeRowEnd();
+            yielder = yielder.next(null);
+          }
+
+          writer.writeResponseEnd();
+        }
+        catch (Exception e) {
+          log.error(e, "Unable to send SQL response [%s]", sqlQueryId);
+          throw new RuntimeException(e);
+        }
+        finally {
+          yielder.close();
+        }
+      });
+
+      if (sqlQuery.includeHeader()) {
+        responseBuilder.header(SqlResource.SQL_HEADER_RESPONSE_HEADER, SqlResource.SQL_HEADER_VALUE);
+      }
+
+      return responseBuilder.build();
+    }
+    catch (Throwable e) {
+      // make sure to close yielder if anything happened before starting to serialize the response.
+      yielder0.close();
+      throw e;
+    }
+  }
+
+  private Response buildTaskResponse(Sequence<Object[]> sequence, String user)
+  {
+    List<Object[]> rows = sequence.toList();
+    int numRows = rows.size();
+    if (numRows != 1) {
+      throw new RE("Expected a single row but got [%d] rows. Please check broker logs for more information.", numRows);
+    }
+    String taskId = (String) rows.get(0)[0];
+    try {
+      Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, user, true);
+
+      if (statementResult.isPresent()) {
+        return Response.status(Response.Status.OK).entity(statementResult.get()).build();
+      } else {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(
+                null,
+                StringUtils.format(
+                    "Unable to find associated task for query id %s",
+                    taskId
+                ),
+                null,
+                null,
+                null
+            ),
+            taskId
+        );
+      }
+    }
+    catch (QueryException e) {
+      return buildNonOkResponse(Response.Status.OK.getStatusCode(), e, taskId);
+    }
+  }
+
+  private Response buildNonOkResponse(int status, SanitizableException e, String sqlQueryId)
+  {
+    // Though transformIfNeeded returns an exception, its purpose is to return
+    // a QueryException
+    Exception cleaned = serverConfig.getErrorResponseTransformStrategy().transformIfNeeded(e);
+    return Response.status(status).entity(cleaned).build();
+  }
+
+
+  private static Optional<List<ColNameAndType>> getSignature(
+      MSQControllerTask msqControllerTask
+  )
+  {
+    // only populate signature for select q's
+    if (msqControllerTask.getQuerySpec().getDestination().getClass() == TaskReportMSQDestination.class) {
+      ColumnMappings columnMappings = msqControllerTask.getQuerySpec().getColumnMappings();
+      List<SqlTypeName> sqlTypeNames = msqControllerTask.getSqlTypeNames();
+      if (sqlTypeNames == null || sqlTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColumnType> nativeTypeNames = msqControllerTask.getNativeTypeNames();
+      if (nativeTypeNames == null || nativeTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColNameAndType> signature = new ArrayList<>(columnMappings.size());
+      int index = 0;
+      for (String colName : columnMappings.getOutputColumnNames()) {
+        signature.add(new ColNameAndType(
+            colName,
+            sqlTypeNames.get(index).getName(),
+            nativeTypeNames.get(index).asTypeString()
+        ));
+        index++;
+      }
+      return Optional.of(signature);
+    }
+    return Optional.empty();
+  }
+
+
+  private void checkTaskPayloadOrThrow(TaskPayloadResponse taskPayloadResponse, String queryId) throws QueryException
+  {
+    if (taskPayloadResponse == null || taskPayloadResponse.getPayload() == null) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Could not get payload details for query[%s] from the overlord",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+
+    if (MSQControllerTask.class != taskPayloadResponse.getPayload().getClass()) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Fetched an unexpected payload for query[%s] from the overlord.",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+  }
+
+  private Optional<ResultSetInformation> getSampleResults(
+      String asyncResultId,
+      boolean isSelectQuery,
+      String dataSource,
+      SqlStatementState sqlStatementState
+  )
+  {
+    if (sqlStatementState == SqlStatementState.SUCCESS) {
+      Map<String, Object> payload = getPayload(overlordWork(overlordClient.taskReportAsMap(asyncResultId)));
+      Optional<Pair<Long, Long>> rowsAndSize = getRowsAndSizeFromPayload(payload, isSelectQuery);
+      return Optional.of(new ResultSetInformation(
+          null,
+          // since the rows can be sampled, get the number of rows from counters
+          rowsAndSize.orElse(new Pair<>(null, null)).lhs,
+          rowsAndSize.orElse(new Pair<>(null, null)).rhs,
+          dataSource,
+          // only populate sample results in case a select query is successful
+          isSelectQuery ? getResults(payload).orElse(null) : null
+      ));
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  private Optional<Pair<Long, Long>> getRowsAndSizeFromPayload(Map<String, Object> payload, boolean isSelectQuery)
+  {
+    List stages = getList(payload, "stages");
+    if (stages == null || stages.isEmpty()) {
+      return Optional.empty();
+    } else {
+      int maxStage = stages.size() - 1; // Last stage output is the total number of rows returned to the end user.
+      Map<String, Object> counterMap = getMap(getMap(payload, "counters"), String.valueOf(maxStage));
+      Long rows = -1L;
+      Long sizeInBytes = -1L;
+      if (counterMap == null) {
+        return Optional.empty();
+      }
+      for (Map.Entry<String, Object> worker : counterMap.entrySet()) {
+        Object workerChannels = worker.getValue();
+        if (workerChannels == null || !(workerChannels instanceof Map)) {
+          return Optional.empty();

Review Comment:
   No. Just defensive coding. 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1244660295


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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);

Review Comment:
   Yes We needed to. Updated the checks. 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1244664179


##########
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))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColNameAndType> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {
+              writer.writeRowStart();
+              for (int i = 0; i < rowSignature.size(); i++) {
+                writer.writeRowField(rowSignature.get(i).getColName(), ((List) results.get().get(k)).get(i));
+              }
+              writer.writeRowEnd();
+            }
+
+            writer.writeResponseEnd();
+          }
+          catch (Exception e) {
+            log.error(e, "Unable to stream results back for query[%s]", queryId);
+            throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+          }
+        }).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;
+    }
+  }
+
+  /**
+   * Queries can be canceled while in anystate. Canceling a query that has already completed will remove its results.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@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(),
+          false
+      );
+      if (sqlStatementResult.isPresent()) {
+        switch (sqlStatementResult.get().getState()) {
+          case ACCEPTED:
+          case RUNNING:
+            overlordClient.cancelTask(queryId);
+            return Response.status(Response.Status.ACCEPTED).build();
+          case SUCCESS:
+          case FAILED:
+            // we would also want to clean up the
+            return Response.ok().build();
+          default:
+            throw new ISE("Illegal State[%s] encountered", sqlStatementResult.get().getState());
+        }
+
+      } 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;
+    }
+  }
+
+  private Response buildStandardResponse(
+      Sequence<Object[]> sequence,
+      SqlQuery sqlQuery,
+      String sqlQueryId,
+      SqlRowTransformer rowTransformer
+  ) throws IOException
+  {
+    final Yielder<Object[]> yielder0 = Yielders.each(sequence);
+
+    try {
+      final Response.ResponseBuilder responseBuilder = Response.ok((StreamingOutput) outputStream -> {
+        CountingOutputStream os = new CountingOutputStream(outputStream);
+        Yielder<Object[]> yielder = yielder0;
+
+        try (final ResultFormat.Writer writer = sqlQuery.getResultFormat().createFormatter(os, jsonMapper)) {
+          writer.writeResponseStart();
+
+          if (sqlQuery.includeHeader()) {
+            writer.writeHeader(
+                rowTransformer.getRowType(),
+                sqlQuery.includeTypesHeader(),
+                sqlQuery.includeSqlTypesHeader()
+            );
+          }
+
+          while (!yielder.isDone()) {
+            final Object[] row = yielder.get();
+            writer.writeRowStart();
+            for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
+              final Object value = rowTransformer.transform(row, i);
+              writer.writeRowField(rowTransformer.getFieldList().get(i), value);
+            }
+            writer.writeRowEnd();
+            yielder = yielder.next(null);
+          }
+
+          writer.writeResponseEnd();
+        }
+        catch (Exception e) {
+          log.error(e, "Unable to send SQL response [%s]", sqlQueryId);
+          throw new RuntimeException(e);
+        }
+        finally {
+          yielder.close();
+        }
+      });
+
+      if (sqlQuery.includeHeader()) {
+        responseBuilder.header(SqlResource.SQL_HEADER_RESPONSE_HEADER, SqlResource.SQL_HEADER_VALUE);
+      }
+
+      return responseBuilder.build();
+    }
+    catch (Throwable e) {
+      // make sure to close yielder if anything happened before starting to serialize the response.
+      yielder0.close();
+      throw e;
+    }
+  }
+
+  private Response buildTaskResponse(Sequence<Object[]> sequence, String user)
+  {
+    List<Object[]> rows = sequence.toList();
+    int numRows = rows.size();
+    if (numRows != 1) {
+      throw new RE("Expected a single row but got [%d] rows. Please check broker logs for more information.", numRows);
+    }
+    String taskId = (String) rows.get(0)[0];
+    try {
+      Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, user, true);
+
+      if (statementResult.isPresent()) {
+        return Response.status(Response.Status.OK).entity(statementResult.get()).build();
+      } else {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(
+                null,
+                StringUtils.format(
+                    "Unable to find associated task for query id %s",
+                    taskId
+                ),
+                null,
+                null,
+                null
+            ),
+            taskId
+        );
+      }
+    }
+    catch (QueryException e) {
+      return buildNonOkResponse(Response.Status.OK.getStatusCode(), e, taskId);
+    }
+  }
+
+  private Response buildNonOkResponse(int status, SanitizableException e, String sqlQueryId)
+  {
+    // Though transformIfNeeded returns an exception, its purpose is to return
+    // a QueryException
+    Exception cleaned = serverConfig.getErrorResponseTransformStrategy().transformIfNeeded(e);
+    return Response.status(status).entity(cleaned).build();
+  }
+
+
+  private static Optional<List<ColNameAndType>> getSignature(
+      MSQControllerTask msqControllerTask
+  )
+  {
+    // only populate signature for select q's
+    if (msqControllerTask.getQuerySpec().getDestination().getClass() == TaskReportMSQDestination.class) {
+      ColumnMappings columnMappings = msqControllerTask.getQuerySpec().getColumnMappings();
+      List<SqlTypeName> sqlTypeNames = msqControllerTask.getSqlTypeNames();
+      if (sqlTypeNames == null || sqlTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColumnType> nativeTypeNames = msqControllerTask.getNativeTypeNames();
+      if (nativeTypeNames == null || nativeTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColNameAndType> signature = new ArrayList<>(columnMappings.size());
+      int index = 0;
+      for (String colName : columnMappings.getOutputColumnNames()) {
+        signature.add(new ColNameAndType(
+            colName,
+            sqlTypeNames.get(index).getName(),
+            nativeTypeNames.get(index).asTypeString()
+        ));
+        index++;
+      }
+      return Optional.of(signature);
+    }
+    return Optional.empty();
+  }
+
+
+  private void checkTaskPayloadOrThrow(TaskPayloadResponse taskPayloadResponse, String queryId) throws QueryException
+  {
+    if (taskPayloadResponse == null || taskPayloadResponse.getPayload() == null) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Could not get payload details for query[%s] from the overlord",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+
+    if (MSQControllerTask.class != taskPayloadResponse.getPayload().getClass()) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Fetched an unexpected payload for query[%s] from the overlord.",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+  }
+
+  private Optional<ResultSetInformation> getSampleResults(
+      String asyncResultId,
+      boolean isSelectQuery,
+      String dataSource,
+      SqlStatementState sqlStatementState
+  )
+  {
+    if (sqlStatementState == SqlStatementState.SUCCESS) {
+      Map<String, Object> payload = getPayload(overlordWork(overlordClient.taskReportAsMap(asyncResultId)));
+      Optional<Pair<Long, Long>> rowsAndSize = getRowsAndSizeFromPayload(payload, isSelectQuery);
+      return Optional.of(new ResultSetInformation(
+          null,
+          // since the rows can be sampled, get the number of rows from counters
+          rowsAndSize.orElse(new Pair<>(null, null)).lhs,
+          rowsAndSize.orElse(new Pair<>(null, null)).rhs,
+          dataSource,
+          // only populate sample results in case a select query is successful
+          isSelectQuery ? getResults(payload).orElse(null) : null
+      ));
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  private Optional<Pair<Long, Long>> getRowsAndSizeFromPayload(Map<String, Object> payload, boolean isSelectQuery)
+  {
+    List stages = getList(payload, "stages");
+    if (stages == null || stages.isEmpty()) {
+      return Optional.empty();
+    } else {
+      int maxStage = stages.size() - 1; // Last stage output is the total number of rows returned to the end user.
+      Map<String, Object> counterMap = getMap(getMap(payload, "counters"), String.valueOf(maxStage));
+      Long rows = -1L;
+      Long sizeInBytes = -1L;
+      if (counterMap == null) {
+        return Optional.empty();
+      }
+      for (Map.Entry<String, Object> worker : counterMap.entrySet()) {
+        Object workerChannels = worker.getValue();
+        if (workerChannels == null || !(workerChannels instanceof Map)) {
+          return Optional.empty();
+        }
+        if (isSelectQuery) {
+          Object output = ((Map<?, ?>) workerChannels).get("output");
+          if (output != null && output instanceof Map) {
+            List<Integer> rowsPerChannel = (List<Integer>) ((Map<String, Object>) output).get("rows");
+            List<Integer> bytesPerChannel = (List<Integer>) ((Map<String, Object>) output).get("bytes");
+            for (Integer row : rowsPerChannel) {
+              rows = rows + row;
+            }
+            for (Integer bytes : bytesPerChannel) {
+              sizeInBytes = sizeInBytes + bytes;
+            }
+          }
+        } else {
+          Object output = ((Map<?, ?>) workerChannels).get("segmentGenerationProgress");
+          if (output != null && output instanceof Map) {
+            rows += (Integer) ((Map<String, Object>) output).get("rowsPushed");
+          }
+        }
+      }
+
+      return Optional.of(new Pair<>(rows == -1L ? null : rows + 1, sizeInBytes == -1L ? null : sizeInBytes + 1));
+    }
+  }
+
+
+  private Optional<SqlStatementResult> getStatementStatus(String queryId, String currentUser, boolean withResults)
+      throws QueryException, ForbiddenException
+  {
+    TaskStatusResponse taskResponse = overlordWork(overlordClient.taskStatus(queryId));
+    if (taskResponse == null) {
+      return Optional.empty();
+    }
+
+    TaskStatusPlus statusPlus = taskResponse.getStatus();
+    if (statusPlus == null || !MSQControllerTask.TYPE.equals(statusPlus.getType())) {
+      return Optional.empty();
+    }
+
+    // since we need the controller payload for auth checks.
+    MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, currentUser);
+    SqlStatementState sqlStatementState = getSqlStatementState(statusPlus);
+
+    if (SqlStatementState.FAILED == sqlStatementState) {
+      Map<String, Object> exceptionDetails = getQueryExceptionDetails(getPayload(overlordWork(overlordClient.taskReportAsMap(
+          queryId))));
+      Map<String, Object> exception = getMap(exceptionDetails, "error");
+      if (exceptionDetails == null || exception == null) {
+        return Optional.of(new SqlStatementResult(
+            queryId,
+            sqlStatementState,
+            taskResponse.getStatus().getCreatedTime(),
+            null,
+            taskResponse.getStatus().getDuration(),
+            null,
+            new QueryException(null, statusPlus.getErrorMsg(), null, null, null)
+        ));
+      }
+
+      final String errorMessage = String.valueOf(exception.getOrDefault("errorMessage", statusPlus.getErrorMsg()));
+      exception.remove("errorMessage");
+      String host = (String) exceptionDetails.getOrDefault("host", null);
+      Map<String, String> stringException = new HashMap<>();
+      for (Map.Entry<String, Object> exceptionKeys : exception.entrySet()) {
+        stringException.put(exceptionKeys.getKey(), String.valueOf(exceptionKeys.getValue()));
+      }
+      return Optional.of(new SqlStatementResult(
+          queryId,
+          sqlStatementState,
+          taskResponse.getStatus().getCreatedTime(),
+          null,
+          taskResponse.getStatus().getDuration(),
+          null,
+          new QueryException(null, errorMessage, null, host, stringException)
+      ));
+    } else {
+      Optional<List<ColNameAndType>> signature = getSignature(msqControllerTask);
+
+      return Optional.of(new SqlStatementResult(
+          queryId,
+          sqlStatementState,
+          taskResponse.getStatus().getCreatedTime(),
+          signature.orElse(null),
+          taskResponse.getStatus().getDuration(),
+          withResults ? getSampleResults(
+              queryId,
+              signature.isPresent(),
+              msqControllerTask.getDataSource(),
+              sqlStatementState
+          ).orElse(null) : null,
+          null
+      ));
+    }
+  }
+
+  private Map<String, Object> getQueryExceptionDetails(Map<String, Object> payload)
+  {
+    return getMap(getMap(payload, "status"), "errorReport");
+  }
+
+  private MSQControllerTask getMSQControllerTaskOrThrow(String queryId, String currentUser) throws ForbiddenException
+  {
+    TaskPayloadResponse taskPayloadResponse = overlordWork(overlordClient.taskPayload(queryId));
+    checkTaskPayloadOrThrow(taskPayloadResponse, queryId);
+
+    MSQControllerTask msqControllerTask = (MSQControllerTask) taskPayloadResponse.getPayload();
+    if (currentUser == null || !currentUser.equals(msqControllerTask.getQuerySpec()
+                                                                    .getQuery()
+                                                                    .getContext()
+                                                                    .get(MSQTaskQueryMaker.USER_KEY))) {
+      throw new ForbiddenException();
+    }
+    return msqControllerTask;
+  }
+
+  private <T> T overlordWork(final ListenableFuture<T> future)
+  {
+    try {
+      return FutureUtils.getUnchecked(future, true);
+    }
+    catch (RuntimeException e) {
+      throw new QueryException(null, "Unable to contact overlord " + e.getMessage(), null, null, null);

Review Comment:
   This can happen for n reasons. Changed it to a developer exception. 
   ``` 
   DruidException.forPersona(DruidException.Persona.DEVELOPER)
                             .ofCategory(DruidException.Category.UNCATEGORIZED)
                             .build("Unable to contact overlord " + e.getMessage());
       }
   ```



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1245112124


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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);

Review Comment:
   This is an outdated diff. In the current patch, results are checked. 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] github-code-scanning[bot] commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["github-code-scanning[bot] (via GitHub)" <gi...@apache.org>]

github-code-scanning[bot] commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1237925844


##########
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))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColNameAndType> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {

Review Comment:
   ## User-controlled data in numeric cast
   
   This cast to a narrower type depends on a [user-provided value](1), potentially causing truncation.
   
   [Show more details](https://github.com/apache/druid/security/code-scanning/5154)



##########
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))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColNameAndType> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {
+              writer.writeRowStart();
+              for (int i = 0; i < rowSignature.size(); i++) {
+                writer.writeRowField(rowSignature.get(i).getColName(), ((List) results.get().get(k)).get(i));
+              }
+              writer.writeRowEnd();
+            }
+
+            writer.writeResponseEnd();
+          }
+          catch (Exception e) {
+            log.error(e, "Unable to stream results back for query[%s]", queryId);
+            throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+          }
+        }).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;
+    }
+  }
+
+  /**
+   * Queries can be canceled while in anystate. Canceling a query that has already completed will remove its results.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@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(),
+          false
+      );
+      if (sqlStatementResult.isPresent()) {
+        switch (sqlStatementResult.get().getState()) {
+          case ACCEPTED:
+          case RUNNING:
+            overlordClient.cancelTask(queryId);
+            return Response.status(Response.Status.ACCEPTED).build();
+          case SUCCESS:
+          case FAILED:
+            // we would also want to clean up the
+            return Response.ok().build();
+          default:
+            throw new ISE("Illegal State[%s] encountered", sqlStatementResult.get().getState());
+        }
+
+      } 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;
+    }
+  }
+
+  private Response buildStandardResponse(
+      Sequence<Object[]> sequence,
+      SqlQuery sqlQuery,
+      String sqlQueryId,
+      SqlRowTransformer rowTransformer
+  ) throws IOException
+  {
+    final Yielder<Object[]> yielder0 = Yielders.each(sequence);
+
+    try {
+      final Response.ResponseBuilder responseBuilder = Response.ok((StreamingOutput) outputStream -> {
+        CountingOutputStream os = new CountingOutputStream(outputStream);
+        Yielder<Object[]> yielder = yielder0;
+
+        try (final ResultFormat.Writer writer = sqlQuery.getResultFormat().createFormatter(os, jsonMapper)) {
+          writer.writeResponseStart();
+
+          if (sqlQuery.includeHeader()) {
+            writer.writeHeader(
+                rowTransformer.getRowType(),
+                sqlQuery.includeTypesHeader(),
+                sqlQuery.includeSqlTypesHeader()
+            );
+          }
+
+          while (!yielder.isDone()) {
+            final Object[] row = yielder.get();
+            writer.writeRowStart();
+            for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
+              final Object value = rowTransformer.transform(row, i);
+              writer.writeRowField(rowTransformer.getFieldList().get(i), value);
+            }
+            writer.writeRowEnd();
+            yielder = yielder.next(null);
+          }
+
+          writer.writeResponseEnd();
+        }
+        catch (Exception e) {
+          log.error(e, "Unable to send SQL response [%s]", sqlQueryId);
+          throw new RuntimeException(e);
+        }
+        finally {
+          yielder.close();
+        }
+      });
+
+      if (sqlQuery.includeHeader()) {
+        responseBuilder.header(SqlResource.SQL_HEADER_RESPONSE_HEADER, SqlResource.SQL_HEADER_VALUE);
+      }
+
+      return responseBuilder.build();
+    }
+    catch (Throwable e) {
+      // make sure to close yielder if anything happened before starting to serialize the response.
+      yielder0.close();
+      throw e;
+    }
+  }
+
+  private Response buildTaskResponse(Sequence<Object[]> sequence, String user)
+  {
+    List<Object[]> rows = sequence.toList();
+    int numRows = rows.size();
+    if (numRows != 1) {
+      throw new RE("Expected a single row but got [%d] rows. Please check broker logs for more information.", numRows);
+    }
+    String taskId = (String) rows.get(0)[0];
+    try {
+      Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, user, true);
+
+      if (statementResult.isPresent()) {
+        return Response.status(Response.Status.OK).entity(statementResult.get()).build();
+      } else {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(
+                null,
+                StringUtils.format(
+                    "Unable to find associated task for query id %s",
+                    taskId
+                ),
+                null,
+                null,
+                null
+            ),
+            taskId
+        );
+      }
+    }
+    catch (QueryException e) {
+      return buildNonOkResponse(Response.Status.OK.getStatusCode(), e, taskId);
+    }
+  }
+
+  private Response buildNonOkResponse(int status, SanitizableException e, String sqlQueryId)
+  {
+    // Though transformIfNeeded returns an exception, its purpose is to return
+    // a QueryException
+    Exception cleaned = serverConfig.getErrorResponseTransformStrategy().transformIfNeeded(e);
+    return Response.status(status).entity(cleaned).build();
+  }
+
+
+  private static Optional<List<ColNameAndType>> getSignature(
+      MSQControllerTask msqControllerTask
+  )
+  {
+    // only populate signature for select q's
+    if (msqControllerTask.getQuerySpec().getDestination().getClass() == TaskReportMSQDestination.class) {
+      ColumnMappings columnMappings = msqControllerTask.getQuerySpec().getColumnMappings();
+      List<SqlTypeName> sqlTypeNames = msqControllerTask.getSqlTypeNames();
+      if (sqlTypeNames == null || sqlTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColumnType> nativeTypeNames = msqControllerTask.getNativeTypeNames();
+      if (nativeTypeNames == null || nativeTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColNameAndType> signature = new ArrayList<>(columnMappings.size());
+      int index = 0;
+      for (String colName : columnMappings.getOutputColumnNames()) {
+        signature.add(new ColNameAndType(
+            colName,
+            sqlTypeNames.get(index).getName(),
+            nativeTypeNames.get(index).asTypeString()
+        ));
+        index++;
+      }
+      return Optional.of(signature);
+    }
+    return Optional.empty();
+  }
+
+
+  private void checkTaskPayloadOrThrow(TaskPayloadResponse taskPayloadResponse, String queryId) throws QueryException
+  {
+    if (taskPayloadResponse == null || taskPayloadResponse.getPayload() == null) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Could not get payload details for query[%s] from the overlord",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+
+    if (MSQControllerTask.class != taskPayloadResponse.getPayload().getClass()) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Fetched an unexpected payload for query[%s] from the overlord.",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+  }
+
+  private Optional<ResultSetInformation> getSampleResults(
+      String asyncResultId,
+      boolean isSelectQuery,
+      String dataSource,
+      SqlStatementState sqlStatementState
+  )
+  {
+    if (sqlStatementState == SqlStatementState.SUCCESS) {
+      Map<String, Object> payload = getPayload(overlordWork(overlordClient.taskReportAsMap(asyncResultId)));
+      Optional<Pair<Long, Long>> rowsAndSize = getRowsAndSizeFromPayload(payload, isSelectQuery);
+      return Optional.of(new ResultSetInformation(
+          null,
+          // since the rows can be sampled, get the number of rows from counters
+          rowsAndSize.orElse(new Pair<>(null, null)).lhs,
+          rowsAndSize.orElse(new Pair<>(null, null)).rhs,
+          dataSource,
+          // only populate sample results in case a select query is successful
+          isSelectQuery ? getResults(payload).orElse(null) : null
+      ));
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  private Optional<Pair<Long, Long>> getRowsAndSizeFromPayload(Map<String, Object> payload, boolean isSelectQuery)
+  {
+    List stages = getList(payload, "stages");
+    if (stages == null || stages.isEmpty()) {
+      return Optional.empty();
+    } else {
+      int maxStage = stages.size() - 1; // Last stage output is the total number of rows returned to the end user.
+      Map<String, Object> counterMap = getMap(getMap(payload, "counters"), String.valueOf(maxStage));
+      Long rows = -1L;

Review Comment:
   ## Boxed variable is never null
   
   The variable 'rows' is only assigned values of primitive type and is never 'null', but it is declared with the boxed type 'Long'.
   
   [Show more details](https://github.com/apache/druid/security/code-scanning/5152)



##########
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))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColNameAndType> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {
+              writer.writeRowStart();
+              for (int i = 0; i < rowSignature.size(); i++) {
+                writer.writeRowField(rowSignature.get(i).getColName(), ((List) results.get().get(k)).get(i));
+              }
+              writer.writeRowEnd();
+            }
+
+            writer.writeResponseEnd();
+          }
+          catch (Exception e) {
+            log.error(e, "Unable to stream results back for query[%s]", queryId);
+            throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+          }
+        }).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;
+    }
+  }
+
+  /**
+   * Queries can be canceled while in anystate. Canceling a query that has already completed will remove its results.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@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(),
+          false
+      );
+      if (sqlStatementResult.isPresent()) {
+        switch (sqlStatementResult.get().getState()) {
+          case ACCEPTED:
+          case RUNNING:
+            overlordClient.cancelTask(queryId);
+            return Response.status(Response.Status.ACCEPTED).build();
+          case SUCCESS:
+          case FAILED:
+            // we would also want to clean up the
+            return Response.ok().build();
+          default:
+            throw new ISE("Illegal State[%s] encountered", sqlStatementResult.get().getState());
+        }
+
+      } 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;
+    }
+  }
+
+  private Response buildStandardResponse(
+      Sequence<Object[]> sequence,
+      SqlQuery sqlQuery,
+      String sqlQueryId,
+      SqlRowTransformer rowTransformer
+  ) throws IOException
+  {
+    final Yielder<Object[]> yielder0 = Yielders.each(sequence);
+
+    try {
+      final Response.ResponseBuilder responseBuilder = Response.ok((StreamingOutput) outputStream -> {
+        CountingOutputStream os = new CountingOutputStream(outputStream);
+        Yielder<Object[]> yielder = yielder0;
+
+        try (final ResultFormat.Writer writer = sqlQuery.getResultFormat().createFormatter(os, jsonMapper)) {
+          writer.writeResponseStart();
+
+          if (sqlQuery.includeHeader()) {
+            writer.writeHeader(
+                rowTransformer.getRowType(),
+                sqlQuery.includeTypesHeader(),
+                sqlQuery.includeSqlTypesHeader()
+            );
+          }
+
+          while (!yielder.isDone()) {
+            final Object[] row = yielder.get();
+            writer.writeRowStart();
+            for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
+              final Object value = rowTransformer.transform(row, i);
+              writer.writeRowField(rowTransformer.getFieldList().get(i), value);
+            }
+            writer.writeRowEnd();
+            yielder = yielder.next(null);
+          }
+
+          writer.writeResponseEnd();
+        }
+        catch (Exception e) {
+          log.error(e, "Unable to send SQL response [%s]", sqlQueryId);
+          throw new RuntimeException(e);
+        }
+        finally {
+          yielder.close();
+        }
+      });
+
+      if (sqlQuery.includeHeader()) {
+        responseBuilder.header(SqlResource.SQL_HEADER_RESPONSE_HEADER, SqlResource.SQL_HEADER_VALUE);
+      }
+
+      return responseBuilder.build();
+    }
+    catch (Throwable e) {
+      // make sure to close yielder if anything happened before starting to serialize the response.
+      yielder0.close();
+      throw e;
+    }
+  }
+
+  private Response buildTaskResponse(Sequence<Object[]> sequence, String user)
+  {
+    List<Object[]> rows = sequence.toList();
+    int numRows = rows.size();
+    if (numRows != 1) {
+      throw new RE("Expected a single row but got [%d] rows. Please check broker logs for more information.", numRows);
+    }
+    String taskId = (String) rows.get(0)[0];
+    try {
+      Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, user, true);
+
+      if (statementResult.isPresent()) {
+        return Response.status(Response.Status.OK).entity(statementResult.get()).build();
+      } else {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(
+                null,
+                StringUtils.format(
+                    "Unable to find associated task for query id %s",
+                    taskId
+                ),
+                null,
+                null,
+                null
+            ),
+            taskId
+        );
+      }
+    }
+    catch (QueryException e) {
+      return buildNonOkResponse(Response.Status.OK.getStatusCode(), e, taskId);
+    }
+  }
+
+  private Response buildNonOkResponse(int status, SanitizableException e, String sqlQueryId)
+  {
+    // Though transformIfNeeded returns an exception, its purpose is to return
+    // a QueryException
+    Exception cleaned = serverConfig.getErrorResponseTransformStrategy().transformIfNeeded(e);
+    return Response.status(status).entity(cleaned).build();
+  }
+
+
+  private static Optional<List<ColNameAndType>> getSignature(
+      MSQControllerTask msqControllerTask
+  )
+  {
+    // only populate signature for select q's
+    if (msqControllerTask.getQuerySpec().getDestination().getClass() == TaskReportMSQDestination.class) {
+      ColumnMappings columnMappings = msqControllerTask.getQuerySpec().getColumnMappings();
+      List<SqlTypeName> sqlTypeNames = msqControllerTask.getSqlTypeNames();
+      if (sqlTypeNames == null || sqlTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColumnType> nativeTypeNames = msqControllerTask.getNativeTypeNames();
+      if (nativeTypeNames == null || nativeTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColNameAndType> signature = new ArrayList<>(columnMappings.size());
+      int index = 0;
+      for (String colName : columnMappings.getOutputColumnNames()) {
+        signature.add(new ColNameAndType(
+            colName,
+            sqlTypeNames.get(index).getName(),
+            nativeTypeNames.get(index).asTypeString()
+        ));
+        index++;
+      }
+      return Optional.of(signature);
+    }
+    return Optional.empty();
+  }
+
+
+  private void checkTaskPayloadOrThrow(TaskPayloadResponse taskPayloadResponse, String queryId) throws QueryException
+  {
+    if (taskPayloadResponse == null || taskPayloadResponse.getPayload() == null) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Could not get payload details for query[%s] from the overlord",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+
+    if (MSQControllerTask.class != taskPayloadResponse.getPayload().getClass()) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Fetched an unexpected payload for query[%s] from the overlord.",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+  }
+
+  private Optional<ResultSetInformation> getSampleResults(
+      String asyncResultId,
+      boolean isSelectQuery,
+      String dataSource,
+      SqlStatementState sqlStatementState
+  )
+  {
+    if (sqlStatementState == SqlStatementState.SUCCESS) {
+      Map<String, Object> payload = getPayload(overlordWork(overlordClient.taskReportAsMap(asyncResultId)));
+      Optional<Pair<Long, Long>> rowsAndSize = getRowsAndSizeFromPayload(payload, isSelectQuery);
+      return Optional.of(new ResultSetInformation(
+          null,
+          // since the rows can be sampled, get the number of rows from counters
+          rowsAndSize.orElse(new Pair<>(null, null)).lhs,
+          rowsAndSize.orElse(new Pair<>(null, null)).rhs,
+          dataSource,
+          // only populate sample results in case a select query is successful
+          isSelectQuery ? getResults(payload).orElse(null) : null
+      ));
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  private Optional<Pair<Long, Long>> getRowsAndSizeFromPayload(Map<String, Object> payload, boolean isSelectQuery)
+  {
+    List stages = getList(payload, "stages");
+    if (stages == null || stages.isEmpty()) {
+      return Optional.empty();
+    } else {
+      int maxStage = stages.size() - 1; // Last stage output is the total number of rows returned to the end user.
+      Map<String, Object> counterMap = getMap(getMap(payload, "counters"), String.valueOf(maxStage));
+      Long rows = -1L;
+      Long sizeInBytes = -1L;

Review Comment:
   ## Boxed variable is never null
   
   The variable 'sizeInBytes' is only assigned values of primitive type and is never 'null', but it is declared with the boxed type 'Long'.
   
   [Show more details](https://github.com/apache/druid/security/code-scanning/5153)



##########
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))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColNameAndType> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {

Review Comment:
   ## Comparison of narrow type with wide type in loop condition
   
   Comparison between [expression](1) of type int and [expression](2) of wider type long.
   
   [Show more details](https://github.com/apache/druid/security/code-scanning/5155)



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1244892536


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,673 @@
+/*
+ * 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.druid.client.indexing.TaskPayloadResponse;
+import org.apache.druid.client.indexing.TaskStatusResponse;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.error.ErrorResponse;
+import org.apache.druid.error.QueryExceptionCompat;
+import org.apache.druid.guice.annotations.MSQ;
+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.sql.MSQTaskQueryMaker;
+import org.apache.druid.msq.sql.MSQTaskSqlEngine;
+import org.apache.druid.msq.sql.SqlStatementState;
+import org.apache.druid.msq.sql.entity.ColumnNameAndTypes;
+import org.apache.druid.msq.sql.entity.ResultSetInformation;
+import org.apache.druid.msq.sql.entity.SqlStatementResult;
+import org.apache.druid.msq.util.SqlStatementResourceHelper;
+import org.apache.druid.query.ExecutionMode;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.query.QueryException;
+import org.apache.druid.rpc.indexing.OverlordClient;
+import org.apache.druid.server.QueryResponse;
+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.SqlRowTransformer;
+import org.apache.druid.sql.SqlStatementFactory;
+import org.apache.druid.sql.http.ResultFormat;
+import org.apache.druid.sql.http.SqlQuery;
+import org.apache.druid.sql.http.SqlResource;
+
+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.Collections;
+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 AuthorizerMapper authorizerMapper;
+  private final ObjectMapper jsonMapper;
+  private final OverlordClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final OverlordClient overlordClient
+  )
+  {
+    this.msqSqlStatementFactory = msqSqlStatementFactory;
+    this.authorizerMapper = authorizerMapper;
+    this.jsonMapper = jsonMapper;
+    this.overlordClient = overlordClient;
+  }
+
+
+  @GET
+  @Path("/enabled")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response isEnabled(@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(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              StringUtils.format(
+                                  "The statement sql api only supports sync mode[%s]. Please set context parameter [%s=%s] in the context payload",
+                                  ExecutionMode.ASYNC,
+                                  QueryContexts.CTX_EXECUTION_MODE,
+                                  ExecutionMode.ASYNC
+                              )
+                          )
+        );
+      }
+
+
+      Thread.currentThread().setName(StringUtils.format("statement_sql[%s]", sqlQueryId));
+
+      final DirectStatement.ResultSet plan = stmt.plan();
+      // in case the engine is async, the query is not run yet. We just return the taskID in case of non explain queries.
+      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);
+      }
+    }
+    catch (DruidException e) {
+      stmt.reporter().failed(e);
+      return buildNonOkResponse(e);
+    }
+    catch (QueryException queryException) {
+      stmt.reporter().failed(queryException);
+      final DruidException underlyingException = DruidException.fromFailure(new QueryExceptionCompat(queryException));
+      return buildNonOkResponse(underlyingException);
+    }
+    catch (ForbiddenException e) {
+      log.debug("Got forbidden request for reason [%s]", e.getErrorMessage());
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.USER)
+                        .ofCategory(DruidException.Category.FORBIDDEN)
+                        .build(Access.DEFAULT_ERROR_MESSAGE)
+      );
+    }
+    // Calcite throws java.lang.AssertionError at various points in planning/validation.
+    catch (AssertionError | Exception e) {
+      stmt.reporter().failed(e);
+      log.noStackTrace().warn(e, "Failed to handle query: %s", sqlQueryId);
+
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.DEVELOPER)
+                        .ofCategory(DruidException.Category.UNCATEGORIZED)
+                        .build(e.getMessage())
+      );
+    }
+    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 {
+      Access authResult = AuthorizationUtils.authorizeAllResourceActions(
+          req,
+          Collections.emptyList(),
+          authorizerMapper
+      );
+      if (!authResult.isAllowed()) {
+        throw new ForbiddenException(authResult.toString());
+      }
+      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 (DruidException e) {
+      return buildNonOkResponse(e);
+    }
+    catch (ForbiddenException e) {
+      log.debug("Got forbidden request for reason [%s]", e.getErrorMessage());
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.USER)
+                        .ofCategory(DruidException.Category.FORBIDDEN)
+                        .build(Access.DEFAULT_ERROR_MESSAGE)
+      );
+    }
+    catch (Exception e) {
+      log.noStackTrace().warn(e, "Failed to handle query: %s", queryId);
+      return buildNonOkResponse(DruidException.forPersona(DruidException.Persona.DEVELOPER)
+                                              .ofCategory(DruidException.Category.UNCATEGORIZED)
+                                              .build(e, "Failed to handle query: [%s]", queryId));
+    }
+  }
+
+  @GET
+  @Path("/{id}/results")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response doGetResults(
+      @PathParam("id") final String queryId,
+      @QueryParam("offset") Long offset,
+      @QueryParam("numRows") Long numberOfRows,
+      @Context final HttpServletRequest req
+  )
+  {
+    try {
+      Access authResult = AuthorizationUtils.authorizeAllResourceActions(
+          req,
+          Collections.emptyList(),
+          authorizerMapper
+      );
+      if (!authResult.isAllowed()) {
+        throw new ForbiddenException(authResult.toString());
+      }
+      final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req);
+
+      if (offset != null && offset < 0) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "offset cannot be negative. Please pass a positive number."
+                          )
+        );
+      }
+      if (numberOfRows != null && numberOfRows < 0) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "numRows cannot be negative. Please pass a positive number."
+                          )
+        );
+      }
+
+      final long start = offset == null ? 0 : offset;
+      final long last = SqlStatementResourceHelper.getLastIndex(numberOfRows, start);
+
+
+      getStatementStatus(queryId, authenticationResult.getIdentity(), false);

Review Comment:
   Oh yes that needs to be removed. Thanks for the catch. 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] abhishekagarwal87 commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["abhishekagarwal87 (via GitHub)" <gi...@apache.org>]

abhishekagarwal87 commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1245055112


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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);

Review Comment:
   it's still ignoring the results. 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] abhishekagarwal87 commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["abhishekagarwal87 (via GitHub)" <gi...@apache.org>]

abhishekagarwal87 commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1239482754


##########
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))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColNameAndType> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {
+              writer.writeRowStart();
+              for (int i = 0; i < rowSignature.size(); i++) {
+                writer.writeRowField(rowSignature.get(i).getColName(), ((List) results.get().get(k)).get(i));
+              }
+              writer.writeRowEnd();
+            }
+
+            writer.writeResponseEnd();
+          }
+          catch (Exception e) {
+            log.error(e, "Unable to stream results back for query[%s]", queryId);
+            throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+          }
+        }).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;
+    }
+  }
+
+  /**
+   * Queries can be canceled while in anystate. Canceling a query that has already completed will remove its results.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@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(),
+          false
+      );
+      if (sqlStatementResult.isPresent()) {
+        switch (sqlStatementResult.get().getState()) {
+          case ACCEPTED:
+          case RUNNING:
+            overlordClient.cancelTask(queryId);
+            return Response.status(Response.Status.ACCEPTED).build();
+          case SUCCESS:
+          case FAILED:
+            // we would also want to clean up the
+            return Response.ok().build();
+          default:
+            throw new ISE("Illegal State[%s] encountered", sqlStatementResult.get().getState());
+        }
+
+      } 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;
+    }
+  }
+
+  private Response buildStandardResponse(
+      Sequence<Object[]> sequence,
+      SqlQuery sqlQuery,
+      String sqlQueryId,
+      SqlRowTransformer rowTransformer
+  ) throws IOException
+  {
+    final Yielder<Object[]> yielder0 = Yielders.each(sequence);
+
+    try {
+      final Response.ResponseBuilder responseBuilder = Response.ok((StreamingOutput) outputStream -> {
+        CountingOutputStream os = new CountingOutputStream(outputStream);
+        Yielder<Object[]> yielder = yielder0;
+
+        try (final ResultFormat.Writer writer = sqlQuery.getResultFormat().createFormatter(os, jsonMapper)) {
+          writer.writeResponseStart();
+
+          if (sqlQuery.includeHeader()) {
+            writer.writeHeader(
+                rowTransformer.getRowType(),
+                sqlQuery.includeTypesHeader(),
+                sqlQuery.includeSqlTypesHeader()
+            );
+          }
+
+          while (!yielder.isDone()) {
+            final Object[] row = yielder.get();
+            writer.writeRowStart();
+            for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
+              final Object value = rowTransformer.transform(row, i);
+              writer.writeRowField(rowTransformer.getFieldList().get(i), value);
+            }
+            writer.writeRowEnd();
+            yielder = yielder.next(null);
+          }
+
+          writer.writeResponseEnd();
+        }
+        catch (Exception e) {
+          log.error(e, "Unable to send SQL response [%s]", sqlQueryId);
+          throw new RuntimeException(e);
+        }
+        finally {
+          yielder.close();
+        }
+      });
+
+      if (sqlQuery.includeHeader()) {
+        responseBuilder.header(SqlResource.SQL_HEADER_RESPONSE_HEADER, SqlResource.SQL_HEADER_VALUE);
+      }
+
+      return responseBuilder.build();
+    }
+    catch (Throwable e) {
+      // make sure to close yielder if anything happened before starting to serialize the response.
+      yielder0.close();
+      throw e;
+    }
+  }
+
+  private Response buildTaskResponse(Sequence<Object[]> sequence, String user)
+  {
+    List<Object[]> rows = sequence.toList();
+    int numRows = rows.size();
+    if (numRows != 1) {
+      throw new RE("Expected a single row but got [%d] rows. Please check broker logs for more information.", numRows);
+    }
+    String taskId = (String) rows.get(0)[0];
+    try {
+      Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, user, true);
+
+      if (statementResult.isPresent()) {
+        return Response.status(Response.Status.OK).entity(statementResult.get()).build();
+      } else {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(
+                null,
+                StringUtils.format(
+                    "Unable to find associated task for query id %s",
+                    taskId
+                ),
+                null,
+                null,
+                null
+            ),
+            taskId
+        );
+      }
+    }
+    catch (QueryException e) {
+      return buildNonOkResponse(Response.Status.OK.getStatusCode(), e, taskId);
+    }
+  }
+
+  private Response buildNonOkResponse(int status, SanitizableException e, String sqlQueryId)
+  {
+    // Though transformIfNeeded returns an exception, its purpose is to return
+    // a QueryException
+    Exception cleaned = serverConfig.getErrorResponseTransformStrategy().transformIfNeeded(e);
+    return Response.status(status).entity(cleaned).build();
+  }
+
+
+  private static Optional<List<ColNameAndType>> getSignature(
+      MSQControllerTask msqControllerTask
+  )
+  {
+    // only populate signature for select q's
+    if (msqControllerTask.getQuerySpec().getDestination().getClass() == TaskReportMSQDestination.class) {
+      ColumnMappings columnMappings = msqControllerTask.getQuerySpec().getColumnMappings();
+      List<SqlTypeName> sqlTypeNames = msqControllerTask.getSqlTypeNames();
+      if (sqlTypeNames == null || sqlTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColumnType> nativeTypeNames = msqControllerTask.getNativeTypeNames();
+      if (nativeTypeNames == null || nativeTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColNameAndType> signature = new ArrayList<>(columnMappings.size());
+      int index = 0;
+      for (String colName : columnMappings.getOutputColumnNames()) {
+        signature.add(new ColNameAndType(
+            colName,
+            sqlTypeNames.get(index).getName(),
+            nativeTypeNames.get(index).asTypeString()
+        ));
+        index++;
+      }
+      return Optional.of(signature);
+    }
+    return Optional.empty();
+  }
+
+
+  private void checkTaskPayloadOrThrow(TaskPayloadResponse taskPayloadResponse, String queryId) throws QueryException
+  {
+    if (taskPayloadResponse == null || taskPayloadResponse.getPayload() == null) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Could not get payload details for query[%s] from the overlord",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+
+    if (MSQControllerTask.class != taskPayloadResponse.getPayload().getClass()) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Fetched an unexpected payload for query[%s] from the overlord.",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+  }
+
+  private Optional<ResultSetInformation> getSampleResults(
+      String asyncResultId,
+      boolean isSelectQuery,
+      String dataSource,
+      SqlStatementState sqlStatementState
+  )
+  {
+    if (sqlStatementState == SqlStatementState.SUCCESS) {
+      Map<String, Object> payload = getPayload(overlordWork(overlordClient.taskReportAsMap(asyncResultId)));
+      Optional<Pair<Long, Long>> rowsAndSize = getRowsAndSizeFromPayload(payload, isSelectQuery);
+      return Optional.of(new ResultSetInformation(
+          null,
+          // since the rows can be sampled, get the number of rows from counters
+          rowsAndSize.orElse(new Pair<>(null, null)).lhs,
+          rowsAndSize.orElse(new Pair<>(null, null)).rhs,
+          dataSource,
+          // only populate sample results in case a select query is successful
+          isSelectQuery ? getResults(payload).orElse(null) : null
+      ));
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  private Optional<Pair<Long, Long>> getRowsAndSizeFromPayload(Map<String, Object> payload, boolean isSelectQuery)
+  {
+    List stages = getList(payload, "stages");
+    if (stages == null || stages.isEmpty()) {
+      return Optional.empty();
+    } else {
+      int maxStage = stages.size() - 1; // Last stage output is the total number of rows returned to the end user.
+      Map<String, Object> counterMap = getMap(getMap(payload, "counters"), String.valueOf(maxStage));
+      Long rows = -1L;
+      Long sizeInBytes = -1L;
+      if (counterMap == null) {
+        return Optional.empty();
+      }
+      for (Map.Entry<String, Object> worker : counterMap.entrySet()) {
+        Object workerChannels = worker.getValue();
+        if (workerChannels == null || !(workerChannels instanceof Map)) {
+          return Optional.empty();

Review Comment:
   do you ever expect this though? 



##########
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))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColNameAndType> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {
+              writer.writeRowStart();
+              for (int i = 0; i < rowSignature.size(); i++) {
+                writer.writeRowField(rowSignature.get(i).getColName(), ((List) results.get().get(k)).get(i));
+              }
+              writer.writeRowEnd();
+            }
+
+            writer.writeResponseEnd();
+          }
+          catch (Exception e) {
+            log.error(e, "Unable to stream results back for query[%s]", queryId);
+            throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+          }
+        }).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;
+    }
+  }
+
+  /**
+   * Queries can be canceled while in anystate. Canceling a query that has already completed will remove its results.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@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(),
+          false
+      );
+      if (sqlStatementResult.isPresent()) {
+        switch (sqlStatementResult.get().getState()) {
+          case ACCEPTED:
+          case RUNNING:
+            overlordClient.cancelTask(queryId);
+            return Response.status(Response.Status.ACCEPTED).build();
+          case SUCCESS:
+          case FAILED:
+            // we would also want to clean up the
+            return Response.ok().build();
+          default:
+            throw new ISE("Illegal State[%s] encountered", sqlStatementResult.get().getState());
+        }
+
+      } 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;
+    }
+  }
+
+  private Response buildStandardResponse(
+      Sequence<Object[]> sequence,
+      SqlQuery sqlQuery,
+      String sqlQueryId,
+      SqlRowTransformer rowTransformer
+  ) throws IOException
+  {
+    final Yielder<Object[]> yielder0 = Yielders.each(sequence);
+
+    try {
+      final Response.ResponseBuilder responseBuilder = Response.ok((StreamingOutput) outputStream -> {
+        CountingOutputStream os = new CountingOutputStream(outputStream);
+        Yielder<Object[]> yielder = yielder0;
+
+        try (final ResultFormat.Writer writer = sqlQuery.getResultFormat().createFormatter(os, jsonMapper)) {
+          writer.writeResponseStart();
+
+          if (sqlQuery.includeHeader()) {
+            writer.writeHeader(
+                rowTransformer.getRowType(),
+                sqlQuery.includeTypesHeader(),
+                sqlQuery.includeSqlTypesHeader()
+            );
+          }
+
+          while (!yielder.isDone()) {
+            final Object[] row = yielder.get();
+            writer.writeRowStart();
+            for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
+              final Object value = rowTransformer.transform(row, i);
+              writer.writeRowField(rowTransformer.getFieldList().get(i), value);
+            }
+            writer.writeRowEnd();
+            yielder = yielder.next(null);
+          }
+
+          writer.writeResponseEnd();
+        }
+        catch (Exception e) {
+          log.error(e, "Unable to send SQL response [%s]", sqlQueryId);
+          throw new RuntimeException(e);
+        }
+        finally {
+          yielder.close();
+        }
+      });
+
+      if (sqlQuery.includeHeader()) {
+        responseBuilder.header(SqlResource.SQL_HEADER_RESPONSE_HEADER, SqlResource.SQL_HEADER_VALUE);
+      }
+
+      return responseBuilder.build();
+    }
+    catch (Throwable e) {
+      // make sure to close yielder if anything happened before starting to serialize the response.
+      yielder0.close();
+      throw e;
+    }
+  }
+
+  private Response buildTaskResponse(Sequence<Object[]> sequence, String user)
+  {
+    List<Object[]> rows = sequence.toList();
+    int numRows = rows.size();
+    if (numRows != 1) {
+      throw new RE("Expected a single row but got [%d] rows. Please check broker logs for more information.", numRows);
+    }
+    String taskId = (String) rows.get(0)[0];

Review Comment:
   please add a size check on rows.get(0)



##########
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))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColNameAndType> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {
+              writer.writeRowStart();
+              for (int i = 0; i < rowSignature.size(); i++) {
+                writer.writeRowField(rowSignature.get(i).getColName(), ((List) results.get().get(k)).get(i));
+              }
+              writer.writeRowEnd();
+            }
+
+            writer.writeResponseEnd();
+          }
+          catch (Exception e) {
+            log.error(e, "Unable to stream results back for query[%s]", queryId);
+            throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+          }
+        }).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;
+    }
+  }
+
+  /**
+   * Queries can be canceled while in anystate. Canceling a query that has already completed will remove its results.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@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(),
+          false
+      );
+      if (sqlStatementResult.isPresent()) {
+        switch (sqlStatementResult.get().getState()) {
+          case ACCEPTED:
+          case RUNNING:
+            overlordClient.cancelTask(queryId);
+            return Response.status(Response.Status.ACCEPTED).build();
+          case SUCCESS:
+          case FAILED:
+            // we would also want to clean up the
+            return Response.ok().build();
+          default:
+            throw new ISE("Illegal State[%s] encountered", sqlStatementResult.get().getState());
+        }
+
+      } 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;
+    }
+  }
+
+  private Response buildStandardResponse(
+      Sequence<Object[]> sequence,
+      SqlQuery sqlQuery,
+      String sqlQueryId,
+      SqlRowTransformer rowTransformer
+  ) throws IOException
+  {
+    final Yielder<Object[]> yielder0 = Yielders.each(sequence);
+
+    try {
+      final Response.ResponseBuilder responseBuilder = Response.ok((StreamingOutput) outputStream -> {
+        CountingOutputStream os = new CountingOutputStream(outputStream);
+        Yielder<Object[]> yielder = yielder0;
+
+        try (final ResultFormat.Writer writer = sqlQuery.getResultFormat().createFormatter(os, jsonMapper)) {
+          writer.writeResponseStart();
+
+          if (sqlQuery.includeHeader()) {
+            writer.writeHeader(
+                rowTransformer.getRowType(),
+                sqlQuery.includeTypesHeader(),
+                sqlQuery.includeSqlTypesHeader()
+            );
+          }
+
+          while (!yielder.isDone()) {
+            final Object[] row = yielder.get();
+            writer.writeRowStart();
+            for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
+              final Object value = rowTransformer.transform(row, i);
+              writer.writeRowField(rowTransformer.getFieldList().get(i), value);
+            }
+            writer.writeRowEnd();
+            yielder = yielder.next(null);
+          }
+
+          writer.writeResponseEnd();
+        }
+        catch (Exception e) {
+          log.error(e, "Unable to send SQL response [%s]", sqlQueryId);
+          throw new RuntimeException(e);
+        }
+        finally {
+          yielder.close();
+        }
+      });
+
+      if (sqlQuery.includeHeader()) {
+        responseBuilder.header(SqlResource.SQL_HEADER_RESPONSE_HEADER, SqlResource.SQL_HEADER_VALUE);
+      }
+
+      return responseBuilder.build();
+    }
+    catch (Throwable e) {
+      // make sure to close yielder if anything happened before starting to serialize the response.
+      yielder0.close();
+      throw e;
+    }
+  }
+
+  private Response buildTaskResponse(Sequence<Object[]> sequence, String user)
+  {
+    List<Object[]> rows = sequence.toList();
+    int numRows = rows.size();
+    if (numRows != 1) {
+      throw new RE("Expected a single row but got [%d] rows. Please check broker logs for more information.", numRows);
+    }
+    String taskId = (String) rows.get(0)[0];
+    try {
+      Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, user, true);
+
+      if (statementResult.isPresent()) {
+        return Response.status(Response.Status.OK).entity(statementResult.get()).build();
+      } else {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(
+                null,
+                StringUtils.format(
+                    "Unable to find associated task for query id %s",
+                    taskId
+                ),
+                null,
+                null,
+                null
+            ),
+            taskId
+        );
+      }
+    }
+    catch (QueryException e) {
+      return buildNonOkResponse(Response.Status.OK.getStatusCode(), e, taskId);
+    }
+  }
+
+  private Response buildNonOkResponse(int status, SanitizableException e, String sqlQueryId)
+  {
+    // Though transformIfNeeded returns an exception, its purpose is to return
+    // a QueryException
+    Exception cleaned = serverConfig.getErrorResponseTransformStrategy().transformIfNeeded(e);
+    return Response.status(status).entity(cleaned).build();
+  }
+
+
+  private static Optional<List<ColNameAndType>> getSignature(
+      MSQControllerTask msqControllerTask
+  )
+  {
+    // only populate signature for select q's
+    if (msqControllerTask.getQuerySpec().getDestination().getClass() == TaskReportMSQDestination.class) {
+      ColumnMappings columnMappings = msqControllerTask.getQuerySpec().getColumnMappings();
+      List<SqlTypeName> sqlTypeNames = msqControllerTask.getSqlTypeNames();
+      if (sqlTypeNames == null || sqlTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColumnType> nativeTypeNames = msqControllerTask.getNativeTypeNames();
+      if (nativeTypeNames == null || nativeTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColNameAndType> signature = new ArrayList<>(columnMappings.size());
+      int index = 0;
+      for (String colName : columnMappings.getOutputColumnNames()) {
+        signature.add(new ColNameAndType(
+            colName,
+            sqlTypeNames.get(index).getName(),
+            nativeTypeNames.get(index).asTypeString()
+        ));
+        index++;
+      }
+      return Optional.of(signature);
+    }
+    return Optional.empty();
+  }
+
+
+  private void checkTaskPayloadOrThrow(TaskPayloadResponse taskPayloadResponse, String queryId) throws QueryException
+  {
+    if (taskPayloadResponse == null || taskPayloadResponse.getPayload() == null) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Could not get payload details for query[%s] from the overlord",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+
+    if (MSQControllerTask.class != taskPayloadResponse.getPayload().getClass()) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Fetched an unexpected payload for query[%s] from the overlord.",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+  }
+
+  private Optional<ResultSetInformation> getSampleResults(
+      String asyncResultId,
+      boolean isSelectQuery,
+      String dataSource,
+      SqlStatementState sqlStatementState
+  )
+  {
+    if (sqlStatementState == SqlStatementState.SUCCESS) {
+      Map<String, Object> payload = getPayload(overlordWork(overlordClient.taskReportAsMap(asyncResultId)));
+      Optional<Pair<Long, Long>> rowsAndSize = getRowsAndSizeFromPayload(payload, isSelectQuery);
+      return Optional.of(new ResultSetInformation(
+          null,
+          // since the rows can be sampled, get the number of rows from counters
+          rowsAndSize.orElse(new Pair<>(null, null)).lhs,
+          rowsAndSize.orElse(new Pair<>(null, null)).rhs,
+          dataSource,
+          // only populate sample results in case a select query is successful
+          isSelectQuery ? getResults(payload).orElse(null) : null
+      ));
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  private Optional<Pair<Long, Long>> getRowsAndSizeFromPayload(Map<String, Object> payload, boolean isSelectQuery)
+  {
+    List stages = getList(payload, "stages");
+    if (stages == null || stages.isEmpty()) {
+      return Optional.empty();
+    } else {
+      int maxStage = stages.size() - 1; // Last stage output is the total number of rows returned to the end user.
+      Map<String, Object> counterMap = getMap(getMap(payload, "counters"), String.valueOf(maxStage));
+      Long rows = -1L;
+      Long sizeInBytes = -1L;
+      if (counterMap == null) {
+        return Optional.empty();
+      }
+      for (Map.Entry<String, Object> worker : counterMap.entrySet()) {
+        Object workerChannels = worker.getValue();
+        if (workerChannels == null || !(workerChannels instanceof Map)) {
+          return Optional.empty();
+        }
+        if (isSelectQuery) {
+          Object output = ((Map<?, ?>) workerChannels).get("output");
+          if (output != null && output instanceof Map) {
+            List<Integer> rowsPerChannel = (List<Integer>) ((Map<String, Object>) output).get("rows");
+            List<Integer> bytesPerChannel = (List<Integer>) ((Map<String, Object>) output).get("bytes");
+            for (Integer row : rowsPerChannel) {
+              rows = rows + row;
+            }
+            for (Integer bytes : bytesPerChannel) {
+              sizeInBytes = sizeInBytes + bytes;
+            }
+          }
+        } else {
+          Object output = ((Map<?, ?>) workerChannels).get("segmentGenerationProgress");
+          if (output != null && output instanceof Map) {
+            rows += (Integer) ((Map<String, Object>) output).get("rowsPushed");
+          }
+        }
+      }
+
+      return Optional.of(new Pair<>(rows == -1L ? null : rows + 1, sizeInBytes == -1L ? null : sizeInBytes + 1));
+    }
+  }
+
+
+  private Optional<SqlStatementResult> getStatementStatus(String queryId, String currentUser, boolean withResults)
+      throws QueryException, ForbiddenException
+  {
+    TaskStatusResponse taskResponse = overlordWork(overlordClient.taskStatus(queryId));
+    if (taskResponse == null) {
+      return Optional.empty();
+    }
+
+    TaskStatusPlus statusPlus = taskResponse.getStatus();
+    if (statusPlus == null || !MSQControllerTask.TYPE.equals(statusPlus.getType())) {
+      return Optional.empty();
+    }
+
+    // since we need the controller payload for auth checks.
+    MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, currentUser);
+    SqlStatementState sqlStatementState = getSqlStatementState(statusPlus);
+
+    if (SqlStatementState.FAILED == sqlStatementState) {
+      Map<String, Object> exceptionDetails = getQueryExceptionDetails(getPayload(overlordWork(overlordClient.taskReportAsMap(
+          queryId))));
+      Map<String, Object> exception = getMap(exceptionDetails, "error");
+      if (exceptionDetails == null || exception == null) {
+        return Optional.of(new SqlStatementResult(
+            queryId,
+            sqlStatementState,
+            taskResponse.getStatus().getCreatedTime(),
+            null,
+            taskResponse.getStatus().getDuration(),
+            null,
+            new QueryException(null, statusPlus.getErrorMsg(), null, null, null)
+        ));
+      }
+
+      final String errorMessage = String.valueOf(exception.getOrDefault("errorMessage", statusPlus.getErrorMsg()));
+      exception.remove("errorMessage");
+      String host = (String) exceptionDetails.getOrDefault("host", null);
+      Map<String, String> stringException = new HashMap<>();
+      for (Map.Entry<String, Object> exceptionKeys : exception.entrySet()) {
+        stringException.put(exceptionKeys.getKey(), String.valueOf(exceptionKeys.getValue()));
+      }
+      return Optional.of(new SqlStatementResult(
+          queryId,
+          sqlStatementState,
+          taskResponse.getStatus().getCreatedTime(),
+          null,
+          taskResponse.getStatus().getDuration(),
+          null,
+          new QueryException(null, errorMessage, null, host, stringException)
+      ));
+    } else {
+      Optional<List<ColNameAndType>> signature = getSignature(msqControllerTask);
+
+      return Optional.of(new SqlStatementResult(
+          queryId,
+          sqlStatementState,
+          taskResponse.getStatus().getCreatedTime(),
+          signature.orElse(null),
+          taskResponse.getStatus().getDuration(),
+          withResults ? getSampleResults(
+              queryId,
+              signature.isPresent(),
+              msqControllerTask.getDataSource(),
+              sqlStatementState
+          ).orElse(null) : null,
+          null
+      ));
+    }
+  }
+
+  private Map<String, Object> getQueryExceptionDetails(Map<String, Object> payload)
+  {
+    return getMap(getMap(payload, "status"), "errorReport");
+  }
+
+  private MSQControllerTask getMSQControllerTaskOrThrow(String queryId, String currentUser) throws ForbiddenException
+  {
+    TaskPayloadResponse taskPayloadResponse = overlordWork(overlordClient.taskPayload(queryId));
+    checkTaskPayloadOrThrow(taskPayloadResponse, queryId);
+
+    MSQControllerTask msqControllerTask = (MSQControllerTask) taskPayloadResponse.getPayload();
+    if (currentUser == null || !currentUser.equals(msqControllerTask.getQuerySpec()
+                                                                    .getQuery()
+                                                                    .getContext()
+                                                                    .get(MSQTaskQueryMaker.USER_KEY))) {
+      throw new ForbiddenException();
+    }
+    return msqControllerTask;
+  }
+
+  private <T> T overlordWork(final ListenableFuture<T> future)

Review Comment:
   not sure why you named it overlordWork but maybe blockOnOverlordResponse is better? 



##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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)

Review Comment:
   this context will be helpful in the javadocs. 



##########
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:
   so the results are in the task report? 



##########
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();

Review Comment:
   hmm. what does it mean for signature to be absent? 



##########
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.",

Review Comment:
   It will be good to include queryId in these messages. 



##########
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))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColNameAndType> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {
+              writer.writeRowStart();
+              for (int i = 0; i < rowSignature.size(); i++) {
+                writer.writeRowField(rowSignature.get(i).getColName(), ((List) results.get().get(k)).get(i));
+              }
+              writer.writeRowEnd();
+            }
+
+            writer.writeResponseEnd();
+          }
+          catch (Exception e) {
+            log.error(e, "Unable to stream results back for query[%s]", queryId);
+            throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+          }
+        }).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;
+    }
+  }
+
+  /**
+   * Queries can be canceled while in anystate. Canceling a query that has already completed will remove its results.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@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(),
+          false
+      );
+      if (sqlStatementResult.isPresent()) {
+        switch (sqlStatementResult.get().getState()) {
+          case ACCEPTED:
+          case RUNNING:
+            overlordClient.cancelTask(queryId);
+            return Response.status(Response.Status.ACCEPTED).build();
+          case SUCCESS:
+          case FAILED:
+            // we would also want to clean up the
+            return Response.ok().build();
+          default:
+            throw new ISE("Illegal State[%s] encountered", sqlStatementResult.get().getState());
+        }
+
+      } 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;
+    }
+  }
+
+  private Response buildStandardResponse(
+      Sequence<Object[]> sequence,
+      SqlQuery sqlQuery,
+      String sqlQueryId,
+      SqlRowTransformer rowTransformer
+  ) throws IOException
+  {
+    final Yielder<Object[]> yielder0 = Yielders.each(sequence);
+
+    try {
+      final Response.ResponseBuilder responseBuilder = Response.ok((StreamingOutput) outputStream -> {
+        CountingOutputStream os = new CountingOutputStream(outputStream);
+        Yielder<Object[]> yielder = yielder0;
+
+        try (final ResultFormat.Writer writer = sqlQuery.getResultFormat().createFormatter(os, jsonMapper)) {
+          writer.writeResponseStart();
+
+          if (sqlQuery.includeHeader()) {
+            writer.writeHeader(
+                rowTransformer.getRowType(),
+                sqlQuery.includeTypesHeader(),
+                sqlQuery.includeSqlTypesHeader()
+            );
+          }
+
+          while (!yielder.isDone()) {
+            final Object[] row = yielder.get();
+            writer.writeRowStart();
+            for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
+              final Object value = rowTransformer.transform(row, i);
+              writer.writeRowField(rowTransformer.getFieldList().get(i), value);
+            }
+            writer.writeRowEnd();
+            yielder = yielder.next(null);
+          }
+
+          writer.writeResponseEnd();
+        }
+        catch (Exception e) {
+          log.error(e, "Unable to send SQL response [%s]", sqlQueryId);
+          throw new RuntimeException(e);
+        }
+        finally {
+          yielder.close();
+        }
+      });
+
+      if (sqlQuery.includeHeader()) {
+        responseBuilder.header(SqlResource.SQL_HEADER_RESPONSE_HEADER, SqlResource.SQL_HEADER_VALUE);
+      }
+
+      return responseBuilder.build();
+    }
+    catch (Throwable e) {
+      // make sure to close yielder if anything happened before starting to serialize the response.
+      yielder0.close();
+      throw e;
+    }
+  }
+
+  private Response buildTaskResponse(Sequence<Object[]> sequence, String user)
+  {
+    List<Object[]> rows = sequence.toList();
+    int numRows = rows.size();
+    if (numRows != 1) {
+      throw new RE("Expected a single row but got [%d] rows. Please check broker logs for more information.", numRows);
+    }
+    String taskId = (String) rows.get(0)[0];
+    try {
+      Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, user, true);
+
+      if (statementResult.isPresent()) {
+        return Response.status(Response.Status.OK).entity(statementResult.get()).build();
+      } else {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(
+                null,
+                StringUtils.format(
+                    "Unable to find associated task for query id %s",
+                    taskId
+                ),
+                null,
+                null,
+                null
+            ),
+            taskId
+        );
+      }
+    }
+    catch (QueryException e) {
+      return buildNonOkResponse(Response.Status.OK.getStatusCode(), e, taskId);
+    }
+  }
+
+  private Response buildNonOkResponse(int status, SanitizableException e, String sqlQueryId)
+  {
+    // Though transformIfNeeded returns an exception, its purpose is to return
+    // a QueryException
+    Exception cleaned = serverConfig.getErrorResponseTransformStrategy().transformIfNeeded(e);
+    return Response.status(status).entity(cleaned).build();
+  }
+
+
+  private static Optional<List<ColNameAndType>> getSignature(
+      MSQControllerTask msqControllerTask
+  )
+  {
+    // only populate signature for select q's
+    if (msqControllerTask.getQuerySpec().getDestination().getClass() == TaskReportMSQDestination.class) {
+      ColumnMappings columnMappings = msqControllerTask.getQuerySpec().getColumnMappings();
+      List<SqlTypeName> sqlTypeNames = msqControllerTask.getSqlTypeNames();
+      if (sqlTypeNames == null || sqlTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColumnType> nativeTypeNames = msqControllerTask.getNativeTypeNames();
+      if (nativeTypeNames == null || nativeTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColNameAndType> signature = new ArrayList<>(columnMappings.size());
+      int index = 0;
+      for (String colName : columnMappings.getOutputColumnNames()) {
+        signature.add(new ColNameAndType(
+            colName,
+            sqlTypeNames.get(index).getName(),
+            nativeTypeNames.get(index).asTypeString()
+        ));
+        index++;
+      }
+      return Optional.of(signature);
+    }
+    return Optional.empty();
+  }
+
+
+  private void checkTaskPayloadOrThrow(TaskPayloadResponse taskPayloadResponse, String queryId) throws QueryException
+  {
+    if (taskPayloadResponse == null || taskPayloadResponse.getPayload() == null) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Could not get payload details for query[%s] from the overlord",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+
+    if (MSQControllerTask.class != taskPayloadResponse.getPayload().getClass()) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Fetched an unexpected payload for query[%s] from the overlord.",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+  }
+
+  private Optional<ResultSetInformation> getSampleResults(
+      String asyncResultId,
+      boolean isSelectQuery,
+      String dataSource,
+      SqlStatementState sqlStatementState
+  )
+  {
+    if (sqlStatementState == SqlStatementState.SUCCESS) {
+      Map<String, Object> payload = getPayload(overlordWork(overlordClient.taskReportAsMap(asyncResultId)));
+      Optional<Pair<Long, Long>> rowsAndSize = getRowsAndSizeFromPayload(payload, isSelectQuery);
+      return Optional.of(new ResultSetInformation(
+          null,
+          // since the rows can be sampled, get the number of rows from counters
+          rowsAndSize.orElse(new Pair<>(null, null)).lhs,
+          rowsAndSize.orElse(new Pair<>(null, null)).rhs,
+          dataSource,
+          // only populate sample results in case a select query is successful
+          isSelectQuery ? getResults(payload).orElse(null) : null
+      ));
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  private Optional<Pair<Long, Long>> getRowsAndSizeFromPayload(Map<String, Object> payload, boolean isSelectQuery)
+  {
+    List stages = getList(payload, "stages");
+    if (stages == null || stages.isEmpty()) {
+      return Optional.empty();
+    } else {
+      int maxStage = stages.size() - 1; // Last stage output is the total number of rows returned to the end user.
+      Map<String, Object> counterMap = getMap(getMap(payload, "counters"), String.valueOf(maxStage));
+      Long rows = -1L;
+      Long sizeInBytes = -1L;
+      if (counterMap == null) {
+        return Optional.empty();
+      }
+      for (Map.Entry<String, Object> worker : counterMap.entrySet()) {
+        Object workerChannels = worker.getValue();
+        if (workerChannels == null || !(workerChannels instanceof Map)) {
+          return Optional.empty();
+        }
+        if (isSelectQuery) {
+          Object output = ((Map<?, ?>) workerChannels).get("output");
+          if (output != null && output instanceof Map) {
+            List<Integer> rowsPerChannel = (List<Integer>) ((Map<String, Object>) output).get("rows");
+            List<Integer> bytesPerChannel = (List<Integer>) ((Map<String, Object>) output).get("bytes");
+            for (Integer row : rowsPerChannel) {
+              rows = rows + row;
+            }
+            for (Integer bytes : bytesPerChannel) {
+              sizeInBytes = sizeInBytes + bytes;
+            }
+          }
+        } else {
+          Object output = ((Map<?, ?>) workerChannels).get("segmentGenerationProgress");
+          if (output != null && output instanceof Map) {
+            rows += (Integer) ((Map<String, Object>) output).get("rowsPushed");
+          }
+        }
+      }
+
+      return Optional.of(new Pair<>(rows == -1L ? null : rows + 1, sizeInBytes == -1L ? null : sizeInBytes + 1));
+    }
+  }
+
+
+  private Optional<SqlStatementResult> getStatementStatus(String queryId, String currentUser, boolean withResults)
+      throws QueryException, ForbiddenException
+  {
+    TaskStatusResponse taskResponse = overlordWork(overlordClient.taskStatus(queryId));
+    if (taskResponse == null) {
+      return Optional.empty();
+    }
+
+    TaskStatusPlus statusPlus = taskResponse.getStatus();
+    if (statusPlus == null || !MSQControllerTask.TYPE.equals(statusPlus.getType())) {
+      return Optional.empty();
+    }
+
+    // since we need the controller payload for auth checks.
+    MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, currentUser);
+    SqlStatementState sqlStatementState = getSqlStatementState(statusPlus);
+
+    if (SqlStatementState.FAILED == sqlStatementState) {
+      Map<String, Object> exceptionDetails = getQueryExceptionDetails(getPayload(overlordWork(overlordClient.taskReportAsMap(
+          queryId))));
+      Map<String, Object> exception = getMap(exceptionDetails, "error");
+      if (exceptionDetails == null || exception == null) {
+        return Optional.of(new SqlStatementResult(
+            queryId,
+            sqlStatementState,
+            taskResponse.getStatus().getCreatedTime(),
+            null,
+            taskResponse.getStatus().getDuration(),
+            null,
+            new QueryException(null, statusPlus.getErrorMsg(), null, null, null)
+        ));
+      }
+
+      final String errorMessage = String.valueOf(exception.getOrDefault("errorMessage", statusPlus.getErrorMsg()));
+      exception.remove("errorMessage");
+      String host = (String) exceptionDetails.getOrDefault("host", null);
+      Map<String, String> stringException = new HashMap<>();
+      for (Map.Entry<String, Object> exceptionKeys : exception.entrySet()) {
+        stringException.put(exceptionKeys.getKey(), String.valueOf(exceptionKeys.getValue()));
+      }
+      return Optional.of(new SqlStatementResult(
+          queryId,
+          sqlStatementState,
+          taskResponse.getStatus().getCreatedTime(),
+          null,
+          taskResponse.getStatus().getDuration(),
+          null,
+          new QueryException(null, errorMessage, null, host, stringException)
+      ));
+    } else {
+      Optional<List<ColNameAndType>> signature = getSignature(msqControllerTask);
+
+      return Optional.of(new SqlStatementResult(
+          queryId,
+          sqlStatementState,
+          taskResponse.getStatus().getCreatedTime(),
+          signature.orElse(null),
+          taskResponse.getStatus().getDuration(),
+          withResults ? getSampleResults(
+              queryId,
+              signature.isPresent(),
+              msqControllerTask.getDataSource(),
+              sqlStatementState
+          ).orElse(null) : null,
+          null
+      ));
+    }
+  }
+
+  private Map<String, Object> getQueryExceptionDetails(Map<String, Object> payload)
+  {
+    return getMap(getMap(payload, "status"), "errorReport");
+  }
+
+  private MSQControllerTask getMSQControllerTaskOrThrow(String queryId, String currentUser) throws ForbiddenException
+  {
+    TaskPayloadResponse taskPayloadResponse = overlordWork(overlordClient.taskPayload(queryId));
+    checkTaskPayloadOrThrow(taskPayloadResponse, queryId);
+
+    MSQControllerTask msqControllerTask = (MSQControllerTask) taskPayloadResponse.getPayload();
+    if (currentUser == null || !currentUser.equals(msqControllerTask.getQuerySpec()
+                                                                    .getQuery()
+                                                                    .getContext()
+                                                                    .get(MSQTaskQueryMaker.USER_KEY))) {
+      throw new ForbiddenException();
+    }
+    return msqControllerTask;
+  }
+
+  private <T> T overlordWork(final ListenableFuture<T> future)
+  {
+    try {
+      return FutureUtils.getUnchecked(future, true);
+    }
+    catch (RuntimeException e) {
+      throw new QueryException(null, "Unable to contact overlord " + e.getMessage(), null, null, null);

Review Comment:
   whats the expected action associated with this error message? 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1239371074


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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);

Review Comment:
   Those are done as part of stmt.plan(). 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1239350411


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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)

Review Comment:
   So that the console can decide if the MSQ extension is enabled or not. 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1244664179


##########
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))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColNameAndType> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {
+              writer.writeRowStart();
+              for (int i = 0; i < rowSignature.size(); i++) {
+                writer.writeRowField(rowSignature.get(i).getColName(), ((List) results.get().get(k)).get(i));
+              }
+              writer.writeRowEnd();
+            }
+
+            writer.writeResponseEnd();
+          }
+          catch (Exception e) {
+            log.error(e, "Unable to stream results back for query[%s]", queryId);
+            throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+          }
+        }).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;
+    }
+  }
+
+  /**
+   * Queries can be canceled while in anystate. Canceling a query that has already completed will remove its results.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@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(),
+          false
+      );
+      if (sqlStatementResult.isPresent()) {
+        switch (sqlStatementResult.get().getState()) {
+          case ACCEPTED:
+          case RUNNING:
+            overlordClient.cancelTask(queryId);
+            return Response.status(Response.Status.ACCEPTED).build();
+          case SUCCESS:
+          case FAILED:
+            // we would also want to clean up the
+            return Response.ok().build();
+          default:
+            throw new ISE("Illegal State[%s] encountered", sqlStatementResult.get().getState());
+        }
+
+      } 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;
+    }
+  }
+
+  private Response buildStandardResponse(
+      Sequence<Object[]> sequence,
+      SqlQuery sqlQuery,
+      String sqlQueryId,
+      SqlRowTransformer rowTransformer
+  ) throws IOException
+  {
+    final Yielder<Object[]> yielder0 = Yielders.each(sequence);
+
+    try {
+      final Response.ResponseBuilder responseBuilder = Response.ok((StreamingOutput) outputStream -> {
+        CountingOutputStream os = new CountingOutputStream(outputStream);
+        Yielder<Object[]> yielder = yielder0;
+
+        try (final ResultFormat.Writer writer = sqlQuery.getResultFormat().createFormatter(os, jsonMapper)) {
+          writer.writeResponseStart();
+
+          if (sqlQuery.includeHeader()) {
+            writer.writeHeader(
+                rowTransformer.getRowType(),
+                sqlQuery.includeTypesHeader(),
+                sqlQuery.includeSqlTypesHeader()
+            );
+          }
+
+          while (!yielder.isDone()) {
+            final Object[] row = yielder.get();
+            writer.writeRowStart();
+            for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
+              final Object value = rowTransformer.transform(row, i);
+              writer.writeRowField(rowTransformer.getFieldList().get(i), value);
+            }
+            writer.writeRowEnd();
+            yielder = yielder.next(null);
+          }
+
+          writer.writeResponseEnd();
+        }
+        catch (Exception e) {
+          log.error(e, "Unable to send SQL response [%s]", sqlQueryId);
+          throw new RuntimeException(e);
+        }
+        finally {
+          yielder.close();
+        }
+      });
+
+      if (sqlQuery.includeHeader()) {
+        responseBuilder.header(SqlResource.SQL_HEADER_RESPONSE_HEADER, SqlResource.SQL_HEADER_VALUE);
+      }
+
+      return responseBuilder.build();
+    }
+    catch (Throwable e) {
+      // make sure to close yielder if anything happened before starting to serialize the response.
+      yielder0.close();
+      throw e;
+    }
+  }
+
+  private Response buildTaskResponse(Sequence<Object[]> sequence, String user)
+  {
+    List<Object[]> rows = sequence.toList();
+    int numRows = rows.size();
+    if (numRows != 1) {
+      throw new RE("Expected a single row but got [%d] rows. Please check broker logs for more information.", numRows);
+    }
+    String taskId = (String) rows.get(0)[0];
+    try {
+      Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, user, true);
+
+      if (statementResult.isPresent()) {
+        return Response.status(Response.Status.OK).entity(statementResult.get()).build();
+      } else {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(
+                null,
+                StringUtils.format(
+                    "Unable to find associated task for query id %s",
+                    taskId
+                ),
+                null,
+                null,
+                null
+            ),
+            taskId
+        );
+      }
+    }
+    catch (QueryException e) {
+      return buildNonOkResponse(Response.Status.OK.getStatusCode(), e, taskId);
+    }
+  }
+
+  private Response buildNonOkResponse(int status, SanitizableException e, String sqlQueryId)
+  {
+    // Though transformIfNeeded returns an exception, its purpose is to return
+    // a QueryException
+    Exception cleaned = serverConfig.getErrorResponseTransformStrategy().transformIfNeeded(e);
+    return Response.status(status).entity(cleaned).build();
+  }
+
+
+  private static Optional<List<ColNameAndType>> getSignature(
+      MSQControllerTask msqControllerTask
+  )
+  {
+    // only populate signature for select q's
+    if (msqControllerTask.getQuerySpec().getDestination().getClass() == TaskReportMSQDestination.class) {
+      ColumnMappings columnMappings = msqControllerTask.getQuerySpec().getColumnMappings();
+      List<SqlTypeName> sqlTypeNames = msqControllerTask.getSqlTypeNames();
+      if (sqlTypeNames == null || sqlTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColumnType> nativeTypeNames = msqControllerTask.getNativeTypeNames();
+      if (nativeTypeNames == null || nativeTypeNames.size() != columnMappings.size()) {
+        return Optional.empty();
+      }
+      List<ColNameAndType> signature = new ArrayList<>(columnMappings.size());
+      int index = 0;
+      for (String colName : columnMappings.getOutputColumnNames()) {
+        signature.add(new ColNameAndType(
+            colName,
+            sqlTypeNames.get(index).getName(),
+            nativeTypeNames.get(index).asTypeString()
+        ));
+        index++;
+      }
+      return Optional.of(signature);
+    }
+    return Optional.empty();
+  }
+
+
+  private void checkTaskPayloadOrThrow(TaskPayloadResponse taskPayloadResponse, String queryId) throws QueryException
+  {
+    if (taskPayloadResponse == null || taskPayloadResponse.getPayload() == null) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Could not get payload details for query[%s] from the overlord",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+
+    if (MSQControllerTask.class != taskPayloadResponse.getPayload().getClass()) {
+      throw new QueryException(
+          QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+          StringUtils.format(
+              "Fetched an unexpected payload for query[%s] from the overlord.",
+              queryId
+          ),
+          null,
+          null,
+          null
+      );
+    }
+  }
+
+  private Optional<ResultSetInformation> getSampleResults(
+      String asyncResultId,
+      boolean isSelectQuery,
+      String dataSource,
+      SqlStatementState sqlStatementState
+  )
+  {
+    if (sqlStatementState == SqlStatementState.SUCCESS) {
+      Map<String, Object> payload = getPayload(overlordWork(overlordClient.taskReportAsMap(asyncResultId)));
+      Optional<Pair<Long, Long>> rowsAndSize = getRowsAndSizeFromPayload(payload, isSelectQuery);
+      return Optional.of(new ResultSetInformation(
+          null,
+          // since the rows can be sampled, get the number of rows from counters
+          rowsAndSize.orElse(new Pair<>(null, null)).lhs,
+          rowsAndSize.orElse(new Pair<>(null, null)).rhs,
+          dataSource,
+          // only populate sample results in case a select query is successful
+          isSelectQuery ? getResults(payload).orElse(null) : null
+      ));
+    } else {
+      return Optional.empty();
+    }
+  }
+
+  private Optional<Pair<Long, Long>> getRowsAndSizeFromPayload(Map<String, Object> payload, boolean isSelectQuery)
+  {
+    List stages = getList(payload, "stages");
+    if (stages == null || stages.isEmpty()) {
+      return Optional.empty();
+    } else {
+      int maxStage = stages.size() - 1; // Last stage output is the total number of rows returned to the end user.
+      Map<String, Object> counterMap = getMap(getMap(payload, "counters"), String.valueOf(maxStage));
+      Long rows = -1L;
+      Long sizeInBytes = -1L;
+      if (counterMap == null) {
+        return Optional.empty();
+      }
+      for (Map.Entry<String, Object> worker : counterMap.entrySet()) {
+        Object workerChannels = worker.getValue();
+        if (workerChannels == null || !(workerChannels instanceof Map)) {
+          return Optional.empty();
+        }
+        if (isSelectQuery) {
+          Object output = ((Map<?, ?>) workerChannels).get("output");
+          if (output != null && output instanceof Map) {
+            List<Integer> rowsPerChannel = (List<Integer>) ((Map<String, Object>) output).get("rows");
+            List<Integer> bytesPerChannel = (List<Integer>) ((Map<String, Object>) output).get("bytes");
+            for (Integer row : rowsPerChannel) {
+              rows = rows + row;
+            }
+            for (Integer bytes : bytesPerChannel) {
+              sizeInBytes = sizeInBytes + bytes;
+            }
+          }
+        } else {
+          Object output = ((Map<?, ?>) workerChannels).get("segmentGenerationProgress");
+          if (output != null && output instanceof Map) {
+            rows += (Integer) ((Map<String, Object>) output).get("rowsPushed");
+          }
+        }
+      }
+
+      return Optional.of(new Pair<>(rows == -1L ? null : rows + 1, sizeInBytes == -1L ? null : sizeInBytes + 1));
+    }
+  }
+
+
+  private Optional<SqlStatementResult> getStatementStatus(String queryId, String currentUser, boolean withResults)
+      throws QueryException, ForbiddenException
+  {
+    TaskStatusResponse taskResponse = overlordWork(overlordClient.taskStatus(queryId));
+    if (taskResponse == null) {
+      return Optional.empty();
+    }
+
+    TaskStatusPlus statusPlus = taskResponse.getStatus();
+    if (statusPlus == null || !MSQControllerTask.TYPE.equals(statusPlus.getType())) {
+      return Optional.empty();
+    }
+
+    // since we need the controller payload for auth checks.
+    MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, currentUser);
+    SqlStatementState sqlStatementState = getSqlStatementState(statusPlus);
+
+    if (SqlStatementState.FAILED == sqlStatementState) {
+      Map<String, Object> exceptionDetails = getQueryExceptionDetails(getPayload(overlordWork(overlordClient.taskReportAsMap(
+          queryId))));
+      Map<String, Object> exception = getMap(exceptionDetails, "error");
+      if (exceptionDetails == null || exception == null) {
+        return Optional.of(new SqlStatementResult(
+            queryId,
+            sqlStatementState,
+            taskResponse.getStatus().getCreatedTime(),
+            null,
+            taskResponse.getStatus().getDuration(),
+            null,
+            new QueryException(null, statusPlus.getErrorMsg(), null, null, null)
+        ));
+      }
+
+      final String errorMessage = String.valueOf(exception.getOrDefault("errorMessage", statusPlus.getErrorMsg()));
+      exception.remove("errorMessage");
+      String host = (String) exceptionDetails.getOrDefault("host", null);
+      Map<String, String> stringException = new HashMap<>();
+      for (Map.Entry<String, Object> exceptionKeys : exception.entrySet()) {
+        stringException.put(exceptionKeys.getKey(), String.valueOf(exceptionKeys.getValue()));
+      }
+      return Optional.of(new SqlStatementResult(
+          queryId,
+          sqlStatementState,
+          taskResponse.getStatus().getCreatedTime(),
+          null,
+          taskResponse.getStatus().getDuration(),
+          null,
+          new QueryException(null, errorMessage, null, host, stringException)
+      ));
+    } else {
+      Optional<List<ColNameAndType>> signature = getSignature(msqControllerTask);
+
+      return Optional.of(new SqlStatementResult(
+          queryId,
+          sqlStatementState,
+          taskResponse.getStatus().getCreatedTime(),
+          signature.orElse(null),
+          taskResponse.getStatus().getDuration(),
+          withResults ? getSampleResults(
+              queryId,
+              signature.isPresent(),
+              msqControllerTask.getDataSource(),
+              sqlStatementState
+          ).orElse(null) : null,
+          null
+      ));
+    }
+  }
+
+  private Map<String, Object> getQueryExceptionDetails(Map<String, Object> payload)
+  {
+    return getMap(getMap(payload, "status"), "errorReport");
+  }
+
+  private MSQControllerTask getMSQControllerTaskOrThrow(String queryId, String currentUser) throws ForbiddenException
+  {
+    TaskPayloadResponse taskPayloadResponse = overlordWork(overlordClient.taskPayload(queryId));
+    checkTaskPayloadOrThrow(taskPayloadResponse, queryId);
+
+    MSQControllerTask msqControllerTask = (MSQControllerTask) taskPayloadResponse.getPayload();
+    if (currentUser == null || !currentUser.equals(msqControllerTask.getQuerySpec()
+                                                                    .getQuery()
+                                                                    .getContext()
+                                                                    .get(MSQTaskQueryMaker.USER_KEY))) {
+      throw new ForbiddenException();
+    }
+    return msqControllerTask;
+  }
+
+  private <T> T overlordWork(final ListenableFuture<T> future)
+  {
+    try {
+      return FutureUtils.getUnchecked(future, true);
+    }
+    catch (RuntimeException e) {
+      throw new QueryException(null, "Unable to contact overlord " + e.getMessage(), null, null, null);

Review Comment:
   ``` DruidException.forPersona(DruidException.Persona.DEVELOPER)
                             .ofCategory(DruidException.Category.UNCATEGORIZED)
                             .build("Unable to contact overlord " + e.getMessage());
       }
       ```



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1244661992


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          queryId
+      );
+    }
+
+    final long start = offset == null ? 0 : offset;
+    final long last = getLastIndex(numberOfRows, start);
+
+    try {
+      AuthorizationUtils.authorizeAllResourceActions(req, Collections.emptyList(), authorizerMapper);

Review Comment:
   All authenticated users are authorized for this API hence we check an empty resource.



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe merged pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe merged PR #14416:
URL: https://github.com/apache/druid/pull/14416


-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

Re: [PR] Adding Interactive API's for MSQ engine (druid)

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1299395739


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,670 @@
+/*
+ * 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.druid.client.indexing.TaskPayloadResponse;
+import org.apache.druid.client.indexing.TaskStatusResponse;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.error.ErrorResponse;
+import org.apache.druid.error.QueryExceptionCompat;
+import org.apache.druid.guice.annotations.MSQ;
+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.sql.MSQTaskQueryMaker;
+import org.apache.druid.msq.sql.MSQTaskSqlEngine;
+import org.apache.druid.msq.sql.SqlStatementState;
+import org.apache.druid.msq.sql.entity.ColumnNameAndTypes;
+import org.apache.druid.msq.sql.entity.ResultSetInformation;
+import org.apache.druid.msq.sql.entity.SqlStatementResult;
+import org.apache.druid.msq.util.SqlStatementResourceHelper;
+import org.apache.druid.query.ExecutionMode;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.query.QueryException;
+import org.apache.druid.rpc.indexing.OverlordClient;
+import org.apache.druid.server.QueryResponse;
+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.SqlRowTransformer;
+import org.apache.druid.sql.SqlStatementFactory;
+import org.apache.druid.sql.http.ResultFormat;
+import org.apache.druid.sql.http.SqlQuery;
+import org.apache.druid.sql.http.SqlResource;
+
+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.Collections;
+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 AuthorizerMapper authorizerMapper;
+  private final ObjectMapper jsonMapper;
+  private final OverlordClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final OverlordClient overlordClient
+  )
+  {
+    this.msqSqlStatementFactory = msqSqlStatementFactory;
+    this.authorizerMapper = authorizerMapper;
+    this.jsonMapper = jsonMapper;
+    this.overlordClient = overlordClient;
+  }
+
+
+  @GET
+  @Path("/enabled")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response isEnabled(@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(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              StringUtils.format(
+                                  "The statement sql api only supports sync mode[%s]. Please set context parameter [%s=%s] in the context payload",
+                                  ExecutionMode.ASYNC,
+                                  QueryContexts.CTX_EXECUTION_MODE,
+                                  ExecutionMode.ASYNC
+                              )
+                          )
+        );
+      }
+
+
+      Thread.currentThread().setName(StringUtils.format("statement_sql[%s]", sqlQueryId));
+
+      final DirectStatement.ResultSet plan = stmt.plan();
+      // in case the engine is async, the query is not run yet. We just return the taskID in case of non explain queries.
+      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);
+      }
+    }
+    catch (DruidException e) {
+      stmt.reporter().failed(e);
+      return buildNonOkResponse(e);
+    }
+    catch (QueryException queryException) {
+      stmt.reporter().failed(queryException);
+      final DruidException underlyingException = DruidException.fromFailure(new QueryExceptionCompat(queryException));
+      return buildNonOkResponse(underlyingException);
+    }
+    catch (ForbiddenException e) {
+      log.debug("Got forbidden request for reason [%s]", e.getErrorMessage());
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.USER)
+                        .ofCategory(DruidException.Category.FORBIDDEN)
+                        .build(Access.DEFAULT_ERROR_MESSAGE)
+      );
+    }
+    // Calcite throws java.lang.AssertionError at various points in planning/validation.
+    catch (AssertionError | Exception e) {
+      stmt.reporter().failed(e);
+      log.noStackTrace().warn(e, "Failed to handle query: %s", sqlQueryId);
+
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.DEVELOPER)
+                        .ofCategory(DruidException.Category.UNCATEGORIZED)
+                        .build(e.getMessage())
+      );
+    }
+    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 {
+      Access authResult = AuthorizationUtils.authorizeAllResourceActions(
+          req,
+          Collections.emptyList(),
+          authorizerMapper
+      );
+      if (!authResult.isAllowed()) {
+        throw new ForbiddenException(authResult.toString());
+      }
+      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 (DruidException e) {
+      return buildNonOkResponse(e);
+    }
+    catch (ForbiddenException e) {
+      log.debug("Got forbidden request for reason [%s]", e.getErrorMessage());
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.USER)
+                        .ofCategory(DruidException.Category.FORBIDDEN)
+                        .build(Access.DEFAULT_ERROR_MESSAGE)
+      );
+    }
+    catch (Exception e) {
+      log.noStackTrace().warn(e, "Failed to handle query: %s", queryId);
+      return buildNonOkResponse(DruidException.forPersona(DruidException.Persona.DEVELOPER)
+                                              .ofCategory(DruidException.Category.UNCATEGORIZED)
+                                              .build(e, "Failed to handle query: [%s]", queryId));
+    }
+  }
+
+  @GET
+  @Path("/{id}/results")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response doGetResults(
+      @PathParam("id") final String queryId,
+      @QueryParam("offset") Long offset,
+      @QueryParam("numRows") Long numberOfRows,
+      @Context final HttpServletRequest req
+  )
+  {
+    try {
+      Access authResult = AuthorizationUtils.authorizeAllResourceActions(
+          req,
+          Collections.emptyList(),
+          authorizerMapper
+      );
+      if (!authResult.isAllowed()) {
+        throw new ForbiddenException(authResult.toString());
+      }
+      final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req);
+
+      if (offset != null && offset < 0) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "offset cannot be negative. Please pass a positive number."
+                          )
+        );
+      }
+      if (numberOfRows != null && numberOfRows < 0) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "numRows cannot be negative. Please pass a positive number."
+                          )
+        );
+      }
+
+      final long start = offset == null ? 0 : offset;
+      final long last = SqlStatementResourceHelper.getLastIndex(numberOfRows, start);
+
+      TaskStatusResponse taskResponse = contactOverlord(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();
+      }
+
+      MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, authenticationResult.getIdentity());
+      SqlStatementState sqlStatementState = SqlStatementResourceHelper.getSqlStatementState(statusPlus);
+
+      if (sqlStatementState == SqlStatementState.RUNNING || sqlStatementState == SqlStatementState.ACCEPTED) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "Query[%s] is currently in [%s] state. Please wait for it to complete.",
+                              queryId,
+                              sqlStatementState
+                          )
+        );
+      } else if (sqlStatementState == SqlStatementState.FAILED) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "Query[%s] failed. Hit status api for more details.",
+                              queryId
+                          )
+        );
+      } else {
+        Optional<List<ColumnNameAndTypes>> signature = SqlStatementResourceHelper.getSignature(msqControllerTask);
+        if (!signature.isPresent()) {
+          return Response.ok().build();
+        }
+        Optional<List<Object>> results = SqlStatementResourceHelper.getResults(SqlStatementResourceHelper.getPayload(
+            contactOverlord(overlordClient.taskReportAsMap(queryId))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColumnNameAndTypes> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (long k = start; k < Math.min(last, results.get().size()); k++) {
+              writer.writeRowStart();
+              for (int i = 0; i < rowSignature.size(); i++) {
+                writer.writeRowField(
+                    rowSignature.get(i).getColName(),
+                    ((List) results.get().get(Math.toIntExact(k))).get(i)
+                );
+              }
+              writer.writeRowEnd();
+            }
+
+            writer.writeResponseEnd();
+          }
+          catch (Exception e) {
+            log.error(e, "Unable to stream results back for query[%s]", queryId);
+            throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+          }
+        }).build();
+
+      }
+    }
+    catch (DruidException e) {
+      return buildNonOkResponse(e);
+    }
+    catch (ForbiddenException e) {
+      log.debug("Got forbidden request for reason [%s]", e.getErrorMessage());
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.USER)
+                        .ofCategory(DruidException.Category.FORBIDDEN)
+                        .build(Access.DEFAULT_ERROR_MESSAGE)
+      );
+    }
+    catch (Exception e) {
+      log.noStackTrace().warn(e, "Failed to handle query: %s", queryId);
+      return buildNonOkResponse(DruidException.forPersona(DruidException.Persona.DEVELOPER)
+                                              .ofCategory(DruidException.Category.UNCATEGORIZED)
+                                              .build(e, "Failed to handle query: [%s]", queryId));
+    }
+  }
+
+  /**
+   * Queries can be canceled while in any {@link SqlStatementState}. Canceling a query that has already completed will be a no-op.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@PathParam("id") final String queryId, @Context final HttpServletRequest req)
+  {
+
+    try {
+      Access authResult = AuthorizationUtils.authorizeAllResourceActions(

Review Comment:
   https://github.com/apache/druid/pull/14878



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] abhishekagarwal87 commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["abhishekagarwal87 (via GitHub)" <gi...@apache.org>]

abhishekagarwal87 commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1245054469


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          queryId
+      );
+    }
+
+    final long start = offset == null ? 0 : offset;
+    final long last = getLastIndex(numberOfRows, start);
+
+    try {
+      AuthorizationUtils.authorizeAllResourceActions(req, Collections.emptyList(), authorizerMapper);

Review Comment:
   if all users are authorized, why do we even call the authorize call? And an authenticated user will still require some permission to call the API. 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on PR #14416:
URL: https://github.com/apache/druid/pull/14416#issuecomment-1611305339

   Thanks @abhishekagarwal87 @adarshsanjeev for the reviews. User Facing docs would be coming soon. 


-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1244660295


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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);

Review Comment:
   Yes We needed to. 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1244665872


##########
server/src/main/java/org/apache/druid/rpc/indexing/OverlordClient.java:
##########
@@ -52,5 +53,7 @@ public interface OverlordClient
 
   ListenableFuture<Map<String, Object>> taskReportAsMap(String taskId);
 
+  ListenableFuture<TaskPayloadResponse> taskPayload(String taskId);

Review Comment:
   This is borrowed for the existing indexingServiceClient hence I did not want to change the name. 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1244745645


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,767 @@
+/*
+ * 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.inject.Inject;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.druid.client.indexing.IndexingServiceClient;
+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.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.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.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.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 IndexingServiceClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final ServerConfig serverConfig,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final IndexingServiceClient 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
+            ),
+            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),
+          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),
+          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 = overlordClient.getTaskStatus(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
+            ),
+            queryId
+        );
+      } else if (sqlStatementState == SqlStatementState.FAILED) {
+        return buildNonOkResponse(
+            Response.Status.NOT_FOUND.getStatusCode(),
+            new QueryException(null, statusPlus.getErrorMsg(), null, null),
+            queryId
+        );
+      } else {
+        MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, authenticationResult.getIdentity());
+        Optional<List<ColNameAndType>> signature = getSignature(msqControllerTask);
+        Optional<List<Object>> results = getResults(overlordClient.getTaskReport(queryId));
+
+        return Response.ok(
+            (StreamingOutput) outputStream -> {
+              CountingOutputStream os = new CountingOutputStream(outputStream);
+
+              try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+                List<ColNameAndType> rowSignature = signature.get();
+                writer.writeResponseStart();
+
+                for (int k = (int) start; k < Math.min(last, results.get().size()); k++) {
+                  writer.writeRowStart();
+                  for (int i = 0; i < rowSignature.size(); i++) {
+                    writer.writeRowField(rowSignature.get(i).getColName(), ((List) results.get().get(k)).get(i));
+                  }
+                  writer.writeRowEnd();
+                }
+
+                writer.writeResponseEnd();
+              }
+              catch (Exception e) {
+                log.error(e, "Unable to stream results back for query[%s]", queryId);
+                throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+              }
+            }
+        ).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;
+    }
+  }
+
+  /**
+   * Queries can be canceled while in anystate. Canceling a query that has already completed will remove its results.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@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(),
+          false
+      );
+      if (sqlStatementResult.isPresent()) {
+        switch (sqlStatementResult.get().getState()) {
+          case ACCEPTED:
+          case RUNNING:
+            overlordClient.cancelTask(queryId);
+            return Response.status(Response.Status.ACCEPTED).build();
+          case SUCCESS:
+          case FAILED:
+            // we would also want to clean up the
+            return Response.ok().build();
+          default:
+            throw new ISE("Illegal State[%s] encountered", sqlStatementResult.get().getState());
+        }
+
+      } 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;
+    }
+  }
+
+  private Response buildStandardResponse(
+      Sequence<Object[]> sequence,
+      SqlQuery sqlQuery,
+      String sqlQueryId,
+      SqlRowTransformer rowTransformer
+  ) throws IOException
+  {
+    final Yielder<Object[]> yielder0 = Yielders.each(sequence);
+
+    try {
+      final Response.ResponseBuilder responseBuilder = Response
+          .ok(
+              (StreamingOutput) outputStream -> {
+                CountingOutputStream os = new CountingOutputStream(outputStream);
+                Yielder<Object[]> yielder = yielder0;
+
+                try (final ResultFormat.Writer writer = sqlQuery.getResultFormat()
+                                                                .createFormatter(os, jsonMapper)) {
+                  writer.writeResponseStart();
+
+                  if (sqlQuery.includeHeader()) {
+                    writer.writeHeader(
+                        rowTransformer.getRowType(),
+                        sqlQuery.includeTypesHeader(),
+                        sqlQuery.includeSqlTypesHeader()
+                    );
+                  }
+
+                  while (!yielder.isDone()) {
+                    final Object[] row = yielder.get();
+                    writer.writeRowStart();
+                    for (int i = 0; i < rowTransformer.getFieldList().size(); i++) {
+                      final Object value = rowTransformer.transform(row, i);
+                      writer.writeRowField(rowTransformer.getFieldList().get(i), value);
+                    }
+                    writer.writeRowEnd();
+                    yielder = yielder.next(null);
+                  }
+
+                  writer.writeResponseEnd();
+                }
+                catch (Exception e) {
+                  log.error(e, "Unable to send SQL response [%s]", sqlQueryId);
+                  throw new RuntimeException(e);
+                }
+                finally {
+                  yielder.close();
+                }
+              }
+          );
+
+      if (sqlQuery.includeHeader()) {
+        responseBuilder.header(SqlResource.SQL_HEADER_RESPONSE_HEADER, SqlResource.SQL_HEADER_VALUE);
+      }
+
+      return responseBuilder.build();
+    }
+    catch (Throwable e) {
+      // make sure to close yielder if anything happened before starting to serialize the response.
+      yielder0.close();
+      throw e;
+    }
+  }
+
+  private Response buildTaskResponse(Sequence<Object[]> sequence, String user)
+  {
+    List<Object[]> rows = sequence.toList();
+    int numRows = rows.size();
+    if (numRows != 1) {
+      throw new RE("Expected a single row but got [%d] rows. Please check broker logs for more information.", numRows);
+    }
+    String taskId = (String) rows.get(0)[0];
+    try {
+      Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, user, true);
+
+      if (statementResult.isPresent()) {
+        return Response.status(Response.Status.OK).entity(statementResult.get()).build();
+      } else {
+        return buildNonOkResponse(
+            Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(),
+            new QueryException(
+                null,
+                StringUtils.format(
+                    "Unable to find associated task for query id %s",
+                    taskId
+                ),
+                null,
+                null
+            ),
+            taskId
+        );
+      }
+    }
+    catch (QueryException e) {
+      return buildNonOkResponse(Response.Status.OK.getStatusCode(), e, taskId);
+    }
+  }
+
+  private Response buildNonOkResponse(int status, SanitizableException e, String sqlQueryId)
+  {
+    // Though transformIfNeeded returns an exception, its purpose is to return
+    // a QueryException
+    Exception cleaned = serverConfig
+        .getErrorResponseTransformStrategy()
+        .transformIfNeeded(e);
+    return Response
+        .status(status)
+        .entity(cleaned)
+        .build();
+  }
+
+
+  private static Optional<List<ColNameAndType>> getSignature(
+      MSQControllerTask msqControllerTask
+  )
+  {
+    // only populate signature for select q's
+    if (msqControllerTask.getQuerySpec().getDestination().getClass() == TaskReportMSQDestination.class) {
+      ColumnMappings columnMappings = msqControllerTask.getQuerySpec().getColumnMappings();
+      List<SqlTypeName> sqlTypeNames = msqControllerTask.getSqlTypeNames();
+      if (sqlTypeNames == null) {
+        return Optional.empty();
+      }
+      List<ColNameAndType> signature = new ArrayList<>(columnMappings.size());
+      int index = 0;
+      for (String colName : columnMappings.getOutputColumnNames()) {
+        signature.add(new ColNameAndType(colName, sqlTypeNames.get(index).getName()));
+        index++;
+      }
+      return Optional.of(signature);
+    }
+    return Optional.empty();
+  }
+
+
+  private void checkTaskPayloadOrThrow(TaskPayloadResponse taskPayloadResponse, String queryId) throws QueryException
+  {
+    if (taskPayloadResponse == null || taskPayloadResponse.getPayload() == null) {
+      throw new QueryException(QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+                               StringUtils.format(
+                                   "Could not get payload details of query[%s] from the overlord",
+                                   queryId
+                               ), null, null
+      );
+    }
+
+    if (MSQControllerTask.class != taskPayloadResponse.getPayload().getClass()) {
+      throw new QueryException(QueryException.UNKNOWN_EXCEPTION_ERROR_CODE,
+                               StringUtils.format(
+                                   "Fetched an unexpected payload details of query[%s] from the overlord.",
+                                   queryId
+                               ), null, null
+      );
+    }
+  }
+
+  private Optional<ResultSetInformation> getSampleResults(
+      String asyncResultId,
+      boolean isSelectQuery,
+      SqlStatementState sqlStatementState
+  )
+  {
+    // only populate sample results in case a select query is successful
+    if (isSelectQuery && sqlStatementState == SqlStatementState.SUCCESS) {
+      Map<String, Object> report = overlordClient.getTaskReport(asyncResultId);
+      Optional<List<Object>> rows = getResults(report);
+
+      if (rows.isPresent()) {
+        return Optional.of(new ResultSetInformation(
+            null,
+            false,
+            (long) rows.get().size(),
+            null,
+            rows.get()
+        ));
+      } else {
+        return Optional.empty();
+      }
+
+    } else {
+      return Optional.empty();
+    }
+  }
+
+
+  private Optional<SqlStatementResult> getStatementStatus(String queryId, String currentUser, boolean withResults)
+      throws QueryException, ForbiddenException
+  {
+    TaskStatusResponse taskResponse = overlordClient.getTaskStatus(queryId);
+    if (taskResponse == null) {
+      return Optional.empty();
+    }
+
+    TaskStatusPlus statusPlus = taskResponse.getStatus();
+    if (statusPlus == null || !MSQControllerTask.TYPE.equals(statusPlus.getType())) {
+      return Optional.empty();
+    }

Review Comment:
   I think we should not add a warning since its bad user input. 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] adarshsanjeev commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["adarshsanjeev (via GitHub)" <gi...@apache.org>]

adarshsanjeev commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1244861332


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,673 @@
+/*
+ * 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.druid.client.indexing.TaskPayloadResponse;
+import org.apache.druid.client.indexing.TaskStatusResponse;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.error.ErrorResponse;
+import org.apache.druid.error.QueryExceptionCompat;
+import org.apache.druid.guice.annotations.MSQ;
+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.sql.MSQTaskQueryMaker;
+import org.apache.druid.msq.sql.MSQTaskSqlEngine;
+import org.apache.druid.msq.sql.SqlStatementState;
+import org.apache.druid.msq.sql.entity.ColumnNameAndTypes;
+import org.apache.druid.msq.sql.entity.ResultSetInformation;
+import org.apache.druid.msq.sql.entity.SqlStatementResult;
+import org.apache.druid.msq.util.SqlStatementResourceHelper;
+import org.apache.druid.query.ExecutionMode;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.query.QueryException;
+import org.apache.druid.rpc.indexing.OverlordClient;
+import org.apache.druid.server.QueryResponse;
+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.SqlRowTransformer;
+import org.apache.druid.sql.SqlStatementFactory;
+import org.apache.druid.sql.http.ResultFormat;
+import org.apache.druid.sql.http.SqlQuery;
+import org.apache.druid.sql.http.SqlResource;
+
+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.Collections;
+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 AuthorizerMapper authorizerMapper;
+  private final ObjectMapper jsonMapper;
+  private final OverlordClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final OverlordClient overlordClient
+  )
+  {
+    this.msqSqlStatementFactory = msqSqlStatementFactory;
+    this.authorizerMapper = authorizerMapper;
+    this.jsonMapper = jsonMapper;
+    this.overlordClient = overlordClient;
+  }
+
+
+  @GET
+  @Path("/enabled")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response isEnabled(@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(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              StringUtils.format(
+                                  "The statement sql api only supports sync mode[%s]. Please set context parameter [%s=%s] in the context payload",
+                                  ExecutionMode.ASYNC,
+                                  QueryContexts.CTX_EXECUTION_MODE,
+                                  ExecutionMode.ASYNC
+                              )
+                          )
+        );
+      }
+
+
+      Thread.currentThread().setName(StringUtils.format("statement_sql[%s]", sqlQueryId));
+
+      final DirectStatement.ResultSet plan = stmt.plan();
+      // in case the engine is async, the query is not run yet. We just return the taskID in case of non explain queries.
+      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);
+      }
+    }
+    catch (DruidException e) {
+      stmt.reporter().failed(e);
+      return buildNonOkResponse(e);
+    }
+    catch (QueryException queryException) {
+      stmt.reporter().failed(queryException);
+      final DruidException underlyingException = DruidException.fromFailure(new QueryExceptionCompat(queryException));
+      return buildNonOkResponse(underlyingException);
+    }
+    catch (ForbiddenException e) {
+      log.debug("Got forbidden request for reason [%s]", e.getErrorMessage());
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.USER)
+                        .ofCategory(DruidException.Category.FORBIDDEN)
+                        .build(Access.DEFAULT_ERROR_MESSAGE)
+      );
+    }
+    // Calcite throws java.lang.AssertionError at various points in planning/validation.
+    catch (AssertionError | Exception e) {
+      stmt.reporter().failed(e);
+      log.noStackTrace().warn(e, "Failed to handle query: %s", sqlQueryId);
+
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.DEVELOPER)
+                        .ofCategory(DruidException.Category.UNCATEGORIZED)
+                        .build(e.getMessage())
+      );
+    }
+    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 {
+      Access authResult = AuthorizationUtils.authorizeAllResourceActions(
+          req,
+          Collections.emptyList(),
+          authorizerMapper
+      );
+      if (!authResult.isAllowed()) {
+        throw new ForbiddenException(authResult.toString());
+      }
+      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 (DruidException e) {
+      return buildNonOkResponse(e);
+    }
+    catch (ForbiddenException e) {
+      log.debug("Got forbidden request for reason [%s]", e.getErrorMessage());
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.USER)
+                        .ofCategory(DruidException.Category.FORBIDDEN)
+                        .build(Access.DEFAULT_ERROR_MESSAGE)
+      );
+    }
+    catch (Exception e) {
+      log.noStackTrace().warn(e, "Failed to handle query: %s", queryId);
+      return buildNonOkResponse(DruidException.forPersona(DruidException.Persona.DEVELOPER)
+                                              .ofCategory(DruidException.Category.UNCATEGORIZED)
+                                              .build(e, "Failed to handle query: [%s]", queryId));
+    }
+  }
+
+  @GET
+  @Path("/{id}/results")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response doGetResults(
+      @PathParam("id") final String queryId,
+      @QueryParam("offset") Long offset,
+      @QueryParam("numRows") Long numberOfRows,
+      @Context final HttpServletRequest req
+  )
+  {
+    try {
+      Access authResult = AuthorizationUtils.authorizeAllResourceActions(
+          req,
+          Collections.emptyList(),
+          authorizerMapper
+      );
+      if (!authResult.isAllowed()) {
+        throw new ForbiddenException(authResult.toString());
+      }
+      final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req);
+
+      if (offset != null && offset < 0) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "offset cannot be negative. Please pass a positive number."
+                          )
+        );
+      }
+      if (numberOfRows != null && numberOfRows < 0) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "numRows cannot be negative. Please pass a positive number."
+                          )
+        );
+      }
+
+      final long start = offset == null ? 0 : offset;
+      final long last = SqlStatementResourceHelper.getLastIndex(numberOfRows, start);
+
+
+      getStatementStatus(queryId, authenticationResult.getIdentity(), false);

Review Comment:
   Is this call required? It seems to be checked after this already.



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org

[GitHub] [druid] cryptoe commented on a diff in pull request #14416: Adding Interactive API's for MSQ engine

["cryptoe (via GitHub)" <gi...@apache.org>]

cryptoe commented on code in PR #14416:
URL: https://github.com/apache/druid/pull/14416#discussion_r1245103222


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java:
##########
@@ -0,0 +1,670 @@
+/*
+ * 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.druid.client.indexing.TaskPayloadResponse;
+import org.apache.druid.client.indexing.TaskStatusResponse;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.error.ErrorResponse;
+import org.apache.druid.error.QueryExceptionCompat;
+import org.apache.druid.guice.annotations.MSQ;
+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.sql.MSQTaskQueryMaker;
+import org.apache.druid.msq.sql.MSQTaskSqlEngine;
+import org.apache.druid.msq.sql.SqlStatementState;
+import org.apache.druid.msq.sql.entity.ColumnNameAndTypes;
+import org.apache.druid.msq.sql.entity.ResultSetInformation;
+import org.apache.druid.msq.sql.entity.SqlStatementResult;
+import org.apache.druid.msq.util.SqlStatementResourceHelper;
+import org.apache.druid.query.ExecutionMode;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.query.QueryException;
+import org.apache.druid.rpc.indexing.OverlordClient;
+import org.apache.druid.server.QueryResponse;
+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.SqlRowTransformer;
+import org.apache.druid.sql.SqlStatementFactory;
+import org.apache.druid.sql.http.ResultFormat;
+import org.apache.druid.sql.http.SqlQuery;
+import org.apache.druid.sql.http.SqlResource;
+
+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.Collections;
+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 AuthorizerMapper authorizerMapper;
+  private final ObjectMapper jsonMapper;
+  private final OverlordClient overlordClient;
+
+
+  @Inject
+  public SqlStatementResource(
+      final @MSQ SqlStatementFactory msqSqlStatementFactory,
+      final AuthorizerMapper authorizerMapper,
+      final ObjectMapper jsonMapper,
+      final OverlordClient overlordClient
+  )
+  {
+    this.msqSqlStatementFactory = msqSqlStatementFactory;
+    this.authorizerMapper = authorizerMapper;
+    this.jsonMapper = jsonMapper;
+    this.overlordClient = overlordClient;
+  }
+
+
+  @GET
+  @Path("/enabled")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response isEnabled(@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(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              StringUtils.format(
+                                  "The statement sql api only supports sync mode[%s]. Please set context parameter [%s=%s] in the context payload",
+                                  ExecutionMode.ASYNC,
+                                  QueryContexts.CTX_EXECUTION_MODE,
+                                  ExecutionMode.ASYNC
+                              )
+                          )
+        );
+      }
+
+
+      Thread.currentThread().setName(StringUtils.format("statement_sql[%s]", sqlQueryId));
+
+      final DirectStatement.ResultSet plan = stmt.plan();
+      // in case the engine is async, the query is not run yet. We just return the taskID in case of non explain queries.
+      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);
+      }
+    }
+    catch (DruidException e) {
+      stmt.reporter().failed(e);
+      return buildNonOkResponse(e);
+    }
+    catch (QueryException queryException) {
+      stmt.reporter().failed(queryException);
+      final DruidException underlyingException = DruidException.fromFailure(new QueryExceptionCompat(queryException));
+      return buildNonOkResponse(underlyingException);
+    }
+    catch (ForbiddenException e) {
+      log.debug("Got forbidden request for reason [%s]", e.getErrorMessage());
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.USER)
+                        .ofCategory(DruidException.Category.FORBIDDEN)
+                        .build(Access.DEFAULT_ERROR_MESSAGE)
+      );
+    }
+    // Calcite throws java.lang.AssertionError at various points in planning/validation.
+    catch (AssertionError | Exception e) {
+      stmt.reporter().failed(e);
+      log.noStackTrace().warn(e, "Failed to handle query: %s", sqlQueryId);
+
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.DEVELOPER)
+                        .ofCategory(DruidException.Category.UNCATEGORIZED)
+                        .build(e.getMessage())
+      );
+    }
+    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 {
+      Access authResult = AuthorizationUtils.authorizeAllResourceActions(
+          req,
+          Collections.emptyList(),
+          authorizerMapper
+      );
+      if (!authResult.isAllowed()) {
+        throw new ForbiddenException(authResult.toString());
+      }
+      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 (DruidException e) {
+      return buildNonOkResponse(e);
+    }
+    catch (ForbiddenException e) {
+      log.debug("Got forbidden request for reason [%s]", e.getErrorMessage());
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.USER)
+                        .ofCategory(DruidException.Category.FORBIDDEN)
+                        .build(Access.DEFAULT_ERROR_MESSAGE)
+      );
+    }
+    catch (Exception e) {
+      log.noStackTrace().warn(e, "Failed to handle query: %s", queryId);
+      return buildNonOkResponse(DruidException.forPersona(DruidException.Persona.DEVELOPER)
+                                              .ofCategory(DruidException.Category.UNCATEGORIZED)
+                                              .build(e, "Failed to handle query: [%s]", queryId));
+    }
+  }
+
+  @GET
+  @Path("/{id}/results")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response doGetResults(
+      @PathParam("id") final String queryId,
+      @QueryParam("offset") Long offset,
+      @QueryParam("numRows") Long numberOfRows,
+      @Context final HttpServletRequest req
+  )
+  {
+    try {
+      Access authResult = AuthorizationUtils.authorizeAllResourceActions(
+          req,
+          Collections.emptyList(),
+          authorizerMapper
+      );
+      if (!authResult.isAllowed()) {
+        throw new ForbiddenException(authResult.toString());
+      }
+      final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req);
+
+      if (offset != null && offset < 0) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "offset cannot be negative. Please pass a positive number."
+                          )
+        );
+      }
+      if (numberOfRows != null && numberOfRows < 0) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "numRows cannot be negative. Please pass a positive number."
+                          )
+        );
+      }
+
+      final long start = offset == null ? 0 : offset;
+      final long last = SqlStatementResourceHelper.getLastIndex(numberOfRows, start);
+
+      TaskStatusResponse taskResponse = contactOverlord(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();
+      }
+
+      MSQControllerTask msqControllerTask = getMSQControllerTaskOrThrow(queryId, authenticationResult.getIdentity());
+      SqlStatementState sqlStatementState = SqlStatementResourceHelper.getSqlStatementState(statusPlus);
+
+      if (sqlStatementState == SqlStatementState.RUNNING || sqlStatementState == SqlStatementState.ACCEPTED) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "Query[%s] is currently in [%s] state. Please wait for it to complete.",
+                              queryId,
+                              sqlStatementState
+                          )
+        );
+      } else if (sqlStatementState == SqlStatementState.FAILED) {
+        return buildNonOkResponse(
+            DruidException.forPersona(DruidException.Persona.USER)
+                          .ofCategory(DruidException.Category.INVALID_INPUT)
+                          .build(
+                              "Query[%s] failed. Hit status api for more details.",
+                              queryId
+                          )
+        );
+      } else {
+        Optional<List<ColumnNameAndTypes>> signature = SqlStatementResourceHelper.getSignature(msqControllerTask);
+        if (!signature.isPresent()) {
+          return Response.ok().build();
+        }
+        Optional<List<Object>> results = SqlStatementResourceHelper.getResults(SqlStatementResourceHelper.getPayload(
+            contactOverlord(overlordClient.taskReportAsMap(queryId))));
+
+        return Response.ok((StreamingOutput) outputStream -> {
+          CountingOutputStream os = new CountingOutputStream(outputStream);
+
+          try (final ResultFormat.Writer writer = ResultFormat.OBJECT.createFormatter(os, jsonMapper)) {
+            List<ColumnNameAndTypes> rowSignature = signature.get();
+            writer.writeResponseStart();
+
+            for (long k = start; k < Math.min(last, results.get().size()); k++) {
+              writer.writeRowStart();
+              for (int i = 0; i < rowSignature.size(); i++) {
+                writer.writeRowField(
+                    rowSignature.get(i).getColName(),
+                    ((List) results.get().get(Math.toIntExact(k))).get(i)
+                );
+              }
+              writer.writeRowEnd();
+            }
+
+            writer.writeResponseEnd();
+          }
+          catch (Exception e) {
+            log.error(e, "Unable to stream results back for query[%s]", queryId);
+            throw new ISE(e, "Unable to stream results back for query[%s]", queryId);
+          }
+        }).build();
+
+      }
+    }
+    catch (DruidException e) {
+      return buildNonOkResponse(e);
+    }
+    catch (ForbiddenException e) {
+      log.debug("Got forbidden request for reason [%s]", e.getErrorMessage());
+      return buildNonOkResponse(
+          DruidException.forPersona(DruidException.Persona.USER)
+                        .ofCategory(DruidException.Category.FORBIDDEN)
+                        .build(Access.DEFAULT_ERROR_MESSAGE)
+      );
+    }
+    catch (Exception e) {
+      log.noStackTrace().warn(e, "Failed to handle query: %s", queryId);
+      return buildNonOkResponse(DruidException.forPersona(DruidException.Persona.DEVELOPER)
+                                              .ofCategory(DruidException.Category.UNCATEGORIZED)
+                                              .build(e, "Failed to handle query: [%s]", queryId));
+    }
+  }
+
+  /**
+   * Queries can be canceled while in any {@link SqlStatementState}. Canceling a query that has already completed will be a no-op.
+   *
+   * @param queryId queryId
+   * @param req     httpServletRequest
+   * @return HTTP 404 if the query ID does not exist,expired or originated by different user. HTTP 202 if the deletion
+   * request has been accepted.
+   */
+  @DELETE
+  @Path("/{id}")
+  @Produces(MediaType.APPLICATION_JSON)
+  public Response deleteQuery(@PathParam("id") final String queryId, @Context final HttpServletRequest req)
+  {
+
+    try {
+      Access authResult = AuthorizationUtils.authorizeAllResourceActions(

Review Comment:
   @abhishekagarwal87 I am following the pattern from https://github.com/apache/druid/blob/baa64e6d8a71b1d3816dea6144ca4322d0432e53/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/SqlTaskResource.java#L110. 
   
   We might need to add a new method to AuthorizationUtilsto set the req header but that can be done as part of a follow up PR. 



-- 
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: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org