Planner test framework

core/src/main/java/org/apache/druid/math/expr/ExpressionProcessing.java

@@ -39,7 +39,6 @@ public class ExpressionProcessing
   @Inject
   private static ExpressionProcessingConfig INSTANCE;
 
-
   /**
    * Many unit tests do not setup modules for this value to be injected, this method provides a manual way to initialize
    * {@link #INSTANCE}
@@ -63,6 +62,24 @@ public static void initializeForHomogenizeNullMultiValueStrings()
     INSTANCE = new ExpressionProcessingConfig(null, null, null, true);
   }
 
+  /**
+   * Get the current configuration when tests want to change it dynamically.
+   */
+  @VisibleForTesting
+  public static ExpressionProcessingConfig currentConfig()
+  {
+    return INSTANCE;
+  }
+
+  /**
+   * Restore a previous config after tests change it.
+   */
+  @VisibleForTesting
+  public static void restoreConfig(ExpressionProcessingConfig config)
+  {
+    INSTANCE = config;
+  }
+
   /**
    * [['is expression support for'],['nested arrays'],['enabled?']]
    */

sql/pom.xml

@@ -243,8 +243,13 @@
       <groupId>org.mockito</groupId>
       <artifactId>mockito-core</artifactId>
       <version>${mockito.version}</version>
-      <scope>test</scope>
     </dependency>
+	<dependency>
+	  <groupId>org.apache.commons</groupId>
+	  <artifactId>commons-text</artifactId>
+	  <version>1.3</version>
+      <scope>test</scope>
+	</dependency>
   </dependencies>
 
   <build>

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

@@ -106,7 +106,7 @@ public class SqlLifecycle
   @GuardedBy("stateLock")
   private State state = State.NEW;
 
-  // init during intialize
+  // init during initialize
   private String sql;
   private QueryContext queryContext;
   private List<TypedValue> parameters;
@@ -170,7 +170,7 @@ private String sqlQueryId()
   }
 
   /**
-   * Assign dynamic parameters to be used to substitute values during query exection. This can be performed at any
+   * Assign dynamic parameters to be used to substitute values during query execution. This can be performed at any
    * part of the lifecycle.
    */
   public void setParameters(List<TypedValue> parameters)
@@ -269,7 +269,7 @@ private void checkAccess(Access access)
 
   /**
    * Prepare the query lifecycle for execution, without completely planning into something that is executable, but
-   * including some initial parsing and validation and any dyanmic parameter type resolution, to support prepared
+   * including some initial parsing and validation and any dynamic parameter type resolution, to support prepared
    * statements via JDBC.
    */
   public PrepareResult prepare() throws RelConversionException
@@ -314,6 +314,12 @@ public void plan() throws RelConversionException
     }
   }
 
+  @VisibleForTesting
+  public PlannerResult plannerResult()
+  {
+    return plannerResult;
+  }
+
   /**
    * This method must be called after {@link #plan()}.
    */
@@ -376,7 +382,6 @@ public void after(boolean isDone, Throwable thrown)
     });
   }
 
