fix(query): forbiden explain explain statement

databendlabs · Oct 22, 2024 · b3cd13e · b3cd13e
1 parent fcf8e09
commit b3cd13e
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 40 deletions.
diff --git a/src/query/ast/src/parser/statement.rs b/src/query/ast/src/parser/statement.rs
@@ -56,55 +56,71 @@ pub enum CreateDatabaseOption {
 pub fn statement_body(i: Input) -> IResult<Statement> {
     let explain = map_res(
         rule! {
-            EXPLAIN ~ ( "(" ~ #comma_separated_list1(explain_option) ~ ")" )? ~ ( AST | SYNTAX | PIPELINE | JOIN | GRAPH | FRAGMENTS | RAW | OPTIMIZED | MEMO )? ~ #statement
+            EXPLAIN ~ ( "(" ~ #comma_separated_list1(explain_option) ~ ")" )? ~ ( AST | SYNTAX | PIPELINE | JOIN | GRAPH | FRAGMENTS | RAW | OPTIMIZED | MEMO )? ~ #statement_body
         },
         |(_, options, opt_kind, statement)| {
-            Ok(Statement::Explain {
-                kind: match opt_kind.map(|token| token.kind) {
-                    Some(TokenKind::SYNTAX) | Some(TokenKind::AST) => {
-                        let pretty_stmt =
-                            pretty_statement(statement.stmt.clone(), 10).map_err(|_| {
-                                nom::Err::Failure(ErrorKind::Other("invalid statement"))
-                            })?;
-                        ExplainKind::Syntax(pretty_stmt)
-                    }
-                    Some(TokenKind::PIPELINE) => ExplainKind::Pipeline,
-                    Some(TokenKind::JOIN) => ExplainKind::Join,
-                    Some(TokenKind::GRAPH) => ExplainKind::Graph,
-                    Some(TokenKind::FRAGMENTS) => ExplainKind::Fragments,
-                    Some(TokenKind::RAW) => ExplainKind::Raw,
-                    Some(TokenKind::OPTIMIZED) => ExplainKind::Optimized,
-                    Some(TokenKind::MEMO) => ExplainKind::Memo("".to_string()),
-                    Some(TokenKind::GRAPHICAL) => ExplainKind::Graphical,
-                    None => ExplainKind::Plan,
-                    _ => unreachable!(),
+            match statement {
+                Statement::Explain {..}=> {
+                    return Err(nom::Err::Failure(ErrorKind::Other("invalid statement")));
                 },
-                options: options.as_ref().map_or(vec![], |(_, opts, _)| opts.clone()),
-                query: Box::new(statement.stmt),
-            })
+                _ => {
+                    Ok(Statement::Explain {
+                        kind: match opt_kind.map(|token| token.kind) {
+                            Some(SYNTAX) | Some(AST) => {
+                                let pretty_stmt =
+                                    pretty_statement(statement.clone(), 10).map_err(|_| {
+                                        nom::Err::Failure(ErrorKind::Other("invalid statement"))
+                                    })?;
+                                ExplainKind::Syntax(pretty_stmt)
+                            }
+                            Some(PIPELINE) => ExplainKind::Pipeline,
+                            Some(JOIN) => ExplainKind::Join,
+                            Some(GRAPH) => ExplainKind::Graph,
+                            Some(FRAGMENTS) => ExplainKind::Fragments,
+                            Some(RAW) => ExplainKind::Raw,
+                            Some(OPTIMIZED) => ExplainKind::Optimized,
+                            Some(MEMO) => ExplainKind::Memo("".to_string()),
+                            Some(GRAPHICAL) => ExplainKind::Graphical,
+                            None => ExplainKind::Plan,
+                            _ => unreachable!(),
+                        },
+                        options: options.as_ref().map_or(vec![], |(_, opts, _)| opts.clone()),
+                        query: Box::new(statement),
+                    })
+                }
+            }
+
         },
     );
-    let explain_analyze = map(
+    let explain_analyze = map_res(
         rule! {
-            EXPLAIN ~ ANALYZE ~ (PARTIAL|GRAPHICAL)? ~ #statement
+            EXPLAIN ~ ANALYZE ~ (PARTIAL|GRAPHICAL)? ~ #statement_body
         },
         |(_, _, opt_partial_or_graphical, statement)| {
-            let (partial, graphical) = match opt_partial_or_graphical {
-                Some(Token {
-                    kind: TokenKind::PARTIAL,
-                    ..
-                }) => (true, false),
-                Some(Token {
-                    kind: TokenKind::GRAPHICAL,
-                    ..
-                }) => (false, true),
-                _ => (false, false),
-            };
-            Statement::ExplainAnalyze {
-                partial,
-                graphical,
-                query: Box::new(statement.stmt),
+            match statement {
+                Statement::ExplainAnalyze { .. } => {
+                    return Err(nom::Err::Failure(ErrorKind::Other("invalid statement")));
+                }
+                _ => {
+                    let (partial, graphical) = match opt_partial_or_graphical {
+                        Some(Token {
+                                 kind: PARTIAL,
+                                 ..
+                             }) => (true, false),
+                        Some(Token {
+                                 kind: GRAPHICAL,
+                                 ..
+                             }) => (false, true),
+                        _ => (false, false),
+                    };
+                    Ok(Statement::ExplainAnalyze {
+                        partial,
+                        graphical,
+                        query: Box::new(statement),
+                    })
+                }
             }
+
         },
     );
 

diff --git a/tests/sqllogictests/suites/mode/standalone/explain/explain.test b/tests/sqllogictests/suites/mode/standalone/explain/explain.test
@@ -10,6 +10,12 @@ create table t1 as select number as a, number as b from numbers(1)
 statement ok
 create table t2 as select number as a, number as b from numbers(5)
 
+statement error 1005
+explain explain select t1.a from t1 where a > 0
+
+statement error 1005
+explain explain analyze select t1.a from t1 where a > 0
+
 query T
 explain select t1.a from t1 where a > 0
 ----