feat(query): better script engine with dynamic schema support

src/query/ast/src/ast/statements/call.rs

@@ -19,10 +19,11 @@ use derive_visitor::Drive;
 use derive_visitor::DriveMut;
 
 use crate::ast::write_comma_separated_string_list;
+use crate::ast::Identifier;
 
 #[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
 pub struct CallStmt {
-    pub name: String,
+    pub name: Identifier,
     pub args: Vec<String>,
 }
 

src/query/ast/src/ast/statements/procedure.rs

@@ -21,16 +21,17 @@ use derive_visitor::DriveMut;
 use crate::ast::write_comma_separated_list;
 use crate::ast::CreateOption;
 use crate::ast::Expr;
+use crate::ast::Identifier;
 use crate::ast::TypeName;
 
-#[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub struct ExecuteImmediateStmt {
-    pub script: String,
+    pub script: Expr,
 }
 
 impl Display for ExecuteImmediateStmt {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
-        write!(f, "EXECUTE IMMEDIATE $$\n{}\n$$", self.script)?;
+        write!(f, "EXECUTE IMMEDIATE {}", self.script)?;
         Ok(())
     }
 }
@@ -170,7 +171,7 @@ impl Display for DescProcedureStmt {
 
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub struct CallProcedureStmt {
-    pub name: String,
+    pub name: Identifier,
     pub args: Vec<Expr>,
 }
 

src/query/ast/src/ast/statements/task.rs

@@ -24,6 +24,7 @@ use crate::ast::quote::QuotedString;
 use crate::ast::write_comma_separated_string_list;
 use crate::ast::write_comma_separated_string_map;
 use crate::ast::Expr;
+use crate::ast::Identifier;
 use crate::ast::ShowLimit;
 
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
@@ -370,7 +371,7 @@ impl Display for ShowTasksStmt {
 
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub struct ExecuteTaskStmt {
-    pub name: String,
+    pub name: Identifier,
 }
 
 impl Display for ExecuteTaskStmt {
@@ -381,7 +382,7 @@ impl Display for ExecuteTaskStmt {
 
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub struct DescribeTaskStmt {
-    pub name: String,
+    pub name: Identifier,
 }
 
 impl Display for DescribeTaskStmt {

src/query/ast/src/parser/script.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use nom::branch::alt;
 use nom::combinator::consumed;
 use nom::combinator::map;
 use nom_rule::rule;
@@ -23,6 +24,25 @@ use crate::parser::input::Input;
 use crate::parser::statement::*;
 use crate::parser::token::*;
 
+#[allow(clippy::large_enum_variant)]
+#[derive(Debug, Clone, PartialEq)]
+pub enum ScriptBlockOrStmt {
+    ScriptBlock(ScriptBlock),
+    Statement(Statement),
+}
+
+pub fn script_block_or_stmt(i: Input) -> IResult<ScriptBlockOrStmt> {
+    alt((
+        map(script_block, ScriptBlockOrStmt::ScriptBlock),
+        map(
+            consumed(rule! {
+                #statement
+            }),
+            |(_, stmt)| ScriptBlockOrStmt::Statement(stmt.stmt),
+        ),
+    ))(i)
+}
+
 pub fn script_block(i: Input) -> IResult<ScriptBlock> {
     map(
         consumed(rule! {

src/query/ast/src/parser/statement.rs

@@ -226,22 +226,14 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         rule! {
             EXECUTE ~ TASK ~ #ident
         },
-        |(_, _, task)| {
-            Statement::ExecuteTask(ExecuteTaskStmt {
-                name: task.to_string(),
-            })
-        },
+        |(_, _, task)| Statement::ExecuteTask(ExecuteTaskStmt { name: task }),
     );
 
     let desc_task = map(
         rule! {
             ( DESC | DESCRIBE ) ~ TASK ~ #ident
         },
-        |(_, _, task)| {
-            Statement::DescribeTask(DescribeTaskStmt {
-                name: task.to_string(),
-            })
-        },
+        |(_, _, task)| Statement::DescribeTask(DescribeTaskStmt { name: task }),
     );
 
     let merge = map(
@@ -2005,12 +1997,7 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         rule! {
             CALL ~ #ident ~ "(" ~ #comma_separated_list0(parameter_to_string) ~ ")"
         },
-        |(_, name, _, args, _)| {
-            Statement::Call(CallStmt {
-                name: name.to_string(),
-                args,
-            })
-        },
+        |(_, name, _, args, _)| Statement::Call(CallStmt { name, args }),
     );
 
     let vacuum_temporary_tables = map(
@@ -2019,7 +2006,7 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         },
         |(_, _, _, opt_limit)| {
             Statement::Call(CallStmt {
-                name: "fuse_vacuum_temporary_table".to_string(),
+                name: Identifier::from_name(None, "fuse_vacuum_temporary_table"),
                 args: opt_limit.map(|v| v.1.to_string()).into_iter().collect(),
             })
         },
@@ -2418,7 +2405,7 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
 
     let execute_immediate = map(
         rule! {
-            EXECUTE ~ IMMEDIATE ~ #code_string
+            EXECUTE ~ IMMEDIATE ~ #expr
         },
         |(_, _, script)| Statement::ExecuteImmediate(ExecuteImmediateStmt { script }),
     );
@@ -2556,7 +2543,7 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         },
         |(_, _, name, _, opt_args, _)| {
             Statement::CallProcedure(CallProcedureStmt {
-                name: name.to_string(),
+                name,
                 args: opt_args.unwrap_or_default(),
             })
         },

src/query/ast/tests/it/parser.rs

@@ -1579,6 +1579,10 @@ fn test_script() {
         r#"select :a + 1"#,
         r#"select IDENTIFIER(:b)"#,
         r#"select a.IDENTIFIER(:b).c + minus(:d)"#,
+        r#"EXECUTE TASK IDENTIFIER(:my_task)"#,
+        r#"DESC TASK IDENTIFIER(:my_task)"#,
+        r#"CALL IDENTIFIER(:test)(a)"#,
+        r#"call PROCEDURE IDENTIFIER(:proc_name)()"#,
     ];
 
     for case in cases {

src/query/ast/tests/it/testdata/script.txt

@@ -1917,6 +1917,106 @@ RunStatement {
 }
 
 
+---------- Input ----------
+EXECUTE TASK IDENTIFIER(:my_task)
+---------- Output ---------
+EXECUTE TASK IDENTIFIER(:my_task)
+---------- AST ------------
+RunStatement {
+    span: Some(
+        0..33,
+    ),
+    stmt: ExecuteTask(
+        ExecuteTaskStmt {
+            name: Identifier {
+                span: Some(
+                    13..33,
+                ),
+                name: "my_task",
+                quote: None,
+                ident_type: Hole,
+            },
+        },
+    ),
+}
+
+
+---------- Input ----------
+DESC TASK IDENTIFIER(:my_task)
+---------- Output ---------
+DESCRIBE TASK IDENTIFIER(:my_task)
+---------- AST ------------
+RunStatement {
+    span: Some(
+        0..30,
+    ),
+    stmt: DescribeTask(
+        DescribeTaskStmt {
+            name: Identifier {
+                span: Some(
+                    10..30,
+                ),
+                name: "my_task",
+                quote: None,
+                ident_type: Hole,
+            },
+        },
+    ),
+}
+
+
+---------- Input ----------
+CALL IDENTIFIER(:test)(a)
+---------- Output ---------
+CALL IDENTIFIER(:test)('a')
+---------- AST ------------
+RunStatement {
+    span: Some(
+        0..25,
+    ),
+    stmt: Call(
+        CallStmt {
+            name: Identifier {
+                span: Some(
+                    5..22,
+                ),
+                name: "test",
+                quote: None,
+                ident_type: Hole,
+            },
+            args: [
+                "a",
+            ],
+        },
+    ),
+}
+
+
+---------- Input ----------
+call PROCEDURE IDENTIFIER(:proc_name)()
+---------- Output ---------
+CALL PROCEDURE IDENTIFIER(:proc_name)()
+---------- AST ------------
+RunStatement {
+    span: Some(
+        0..39,
+    ),
+    stmt: CallProcedure(
+        CallProcedureStmt {
+            name: Identifier {
+                span: Some(
+                    15..37,
+                ),
+                name: "proc_name",
+                quote: None,
+                ident_type: Hole,
+            },
+            args: [],
+        },
+    ),
+}
+
+
 ---------- Input ----------
 BEGIN
     LOOP