feat(script): add script compiler (#15113)

* feat(script): add script compiler

* fix

* fix

* fix

* fix

* fix

* add executor

* add timeout check

* fix

Author: andylokandy

Cargo.lock

@@ -46,6 +46,7 @@ members = [
     "src/query/pipeline/sinks",
     "src/query/pipeline/sources",
     "src/query/pipeline/transforms",
+    "src/query/script",
     "src/query/settings",
     "src/query/sql",
     "src/query/storages/common/blocks",

Cargo.toml

@@ -352,6 +352,9 @@ build_exceptions! {
     TenantQuotaUnknown(2902),
     TenantQuotaExceeded(2903),
 
+    // Script error codes.
+    ScriptSemanticError(3001),
+    ScriptExecutionError(3002),
 }
 
 // Storage errors [3001, 4000].

src/common/exception/src/exception_code.rs

@@ -320,10 +320,10 @@ pub enum SetOperator {
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub struct OrderByExpr {
     pub expr: Expr,
-    // Optional `ASC` or `DESC`
+    /// `ASC` or `DESC`
     #[drive(skip)]
     pub asc: Option<bool>,
-    // Optional `NULLS FIRST` or `NULLS LAST`
+    /// `NULLS FIRST` or `NULLS LAST`
     #[drive(skip)]
     pub nulls_first: Option<bool>,
 }
@@ -352,16 +352,16 @@ impl Display for OrderByExpr {
 /// One item of the comma-separated list following `SELECT`
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub enum SelectTarget {
-    // Expression with alias, e.g. `SELECT t.a, b AS a, a+1 AS b FROM t`
+    /// Expression with alias, e.g. `SELECT t.a, b AS a, a+1 AS b FROM t`
     AliasedExpr {
         expr: Box<Expr>,
         alias: Option<Identifier>,
     },
 
-    // Qualified star name, e.g. `SELECT t.*  exclude a, columns(expr) FROM t`.
-    // Columns("pattern_str")
-    // Columns(lambda expression)
-    // For simplicity, star wildcard is involved.
+    /// Qualified star name, e.g. `SELECT t.*  exclude a, columns(expr) FROM t`.
+    /// Columns("pattern_str")
+    /// Columns(lambda expression)
+    /// For simplicity, star wildcard is involved.
     StarColumns {
         qualified: QualifiedName,
         column_filter: Option<ColumnFilter>,

src/query/ast/src/ast/query.rs

@@ -19,7 +19,6 @@ use databend_common_exception::Span;
 
 use crate::ast::Expr;
 use crate::ast::Identifier;
-use crate::ast::Query;
 use crate::ast::Statement;
 use crate::ast::TypeName;
 
@@ -64,15 +63,15 @@ impl Display for VariableDeclare {
 }
 
 #[derive(Debug, Clone, PartialEq)]
-pub struct QueryDeclare {
+pub struct StatementDeclare {
     pub name: Identifier,
-    pub query: Query,
+    pub stmt: Statement,
 }
 
-impl Display for QueryDeclare {
+impl Display for StatementDeclare {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        let QueryDeclare { name, query } = self;
-        write!(f, "{name} RESULTSET := {query}")
+        let StatementDeclare { name, stmt } = self;
+        write!(f, "{name} RESULTSET := {stmt}")
     }
 }
 
@@ -82,9 +81,13 @@ pub enum ScriptStatement {
         span: Span,
         declare: VariableDeclare,
     },
-    LetQuery {
+    LetStatement {
         span: Span,
-        declare: QueryDeclare,
+        declare: StatementDeclare,
+    },
+    RunStatement {
+        span: Span,
+        stmt: Statement,
     },
     Assign {
         span: Span,
@@ -149,17 +152,14 @@ pub enum ScriptStatement {
         results: Vec<Vec<ScriptStatement>>,
         else_result: Option<Vec<ScriptStatement>>,
     },
-    SQLStatement {
-        span: Span,
-        stmt: Statement,
-    },
 }
 
 impl Display for ScriptStatement {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
             ScriptStatement::LetVar { declare, .. } => write!(f, "LET {declare}"),
-            ScriptStatement::LetQuery { declare, .. } => write!(f, "LET {declare}"),
+            ScriptStatement::LetStatement { declare, .. } => write!(f, "LET {declare}"),
+            ScriptStatement::RunStatement { stmt, .. } => write!(f, "{stmt}"),
             ScriptStatement::Assign { name, value, .. } => write!(f, "{name} := {value}"),
             ScriptStatement::Return { value, .. } => {
                 if let Some(value) = value {
@@ -352,7 +352,6 @@ impl Display for ScriptStatement {
                 }
                 write!(f, "END IF")
             }
-            ScriptStatement::SQLStatement { stmt, .. } => write!(f, "{stmt}"),
         }
     }
 }

src/query/ast/src/ast/statements/script.rs

@@ -1480,21 +1480,17 @@ pub fn literal_string(i: Input) -> IResult<String> {
             QuotedString
         },
         |token| {
-            if token
-                .text()
-                .chars()
-                .next()
-                .filter(|c| i.dialect.is_string_quote(*c))
-                .is_some()
-            {
-                let str = &token.text()[1..token.text().len() - 1];
-                let unescaped = unescape_string(str, '\'').ok_or(nom::Err::Failure(
-                    ErrorKind::Other("invalid escape or unicode"),
-                ))?;
-                Ok(unescaped)
-            } else {
-                Err(nom::Err::Error(ErrorKind::ExpectToken(QuotedString)))
+            if let Some(quote) = token.text().chars().next() {
+                if i.dialect.is_string_quote(quote) {
+                    let str = &token.text()[1..token.text().len() - 1];
+                    let unescaped = unescape_string(str, quote).ok_or(nom::Err::Failure(
+                        ErrorKind::Other("invalid escape or unicode"),
+                    ))?;
+                    return Ok(unescaped);
+                }
             }
+
+            Err(nom::Err::Error(ErrorKind::ExpectToken(QuotedString)))
         },
     )(i)
 }

src/query/ast/src/parser/expr.rs

@@ -19,19 +19,22 @@ use crate::ast::*;
 use crate::parser::common::*;
 use crate::parser::expr::*;
 use crate::parser::input::Input;
-use crate::parser::query::*;
 use crate::parser::statement::*;
 use crate::parser::token::*;
 use crate::rule;
 
+pub fn script_stmts(i: Input) -> IResult<Vec<ScriptStatement>> {
+    semicolon_terminated_list1(script_stmt)(i)
+}
+
 pub fn script_stmt(i: Input) -> IResult<ScriptStatement> {
-    let let_query_stmt = map(
+    let let_stmt_stmt = map(
         consumed(rule! {
-            LET ~^#ident ~ RESULTSET ~ ^":=" ~ ^#query
+            LET ~^#ident ~ RESULTSET ~ ^":=" ~ ^#statement_body
         }),
-        |(span, (_, name, _, _, query))| ScriptStatement::LetQuery {
+        |(span, (_, name, _, _, stmt))| ScriptStatement::LetStatement {
             span: transform_span(span.tokens),
-            declare: QueryDeclare { name, query },
+            declare: StatementDeclare { name, stmt },
         },
     );
     let let_var_stmt = map(
@@ -47,6 +50,15 @@ pub fn script_stmt(i: Input) -> IResult<ScriptStatement> {
             },
         },
     );
+    let run_stmt = map(
+        consumed(rule! {
+            #statement_body
+        }),
+        |(span, stmt)| ScriptStatement::RunStatement {
+            span: transform_span(span.tokens),
+            stmt,
+        },
+    );
     let assign_stmt = map(
         consumed(rule! {
             #ident ~ ":=" ~ ^#expr
@@ -199,19 +211,11 @@ pub fn script_stmt(i: Input) -> IResult<ScriptStatement> {
             }
         },
     );
-    let sql_stmt = map(
-        consumed(rule! {
-            #statement_body
-        }),
-        |(span, stmt)| ScriptStatement::SQLStatement {
-            span: transform_span(span.tokens),
-            stmt,
-        },
-    );
 
     rule!(
-        #let_query_stmt
+        #let_stmt_stmt
         | #let_var_stmt
+        | #run_stmt
         | #assign_stmt
         | #return_stmt
         | #for_loop_stmt
@@ -223,6 +227,5 @@ pub fn script_stmt(i: Input) -> IResult<ScriptStatement> {
         | #continue_stmt
         | #case_stmt
         | #if_stmt
-        | #sql_stmt
     )(i)
 }

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

@@ -28,6 +28,7 @@ pub fn unescape_string(s: &str, quote: char) -> Option<String> {
                 Some('r') => s.push('\r'),
                 Some('t') => s.push('\t'),
                 Some('\'') => s.push('\''),
+                Some('"') => s.push('"'),
                 Some('\\') => s.push('\\'),
                 Some('u') => s.push(unescape_unicode(&mut chars)?),
                 Some('x') => s.push(unescape_byte(&mut chars)?),

src/query/ast/src/parser/unescape.rs

@@ -782,7 +782,7 @@ fn test_statement_error() {
 #[test]
 fn test_raw_insert_stmt() {
     let mut mint = Mint::new("tests/it/testdata");
-    let file = &mut mint.new_goldenfile("raw_insert.txt").unwrap();
+    let file = &mut mint.new_goldenfile("raw-insert.txt").unwrap();
     let cases = &[
         r#"insert into t (c1, c2) values (1, 2), (3, 4);"#,
         r#"insert into t (c1, c2) values (1, 2);"#,
@@ -978,48 +978,72 @@ fn test_expr() {
 }
 
 #[test]
-fn test_experimental_expr() {
+fn test_expr_error() {
     let mut mint = Mint::new("tests/it/testdata");
-    let file = &mut mint.new_goldenfile("experimental-expr.txt").unwrap();
+    let file = &mut mint.new_goldenfile("expr-error.txt").unwrap();
 
     let cases = &[
-        r#"a"#,
+        r#"5 * (a and ) 1"#,
+        r#"a + +"#,
+        r#"CAST(col1 AS foo)"#,
+        r#"1 a"#,
+        r#"CAST(col1)"#,
         r#"a.add(b)"#,
-        r#"a.sub(b).add(e)"#,
-        r#"a.sub(b).add(e)"#,
-        r#"1 + {'k1': 4}.k1"#,
-        r#"'3'.plus(4)"#,
-        r#"(3).add({'k1': 4 }.k1)"#,
         r#"[ x * 100 FOR x in [1,2,3] if x % 2 = 0 ]"#,
+        r#"
+            G.E.B IS NOT NULL
+            AND col1 NOT BETWEEN col2 AND
+            AND 1 + col3 DIV sum(col4)
+        "#,
     ];
 
     for case in cases {
-        run_parser_with_dialect(file, expr, Dialect::Experimental, ParseMode::Default, case);
+        run_parser(file, expr, case);
     }
 }
 
 #[test]
-fn test_expr_error() {
+fn test_dialect() {
     let mut mint = Mint::new("tests/it/testdata");
-    let file = &mut mint.new_goldenfile("expr-error.txt").unwrap();
+    let file = &mut mint.new_goldenfile("dialect.txt").unwrap();
 
     let cases = &[
-        r#"5 * (a and ) 1"#,
-        r#"a + +"#,
-        r#"CAST(col1 AS foo)"#,
-        r#"1 a"#,
-        r#"CAST(col1)"#,
+        r#"'a'"#,
+        r#""a""#,
+        r#"`a`"#,
+        r#"'a''b'"#,
+        r#"'a""b'"#,
+        r#"'a\'b'"#,
+        r#"'a"b'"#,
+        r#"'a`b'"#,
+        r#""a''b""#,
+        r#""a""b""#,
+        r#""a'b""#,
+        r#""a\"b""#,
+        r#""a`b""#,
+    ];
+
+    for case in cases {
+        run_parser_with_dialect(file, expr, Dialect::PostgreSQL, ParseMode::Default, case);
+    }
+
+    for case in cases {
+        run_parser_with_dialect(file, expr, Dialect::MySQL, ParseMode::Default, case);
+    }
+
+    let cases = &[
+        r#"a"#,
         r#"a.add(b)"#,
+        r#"a.sub(b).add(e)"#,
+        r#"a.sub(b).add(e)"#,
+        r#"1 + {'k1': 4}.k1"#,
+        r#"'3'.plus(4)"#,
+        r#"(3).add({'k1': 4 }.k1)"#,
         r#"[ x * 100 FOR x in [1,2,3] if x % 2 = 0 ]"#,
-        r#"
-            G.E.B IS NOT NULL AND
-                col1 NOT BETWEEN col2 AND
-                    AND 1 + col3 DIV sum(col4)
-        "#,
     ];
 
     for case in cases {
-        run_parser(file, expr, case);
+        run_parser_with_dialect(file, expr, Dialect::Experimental, ParseMode::Default, case);
     }
 }