feat(query): support procedure

1. support parse

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

@@ -18,6 +18,10 @@ use std::fmt::Formatter;
 use derive_visitor::Drive;
 use derive_visitor::DriveMut;
 
+use crate::ast::write_comma_separated_list;
+use crate::ast::CreateOption;
+use crate::ast::TypeName;
+
 #[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
 pub struct ExecuteImmediateStmt {
     pub script: String,
@@ -29,3 +33,143 @@ impl Display for ExecuteImmediateStmt {
         Ok(())
     }
 }
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct ProcedureReturnType {
+    pub name: Option<String>,
+    pub data_type: TypeName,
+}
+
+impl Display for ProcedureReturnType {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        if let Some(name) = &self.name {
+            write!(f, "{} {}", name, self.data_type)
+        } else {
+            write!(f, "{}", self.data_type)
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub enum ProcedureLanguage {
+    SQL,
+}
+
+impl Display for ProcedureLanguage {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ProcedureLanguage::SQL => write!(f, "LANGUAGE SQL "),
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct CreateProcedureStmt {
+    pub create_option: CreateOption,
+    pub name: String,
+    pub language: ProcedureLanguage,
+    pub args: Option<Vec<TypeName>>,
+    pub return_type: Vec<ProcedureReturnType>,
+    pub comment: Option<String>,
+    pub script: String,
+}
+
+impl Display for CreateProcedureStmt {
+    // CREATE [ OR REPLACE ] PROCEDURE <name> ()
+    // RETURNS { <result_data_type> }[ NOT NULL ]
+    // LANGUAGE SQL
+    // [ COMMENT = '<string_literal>' ] AS <procedure_definition>
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        write!(f, "CREATE ")?;
+        if let CreateOption::CreateOrReplace = self.create_option {
+            write!(f, "OR REPLACE ")?;
+        }
+        write!(f, "PROCEDURE {}", self.name)?;
+        if let Some(args) = &self.args {
+            if args.is_empty() {
+                write!(f, "() ")?;
+            } else {
+                write!(f, "(")?;
+                for arg in args {
+                    write!(f, "{}", arg)?;
+                }
+                write!(f, ") ")?;
+            }
+        } else {
+            write!(f, "() ")?;
+        }
+        if self.return_type.len() == 1 {
+            if let Some(name) = &self.return_type[0].name {
+                write!(
+                    f,
+                    "RETURNS TABLE({} {}) ",
+                    name, self.return_type[0].data_type
+                )?;
+            } else {
+                write!(f, "RETURNS {} ", self.return_type[0].data_type)?;
+            }
+        } else {
+            write!(f, "RETURNS TABLE(")?;
+            write_comma_separated_list(f, self.return_type.clone())?;
+            write!(f, ") ")?;
+        }
+
+        write!(f, "{}", self.language)?;
+        if let Some(comment) = &self.comment {
+            write!(f, "COMMENT='{}' ", comment)?;
+        }
+        write!(f, "AS $$\n{}\n$$", self.script)?;
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct DropProcedureStmt {
+    pub name: String,
+    pub args: Option<Vec<TypeName>>,
+}
+
+impl Display for DropProcedureStmt {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "DROP PROCEDURE {}", self.name)?;
+        if let Some(args) = &self.args {
+            if args.is_empty() {
+                write!(f, "() ")?;
+            } else {
+                write!(f, "(")?;
+                for arg in args {
+                    write!(f, "{}", arg)?;
+                }
+                write!(f, ") ")?;
+            }
+        } else {
+            write!(f, "() ")?;
+        }
+        Ok(())
+    }
+}
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct DescProcedureStmt {
+    pub name: String,
+    pub args: Option<Vec<TypeName>>,
+}
+
+impl Display for DescProcedureStmt {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "DESCRIBE PROCEDURE {}", self.name)?;
+        if let Some(args) = &self.args {
+            if args.is_empty() {
+                write!(f, "() ")?;
+            } else {
+                write!(f, "(")?;
+                for arg in args {
+                    write!(f, "{}", arg)?;
+                }
+                write!(f, ") ")?;
+            }
+        } else {
+            write!(f, "() ")?;
+        }
+        Ok(())
+    }
+}

src/query/ast/src/ast/statements/statement.rs

