feat(query): add THROW support to SQL procedures (#19067)

Author: TCeason

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

@@ -207,6 +207,10 @@ pub enum ScriptStatement {
         span: Span,
         value: Option<ReturnItem>,
     },
+    Throw {
+        span: Span,
+        message: Option<Expr>,
+    },
     ForLoop {
         span: Span,
         variable: Identifier,
@@ -290,6 +294,13 @@ impl Display for ScriptStatement {
                     write!(f, "RETURN")
                 }
             }
+            ScriptStatement::Throw { message, .. } => {
+                if let Some(message) = message {
+                    write!(f, "THROW {message}")
+                } else {
+                    write!(f, "THROW")
+                }
+            }
             ScriptStatement::ForLoop {
                 variable,
                 is_reverse,

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

@@ -269,6 +269,15 @@ pub fn script_stmt(i: Input) -> IResult<ScriptStatement> {
             value: None,
         },
     );
+    let throw_stmt = map(
+        consumed(rule! {
+            THROW ~ #expr?
+        }),
+        |(span, (_, message))| ScriptStatement::Throw {
+            span: transform_span(span.tokens),
+            message,
+        },
+    );
     let for_loop_stmt = map(
         consumed(rule! {
             FOR ~ ^#ident ~ ^IN ~ REVERSE?
@@ -439,6 +448,7 @@ pub fn script_stmt(i: Input) -> IResult<ScriptStatement> {
         | #return_stmt_stmt
         | #return_var_stmt
         | #return_stmt
+        | #throw_stmt
         | #break_stmt
         | #continue_stmt
     );

src/query/ast/src/parser/token.rs

@@ -1240,6 +1240,8 @@ pub enum TokenKind {
     TENANT,
     #[token("THEN", ignore(ascii_case))]
     THEN,
+    #[token("THROW", ignore(ascii_case))]
+    THROW,
     #[token("THURSDAY", ignore(ascii_case))]
     THURSDAY,
     #[token("TIME", ignore(ascii_case))]

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

@@ -1555,6 +1555,7 @@ fn test_script() {
         r#"RETURN profit"#,
         r#"RETURN TABLE(t1)"#,
         r#"RETURN TABLE(select count(*) from t1)"#,
+        r#"THROW 'Email already exists.'"#,
         r#"
             FOR i IN REVERSE 1 TO maximum_count DO
                 counter := counter + 1;

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

@@ -536,6 +536,28 @@ Return {
 }
 
 
+---------- Input ----------
+THROW 'Email already exists.'
+---------- Output ---------
+THROW 'Email already exists.'
+---------- AST ------------
+Throw {
+    span: Some(
+        0..29,
+    ),
+    message: Some(
+        Literal {
+            span: Some(
+                6..29,
+            ),
+            value: String(
+                "Email already exists.",
+            ),
+        },
+    ),
+}
+
+
 ---------- Input ----------
 FOR i IN REVERSE 1 TO maximum_count DO
     counter := counter + 1;

src/query/script/src/compiler.rs

@@ -193,6 +193,22 @@ impl Compiler {
                     output.append(&mut self.compile_sql_statement(*span, stmt, to_set.clone())?);
                     output.push(ScriptIR::ReturnSet { set: to_set });
                 }
+                ScriptStatement::Throw { span, message } => {
+                    let mut throw_value = None;
+                    if let Some(expr) = message {
+                        let to_var = VarRef::new_internal(
+                            expr.span(),
+                            "throw_value",
+                            &mut self.ref_allocator,
+                        );
+                        output.append(&mut self.compile_expr(expr, to_var.clone())?);
+                        throw_value = Some(to_var);
+                    }
+                    output.push(ScriptIR::Throw {
+                        span: *span,
+                        message: throw_value,
+                    });
+                }
                 ScriptStatement::ForLoop {
                     span,
                     variable,

src/query/script/src/executor.rs

@@ -37,6 +37,7 @@ pub trait Client {
         -> Result<Self::Var>;
     fn num_rows(&self, block: &Self::Set) -> usize;
     fn is_true(&self, scalar: &Self::Var) -> Result<bool>;
+    fn format_error(&self, value: &Self::Var) -> Result<String>;
 }
 
 #[derive(Debug, Clone)]
@@ -181,6 +182,15 @@ impl<C: Client> Executor<C> {
                 self.return_value = Some(ReturnValue::Set(self.get_set(set)?.clone()));
                 self.goto_end();
             }
+            ScriptIR::Throw { span, message } => {
+                let msg = if let Some(var) = message {
+                    let value = self.get_var(var)?;
+                    self.client.format_error(value)?
+                } else {
+                    "Script threw an error".to_string()
+                };
+                return Err(ErrorCode::ScriptExecutionError(msg).set_span(*span));
+            }
         }
 
         self.pc += 1;

src/query/script/src/ir.rs

@@ -93,6 +93,8 @@ pub enum ScriptIR {
     ReturnVar { var: VarRef },
     /// Returns a result set from the script.
     ReturnSet { set: SetRef },
+    /// Throws an error from the script.
+    Throw { span: Span, message: Option<VarRef> },
 }
 
 impl Display for ScriptIR {
@@ -123,6 +125,13 @@ impl Display for ScriptIR {
             ScriptIR::Return => write!(f, "RETURN")?,
             ScriptIR::ReturnVar { var } => write!(f, "RETURN {var}")?,
             ScriptIR::ReturnSet { set } => write!(f, "RETURN {set}")?,
+            ScriptIR::Throw { message, .. } => {
+                if let Some(message) = message {
+                    write!(f, "THROW {message}")?;
+                } else {
+                    write!(f, "THROW")?;
+                }
+            }
         };
         Ok(())
     }

src/query/script/tests/it/main.rs

@@ -696,6 +696,13 @@ impl Client for MockClient {
     fn is_true(&self, scalar: &Self::Var) -> Result<bool> {
         Ok(*scalar == Literal::Boolean(true))
     }
+
+    fn format_error(&self, value: &Self::Var) -> Result<String> {
+        Ok(match value {
+            Literal::String(s) => s.clone(),
+            _ => value.to_string(),
+        })
+    }
 }
 
 #[derive(Debug, Clone)]

src/query/service/src/interpreters/util.rs

@@ -20,6 +20,7 @@ use databend_common_ast::parser::tokenize_sql;
 use databend_common_ast::parser::Dialect;
 use databend_common_catalog::catalog::Catalog;
 use databend_common_exception::ErrorCode;
+use databend_common_expression::display::scalar_ref_to_string;
 use databend_common_expression::ComputedExpr;
 use databend_common_expression::DataBlock;
 use databend_common_expression::DataSchemaRef;
@@ -217,6 +218,13 @@ impl Client for ScriptClient {
             ))),
         }
     }
+
+    fn format_error(&self, value: &Self::Var) -> databend_common_exception::Result<String> {
+        Ok(match value {
+            Scalar::String(s) => s.clone(),
+            _ => scalar_ref_to_string(&value.as_ref()),
+        })
+    }
 }
 
 #[derive(serde::Serialize)]