Include statement attributes in EXPLAIN PLAN output

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

@@ -135,6 +135,8 @@ public void validate() throws SqlParseException, ValidationException
     try {
       handler.validate();
       plannerContext.setResourceActions(handler.resourceActions());
+      plannerContext.setStatementKind(handler.statementKind());
+      plannerContext.setTargetDataSource(handler.targetDataSource());
     }
     catch (RuntimeException e) {
       throw new ValidationException(e);

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

@@ -49,6 +49,7 @@
 import org.apache.druid.sql.calcite.run.EngineFeature;
 import org.apache.druid.sql.calcite.run.QueryMaker;
 
+import javax.annotation.Nullable;
 import java.util.List;
 import java.util.regex.Pattern;
 
@@ -273,6 +274,19 @@ public void validate() throws ValidationException
       }
       super.validate();
     }
+
+    @Override
+    public String statementKind()
+    {
+      return DruidSqlInsert.OPERATOR.getName();
+    }
+
+    @Nullable
+    @Override
+    public SqlNode targetDataSource()
+    {
+      return sqlNode.getTargetTable();
+    }
   }
 
   /**
@@ -331,5 +345,18 @@ public void validate() throws ValidationException
         );
       }
     }
+
+    @Override
+    public String statementKind()
+    {
+      return DruidSqlReplace.OPERATOR.getName();
+    }
+
+    @Nullable
+    @Override
+    public SqlNode targetDataSource()
+    {
+      return sqlNode.getTargetTable();
+    }
   }
 }

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

@@ -28,6 +28,7 @@
 import org.apache.calcite.avatica.remote.TypedValue;
 import org.apache.calcite.linq4j.QueryProvider;
 import org.apache.calcite.schema.SchemaPlus;
+import org.apache.calcite.sql.SqlNode;
 import org.apache.druid.java.util.common.DateTimes;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.Numbers;
@@ -111,6 +112,10 @@ public class PlannerContext
   private String planningError;
   private QueryMaker queryMaker;
   private VirtualColumnRegistry joinExpressionVirtualColumnRegistry;
+  private String statementKind;
+  // Only valid for INSERT/REPLACE queries
+  @Nullable
+  private SqlNode targetDataSource;
 
   private PlannerContext(
       final PlannerToolbox plannerToolbox,
@@ -502,4 +507,31 @@ public void setJoinExpressionVirtualColumnRegistry(VirtualColumnRegistry joinExp
   {
     this.joinExpressionVirtualColumnRegistry = joinExpressionVirtualColumnRegistry;
   }
+
+  public String getStatementKind()
+  {
+    return this.statementKind;
+  }
+
+  public void setStatementKind(String statementKind)
+  {
+    if (this.statementKind != null) {
+      throw new ISE("StatementKind has already been set");
+    }
+    this.statementKind = statementKind;
+  }
+
+  @Nullable
+  public SqlNode getTargetDataSource()
+  {
+    return this.targetDataSource;
+  }
+
+  public void setTargetDataSource(SqlNode targetDataSource)
+  {
+    if (this.targetDataSource != null) {
+      throw new ISE("TargetDataSource has already been set");
+    }
+    this.targetDataSource = targetDataSource;
+  }
 }

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

@@ -228,6 +228,22 @@ public PlannerResult plan() throws ValidationException
       throw new UnsupportedSQLQueryException(errorMessage);
     }
   }
+  @Override
+  public String statementKind()
+  {
+    return "SELECT";
+  }
+
+  /**
+   * No target datasource for SELECT/non-DML statements.
+   */
+  @Nullable
+  @Override
+  public SqlNode targetDataSource()
+  {
+    // nothing for SELECT queries.
+    return null;
+  }
 
   private static Set<RelOptTable> getBindableTables(final RelNode relNode)
   {
@@ -384,7 +400,7 @@ protected 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
+   * 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 implicitly cast it
    *
    * @param rel Instance of the root {@link DruidRel} which is formed by running the planner transformations on it
@@ -414,6 +430,10 @@ private String explainSqlPlanAsNativeQueries(DruidRel<?> rel) throws JsonProcess
       nativeQueriesArrayNode.add(objectNode);
     }
 