@@ -323,6 +323,12 @@ pub enum Statement {
 
     // Stored procedures
     ExecuteImmediate(ExecuteImmediateStmt),
+    CreateProcedure(CreateProcedureStmt),
+    DropProcedure(DropProcedureStmt),
+    ShowProcedures {
+        show_options: Option<ShowOptions>,
+    },
+    DescProcedure(DescProcedureStmt),
 
     // Sequence
     CreateSequence(CreateSequenceStmt),
@@ -795,6 +801,15 @@ impl Display for Statement {
             Statement::DropNotification(stmt) => write!(f, "{stmt}")?,
             Statement::DescribeNotification(stmt) => write!(f, "{stmt}")?,
             Statement::ExecuteImmediate(stmt) => write!(f, "{stmt}")?,
+            Statement::CreateProcedure(stmt) => write!(f, "{stmt}")?,
+            Statement::DropProcedure(stmt) => write!(f, "{stmt}")?,
+            Statement::DescProcedure(stmt) => write!(f, "{stmt}")?,
+            Statement::ShowProcedures { show_options } => {
+                write!(f, "SHOW PROCEDURES")?;
+                if let Some(show_options) = show_options {
+                    write!(f, " {show_options}")?;
+                }
+            }
             Statement::CreateSequence(stmt) => write!(f, "{stmt}")?,
             Statement::DropSequence(stmt) => write!(f, "{stmt}")?,
             Statement::CreateDynamicTable(stmt) => write!(f, "{stmt}")?,

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

@@ -2177,6 +2177,88 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         |(_, action)| Statement::System(SystemStmt { action }),
     );
 
+    pub fn procedure_table_return(i: Input) -> IResult<ProcedureReturnType> {
+        map(rule! { #ident ~ #type_name }, |(name, data_type)| {
+            ProcedureReturnType {
+                name: Some(name.to_string()),
+                data_type,
+            }
+        })(i)
+    }
+
+    fn procedure_return(i: Input) -> IResult<Vec<ProcedureReturnType>> {
+        let procedure_table_return = map(
+            rule! {
+                TABLE ~ "(" ~ #comma_separated_list1(procedure_table_return) ~ ")"
+            },
+            |(_, _, test, _)| test,
+        );
+        let procedure_single_return = map(rule! { #type_name }, |data_type| {
+            vec![ProcedureReturnType {
+                name: None,
+                data_type,
+            }]
+        });
+        rule!(#procedure_single_return: "<type_name>"
+            | #procedure_table_return: "TABLE(<var_name> <type_name>, ...)")(i)
+    }
+
+    // CREATE [ OR REPLACE ] PROCEDURE <name> ()
+    // RETURNS { <result_data_type> }[ NOT NULL ]
+    // LANGUAGE SQL
+    // [ COMMENT = '<string_literal>' ] AS <procedure_definition>
+    let create_procedure = map_res(
+        rule! {
+            CREATE ~ ( OR ~ ^REPLACE )? ~ PROCEDURE ~ #ident ~ "(" ~ ")" ~ RETURNS ~ #procedure_return ~ LANGUAGE ~ SQL  ~ (COMMENT ~ "=" ~ #literal_string)? ~ AS ~ #code_string
+        },
+        |(_, opt_or_replace, _, name, _, _, _, return_type, _, _, opt_comment, _, script)| {
+            let create_option = parse_create_option(opt_or_replace.is_some(), false)?;
+            let stmt = CreateProcedureStmt {
+                create_option,
+                name: name.to_string(),
+                args: None,
+                return_type,
+                language: ProcedureLanguage::SQL,
+                comment: match opt_comment {
+                    Some(opt) => Some(opt.2),
+                    None => None,
+                },
+                script,
+            };
+            Ok(Statement::CreateProcedure(stmt))
+        },
+    );
+
+    let drop_procedure = map(
+        rule! {
+            DROP ~ PROCEDURE ~ #ident ~ "(" ~ ")"
+        },
+        |(_, _, name, _, _)| {
+            Statement::DropProcedure(DropProcedureStmt {
+                name: name.to_string(),
+                args: None,
+            })
+        },
+    );
+    let show_procedures = map(
+        rule! {
+            SHOW ~ PROCEDURES ~ #show_options?
+        },
+        |(_, _, show_options)| Statement::ShowProcedures { show_options },
+    );
+
+    let describe_procedure = map(
+        rule! {
+            ( DESC | DESCRIBE ) ~ PROCEDURE ~ #ident ~ "(" ~ ")"
+        },
+        |(_, _, name, _, _)| {
+            Statement::DescProcedure(DescProcedureStmt {
+                name: name.to_string(),
+                args: None,
+            })
+        },
+    );
+
     alt((
         // query, explain,show
         rule!(
@@ -2393,6 +2475,10 @@ AS
             | #desc_connection: "`DESC | DESCRIBE CONNECTION  <connection_name>`"
             | #show_connections: "`SHOW CONNECTIONS`"
             | #execute_immediate : "`EXECUTE IMMEDIATE $$ <script> $$`"
+            | #create_procedure : "`CREATE [ OR REPLACE ] PROCEDURE <procedure_name>() RETURNS { <result_data_type> [ NOT NULL ] | TABLE(<var_name> <data_type>, ...)} LANGUAGE SQL [ COMMENT = '<string_literal>' ] AS <procedure_definition>`"
+            | #drop_procedure : "`DROP PROCEDURE <procedure_name>()`"
+            | #show_procedures : "`SHOW PROCEDURES [<show_options>]()`"
+            | #describe_procedure : "`DESC PROCEDURE <procedure_name>()`"
         ),
     ))(i)
 }

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

@@ -706,6 +706,10 @@ pub enum TokenKind {
     INTO,
     #[token("INVERTED", ignore(ascii_case))]
     INVERTED,
+    #[token("PROCEDURE", ignore(ascii_case))]
+    PROCEDURE,
+    #[token("PROCEDURES", ignore(ascii_case))]
+    PROCEDURES,
     #[token("IMMEDIATE", ignore(ascii_case))]
     IMMEDIATE,
     #[token("IS", ignore(ascii_case))]
@@ -1319,6 +1323,8 @@ pub enum TokenKind {
     PRIMARY,
     #[token("SOURCE", ignore(ascii_case))]
     SOURCE,
+    #[token("SQL", ignore(ascii_case))]
+    SQL,
 }
 
 // Reference: https://www.postgresql.org/docs/current/sql-keywords-appendix.html
@@ -1569,6 +1575,7 @@ impl TokenKind {
             | TokenKind::AND
             | TokenKind::ANY
             | TokenKind::FUNCTION
+            | TokenKind::PROCEDURE
             | TokenKind::ASC
             | TokenKind::ANTI
             // | TokenKind::ASYMMETRIC

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

@@ -841,6 +841,28 @@ fn test_statement() {
             PRIMARY KEY username
             SOURCE (mysql(host='localhost' username='root' password='1234'))
             COMMENT 'This is a comment';"#,
+        // Stored Procedure
+        r#"describe PROCEDURE p1()"#,
+        r#"drop PROCEDURE p1()"#,
+        r#"show PROCEDURES like 'p1%'"#,
+        r#"create PROCEDURE p1() returns string not null language sql comment = 'test' as $$
+            BEGIN
+                LET sum := 0;
+                FOR x IN SELECT * FROM numbers(100) DO
+                    sum := sum + x.number;
+                END FOR;
+                RETURN sum;
+            END;
+            $$;"#,
+        r#"create PROCEDURE p1() returns table(a string not null, b int null) language sql comment = 'test' as $$
+            BEGIN
+                LET sum := 0;
+                FOR x IN SELECT * FROM numbers(100) DO
+                    sum := sum + x.number;
+                END FOR;
+                RETURN sum;
+            END;
+            $$;"#,
     ];
 
     for case in cases {
@@ -959,6 +981,18 @@ fn test_statement_error() {
         PRIMARY KEY username
         SOURCE ()
         COMMENT 'This is a comment';"#,
+        // Stored Procedure
+        r#"desc procedure p1"#,
+        r#"drop procedure p1"#,
+        r#"create PROCEDURE p1() returns table(string not null, int null) language sql comment = 'test' as $$
+            BEGIN
+                LET sum := 0;
+                FOR x IN SELECT * FROM numbers(100) DO
+                    sum := sum + x.number;
+                END FOR;
+                RETURN sum;
+            END;
+            $$;"#,
     ];
 
     for case in cases {