-
   @VisibleForTesting
   public ValidationResult runAnalyzeResources(AuthenticationResult authenticationResult)
   {

sql/src/main/java/org/apache/druid/sql/calcite/planner/CapturedState.java

@@ -0,0 +1,111 @@
+/*
+ * 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.calcite.planner;
+
+import org.apache.calcite.interpreter.BindableRel;
+import org.apache.calcite.rel.RelRoot;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.sql.SqlInsert;
+import org.apache.calcite.sql.SqlNode;
+import org.apache.druid.sql.calcite.rel.DruidRel;
+
+/**
+ * Planner state capture for tests. The captured objects are available as
+ * public fields since this is only ever meant to be used in tests, and
+ * tests are already tightly coupled to the planner.
+ * <p>
+ * Spotbugs really doesn't like public fields referenced in another (test)
+ * package. This file appears in spotbugs-exclude.xml to avoid the issue.
+ */
+public class CapturedState implements PlannerStateCapture
+{
+  public String sql;
+  public SqlNode sqlNode;
+  public RelRoot relRoot;
+  public DruidRel<?> druidRel;
+  public RelDataType parameterTypes;
+  public PlannerContext plannerContext;
+  public ValidationResult validationResult;
+  public SqlNode queryNode;
+  public SqlInsert insertNode;
+  public BindableRel bindableRel;
+  public Object execPlan;
+
+  @Override
+  public void capturePlannerContext(PlannerContext plannerContext)
+  {
+    this.plannerContext = plannerContext;
+  }
+
+  @Override
+  public void captureSql(String sql)
+  {
+    this.sql = sql;
+  }
+
+  @Override
+  public void captureParse(SqlNode root)
+  {
+    this.sqlNode = root;
+  }
+
+  @Override
+  public void captureQueryRel(RelRoot rootQueryRel)
+  {
+    this.relRoot = rootQueryRel;
+  }
+
+  @Override
+  public void captureDruidRel(DruidRel<?> druidRel)
+  {
+    this.druidRel = druidRel;
+    this.execPlan = null;
+  }
+
+  @Override
+  public void captureParameterTypes(RelDataType parameterTypes)
+  {
+    this.parameterTypes = parameterTypes;
+  }
+
+  @Override
+  public void captureValidationResult(ValidationResult validationResult)
+  {
+    this.validationResult = validationResult;
+  }
+
+  @Override
+  public void captureQuery(SqlNode query)
+  {
+    this.queryNode = query;
+  }
+
+  @Override
+  public void captureInsert(SqlInsert insert)
+  {
+    this.insertNode = insert;
+  }
+
+  @Override
+  public void captureBindableRel(BindableRel bindableRel)
+  {
+    this.bindableRel = bindableRel;
+  }
+}

sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidPlanner.java

@@ -115,6 +115,7 @@ public class DruidPlanner implements Closeable
   private final Planner planner;
   private final PlannerContext plannerContext;
   private final QueryMakerFactory queryMakerFactory;
+  private PlannerStateCapture stateCapture;
 
   private RexBuilder rexBuilder;
 
@@ -128,6 +129,23 @@ public class DruidPlanner implements Closeable
     this.planner = Frameworks.getPlanner(frameworkConfig);
     this.plannerContext = plannerContext;
     this.queryMakerFactory = queryMakerFactory;
+    this.stateCapture = new NoOpCapture();
+  }
+
+  public void captureState(PlannerStateCapture capture)
+  {
+    this.stateCapture = capture;
+    this.stateCapture.capturePlannerContext(plannerContext);
+  }
+
+  private ParsedNodes parse() throws SqlParseException, ValidationException
+  {
+    resetPlanner();
+    SqlNode root = planner.parse(plannerContext.getSql());
+    stateCapture.captureSql(plannerContext.getSql());
+    final ParsedNodes parsed = ParsedNodes.create(root, plannerContext.getTimeZone());
+    stateCapture.captureParse(root);
+    return parsed;
   }
 
   /**
@@ -137,8 +155,7 @@ public class DruidPlanner implements Closeable
    */
   public ValidationResult validate(boolean authorizeContextParams) throws SqlParseException, ValidationException
   {
-    resetPlanner();
-    final ParsedNodes parsed = ParsedNodes.create(planner.parse(plannerContext.getSql()), plannerContext.getTimeZone());
+    final ParsedNodes parsed = parse();
     final SqlValidator validator = getValidator();
     final SqlNode validatedQueryNode;
 
@@ -165,7 +182,9 @@ public ValidationResult validate(boolean authorizeContextParams) throws SqlParse
     }
 
     plannerContext.setResourceActions(resourceActions);