+    nativeQueriesArrayNode.add(jsonMapper.createObjectNode().put(
+        "statementKind", jsonMapper.convertValue(rel.getPlannerContext().getStatementKind(), String.class)));
+    nativeQueriesArrayNode.add(jsonMapper.createObjectNode().put(
+        "targetDataSource", StringUtils.format("%s", rel.getPlannerContext().getTargetDataSource())));
     return jsonMapper.writeValueAsString(nativeQueriesArrayNode);
   }
 

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

@@ -21,12 +21,14 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.calcite.schema.SchemaPlus;
+import org.apache.calcite.sql.SqlNode;
 import org.apache.calcite.tools.ValidationException;
 import org.apache.druid.query.QueryContext;
 import org.apache.druid.server.security.ResourceAction;
 import org.apache.druid.sql.calcite.run.SqlEngine;
 import org.joda.time.DateTimeZone;
 
+import javax.annotation.Nullable;
 import java.util.Map;
 import java.util.Set;
 
@@ -42,6 +44,18 @@ public interface SqlStatementHandler
   PrepareResult prepareResult();
   PlannerResult plan() throws ValidationException;
 
+  /**
+   * @return the statement kind for a SQL statement. For example, SELECT, INSERT, or REPLACE.
+   */
+  String statementKind();
+
+  /**
+   *
+   * @return the SQL node representing the target table in a SQL statement. Returns
+   * null for SELECT/non-DML statements.
+   */
+  @Nullable SqlNode targetDataSource();
+
   /**
    * Context available to statement handlers.
    */

sql/src/test/java/org/apache/druid/sql/avatica/DruidAvaticaHandlerTest.java

