Expose HTTP Response headers from SqlResource

benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlBenchmark.java

@@ -518,7 +518,7 @@ public void querySql(Blackhole blackhole) throws Exception
     final String sql = QUERIES.get(Integer.parseInt(query));
     try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sql, new QueryContext(context))) {
       final PlannerResult plannerResult = planner.plan();
-      final Sequence<Object[]> resultSequence = plannerResult.run();
+      final Sequence<Object[]> resultSequence = plannerResult.run().getResults();
       final Object[] lastRow = resultSequence.accumulate(null, (accumulated, in) -> in);
       blackhole.consume(lastRow);
     }

benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlExpressionBenchmark.java

@@ -354,7 +354,7 @@ public void querySql(Blackhole blackhole) throws Exception
     final String sql = QUERIES.get(Integer.parseInt(query));
     try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sql, new QueryContext(context))) {
       final PlannerResult plannerResult = planner.plan();
-      final Sequence<Object[]> resultSequence = plannerResult.run();
+      final Sequence<Object[]> resultSequence = plannerResult.run().getResults();
       final Object[] lastRow = resultSequence.accumulate(null, (accumulated, in) -> in);
       blackhole.consume(lastRow);
     }

benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlNestedDataBenchmark.java

@@ -320,7 +320,7 @@ public void querySql(Blackhole blackhole) throws Exception
     final String sql = QUERIES.get(Integer.parseInt(query));
     try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sql, new QueryContext(context))) {
       final PlannerResult plannerResult = planner.plan();
-      final Sequence<Object[]> resultSequence = plannerResult.run();
+      final Sequence<Object[]> resultSequence = plannerResult.run().getResults();
       final Object[] lastRow = resultSequence.accumulate(null, (accumulated, in) -> in);
       blackhole.consume(lastRow);
     }

benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlVsNativeBenchmark.java

@@ -169,7 +169,7 @@ public void queryPlanner(Blackhole blackhole) throws Exception
   {
     try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, sqlQuery, new QueryContext())) {
       final PlannerResult plannerResult = planner.plan();
-      final Sequence<Object[]> resultSequence = plannerResult.run();
+      final Sequence<Object[]> resultSequence = plannerResult.run().getResults();
       final Object[] lastRow = resultSequence.accumulate(null, (accumulated, in) -> in);
       blackhole.consume(lastRow);
     }

core/src/main/java/org/apache/druid/guice/annotations/NativeQuery.java

@@ -0,0 +1,40 @@
+/*
+ * 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.guice.annotations;
+
+import com.google.inject.BindingAnnotation;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Binding annotation for implements of interfaces that are focused on running native queries.  This is generally
+ * contrasted with the MSQ annotation.
+ *
+ * @see Parent
+ */
+@Target({ElementType.FIELD, ElementType.PARAMETER, ElementType.METHOD})
+@Retention(RetentionPolicy.RUNTIME)
+@BindingAnnotation
+public @interface NativeQuery
+{
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/guice/annotations/MSQ.java

@@ -0,0 +1,40 @@
+/*
+ * 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.guice.annotations;
+
+import com.google.inject.BindingAnnotation;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Binding annotation for implements of interfaces that are MSQ (MultiStageQuery) focused.  This is generally
+ * contrasted with the NativeQ annotation.
+ *
+ * @see Parent
+ */
+@Target({ElementType.FIELD, ElementType.PARAMETER, ElementType.METHOD})
+@Retention(RetentionPolicy.RUNTIME)
+@BindingAnnotation
+public @interface MSQ
+{
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQSqlModule.java

@@ -22,10 +22,16 @@
 import com.fasterxml.jackson.databind.Module;
 import com.google.inject.Binder;
 import com.google.inject.Inject;
+import com.google.inject.Provides;
 import org.apache.druid.discovery.NodeRole;
+import org.apache.druid.guice.LazySingleton;
 import org.apache.druid.guice.annotations.LoadScope;
+import org.apache.druid.guice.annotations.MSQ;
 import org.apache.druid.initialization.DruidModule;
 import org.apache.druid.metadata.input.InputSourceModule;
+import org.apache.druid.msq.sql.MSQTaskSqlEngine;
+import org.apache.druid.sql.SqlStatementFactory;
+import org.apache.druid.sql.SqlToolbox;
 import org.apache.druid.sql.calcite.external.ExternalOperatorConversion;
 import org.apache.druid.sql.guice.SqlBindings;
 
@@ -54,7 +60,21 @@ public void configure(Binder binder)
     // We want this module to bring InputSourceModule along for the ride.
     binder.install(new InputSourceModule());
 
+    binder.bind(MSQTaskSqlEngine.class).in(LazySingleton.class);
+
     // Set up the EXTERN macro.
     SqlBindings.addOperatorConversion(binder, ExternalOperatorConversion.class);
   }
+
+
+  @Provides
+  @MSQ
+  @LazySingleton
+  public SqlStatementFactory makeMSQSqlStatementFactory(
+      final MSQTaskSqlEngine engine,
+      SqlToolbox toolbox
+  )
+  {
+    return new SqlStatementFactory(toolbox.withEngine(engine));
+  }
 }

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