-    return new ValidationResult(resourceActions);
+    ValidationResult validationResult = new ValidationResult(resourceActions);
+    stateCapture.captureValidationResult(validationResult);
+    return validationResult;
   }
 
   /**
@@ -177,15 +196,16 @@ public ValidationResult validate(boolean authorizeContextParams) throws SqlParse
    */
   public PrepareResult prepare() throws SqlParseException, ValidationException, RelConversionException
   {
-    resetPlanner();
-
-    final ParsedNodes parsed = ParsedNodes.create(planner.parse(plannerContext.getSql()), plannerContext.getTimeZone());
+    final ParsedNodes parsed = parse();
     final SqlNode validatedQueryNode = planner.validate(parsed.getQueryNode());
+    stateCapture.captureQuery(validatedQueryNode);
     final RelRoot rootQueryRel = planner.rel(validatedQueryNode);
+    stateCapture.captureQueryRel(rootQueryRel);
 
     final SqlValidator validator = getValidator();
     final RelDataTypeFactory typeFactory = rootQueryRel.rel.getCluster().getTypeFactory();
     final RelDataType parameterTypes = validator.getParameterRowType(validator.validate(validatedQueryNode));
+    stateCapture.captureParameterTypes(parameterTypes);
     final RelDataType returnedRowType;
 
     if (parsed.getExplainNode() != null) {
@@ -208,9 +228,7 @@ public PrepareResult prepare() throws SqlParseException, ValidationException, Re
    */
   public PlannerResult plan() throws SqlParseException, ValidationException, RelConversionException
   {
-    resetPlanner();
-
-    final ParsedNodes parsed = ParsedNodes.create(planner.parse(plannerContext.getSql()), plannerContext.getTimeZone());
+    final ParsedNodes parsed = parse();
 
     try {
       if (parsed.getIngestionGranularity() != null) {
@@ -235,6 +253,8 @@ public PlannerResult plan() throws SqlParseException, ValidationException, RelCo
     this.rexBuilder = new RexBuilder(planner.getTypeFactory());
     final SqlNode parameterizedQueryNode = rewriteDynamicParameters(parsed.getQueryNode());
     final SqlNode validatedQueryNode = planner.validate(parameterizedQueryNode);
+    stateCapture.captureQuery(validatedQueryNode);
+    stateCapture.captureInsert(parsed.getInsertOrReplace());
     final RelRoot rootQueryRel = planner.rel(validatedQueryNode);
 
     try {
@@ -243,7 +263,7 @@ public PlannerResult plan() throws SqlParseException, ValidationException, RelCo
     catch (Exception e) {
       Throwable cannotPlanException = Throwables.getCauseOfType(e, RelOptPlanner.CannotPlanException.class);
       if (null == cannotPlanException) {
-        // Not a CannotPlanningException, rethrow without trying with bindable
+        // Not a CannotPlanException, rethrow without trying with bindable
         throw e;
       }
 
@@ -310,6 +330,8 @@ private PlannerResult planWithDruidConvention(
   ) throws ValidationException, RelConversionException
   {
     final RelRoot possiblyLimitedRoot = possiblyWrapRootWithOuterLimitFromContext(root);
+    stateCapture.captureQueryRel(possiblyLimitedRoot);
+
     final QueryMaker queryMaker = buildQueryMaker(root, insertOrReplace);
     plannerContext.setQueryMaker(queryMaker);
 
@@ -321,6 +343,7 @@ private PlannerResult planWithDruidConvention(
                .plus(root.collation),
         parameterized
     );
+    stateCapture.captureDruidRel(druidRel);
 
     if (explain != null) {
       return planExplanation(druidRel, explain, true);
@@ -366,6 +389,7 @@ private PlannerResult planWithBindableConvention(
         planner.getEmptyTraitSet().replace(BindableConvention.INSTANCE).plus(root.collation),
         root.rel
     );
+    stateCapture.captureBindableRel(bindableRel);
 
     if (!root.isRefTrivial()) {
       // Add a projection on top to accommodate root.fields.
@@ -469,7 +493,7 @@ private PlannerResult planExplanation(
   /**
    * This method doesn't utilize the Calcite's internal {@link RelOptUtil#dumpPlan} since that tends to be verbose
    * and not indicative of the native Druid Queries which will get executed
-   * This method assumes that the Planner has converted the RelNodes to DruidRels, and thereby we can implictly cast it
+   * This method assumes that the Planner has converted the RelNodes to DruidRels, and thereby we can implicitly cast it
    *
    * @param rel Instance of the root {@link DruidRel} which is formed by running the planner transformations on it
    * @return A string representing an array of native queries that correspond to the given SQL query, in JSON format