support EXPLAIN and EXPLAIN ANALYZE for EXECUTE prepared statement (#17318)

Author: Young-Leo

integration-test/src/test/java/org/apache/iotdb/relational/it/insertquery/IoTDBInsertQueryIT.java

@@ -459,7 +459,8 @@ public void testExplain() throws SQLException {
       Assert.assertTrue(
           e.getMessage(),
           e.getMessage()
-              .contains("700: line 1:9: mismatched input 'INSERT'. Expecting: 'ANALYZE', <query>"));
+              .contains(
+                  "700: line 1:9: mismatched input 'INSERT'. Expecting: 'ANALYZE', 'EXECUTE', <query>"));
     }
 
     try {
@@ -471,7 +472,7 @@ public void testExplain() throws SQLException {
           e.getMessage(),
           e.getMessage()
               .contains(
-                  "700: line 1:17: mismatched input 'INSERT'. Expecting: 'VERBOSE', <query>"));
+                  "700: line 1:17: mismatched input 'INSERT'. Expecting: 'EXECUTE', 'VERBOSE', <query>"));
     }
   }
 

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/Coordinator.java

@@ -86,6 +86,8 @@
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.DropTable;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Execute;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ExecuteImmediate;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Explain;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ExplainAnalyze;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Expression;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ExtendRegion;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Flush;
@@ -689,39 +691,59 @@ private IQueryExecution createQueryExecutionForTableModel(
     List<Expression> parameters = Collections.emptyList();
     Map<NodeRef<Parameter>, Expression> parameterLookup = Collections.emptyMap();
 
-    if (statement instanceof Execute) {
-      Execute executeStatement = (Execute) statement;
+    // Unwrap Explain/ExplainAnalyze to check for inner Execute/ExecuteImmediate
+    org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Statement innerStatement = statement;
+    if (statement instanceof Explain) {
+      innerStatement = ((Explain) statement).getStatement();
+    } else if (statement instanceof ExplainAnalyze) {
+      innerStatement = ((ExplainAnalyze) statement).getStatement();
+    }
+
+    if (innerStatement instanceof Execute) {
+      Execute executeStatement = (Execute) innerStatement;
       String statementName = executeStatement.getStatementName().getValue();
 
-      // Get prepared statement from session (contains cached AST)
       PreparedStatementInfo preparedInfo = clientSession.getPreparedStatement(statementName);
       if (preparedInfo == null) {
         throw new SemanticException(
             String.format("Prepared statement '%s' does not exist", statementName));
       }
 
-      // Use cached AST
-      statementToUse = preparedInfo.getSql();
-
-      // Bind parameters: create parameterLookup map
-      // Note: bindParameters() internally validates parameter count
+      org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Statement resolvedSql =
+          preparedInfo.getSql();
       parameterLookup =
-          ParameterExtractor.bindParameters(statementToUse, executeStatement.getParameters());
+          ParameterExtractor.bindParameters(resolvedSql, executeStatement.getParameters());
       parameters = new ArrayList<>(executeStatement.getParameters());
 
-    } else if (statement instanceof ExecuteImmediate) {
-      ExecuteImmediate executeImmediateStatement = (ExecuteImmediate) statement;
+      if (statement instanceof Explain) {
+        statementToUse = new Explain(resolvedSql);
+      } else if (statement instanceof ExplainAnalyze) {
+        statementToUse = new ExplainAnalyze(resolvedSql, ((ExplainAnalyze) statement).isVerbose());
+      } else {
+        statementToUse = resolvedSql;
+      }
+
+    } else if (innerStatement instanceof ExecuteImmediate) {
+      ExecuteImmediate executeImmediateStatement = (ExecuteImmediate) innerStatement;
 
-      // EXECUTE IMMEDIATE needs to parse SQL first
       String sql = executeImmediateStatement.getSqlString();
       List<Literal> literalParameters = executeImmediateStatement.getParameters();
 
-      statementToUse = sqlParser.createStatement(sql, clientSession.getZoneId(), clientSession);
+      org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Statement resolvedSql =
+          sqlParser.createStatement(sql, clientSession.getZoneId(), clientSession);
 
       if (!literalParameters.isEmpty()) {
-        parameterLookup = ParameterExtractor.bindParameters(statementToUse, literalParameters);
+        parameterLookup = ParameterExtractor.bindParameters(resolvedSql, literalParameters);
         parameters = new ArrayList<>(literalParameters);
       }
+
+      if (statement instanceof Explain) {
+        statementToUse = new Explain(resolvedSql);
+      } else if (statement instanceof ExplainAnalyze) {
+        statementToUse = new ExplainAnalyze(resolvedSql, ((ExplainAnalyze) statement).isVerbose());
+      } else {
+        statementToUse = resolvedSql;
+      }
     }
 
     if (statement instanceof WrappedInsertStatement) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/parser/AstBuilder.java

@@ -1790,13 +1790,28 @@ public Node visitStatementDefault(RelationalSqlParser.StatementDefaultContext ct
 
   @Override
   public Node visitExplain(RelationalSqlParser.ExplainContext ctx) {
-    return new Explain(getLocation(ctx), (Statement) visit(ctx.query()));
+    Statement innerStatement;
+    if (ctx.query() != null) {
+      innerStatement = (Statement) visit(ctx.query());
+    } else if (ctx.executeStatement() != null) {
+      innerStatement = (Statement) visit(ctx.executeStatement());
+    } else {
+      innerStatement = (Statement) visit(ctx.executeImmediateStatement());
+    }
+    return new Explain(getLocation(ctx), innerStatement);
   }
 
   @Override
   public Node visitExplainAnalyze(RelationalSqlParser.ExplainAnalyzeContext ctx) {
-    return new ExplainAnalyze(
-        getLocation(ctx), ctx.VERBOSE() != null, (Statement) visit(ctx.query()));
+    Statement innerStatement;
+    if (ctx.query() != null) {
+      innerStatement = (Statement) visit(ctx.query());
+    } else if (ctx.executeStatement() != null) {
+      innerStatement = (Statement) visit(ctx.executeStatement());
+    } else {
+      innerStatement = (Statement) visit(ctx.executeImmediateStatement());
+    }
+    return new ExplainAnalyze(getLocation(ctx), ctx.VERBOSE() != null, innerStatement);
   }
 
   // ********************** author expressions ********************

iotdb-core/relational-grammar/src/main/antlr4/org/apache/iotdb/db/relational/grammar/sql/RelationalSql.g4

@@ -911,8 +911,8 @@ deallocateStatement
 // ------------------------------------------- Query Statement ---------------------------------------------------------
 queryStatement
     : query                                                        #statementDefault
-    | EXPLAIN query                                                #explain
-    | EXPLAIN ANALYZE VERBOSE? query                               #explainAnalyze
+    | EXPLAIN (query | executeStatement | executeImmediateStatement) #explain
+    | EXPLAIN ANALYZE VERBOSE? (query | executeStatement | executeImmediateStatement) #explainAnalyze
     ;
 
 query