@@ -31,7 +31,6 @@
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.granularity.Granularities;
 import org.apache.druid.java.util.common.granularity.Granularity;
-import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.guava.Sequences;
 import org.apache.druid.msq.exec.MSQTasks;
 import org.apache.druid.msq.indexing.ColumnMapping;
@@ -49,6 +48,7 @@
 import org.apache.druid.segment.DimensionHandlerUtils;
 import org.apache.druid.segment.column.ColumnHolder;
 import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.server.QueryResponse;
 import org.apache.druid.sql.calcite.parser.DruidSqlInsert;
 import org.apache.druid.sql.calcite.parser.DruidSqlReplace;
 import org.apache.druid.sql.calcite.planner.PlannerContext;
@@ -105,7 +105,7 @@ public class MSQTaskQueryMaker implements QueryMaker
   }
 
   @Override
-  public Sequence<Object[]> runQuery(final DruidQuery druidQuery)
+  public QueryResponse runQuery(final DruidQuery druidQuery)
   {
     String taskId = MSQTasks.controllerTaskId(plannerContext.getSqlQueryId());
 
@@ -259,7 +259,7 @@ public Sequence<Object[]> runQuery(final DruidQuery druidQuery)
     );
 
     FutureUtils.getUnchecked(overlordClient.runTask(taskId, controllerTask), true);
-    return Sequences.simple(Collections.singletonList(new Object[]{taskId}));
+    return QueryResponse.withEmptyContext(Sequences.simple(Collections.singletonList(new Object[]{taskId})));
   }
 
   private static Map<String, ColumnType> buildAggregationIntermediateTypeMap(final DruidQuery druidQuery)

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

@@ -32,7 +32,6 @@
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.tools.ValidationException;
 import org.apache.calcite.util.Pair;
-import org.apache.druid.guice.LazySingleton;
 import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.granularity.Granularities;
@@ -55,7 +54,6 @@
 import java.util.List;
 import java.util.Set;
 
