shared tries

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

@@ -0,0 +1,37 @@
+/*
+ * 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.SharedQuidemTest;
+
+public class MSQSharedQuidemTest extends SharedQuidemTest
+{
+
+  public MSQSharedQuidemTest()
+  {
+    super();
+  }
+
+  @Override
+  protected boolean isNative()
+  {
+    return false;
+  }
+}

sql/src/test/java/org/apache/druid/quidem/DruidQuidemTestBase.java

@@ -49,6 +49,7 @@ import java.io.Writer;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.function.Function;
 
 import static org.junit.jupiter.api.Assertions.fail;
 
@@ -107,7 +108,7 @@ public abstract class DruidQuidemTestBase
       }
       filter = new WildcardFileFilter(filterStr);
     }
-    druidQuidemRunner = new DruidQuidemRunner();
+    druidQuidemRunner = new DruidQuidemRunner(this::getEnv);
   }
 
   /** Creates a command handler. */
@@ -128,8 +129,11 @@ public abstract class DruidQuidemTestBase
 
   public static class DruidQuidemRunner
   {
-    public DruidQuidemRunner()
+    private final Function<String, Object> envFunction;
+
+    public DruidQuidemRunner(Function<String, Object> envFunction)
     {
+      this.envFunction = envFunction;
     }
 
     public void run(File inFile) throws Exception
@@ -149,7 +153,7 @@ public abstract class DruidQuidemTestBase
         ConfigBuilder configBuilder = Quidem.configBuilder()
             .withConnectionFactory(connectionFactory)
             .withPropertyHandler(connectionFactory)
-            .withEnv(connectionFactory::getEnv)
+            .withEnv(envFunction)
             .withCommandHandler(new DruidQuidemCommandHandler());
 
         Config config = configBuilder
@@ -216,4 +220,9 @@ public abstract class DruidQuidemTestBase
   {
     DruidAvaticaTestDriver.CONFIG_STORE.close();
   }
+
+  protected Object getEnv(String name)
+  {
+    return null;
+  }
 }

sql/src/test/java/org/apache/druid/quidem/SharedQuidemTest.java

@@ -0,0 +1,54 @@
+/*
+ * 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.quidem;
+
+import java.io.File;
+
+public class SharedQuidemTest extends DruidQuidemTestBase
+{
+  public SharedQuidemTest()
+  {
+    super();
+  }
+
+  @Override
+  protected final File getTestRoot()
+  {
+    return ProjectPathUtils.getPathFromProjectRoot("sql/src/test/quidem/" + SharedQuidemTest.class.getName());
+  }
+
+  @Override
+  protected final Object getEnv(String name)
+  {
+    if (name.equals("isMSQ")) {
+      return !isNative();
+    }
+    if (name.equals("isNative")) {
+      return isNative();
+    }
+    return super.getEnv(name);
+  }
+
+  protected boolean isNative()
+  {
+    return true;
+  }
+}
+

sql/src/test/java/org/apache/druid/sql/calcite/QTestCase.java

@@ -69,7 +69,7 @@ public class QTestCase
             isValidTestCaseFile(testInfo.getIQFile());
           }
 
-          DruidQuidemRunner runner = new DruidQuidemTestBase.DruidQuidemRunner();
+          DruidQuidemRunner runner = new DruidQuidemTestBase.DruidQuidemRunner(name -> null);
           runner.run(testInfo.getIQFile());
         }
         catch (Exception e) {

sql/src/test/quidem/org.apache.druid.quidem.SharedQuidemTest/shared.iq

@@ -0,0 +1,21 @@
+!set plannerStrategy DECOUPLED
+isMSQ false
+!show isMSQ
+
+!if (isMSQ) {
+!use druidtest://?componentSupplier=DrillWindowQueryMSQComponentSupplier
+!}
+!if (isNative) {
+!use druidtest://?numMergeBuffers=3
+!}
+
+select 1+1;
+
+!nativePlan
+
+!if (isMSQ) {
+!msqPlan
+!}
+
+
+