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

* feat: better script engine

* update

* update

* update

* fix tests

* update call procedure

* support cursor

* support cursor

* support cursor

* support cursor

* support cursor

* update

* update

Author: sundy-li

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/script.rs

@@ -115,6 +115,47 @@ impl Display for ReturnItem {
     }
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct DeclareCursor {
+    pub span: Span,
+    pub name: Identifier,
+    pub stmt: Option<Statement>,
+    pub resultset: Option<Identifier>,
+}
+
+impl Display for DeclareCursor {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        let DeclareCursor {
+            name,
+            stmt,
+            resultset,
+            ..
+        } = self;
+        if let Some(stmt) = stmt {
+            write!(f, "{name} CURSOR FOR {stmt}")
+        } else if let Some(resultset) = resultset {
+            write!(f, "{name} CURSOR FOR {resultset}")
+        } else {
+            write!(f, "{name} CURSOR")
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum IterableItem {
+    Resultset(Identifier),
+    Cursor(Identifier),
+}
+
+impl Display for IterableItem {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        match self {
+            IterableItem::Resultset(name) => write!(f, "{name}"),
+            IterableItem::Cursor(name) => write!(f, "{name}"),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq)]
 pub enum ScriptStatement {
     LetVar {
@@ -123,6 +164,9 @@ pub enum ScriptStatement {
     LetStatement {
         declare: DeclareSet,
     },
+    LetCursor {
+        declare: DeclareCursor,
+    },
     RunStatement {
         span: Span,
         stmt: Statement,
@@ -132,6 +176,19 @@ pub enum ScriptStatement {
         name: Identifier,
         value: Expr,
     },
+    OpenCursor {
+        span: Span,
+        cursor: Identifier,
+    },
+    FetchCursor {
+        span: Span,
+        cursor: Identifier,
+        into_var: Identifier,
+    },
+    CloseCursor {
+        span: Span,
+        cursor: Identifier,
+    },
     Return {
         span: Span,
         value: Option<ReturnItem>,
@@ -148,7 +205,7 @@ pub enum ScriptStatement {
     ForInSet {
         span: Span,
         variable: Identifier,
-        resultset: Identifier,
+        iterable: IterableItem,
         body: Vec<ScriptStatement>,
         label: Option<Identifier>,
     },
@@ -204,8 +261,14 @@ impl Display for ScriptStatement {
         match self {
             ScriptStatement::LetVar { declare, .. } => write!(f, "LET {declare}"),
             ScriptStatement::LetStatement { declare, .. } => write!(f, "LET {declare}"),
+            ScriptStatement::LetCursor { declare, .. } => write!(f, "LET {declare}"),
             ScriptStatement::RunStatement { stmt, .. } => write!(f, "{stmt}"),
             ScriptStatement::Assign { name, value, .. } => write!(f, "{name} := {value}"),
+            ScriptStatement::OpenCursor { cursor, .. } => write!(f, "OPEN {cursor}"),
+            ScriptStatement::FetchCursor {
+                cursor, into_var, ..
+            } => write!(f, "FETCH {cursor} INTO {into_var}"),
+            ScriptStatement::CloseCursor { cursor, .. } => write!(f, "CLOSE {cursor}"),
             ScriptStatement::Return { value, .. } => {
                 if let Some(value) = value {
                     write!(f, "RETURN {value}")
@@ -242,12 +305,12 @@ impl Display for ScriptStatement {
             }
             ScriptStatement::ForInSet {
                 variable,
-                resultset,
+                iterable,
                 body,
                 label,
                 ..
             } => {
-                writeln!(f, "FOR {variable} IN {resultset} DO")?;
+                writeln!(f, "FOR {variable} IN {iterable} DO")?;
                 for stmt in body {
                     writeln!(
                         f,

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! {
@@ -79,6 +99,53 @@ pub fn declare_set(i: Input) -> IResult<DeclareSet> {
     )(i)
 }
 
+pub fn declare_cursor(i: Input) -> IResult<DeclareCursor> {
+    map(
+        consumed(rule! {
+            #ident ~ CURSOR ~ ^FOR ~ ^#cursor_target
+        }),
+        |(span, (name, _, _, target))| match target {
+            CursorTarget::Resultset(resultset) => DeclareCursor {
+                span: transform_span(span.tokens),
+                name,
+                stmt: None,
+                resultset: Some(resultset),
+            },
+            CursorTarget::Statement(stmt) => DeclareCursor {
+                span: transform_span(span.tokens),
+                name,
+                stmt: Some(stmt),
+                resultset: None,
+            },
+        },
+    )(i)
+}
+
+#[allow(clippy::large_enum_variant)]
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) enum CursorTarget {
+    Resultset(Identifier),
+    Statement(Statement),
+}
+
+pub(crate) fn cursor_target(i: Input) -> IResult<CursorTarget> {
+    // Try identifier first, then statement
+    let resultset = map(ident, CursorTarget::Resultset);
+    let statement = map(statement_body, CursorTarget::Statement);
+
+    rule!(
+        #resultset
+        | #statement
+    )(i)
+}
+
+pub(crate) fn iterable_item(i: Input) -> IResult<IterableItem> {
+    // For now, we'll treat all identifiers as potential iterables
+    // The compiler will determine if it's a cursor or resultset
+    // based on what was actually declared
+    map(ident, IterableItem::Resultset)(i)
+}
+
 pub fn script_stmts(i: Input) -> IResult<Vec<ScriptStatement>> {
     semicolon_terminated_list1(script_stmt)(i)
 }
@@ -96,6 +163,40 @@ pub fn script_stmt(i: Input) -> IResult<ScriptStatement> {
         },
         |(_, declare)| ScriptStatement::LetStatement { declare },
     );
+    let let_cursor_stmt = map(
+        rule! {
+            LET ~ #declare_cursor
+        },
+        |(_, declare)| ScriptStatement::LetCursor { declare },
+    );
+    let open_cursor_stmt = map(
+        consumed(rule! {
+            OPEN ~ #ident
+        }),
+        |(span, (_, cursor))| ScriptStatement::OpenCursor {
+            span: transform_span(span.tokens),
+            cursor,
+        },
+    );
+    let fetch_cursor_stmt = map(
+        consumed(rule! {
+            FETCH ~ #ident ~ ^INTO ~ ^#ident
+        }),
+        |(span, (_, cursor, _, into_var))| ScriptStatement::FetchCursor {
+            span: transform_span(span.tokens),
+            cursor,
+            into_var,
+        },
+    );
+    let close_cursor_stmt = map(
+        consumed(rule! {
+            CLOSE ~ #ident
+        }),
+        |(span, (_, cursor))| ScriptStatement::CloseCursor {
+            span: transform_span(span.tokens),
+            cursor,
+        },
+    );
     let run_stmt = map(
         consumed(rule! {
             #statement_body
@@ -173,14 +274,14 @@ pub fn script_stmt(i: Input) -> IResult<ScriptStatement> {
     );
     let for_in_set_stmt = map(
         consumed(rule! {
-            FOR ~ ^#ident ~ ^IN ~ #ident ~ ^DO
+            FOR ~ ^#ident ~ ^IN ~ #iterable_item ~ ^DO
             ~ ^#semicolon_terminated_list1(script_stmt)
             ~ ^END ~ ^FOR ~ #ident?
         }),
-        |(span, (_, variable, _, resultset, _, body, _, _, label))| ScriptStatement::ForInSet {
+        |(span, (_, variable, _, iterable, _, body, _, _, label))| ScriptStatement::ForInSet {
             span: transform_span(span.tokens),
             variable,
-            resultset,
+            iterable,
             body,
             label,
         },
@@ -299,24 +400,48 @@ pub fn script_stmt(i: Input) -> IResult<ScriptStatement> {
         },
     );
 
-    rule!(
+    let cursor_stmts = rule!(
+        #let_cursor_stmt
+        | #open_cursor_stmt
+        | #fetch_cursor_stmt
+        | #close_cursor_stmt
+    );
+
+    let assignment_stmts = rule!(
         #let_stmt_stmt
         | #let_var_stmt
-        | #run_stmt
         | #assign_stmt
-        | #return_set_stmt
+    );
+
+    let control_flow_stmts = rule!(
+        #return_set_stmt
         | #return_stmt_stmt
         | #return_var_stmt
         | #return_stmt
-        | #for_loop_stmt
+        | #break_stmt
+        | #continue_stmt
+    );
+
+    let loop_stmts = rule!(
+        #for_loop_stmt
         | #for_in_set_stmt
         | #for_in_stmt_stmt
         | #while_loop_stmt
         | #repeat_loop_stmt
         | #loop_stmt
-        | #break_stmt
-        | #continue_stmt
-        | #case_stmt
+    );
+
+    let conditional_stmts = rule!(
+        #case_stmt
         | #if_stmt
+    );
+
+    rule!(
+        #assignment_stmts
+        | #cursor_stmts
+        | #control_flow_stmts
+        | #loop_stmts
+        | #conditional_stmts
+        | #run_stmt
     )(i)
 }