-@LazySingleton
 public class MSQTaskSqlEngine implements SqlEngine
 {
   public static final Set<String> SYSTEM_CONTEXT_PARAMETERS =

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

@@ -25,6 +25,7 @@
 import com.google.inject.Inject;
 import org.apache.calcite.plan.RelOptPlanner;
 import org.apache.druid.common.exception.SanitizableException;
+import org.apache.druid.guice.annotations.MSQ;
 import org.apache.druid.indexer.TaskState;
 import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.guava.Yielder;
@@ -37,6 +38,7 @@
 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.AuthorizationUtils;
@@ -47,7 +49,6 @@
 import org.apache.druid.sql.SqlPlanningException;
 import org.apache.druid.sql.SqlRowTransformer;
 import org.apache.druid.sql.SqlStatementFactory;
-import org.apache.druid.sql.SqlStatementFactoryFactory;
 import org.apache.druid.sql.http.ResultFormat;
 import org.apache.druid.sql.http.SqlQuery;
 import org.apache.druid.sql.http.SqlResource;
@@ -91,14 +92,13 @@ public class SqlTaskResource
 
   @Inject
   public SqlTaskResource(
-      final MSQTaskSqlEngine engine,
-      final SqlStatementFactoryFactory sqlStatementFactoryFactory,
+      final @MSQ SqlStatementFactory sqlStatementFactory,
       final ServerConfig serverConfig,
       final AuthorizerMapper authorizerMapper,
       final ObjectMapper jsonMapper
   )
   {
-    this.sqlStatementFactory = sqlStatementFactoryFactory.factorize(engine);
+    this.sqlStatementFactory = sqlStatementFactory;
     this.serverConfig = serverConfig;
     this.authorizerMapper = authorizerMapper;
     this.jsonMapper = jsonMapper;
@@ -147,7 +147,8 @@ public Response doPost(
     final String sqlQueryId = stmt.sqlQueryId();
     try {
       final DirectStatement.ResultSet plan = stmt.plan();
-      final Sequence<Object[]> sequence = plan.run();
+      final QueryResponse response = plan.run();
+      final Sequence sequence = response.getResults();
       final SqlRowTransformer rowTransformer = plan.createRowTransformer();
       final boolean isTaskStruct = MSQTaskSqlEngine.TASK_STRUCT_FIELD_NAMES.equals(rowTransformer.getFieldList());
 

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

@@ -35,6 +35,7 @@
 import com.google.inject.Key;
 import com.google.inject.Module;
 import com.google.inject.TypeLiteral;
+import com.google.inject.util.Modules;
 import com.google.inject.util.Providers;
 import org.apache.druid.common.config.NullHandling;
 import org.apache.druid.data.input.impl.DimensionsSpec;
@@ -47,6 +48,7 @@
 import org.apache.druid.guice.JoinableFactoryModule;
 import org.apache.druid.guice.JsonConfigProvider;
 import org.apache.druid.guice.annotations.EscalatedGlobal;
+import org.apache.druid.guice.annotations.MSQ;
 import org.apache.druid.guice.annotations.Self;
 import org.apache.druid.hll.HyperLogLogCollector;
 import org.apache.druid.indexing.common.SegmentCacheManagerFactory;
@@ -125,6 +127,7 @@
 import org.apache.druid.sql.DirectStatement;
 import org.apache.druid.sql.SqlQueryPlus;
 import org.apache.druid.sql.SqlStatementFactory;
+import org.apache.druid.sql.SqlToolbox;
 import org.apache.druid.sql.calcite.BaseCalciteQueryTest;
 import org.apache.druid.sql.calcite.external.ExternalDataSource;
 import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
@@ -341,7 +344,17 @@ public String getFormatString()
         new JoinableFactoryModule(),
         new IndexingServiceTuningConfigModule(),
         new MSQIndexingModule(),
-        new MSQSqlModule(),
+        Modules.override(new MSQSqlModule()).with(
+            binder -> {
+              // Our Guice configuration currently requires bindings to exist even if they aren't ever used, the
+              // following bindings are overriding other bindings that end up needing a lot more dependencies.
+              // We replace the bindings with something that returns null to make things more brittle in case they
+              // actually are used somewhere in the test.
+              binder.bind(SqlStatementFactory.class).annotatedWith(MSQ.class).toProvider(Providers.of(null));
+              binder.bind(SqlToolbox.class).toProvider(Providers.of(null));
+              binder.bind(MSQTaskSqlEngine.class).toProvider(Providers.of(null));
+            }
+        ),
         new MSQExternalDataSourceModule()
     ));
 
@@ -580,7 +593,7 @@ private String runMultiStageQuery(String query, Map<String, Object> context)
         )
     );
 
-    final List<Object[]> sequence = stmt.execute().toList();
+    final List<Object[]> sequence = stmt.execute().getResults().toList();
     return (String) Iterables.getOnlyElement(sequence)[0];
   }
 

processing/src/main/java/org/apache/druid/query/context/ResponseContext.java

@@ -39,7 +39,6 @@
 import org.joda.time.Interval;
 
 import javax.annotation.Nullable;
-
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Comparator;
@@ -733,6 +732,10 @@ private Object addValue(Key key, Object value)
    */
   public void merge(ResponseContext responseContext)
   {
+    if (responseContext == null) {
+      return;
+    }
+
     responseContext.getDelegate().forEach((key, newValue) -> {
       if (newValue != null) {
         add(key, newValue);