SQL: Parse queries immediately for non-JDBC endpoints.

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java

@@ -85,6 +85,7 @@
 import org.apache.druid.server.security.ResourceAction;
 import org.apache.druid.sql.DirectStatement;
 import org.apache.druid.sql.HttpStatement;
+import org.apache.druid.sql.SqlQueryPlus;
 import org.apache.druid.sql.SqlRowTransformer;
 import org.apache.druid.sql.SqlStatementFactory;
 import org.apache.druid.sql.http.ResultFormat;
@@ -174,17 +175,29 @@ public Response doPost(@Context final HttpServletRequest req,
   }
 
   @VisibleForTesting
-  Response doPost(final SqlQuery sqlQuery,
-                  final HttpServletRequest req)
+  Response doPost(
+      SqlQuery sqlQuery, // Not final: reassigned using createModifiedSqlQuery
+      final HttpServletRequest req
+  )
   {
-    SqlQuery modifiedQuery = createModifiedSqlQuery(sqlQuery);
+    final SqlQueryPlus sqlQueryPlus;
+    final HttpStatement stmt;
+    final QueryContext queryContext;
+
+    try {
+      sqlQuery = createModifiedSqlQuery(sqlQuery);
+      sqlQueryPlus = SqlResource.makeSqlQueryPlus(sqlQuery, req);
+      queryContext = QueryContext.of(sqlQueryPlus.context());
+      stmt = msqSqlStatementFactory.httpStatement(SqlResource.makeSqlQueryPlus(sqlQuery, req), req);
+    }
+    catch (Exception e) {
+      return SqlResource.handleExceptionBeforeStatementCreated(e, sqlQuery.queryContext());
+    }
 
-    final HttpStatement stmt = msqSqlStatementFactory.httpStatement(modifiedQuery, req);
     final String sqlQueryId = stmt.sqlQueryId();
     final String currThreadName = Thread.currentThread().getName();
     boolean isDebug = false;
     try {
-      QueryContext queryContext = QueryContext.of(modifiedQuery.getContext());
       isDebug = queryContext.isDebug();
       contextChecks(queryContext);
 
@@ -202,7 +215,7 @@ Response doPost(final SqlQuery sqlQuery,
         return buildTaskResponse(sequence, stmt.query().authResult());
       } else {
         // Used for EXPLAIN
-        return buildStandardResponse(sequence, modifiedQuery, sqlQueryId, rowTransformer);
+        return buildStandardResponse(sequence, sqlQuery, sqlQueryId, rowTransformer);
       }
     }
     catch (DruidException e) {

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlTaskResource.java

@@ -43,6 +43,7 @@
 import org.apache.druid.server.security.ForbiddenException;
 import org.apache.druid.sql.DirectStatement;
 import org.apache.druid.sql.HttpStatement;
+import org.apache.druid.sql.SqlQueryPlus;
 import org.apache.druid.sql.SqlRowTransformer;
 import org.apache.druid.sql.SqlStatementFactory;
 import org.apache.druid.sql.http.ResultFormat;
@@ -127,7 +128,16 @@ public Response doPost(
   )
   {
     // Queries run as MSQ tasks look like regular queries, but return the task ID as their only output.
-    final HttpStatement stmt = sqlStatementFactory.httpStatement(sqlQuery, req);
+    final SqlQueryPlus sqlQueryPlus;
+    final HttpStatement stmt;
+    try {
+      sqlQueryPlus = SqlResource.makeSqlQueryPlus(sqlQuery, req);
+      stmt = sqlStatementFactory.httpStatement(sqlQueryPlus, req);
+    }
+    catch (Exception e) {
+      return SqlResource.handleExceptionBeforeStatementCreated(e, sqlQuery.queryContext());
+    }
+
     final String sqlQueryId = stmt.sqlQueryId();
     try {
       final DirectStatement.ResultSet plan = stmt.plan();

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/ResultsContextSerdeTest.java

@@ -31,6 +31,7 @@
 import org.apache.druid.query.QuerySegmentWalker;
 import org.apache.druid.query.policy.NoopPolicyEnforcer;
 import org.apache.druid.server.security.AuthConfig;
+import org.apache.druid.sql.calcite.parser.DruidSqlParser;
 import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
 import org.apache.druid.sql.calcite.planner.CatalogResolver;
 import org.apache.druid.sql.calcite.planner.PlannerConfig;
@@ -86,9 +87,11 @@ NamedViewSchema.NAME, new NamedViewSchema(EasyMock.createMock(ViewSchema.class))
         EasyMock.createMock(QueryRunnerFactoryConglomerate.class)
     );
 
+    final String sql = "SELECT 1";
     PlannerContext plannerContext = PlannerContext.create(
         toolbox,
-        "DUMMY",
+        sql,
+        DruidSqlParser.parse(sql, false).getMainStatement(),
         engine,
         Collections.emptyMap(),
         null

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java

@@ -810,12 +810,12 @@ private String runMultiStageQuery(
   )
   {
     final DirectStatement stmt = sqlStatementFactory.directStatement(
-        new SqlQueryPlus(
-            query,
-            context,
-            parameters,
-            authenticationResult
-        )
+        SqlQueryPlus.builder()
+                    .sql(query)
+                    .context(context)
+                    .parameters(parameters)
+                    .auth(authenticationResult)
+                    .build()
     );
 
     final List<Object[]> sequence = stmt.execute().getResults().toList();

server/src/main/java/org/apache/druid/server/QueryResultPusher.java

@@ -21,6 +21,7 @@
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.io.CountingOutputStream;
 import org.apache.druid.client.DirectDruidClient;
 import org.apache.druid.error.DruidException;
@@ -270,17 +271,14 @@ private Response handleDruidException(ResultsWriter resultsWriter, DruidExceptio
     }
 
     if (response == null) {
-      final Response.ResponseBuilder bob = Response
-          .status(e.getStatusCode())
-          .type(contentType)
-          .entity(new ErrorResponse(e));
-
-      bob.header(QueryResource.QUERY_ID_RESPONSE_HEADER, queryId);
-      for (Map.Entry<String, String> entry : extraHeaders.entrySet()) {
-        bob.header(entry.getKey(), entry.getValue());
-      }
-
-      return bob.build();
+      return handleDruidExceptionBeforeResponseStarted(
+          e,
+          contentType,
+          ImmutableMap.<String, String>builder()
+                      .putAll(extraHeaders)
+                      .put(QueryResource.QUERY_ID_RESPONSE_HEADER, queryId)
+                      .build()
+      );
     } else {
       if (response.isCommitted()) {
         QueryResource.NO_STACK_LOGGER.warn(e, "Response was committed without the accumulator writing anything!?");
@@ -302,6 +300,27 @@ private Response handleDruidException(ResultsWriter resultsWriter, DruidExceptio
     }
   }
 
+  /**
+   * Generates a response for a {@link DruidException} that occurs prior to any query results being sent out.
+   */
+  public static Response handleDruidExceptionBeforeResponseStarted(
+      final DruidException e,
+      final MediaType contentType,
+      final Map<String, String> extraHeaders
+  )
+  {
+    final Response.ResponseBuilder bob = Response
+        .status(e.getStatusCode())
+        .type(contentType)
+        .entity(new ErrorResponse(e));
+
+    for (Map.Entry<String, String> entry : extraHeaders.entrySet()) {
+      bob.header(entry.getKey(), entry.getValue());
+    }
+
+    return bob.build();
+  }
+
   public interface ResultsWriter extends Closeable
   {
     /**

sql/src/main/java/org/apache/druid/sql/DirectStatement.java

@@ -241,9 +241,9 @@ protected DruidPlanner createPlanner()
     return sqlToolbox.plannerFactory.createPlanner(
         sqlToolbox.engine,
         queryPlus.sql(),
+        queryPlus.sqlNode(),
         queryContext,
-        hook,
-        false
+        hook
     );
   }
 

sql/src/main/java/org/apache/druid/sql/HttpStatement.java

@@ -23,7 +23,6 @@
 import org.apache.druid.server.security.AuthorizationUtils;
 import org.apache.druid.server.security.ResourceAction;
 import org.apache.druid.sql.calcite.planner.DruidPlanner;
-import org.apache.druid.sql.http.SqlQuery;
 
 import javax.servlet.http.HttpServletRequest;
 import java.util.Set;
@@ -45,15 +44,13 @@ public class HttpStatement extends DirectStatement
 
   public HttpStatement(
       final SqlToolbox lifecycleToolbox,
-      final SqlQuery sqlQuery,
+      final SqlQueryPlus sqlQueryPlus,
       final HttpServletRequest req
   )
   {
     super(
         lifecycleToolbox,
-        SqlQueryPlus.builder(sqlQuery)
-                    .auth(AuthorizationUtils.authenticationResultFromRequest(req))
-                    .build(),
+        sqlQueryPlus,
         req.getRemoteAddr()
     );
     this.req = req;
@@ -65,9 +62,9 @@ protected DruidPlanner createPlanner()
     return sqlToolbox.plannerFactory.createPlanner(
         sqlToolbox.engine,
         queryPlus.sql(),
+        queryPlus.sqlNode(),
         queryContext,
-        hook,
-        true
+        hook
     );
   }
 

sql/src/main/java/org/apache/druid/sql/PreparedStatement.java

@@ -31,7 +31,6 @@
 public class PreparedStatement extends AbstractStatement
 {
   private final SqlQueryPlus originalRequest;
-  private PrepareResult prepareResult;
 
   public PreparedStatement(
       final SqlToolbox lifecycleToolbox,
@@ -70,8 +69,7 @@ public PrepareResult prepare()
       authorize(planner, authorizer());
 
       // Do the prepare step.
-      this.prepareResult = planner.prepare();
-      return prepareResult;
+      return planner.prepare();
     }
     catch (RuntimeException e) {
       reporter.failed(e);
@@ -92,7 +90,7 @@ public DirectStatement execute(List<TypedValue> parameters)
   {
     return new DirectStatement(
         sqlToolbox,
-        originalRequest.withParameters(parameters)
+        originalRequest.freshCopy().withParameters(parameters)
     );
   }
 
@@ -101,9 +99,9 @@ protected DruidPlanner getPlanner()
     return sqlToolbox.plannerFactory.createPlanner(
         sqlToolbox.engine,
         queryPlus.sql(),
+        queryPlus.freshCopy().sqlNode(),
         queryContext,
-        hook,
-        false
+        hook
     );
   }
 }