@@ -465,7 +465,7 @@ public void testExplainSelectCount() throws SQLException
               ImmutableMap.of(
                   "PLAN",
                   StringUtils.format(
-                      "[{\"query\":{\"queryType\":\"timeseries\",\"dataSource\":{\"type\":\"table\",\"name\":\"foo\"},\"intervals\":{\"type\":\"intervals\",\"intervals\":[\"-146136543-09-08T08:23:32.096Z/146140482-04-24T15:36:27.903Z\"]},\"granularity\":{\"type\":\"all\"},\"aggregations\":[{\"type\":\"count\",\"name\":\"a0\"}],\"context\":{\"sqlQueryId\":\"%s\",\"sqlStringifyArrays\":false,\"sqlTimeZone\":\"America/Los_Angeles\"}},\"signature\":[{\"name\":\"a0\",\"type\":\"LONG\"}]}]",
+                      "[{\"query\":{\"queryType\":\"timeseries\",\"dataSource\":{\"type\":\"table\",\"name\":\"foo\"},\"intervals\":{\"type\":\"intervals\",\"intervals\":[\"-146136543-09-08T08:23:32.096Z/146140482-04-24T15:36:27.903Z\"]},\"granularity\":{\"type\":\"all\"},\"aggregations\":[{\"type\":\"count\",\"name\":\"a0\"}],\"context\":{\"sqlQueryId\":\"%s\",\"sqlStringifyArrays\":false,\"sqlTimeZone\":\"America/Los_Angeles\"}},\"signature\":[{\"name\":\"a0\",\"type\":\"LONG\"}]},{\"statementKind\":\"SELECT\"},{\"targetDataSource\":\"null\"}]",
                       DUMMY_SQL_QUERY_ID
                   ),
                   "RESOURCES",

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

@@ -170,6 +170,9 @@ public static void setupNullValues()
   public static final PlannerConfig PLANNER_CONFIG_AUTHORIZE_SYS_TABLES =
       PlannerConfig.builder().authorizeSystemTablesDirectly(true).build();
 
+  public static final PlannerConfig PLANNER_CONFIG_LEGACY_QUERY_EXPLAIN =
+      PlannerConfig.builder().useNativeQueryExplain(false).build();
+
   public static final PlannerConfig PLANNER_CONFIG_NATIVE_QUERY_EXPLAIN =
       PlannerConfig.builder().useNativeQueryExplain(true).build();
 

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

@@ -48,12 +48,14 @@ public void testExplainCountStarOnView()
                                + "\"granularity\":{\"type\":\"all\"},"
                                + "\"aggregations\":[{\"type\":\"count\",\"name\":\"a0\"}],"
                                + "\"context\":{\"defaultTimeout\":300000,\"maxScatterGatherBytes\":9223372036854775807,\"sqlCurrentTimestamp\":\"2000-01-01T00:00:00Z\",\"sqlQueryId\":\"dummy\",\"vectorize\":\"false\",\"vectorizeVirtualColumns\":\"false\"}},"
-                               + "\"signature\":[{\"name\":\"a0\",\"type\":\"LONG\"}]"
-                               + "}]";
+                               + "\"signature\":[{\"name\":\"a0\",\"type\":\"LONG\"}]},"
+                               + "{\"statementKind\":\"SELECT\"},"
+                               + "{\"targetDataSource\":\"null\"}"
+                               + "]";
     final String resources = "[{\"name\":\"aview\",\"type\":\"VIEW\"}]";
 
     testQuery(
-        PlannerConfig.builder().useNativeQueryExplain(false).build(),
+        PLANNER_CONFIG_LEGACY_QUERY_EXPLAIN,
         query,
         CalciteTests.REGULAR_USER_AUTH_RESULT,
         ImmutableList.of(),
@@ -122,8 +124,10 @@ public void testExplainExactCountDistinctOfSemiJoinResult()
                                + "\"aggregations\":[{\"type\":\"count\",\"name\":\"a0\"}],"
                                + "\"limitSpec\":{\"type\":\"NoopLimitSpec\"},"
                                + "\"context\":{\"defaultTimeout\":300000,\"maxScatterGatherBytes\":9223372036854775807,\"sqlCurrentTimestamp\":\"2000-01-01T00:00:00Z\",\"sqlQueryId\":\"dummy\",\"vectorize\":\"false\",\"vectorizeVirtualColumns\":\"false\"}},"
-                               + "\"signature\":[{\"name\":\"a0\",\"type\":\"LONG\"}]"
-                               + "}]";
+                               + "\"signature\":[{\"name\":\"a0\",\"type\":\"LONG\"}]},"
+                               + "{\"statementKind\":\"SELECT\"},"
+                               + "{\"targetDataSource\":\"null\"}"
+                               + "]";
     final String resources = "[{\"name\":\"foo\",\"type\":\"DATASOURCE\"}]";
 
     testQuery(
@@ -133,15 +137,15 @@ public void testExplainExactCountDistinctOfSemiJoinResult()
     );
 
     testQuery(
-        PlannerConfig.builder().useNativeQueryExplain(false).build(),
+        PLANNER_CONFIG_LEGACY_QUERY_EXPLAIN,
         query,
         CalciteTests.REGULAR_USER_AUTH_RESULT,
         ImmutableList.of(),
         ImmutableList.of(new Object[]{legacyExplanation, resources})
     );
   }
 
-  // This testcase has been added here and not in CalciteSelectQueryTests since this checks if the overrides are working
+  // This testcase has been added here and not in CalciteSelectQueryTest since this checks if the overrides are working
   // properly when displaying the output of "EXPLAIN PLAN FOR ..." queries
   @Test
   public void testExplainSelectStarWithOverrides()
@@ -165,8 +169,10 @@ public void testExplainSelectStarWithOverrides()
                          + "\"legacy\":false,"
                          + "\"context\":{\"defaultTimeout\":300000,\"maxScatterGatherBytes\":9223372036854775807,\"sqlCurrentTimestamp\":\"2000-01-01T00:00:00Z\",\"sqlQueryId\":\"dummy\",\"vectorize\":\"false\",\"vectorizeVirtualColumns\":\"false\"},"
                          + "\"granularity\":{\"type\":\"all\"}},"
-                         + "\"signature\":[{\"name\":\"__time\",\"type\":\"LONG\"},{\"name\":\"dim1\",\"type\":\"STRING\"},{\"name\":\"dim2\",\"type\":\"STRING\"},{\"name\":\"dim3\",\"type\":\"STRING\"},{\"name\":\"cnt\",\"type\":\"LONG\"},{\"name\":\"m1\",\"type\":\"FLOAT\"},{\"name\":\"m2\",\"type\":\"DOUBLE\"},{\"name\":\"unique_dim1\",\"type\":\"COMPLEX<hyperUnique>\"}]"
-                         + "}]";
+                         + "\"signature\":[{\"name\":\"__time\",\"type\":\"LONG\"},{\"name\":\"dim1\",\"type\":\"STRING\"},{\"name\":\"dim2\",\"type\":\"STRING\"},{\"name\":\"dim3\",\"type\":\"STRING\"},{\"name\":\"cnt\",\"type\":\"LONG\"},{\"name\":\"m1\",\"type\":\"FLOAT\"},{\"name\":\"m2\",\"type\":\"DOUBLE\"},{\"name\":\"unique_dim1\",\"type\":\"COMPLEX<hyperUnique>\"}]},"
+                         + "{\"statementKind\":\"SELECT\"},"
+                         + "{\"targetDataSource\":\"null\"}"
+                         + "]";
 
     String explanationWithContext = "[{"
                                     + "\"query\":{\"queryType\":\"scan\","
@@ -177,8 +183,10 @@ public void testExplainSelectStarWithOverrides()
                                     + "\"legacy\":false,"
                                     + "\"context\":{\"defaultTimeout\":300000,\"maxScatterGatherBytes\":9223372036854775807,\"sqlCurrentTimestamp\":\"2000-01-01T00:00:00Z\",\"sqlQueryId\":\"dummy\",\"useNativeQueryExplain\":true,\"vectorize\":\"false\",\"vectorizeVirtualColumns\":\"false\"},"
                                     + "\"granularity\":{\"type\":\"all\"}},"
-                                    + "\"signature\":[{\"name\":\"__time\",\"type\":\"LONG\"},{\"name\":\"dim1\",\"type\":\"STRING\"},{\"name\":\"dim2\",\"type\":\"STRING\"},{\"name\":\"dim3\",\"type\":\"STRING\"},{\"name\":\"cnt\",\"type\":\"LONG\"},{\"name\":\"m1\",\"type\":\"FLOAT\"},{\"name\":\"m2\",\"type\":\"DOUBLE\"},{\"name\":\"unique_dim1\",\"type\":\"COMPLEX<hyperUnique>\"}]"
-                                    + "}]";
+                                    + "\"signature\":[{\"name\":\"__time\",\"type\":\"LONG\"},{\"name\":\"dim1\",\"type\":\"STRING\"},{\"name\":\"dim2\",\"type\":\"STRING\"},{\"name\":\"dim3\",\"type\":\"STRING\"},{\"name\":\"cnt\",\"type\":\"LONG\"},{\"name\":\"m1\",\"type\":\"FLOAT\"},{\"name\":\"m2\",\"type\":\"DOUBLE\"},{\"name\":\"unique_dim1\",\"type\":\"COMPLEX<hyperUnique>\"}]},"
+                                    + "{\"statementKind\":\"SELECT\"},"
+                                    + "{\"targetDataSource\":\"null\"}"
+                                    + "]";
     String sql = "EXPLAIN PLAN FOR SELECT * FROM druid.foo";
     String resources = "[{\"name\":\"foo\",\"type\":\"DATASOURCE\"}]";
 
@@ -239,11 +247,17 @@ public void testExplainMultipleTopLevelUnionAllQueries()
                                + "{"
                                + "\"query\":{\"queryType\":\"scan\",\"dataSource\":{\"type\":\"table\",\"name\":\"foo\"},\"intervals\":{\"type\":\"intervals\",\"intervals\":[\"-146136543-09-08T08:23:32.096Z/146140482-04-24T15:36:27.903Z\"]},\"resultFormat\":\"compactedList\",\"filter\":{\"type\":\"selector\",\"dimension\":\"dim1\",\"value\":\"44\"},\"columns\":[\"dim1\"],\"legacy\":false,\"context\":{\"defaultTimeout\":300000,\"maxScatterGatherBytes\":9223372036854775807,\"sqlCurrentTimestamp\":\"2000-01-01T00:00:00Z\",\"sqlQueryId\":\"dummy\",\"vectorize\":\"false\",\"vectorizeVirtualColumns\":\"false\"},\"granularity\":{\"type\":\"all\"}},"
                                + "\"signature\":[{\"name\":\"dim1\",\"type\":\"STRING\"}]"
+                               + "},"
+                               + "{"
+                               + "\"statementKind\":\"SELECT\""
+                               + "},"
+                               + "{"
+                               + "\"targetDataSource\":\"null\""
                                + "}]";
     final String resources = "[{\"name\":\"foo\",\"type\":\"DATASOURCE\"}]";
 
     testQuery(
-        PlannerConfig.builder().useNativeQueryExplain(false).build(),
+        PLANNER_CONFIG_LEGACY_QUERY_EXPLAIN,
         query,
         CalciteTests.REGULAR_USER_AUTH_RESULT,
         ImmutableList.of(),
@@ -291,8 +305,10 @@ public void testExplainSelectMvfilterExpressions()
                                                       + "\"legacy\":false,"
                                                       + "\"context\":{\"defaultTimeout\":300000,\"forceExpressionVirtualColumns\":true,\"maxScatterGatherBytes\":9223372036854775807,\"sqlCurrentTimestamp\":\"2000-01-01T00:00:00Z\",\"sqlQueryId\":\"dummy\",\"useNativeQueryExplain\":true,\"vectorize\":\"false\",\"vectorizeVirtualColumns\":\"false\"},"
                                                       + "\"granularity\":{\"type\":\"all\"}},"
-                                                      + "\"signature\":[{\"name\":\"v0\",\"type\":\"STRING\"},{\"name\":\"v1\",\"type\":\"STRING\"}]"
-                                                      + "}]";
+                                                      + "\"signature\":[{\"name\":\"v0\",\"type\":\"STRING\"},{\"name\":\"v1\",\"type\":\"STRING\"}]},"
+                                                      + "{\"statementKind\":\"SELECT\"},"
+                                                      + "{\"targetDataSource\":\"null\"}"
+                                                      + "]";
     final String expectedResources = "[{\"name\":\"foo\",\"type\":\"DATASOURCE\"}]";
 
     testQuery(
@@ -316,9 +332,10 @@ public void testExplainSelectMvfilterExpressions()
                                               + "\"legacy\":false,"
                                               + "\"context\":{\"defaultTimeout\":300000,\"maxScatterGatherBytes\":9223372036854775807,\"sqlCurrentTimestamp\":\"2000-01-01T00:00:00Z\",\"sqlQueryId\":\"dummy\",\"useNativeQueryExplain\":true,\"vectorize\":\"false\",\"vectorizeVirtualColumns\":\"false\"},"
                                               + "\"granularity\":{\"type\":\"all\"}},"
-                                              + "\"signature\":[{\"name\":\"v0\",\"type\":\"STRING\"},{\"name\":\"v1\",\"type\":\"STRING\"}]"
-                                              + "}]";
-
+                                              + "\"signature\":[{\"name\":\"v0\",\"type\":\"STRING\"},{\"name\":\"v1\",\"type\":\"STRING\"}]},"
+                                              + "{\"statementKind\":\"SELECT\"},"
+                                              + "{\"targetDataSource\":\"null\"}"
+                                              + "]";
     final Map<String, Object> mvFilteredContext = new HashMap<>(QUERY_CONTEXT_DEFAULT);
     mvFilteredContext.put(PlannerConfig.CTX_KEY_USE_NATIVE_QUERY_EXPLAIN, true);
 
@@ -355,8 +372,10 @@ public void testExplainSelectTimestampExpression()
                                 + "\"legacy\":false,"
                                 + "\"context\":{\"defaultTimeout\":300000,\"maxScatterGatherBytes\":9223372036854775807,\"sqlCurrentTimestamp\":\"2000-01-01T00:00:00Z\",\"sqlQueryId\":\"dummy\",\"useNativeQueryExplain\":true,\"vectorize\":\"false\",\"vectorizeVirtualColumns\":\"false\"},"
                                 + "\"granularity\":{\"type\":\"all\"}},"
-                                + "\"signature\":[{\"name\":\"v0\",\"type\":\"LONG\"}]"
-                                + "}]";
+                                + "\"signature\":[{\"name\":\"v0\",\"type\":\"LONG\"}]},"
+                                + "{\"statementKind\":\"SELECT\"},"
+                                + "{\"targetDataSource\":\"null\"}"
+                                + "]";
     final String expectedResources = "[{\"name\":\"foo\",\"type\":\"DATASOURCE\"}]";
 
     // Verify the query plan