diff --git a/extensions-core/multi-stage-query/pom.xml b/extensions-core/multi-stage-query/pom.xml
index ab23a26e36f..22e85f0b0fc 100644
--- a/extensions-core/multi-stage-query/pom.xml
+++ b/extensions-core/multi-stage-query/pom.xml
@@ -213,6 +213,37 @@
junit
test
+
+ net.hydromatic
+ quidem
+ test
+
+
+ org.apache.calcite.avatica
+ avatica-core
+ test
+
+
+ org.apache.calcite
+ calcite-testkit
+ ${calcite.version}
+
+
+ org.junit.jupiter
+ *
+
+
+ org.hamcrest
+ *
+
+
+ test
+
+
+ org.reflections
+ reflections
+ test
+
org.junit.jupiter
junit-jupiter-api
@@ -233,11 +264,6 @@
junit-jupiter-params
test
-
- org.reflections
- reflections
- test
-
org.easymock
easymock
diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/report/MSQResultsReport.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/report/MSQResultsReport.java
index 0479b295955..8405830773d 100644
--- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/report/MSQResultsReport.java
+++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/report/MSQResultsReport.java
@@ -26,6 +26,7 @@ import com.google.common.base.Preconditions;
import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.druid.common.config.Configs;
import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.segment.column.RowSignature;
import javax.annotation.Nullable;
import java.util.List;
@@ -34,7 +35,7 @@ import java.util.Objects;
public class MSQResultsReport
{
/**
- * Like {@link org.apache.druid.segment.column.RowSignature}, but allows duplicate column names for compatibility
+ * Like {@link RowSignature}, but allows duplicate column names for compatibility
* with SQL (which also allows duplicate column names in query results).
*/
private final List signature;
@@ -135,5 +136,18 @@ public class MSQResultsReport
{
return name + ":" + type;
}
+
+ public static RowSignature toRowSignature(List columnAndTypes)
+ {
+ final RowSignature.Builder builder = RowSignature.builder();
+ for (MSQResultsReport.ColumnAndType columnAndType : columnAndTypes) {
+ builder.add(columnAndType.getName(), columnAndType.getType());
+ }
+ RowSignature rowSignature = builder.build();
+ if (rowSignature.size() != columnAndTypes.size()) {
+ throw new IllegalArgumentException("Duplicate column names are not allowed in RowSignature");
+ }
+ return rowSignature;
+ }
}
}
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQDrillWindowQueryTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQDrillWindowQueryTest.java
index 1b2eebe7742..4e9821d1083 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQDrillWindowQueryTest.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQDrillWindowQueryTest.java
@@ -20,7 +20,6 @@
package org.apache.druid.msq.exec;
import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.common.collect.ImmutableList;
import com.google.inject.Injector;
import com.google.inject.Module;
import org.apache.druid.guice.DruidInjectorBuilder;
@@ -28,9 +27,7 @@ import org.apache.druid.msq.exec.MSQDrillWindowQueryTest.DrillWindowQueryMSQComp
import org.apache.druid.msq.sql.MSQTaskSqlEngine;
import org.apache.druid.msq.test.CalciteMSQTestsHelper;
import org.apache.druid.msq.test.ExtractResultsFactory;
-import org.apache.druid.msq.test.MSQTestBase;
import org.apache.druid.msq.test.MSQTestOverlordServiceClient;
-import org.apache.druid.msq.test.MSQTestTaskActionClient;
import org.apache.druid.msq.test.VerifyMSQSupportedNativeQueriesPredicate;
import org.apache.druid.query.groupby.TestGroupByBuffers;
import org.apache.druid.server.QueryLifecycleFactory;
@@ -55,6 +52,7 @@ public class MSQDrillWindowQueryTest extends DrillWindowQueryTest
{
super.configureGuice(builder);
builder.addModules(CalciteMSQTestsHelper.fetchModules(tempDirProducer::newTempFolder, TestGroupByBuffers.createDefault()).toArray(new Module[0]));
+ builder.addModule(new TestMSQSqlModule());
}
@Override
@@ -64,15 +62,7 @@ public class MSQDrillWindowQueryTest extends DrillWindowQueryTest
Injector injector
)
{
- final WorkerMemoryParameters workerMemoryParameters = MSQTestBase.makeTestWorkerMemoryParameters();
- final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
- queryJsonMapper,
- injector,
- new MSQTestTaskActionClient(queryJsonMapper, injector),
- workerMemoryParameters,
- ImmutableList.of()
- );
- return new MSQTaskSqlEngine(indexingServiceClient, queryJsonMapper);
+ return injector.getInstance(MSQTaskSqlEngine.class);
}
}
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java
new file mode 100644
index 00000000000..e5a0f1cb06e
--- /dev/null
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java
@@ -0,0 +1,89 @@
+/*
+ * 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.exec;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableList;
+import com.google.inject.Injector;
+import com.google.inject.Provides;
+import org.apache.druid.guice.LazySingleton;
+import org.apache.druid.initialization.ServerInjectorBuilderTest.TestDruidModule;
+import org.apache.druid.msq.guice.MultiStageQuery;
+import org.apache.druid.msq.sql.MSQTaskSqlEngine;
+import org.apache.druid.msq.test.MSQTestBase;
+import org.apache.druid.msq.test.MSQTestOverlordServiceClient;
+import org.apache.druid.msq.test.MSQTestTaskActionClient;
+import org.apache.druid.sql.SqlStatementFactory;
+import org.apache.druid.sql.SqlToolbox;
+import org.apache.druid.sql.avatica.DruidMeta;
+import org.apache.druid.sql.avatica.MSQDruidMeta;
+
+public class TestMSQSqlModule extends TestDruidModule
+{
+ @Provides
+ @MultiStageQuery
+ @LazySingleton
+ public SqlStatementFactory makeMSQSqlStatementFactory(
+ final MSQTaskSqlEngine sqlEngine,
+ SqlToolbox toolbox)
+ {
+ return new SqlStatementFactory(toolbox.withEngine(sqlEngine));
+ }
+
+ @Provides
+ @LazySingleton
+ public MSQTaskSqlEngine createEngine(
+ ObjectMapper queryJsonMapper,
+ MSQTestOverlordServiceClient indexingServiceClient)
+ {
+ return new MSQTaskSqlEngine(indexingServiceClient, queryJsonMapper);
+ }
+
+ @Provides
+ @LazySingleton
+ private MSQTestOverlordServiceClient makeOverlordServiceClient(
+ ObjectMapper queryJsonMapper,
+ Injector injector,
+ WorkerMemoryParameters workerMemoryParameters)
+ {
+ final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
+ queryJsonMapper,
+ injector,
+ new MSQTestTaskActionClient(queryJsonMapper, injector),
+ workerMemoryParameters,
+ ImmutableList.of()
+ );
+ return indexingServiceClient;
+ }
+
+ @Provides
+ @LazySingleton
+ private WorkerMemoryParameters makeWorkerMemoryParameters()
+ {
+ return MSQTestBase.makeTestWorkerMemoryParameters();
+ }
+
+ @Provides
+ @LazySingleton
+ public DruidMeta createMeta(MSQDruidMeta druidMeta)
+ {
+ return druidMeta;
+ }
+}
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/quidem/MSQQuidemTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/quidem/MSQQuidemTest.java
new file mode 100644
index 00000000000..e5c3a28e5cc
--- /dev/null
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/quidem/MSQQuidemTest.java
@@ -0,0 +1,39 @@
+/*
+ * 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.quidem;
+
+import org.apache.druid.quidem.DruidQuidemTestBase;
+import org.apache.druid.quidem.ProjectPathUtils;
+import java.io.File;
+
+public class MSQQuidemTest extends DruidQuidemTestBase
+{
+
+ public MSQQuidemTest()
+ {
+ super();
+ }
+
+ @Override
+ protected File getTestRoot()
+ {
+ return ProjectPathUtils.getPathFromProjectRoot("extensions-core/multi-stage-query/src/test/quidem/" + getClass().getName());
+ }
+}
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteArraysQueryMSQTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteArraysQueryMSQTest.java
index 124b4fce258..b761e402d35 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteArraysQueryMSQTest.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteArraysQueryMSQTest.java
@@ -19,63 +19,17 @@
package org.apache.druid.msq.test;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.common.collect.ImmutableList;
-import com.google.inject.Injector;
-import com.google.inject.Module;
-import org.apache.druid.guice.DruidInjectorBuilder;
-import org.apache.druid.msq.exec.WorkerMemoryParameters;
import org.apache.druid.msq.sql.MSQTaskSqlEngine;
-import org.apache.druid.msq.test.CalciteArraysQueryMSQTest.ArraysQueryMSQComponentSupplier;
-import org.apache.druid.query.groupby.TestGroupByBuffers;
-import org.apache.druid.server.QueryLifecycleFactory;
import org.apache.druid.sql.calcite.CalciteArraysQueryTest;
import org.apache.druid.sql.calcite.QueryTestBuilder;
import org.apache.druid.sql.calcite.SqlTestFrameworkConfig;
-import org.apache.druid.sql.calcite.TempDirProducer;
-import org.apache.druid.sql.calcite.run.SqlEngine;
/**
* Runs {@link CalciteArraysQueryTest} but with MSQ engine
*/
-@SqlTestFrameworkConfig.ComponentSupplier(ArraysQueryMSQComponentSupplier.class)
+@SqlTestFrameworkConfig.ComponentSupplier(StandardMSQComponentSupplier.class)
public class CalciteArraysQueryMSQTest extends CalciteArraysQueryTest
{
- public static class ArraysQueryMSQComponentSupplier extends ArraysComponentSupplier
- {
- public ArraysQueryMSQComponentSupplier(TempDirProducer tempFolderProducer)
- {
- super(tempFolderProducer);
- }
-
- @Override
- public void configureGuice(DruidInjectorBuilder builder)
- {
- super.configureGuice(builder);
- builder.addModules(
- CalciteMSQTestsHelper.fetchModules(tempDirProducer::newTempFolder, TestGroupByBuffers.createDefault()).toArray(new Module[0])
- );
- }
-
- @Override
- public SqlEngine createEngine(
- QueryLifecycleFactory qlf,
- ObjectMapper queryJsonMapper,
- Injector injector
- )
- {
- final WorkerMemoryParameters workerMemoryParameters = MSQTestBase.makeTestWorkerMemoryParameters();
- final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
- queryJsonMapper,
- injector,
- new MSQTestTaskActionClient(queryJsonMapper, injector),
- workerMemoryParameters,
- ImmutableList.of()
- );
- return new MSQTaskSqlEngine(indexingServiceClient, queryJsonMapper);
- }
- }
-
@Override
protected QueryTestBuilder testBuilder()
{
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteNestedDataQueryMSQTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteNestedDataQueryMSQTest.java
index 5d4c0994ea0..0a2b8955a3d 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteNestedDataQueryMSQTest.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteNestedDataQueryMSQTest.java
@@ -25,7 +25,7 @@ import com.google.inject.Injector;
import com.google.inject.Module;
import org.apache.druid.guice.DruidInjectorBuilder;
import org.apache.druid.java.util.common.ISE;
-import org.apache.druid.msq.exec.WorkerMemoryParameters;
+import org.apache.druid.msq.exec.TestMSQSqlModule;
import org.apache.druid.msq.sql.MSQTaskSqlEngine;
import org.apache.druid.query.groupby.TestGroupByBuffers;
import org.apache.druid.server.QueryLifecycleFactory;
@@ -58,6 +58,7 @@ public class CalciteNestedDataQueryMSQTest extends CalciteNestedDataQueryTest
builder.addModules(
CalciteMSQTestsHelper.fetchModules(tempDirProducer::newTempFolder, TestGroupByBuffers.createDefault()).toArray(new Module[0])
);
+ builder.addModule(new TestMSQSqlModule());
}
@Override
@@ -67,15 +68,7 @@ public class CalciteNestedDataQueryMSQTest extends CalciteNestedDataQueryTest
Injector injector
)
{
- final WorkerMemoryParameters workerMemoryParameters = MSQTestBase.makeTestWorkerMemoryParameters();
- final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
- queryJsonMapper,
- injector,
- new MSQTestTaskActionClient(queryJsonMapper, injector),
- workerMemoryParameters,
- ImmutableList.of()
- );
- return new MSQTaskSqlEngine(indexingServiceClient, queryJsonMapper);
+ return injector.getInstance(MSQTaskSqlEngine.class);
}
}
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectJoinQueryMSQTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectJoinQueryMSQTest.java
index 6bbf9c6da5e..30724dcd4ce 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectJoinQueryMSQTest.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectJoinQueryMSQTest.java
@@ -19,26 +19,15 @@
package org.apache.druid.msq.test;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.common.collect.ImmutableList;
-import com.google.inject.Injector;
-import com.google.inject.Module;
-import org.apache.calcite.rel.RelRoot;
-import org.apache.druid.guice.DruidInjectorBuilder;
-import org.apache.druid.msq.exec.WorkerMemoryParameters;
+import com.google.common.collect.ImmutableMap;
import org.apache.druid.msq.sql.MSQTaskSqlEngine;
-import org.apache.druid.query.groupby.TestGroupByBuffers;
-import org.apache.druid.server.QueryLifecycleFactory;
+import org.apache.druid.sql.calcite.BaseCalciteQueryTest;
import org.apache.druid.sql.calcite.CalciteJoinQueryTest;
import org.apache.druid.sql.calcite.QueryTestBuilder;
import org.apache.druid.sql.calcite.SqlTestFrameworkConfig;
-import org.apache.druid.sql.calcite.TempDirProducer;
import org.apache.druid.sql.calcite.planner.JoinAlgorithm;
import org.apache.druid.sql.calcite.planner.PlannerContext;
-import org.apache.druid.sql.calcite.run.EngineFeature;
-import org.apache.druid.sql.calcite.run.QueryMaker;
-import org.apache.druid.sql.calcite.run.SqlEngine;
-import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier;
+import java.util.Map;
/**
* Runs {@link CalciteJoinQueryTest} but with MSQ engine.
@@ -48,7 +37,6 @@ public class CalciteSelectJoinQueryMSQTest
/**
* Run all tests with {@link JoinAlgorithm#BROADCAST}.
*/
- @SqlTestFrameworkConfig.ComponentSupplier(BroadcastJoinComponentSupplier.class)
public static class BroadcastTest extends Base
{
@Override
@@ -57,12 +45,17 @@ public class CalciteSelectJoinQueryMSQTest
return super.testBuilder()
.verifyNativeQueries(new VerifyMSQSupportedNativeQueriesPredicate());
}
+
+ @Override
+ protected JoinAlgorithm joinAlgorithm()
+ {
+ return JoinAlgorithm.BROADCAST;
+ }
}
/**
* Run all tests with {@link JoinAlgorithm#SORT_MERGE}.
*/
- @SqlTestFrameworkConfig.ComponentSupplier(SortMergeJoinComponentSupplier.class)
public static class SortMergeTest extends Base
{
@Override
@@ -79,86 +72,31 @@ public class CalciteSelectJoinQueryMSQTest
return super.testBuilder()
.verifyNativeQueries(xs -> false);
}
+
+ @Override
+ protected JoinAlgorithm joinAlgorithm()
+ {
+ return JoinAlgorithm.SORT_MERGE;
+ }
}
+ @SqlTestFrameworkConfig.ComponentSupplier(StandardMSQComponentSupplier.class)
public abstract static class Base extends CalciteJoinQueryTest
{
+ protected abstract JoinAlgorithm joinAlgorithm();
+
@Override
protected QueryTestBuilder testBuilder()
{
- return new QueryTestBuilder(new CalciteTestConfig(true))
+ Map defaultCtx = ImmutableMap.builder()
+ .putAll(BaseCalciteQueryTest.QUERY_CONTEXT_DEFAULT)
+ .put(PlannerContext.CTX_SQL_JOIN_ALGORITHM, joinAlgorithm().toString())
+ .build();
+ return new QueryTestBuilder(new CalciteTestConfig(defaultCtx, true))
.addCustomRunner(
new ExtractResultsFactory(
() -> (MSQTestOverlordServiceClient) ((MSQTaskSqlEngine) queryFramework().engine()).overlordClient()))
.skipVectorize(true);
}
}
-
- protected static class SortMergeJoinComponentSupplier extends AbstractJoinComponentSupplier
- {
- public SortMergeJoinComponentSupplier(TempDirProducer tempFolderProducer)
- {
- super(tempFolderProducer, JoinAlgorithm.SORT_MERGE);
- }
- }
-
- protected static class BroadcastJoinComponentSupplier extends AbstractJoinComponentSupplier
- {
- public BroadcastJoinComponentSupplier(TempDirProducer tempFolderProducer)
- {
- super(tempFolderProducer, JoinAlgorithm.BROADCAST);
- }
- }
-
- protected abstract static class AbstractJoinComponentSupplier extends StandardComponentSupplier
- {
- private JoinAlgorithm joinAlgorithm;
-
- public AbstractJoinComponentSupplier(TempDirProducer tempFolderProducer, JoinAlgorithm joinAlgorithm)
- {
- super(tempFolderProducer);
- this.joinAlgorithm = joinAlgorithm;
- }
-
- @Override
- public void configureGuice(DruidInjectorBuilder builder)
- {
- super.configureGuice(builder);
- builder.addModules(
- CalciteMSQTestsHelper.fetchModules(tempDirProducer::newTempFolder, TestGroupByBuffers.createDefault()).toArray(new Module[0])
- );
- }
-
- @Override
- public SqlEngine createEngine(
- QueryLifecycleFactory qlf,
- ObjectMapper queryJsonMapper,
- Injector injector
- )
- {
- final WorkerMemoryParameters workerMemoryParameters = MSQTestBase.makeTestWorkerMemoryParameters();
- final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
- queryJsonMapper,
- injector,
- new MSQTestTaskActionClient(queryJsonMapper, injector),
- workerMemoryParameters,
- ImmutableList.of()
- );
- return new MSQTaskSqlEngine(indexingServiceClient, queryJsonMapper)
- {
- @Override
- public boolean featureAvailable(EngineFeature feature)
- {
- return super.featureAvailable(feature);
- }
-
- @Override
- public QueryMaker buildQueryMakerForSelect(RelRoot relRoot, PlannerContext plannerContext)
- {
- plannerContext.queryContextMap().put(PlannerContext.CTX_SQL_JOIN_ALGORITHM, joinAlgorithm.toString());
- return super.buildQueryMakerForSelect(relRoot, plannerContext);
- }
- };
- }
- }
}
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java
index 2de9229b4ad..6e2032a1095 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java
@@ -19,25 +19,14 @@
package org.apache.druid.msq.test;
-import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
-import com.google.inject.Injector;
-import com.google.inject.Module;
import org.apache.druid.common.config.NullHandling;
-import org.apache.druid.guice.DruidInjectorBuilder;
import org.apache.druid.java.util.common.ISE;
-import org.apache.druid.msq.exec.WorkerMemoryParameters;
import org.apache.druid.msq.sql.MSQTaskSqlEngine;
-import org.apache.druid.msq.test.CalciteSelectQueryMSQTest.SelectMSQComponentSupplier;
-import org.apache.druid.query.groupby.TestGroupByBuffers;
-import org.apache.druid.server.QueryLifecycleFactory;
import org.apache.druid.sql.calcite.CalciteQueryTest;
import org.apache.druid.sql.calcite.QueryTestBuilder;
import org.apache.druid.sql.calcite.SqlTestFrameworkConfig;
-import org.apache.druid.sql.calcite.TempDirProducer;
-import org.apache.druid.sql.calcite.run.SqlEngine;
-import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier;
import org.junit.Assert;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
@@ -48,43 +37,9 @@ import java.util.concurrent.TimeUnit;
/**
* Runs {@link CalciteQueryTest} but with MSQ engine
*/
-@SqlTestFrameworkConfig.ComponentSupplier(SelectMSQComponentSupplier.class)
+@SqlTestFrameworkConfig.ComponentSupplier(StandardMSQComponentSupplier.class)
public class CalciteSelectQueryMSQTest extends CalciteQueryTest
{
- public static class SelectMSQComponentSupplier extends StandardComponentSupplier
- {
- public SelectMSQComponentSupplier(TempDirProducer tempFolderProducer)
- {
- super(tempFolderProducer);
- }
-
- @Override
- public void configureGuice(DruidInjectorBuilder builder)
- {
- super.configureGuice(builder);
- builder.addModules(CalciteMSQTestsHelper.fetchModules(tempDirProducer::newTempFolder, TestGroupByBuffers.createDefault()).toArray(new Module[0]));
- }
-
-
- @Override
- public SqlEngine createEngine(
- QueryLifecycleFactory qlf,
- ObjectMapper queryJsonMapper,
- Injector injector
- )
- {
- final WorkerMemoryParameters workerMemoryParameters = MSQTestBase.makeTestWorkerMemoryParameters();
- final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
- queryJsonMapper,
- injector,
- new MSQTestTaskActionClient(queryJsonMapper, injector),
- workerMemoryParameters,
- ImmutableList.of()
- );
- return new MSQTaskSqlEngine(indexingServiceClient, queryJsonMapper);
- }
- }
-
@Override
protected QueryTestBuilder testBuilder()
{
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteUnionQueryMSQTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteUnionQueryMSQTest.java
index e4b678402a8..f0686498786 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteUnionQueryMSQTest.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteUnionQueryMSQTest.java
@@ -19,14 +19,9 @@
package org.apache.druid.msq.test;
-import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableList;
-import com.google.inject.Injector;
-import com.google.inject.Module;
import org.apache.druid.common.config.NullHandling;
-import org.apache.druid.guice.DruidInjectorBuilder;
import org.apache.druid.java.util.common.granularity.Granularities;
-import org.apache.druid.msq.exec.WorkerMemoryParameters;
import org.apache.druid.msq.sql.MSQTaskSqlEngine;
import org.apache.druid.query.QueryDataSource;
import org.apache.druid.query.TableDataSource;
@@ -35,62 +30,21 @@ import org.apache.druid.query.aggregation.CountAggregatorFactory;
import org.apache.druid.query.aggregation.LongSumAggregatorFactory;
import org.apache.druid.query.dimension.DefaultDimensionSpec;
import org.apache.druid.query.groupby.GroupByQuery;
-import org.apache.druid.query.groupby.TestGroupByBuffers;
-import org.apache.druid.server.QueryLifecycleFactory;
import org.apache.druid.sql.calcite.BaseCalciteQueryTest;
import org.apache.druid.sql.calcite.CalciteUnionQueryTest;
import org.apache.druid.sql.calcite.QueryTestBuilder;
import org.apache.druid.sql.calcite.SqlTestFrameworkConfig;
-import org.apache.druid.sql.calcite.TempDirProducer;
import org.apache.druid.sql.calcite.filtration.Filtration;
-import org.apache.druid.sql.calcite.run.SqlEngine;
import org.apache.druid.sql.calcite.util.CalciteTests;
-import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
/**
* Runs {@link CalciteUnionQueryTest} but with MSQ engine
*/
-@SqlTestFrameworkConfig.ComponentSupplier(CalciteUnionQueryMSQTest.UnionQueryMSQComponentSupplier.class)
+@SqlTestFrameworkConfig.ComponentSupplier(StandardMSQComponentSupplier.class)
public class CalciteUnionQueryMSQTest extends CalciteUnionQueryTest
{
-
- public static class UnionQueryMSQComponentSupplier extends StandardComponentSupplier
- {
- public UnionQueryMSQComponentSupplier(TempDirProducer tempFolderProducer)
- {
- super(tempFolderProducer);
- }
-
- @Override
- public void configureGuice(DruidInjectorBuilder builder)
- {
- super.configureGuice(builder);
- builder.addModules(
- CalciteMSQTestsHelper.fetchModules(tempDirProducer::newTempFolder, TestGroupByBuffers.createDefault()).toArray(new Module[0])
- );
- }
-
- @Override
- public SqlEngine createEngine(
- QueryLifecycleFactory qlf,
- ObjectMapper queryJsonMapper,
- Injector injector
- )
- {
- final WorkerMemoryParameters workerMemoryParameters = MSQTestBase.makeTestWorkerMemoryParameters();
- final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
- queryJsonMapper,
- injector,
- new MSQTestTaskActionClient(queryJsonMapper, injector),
- workerMemoryParameters,
- ImmutableList.of()
- );
- return new MSQTaskSqlEngine(indexingServiceClient, queryJsonMapper);
- }
- }
-
@Override
protected QueryTestBuilder testBuilder()
{
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/ExtractResultsFactory.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/ExtractResultsFactory.java
index 5d3c1d996ac..d2182b429cf 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/ExtractResultsFactory.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/ExtractResultsFactory.java
@@ -20,7 +20,7 @@
package org.apache.druid.msq.test;
import org.apache.druid.java.util.common.ISE;
-import org.apache.druid.msq.indexing.report.MSQResultsReport;
+import org.apache.druid.msq.indexing.report.MSQResultsReport.ColumnAndType;
import org.apache.druid.msq.indexing.report.MSQTaskReport;
import org.apache.druid.msq.indexing.report.MSQTaskReportPayload;
import org.apache.druid.segment.column.RowSignature;
@@ -104,7 +104,7 @@ public class ExtractResultsFactory implements QueryTestRunner.QueryRunStepFactor
}
extractedResults.add(
results.withSignatureAndResults(
- convertColumnAndTypeToRowSignature(payload.getResults().getSignature()), resultRows)
+ ColumnAndType.toRowSignature(payload.getResults().getSignature()), resultRows)
);
}
}
@@ -114,15 +114,6 @@ public class ExtractResultsFactory implements QueryTestRunner.QueryRunStepFactor
{
return extractedResults;
}
-
- private RowSignature convertColumnAndTypeToRowSignature(final List columnAndTypes)
- {
- final RowSignature.Builder builder = RowSignature.builder();
- for (MSQResultsReport.ColumnAndType columnAndType : columnAndTypes) {
- builder.add(columnAndType.getName(), columnAndType.getType());
- }
- return builder.build();
- }
};
}
}
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/StandardMSQComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/StandardMSQComponentSupplier.java
new file mode 100644
index 00000000000..e519a5ca387
--- /dev/null
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/StandardMSQComponentSupplier.java
@@ -0,0 +1,60 @@
+/*
+ * 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.test;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.inject.Injector;
+import com.google.inject.Module;
+import org.apache.druid.guice.DruidInjectorBuilder;
+import org.apache.druid.msq.exec.TestMSQSqlModule;
+import org.apache.druid.msq.sql.MSQTaskSqlEngine;
+import org.apache.druid.query.groupby.TestGroupByBuffers;
+import org.apache.druid.server.QueryLifecycleFactory;
+import org.apache.druid.sql.calcite.TempDirProducer;
+import org.apache.druid.sql.calcite.run.SqlEngine;
+import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier;
+
+public final class StandardMSQComponentSupplier extends StandardComponentSupplier
+{
+ public StandardMSQComponentSupplier(TempDirProducer tempFolderProducer)
+ {
+ super(tempFolderProducer);
+ }
+
+ @Override
+ public void configureGuice(DruidInjectorBuilder builder)
+ {
+ super.configureGuice(builder);
+ builder.addModules(
+ CalciteMSQTestsHelper.fetchModules(tempDirProducer::newTempFolder, TestGroupByBuffers.createDefault())
+ .toArray(new Module[0])
+ );
+ builder.addModule(new TestMSQSqlModule());
+ }
+
+ @Override
+ public SqlEngine createEngine(
+ QueryLifecycleFactory qlf,
+ ObjectMapper queryJsonMapper,
+ Injector injector)
+ {
+ return injector.getInstance(MSQTaskSqlEngine.class);
+ }
+}
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/MSQDruidMeta.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/MSQDruidMeta.java
new file mode 100644
index 00000000000..845d8adb2b3
--- /dev/null
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/MSQDruidMeta.java
@@ -0,0 +1,163 @@
+/*
+ * 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.sql.avatica;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.inject.Inject;
+import org.apache.calcite.avatica.Meta;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.commons.lang3.RegExUtils;
+import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.msq.guice.MultiStageQuery;
+import org.apache.druid.msq.indexing.report.MSQResultsReport.ColumnAndType;
+import org.apache.druid.msq.indexing.report.MSQTaskReport;
+import org.apache.druid.msq.indexing.report.MSQTaskReportPayload;
+import org.apache.druid.msq.test.MSQTestBase;
+import org.apache.druid.msq.test.MSQTestOverlordServiceClient;
+import org.apache.druid.segment.column.RowSignature;
+import org.apache.druid.server.security.AuthenticatorMapper;
+import org.apache.druid.sql.SqlStatementFactory;
+import org.apache.druid.sql.calcite.planner.DruidTypeSystem;
+import org.apache.druid.sql.calcite.table.RowSignatures;
+import org.apache.druid.sql.hook.DruidHook;
+import org.apache.druid.sql.hook.DruidHookDispatcher;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.regex.Pattern;
+
+public class MSQDruidMeta extends DruidMeta
+{
+ protected final MSQTestOverlordServiceClient overlordClient;
+ protected final ObjectMapper objectMapper;
+ protected final DruidHookDispatcher hookDispatcher;
+
+ @Inject
+ public MSQDruidMeta(
+ final @MultiStageQuery SqlStatementFactory sqlStatementFactory,
+ final AvaticaServerConfig config,
+ final ErrorHandler errorHandler,
+ final AuthenticatorMapper authMapper,
+ final MSQTestOverlordServiceClient overlordClient,
+ final ObjectMapper objectMapper,
+ final DruidHookDispatcher hookDispatcher)
+ {
+ super(sqlStatementFactory, config, errorHandler, authMapper);
+ this.overlordClient = overlordClient;
+ this.objectMapper = objectMapper;
+ this.hookDispatcher = hookDispatcher;
+ }
+
+ @Override
+ protected ExecuteResult doFetch(AbstractDruidJdbcStatement druidStatement, int maxRows)
+ {
+ String taskId = extractTaskId(druidStatement);
+
+ MSQTaskReportPayload payload = (MSQTaskReportPayload) overlordClient.getReportForTask(taskId)
+ .get(MSQTaskReport.REPORT_KEY)
+ .getPayload();
+ if (payload.getStatus().getStatus().isFailure()) {
+ throw new ISE(
+ "Query task [%s] failed due to %s",
+ taskId,
+ payload.getStatus().getErrorReport().toString()
+ );
+ }
+
+ if (!payload.getStatus().getStatus().isComplete()) {
+ throw new ISE("Query task [%s] should have finished", taskId);
+ }
+ final List> resultRows = MSQTestBase.getRows(payload.getResults());
+ if (resultRows == null) {
+ throw new ISE("Results report not present in the task's report payload");
+ }
+ try {
+ String str = objectMapper
+ .writerWithDefaultPrettyPrinter()
+ .writeValueAsString(payload.getStages());
+
+ str = maskMSQPlan(str, taskId);
+
+ hookDispatcher.dispatch(DruidHook.MSQ_PLAN, str);
+ }
+ catch (JsonProcessingException e) {
+ hookDispatcher.dispatch(DruidHook.MSQ_PLAN, "error happened during json serialization");
+ }
+
+ Signature signature = makeSignature(druidStatement, payload.getResults().getSignature());
+ @SuppressWarnings("unchecked")
+ Frame firstFrame = Frame.create(0, true, (List