Support basic CREATE PROCEDURE of MSSQL (#900)

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

src/ast/ddl.rs

@@ -489,6 +489,19 @@ impl fmt::Display for IndexType {
         }
     }
 }
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct ProcedureParam {
+    pub name: Ident,
+    pub data_type: DataType,
+}
+
+impl fmt::Display for ProcedureParam {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{} {}", self.name, self.data_type)
+    }
+}
 
 /// SQL column definition
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]

src/ast/mod.rs

@@ -30,8 +30,8 @@ pub use self::data_type::{
 };
 pub use self::ddl::{
     AlterColumnOperation, AlterIndexOperation, AlterTableOperation, ColumnDef, ColumnOption,
-    ColumnOptionDef, GeneratedAs, IndexType, KeyOrIndexDisplay, ReferentialAction, TableConstraint,
+    ColumnOptionDef, GeneratedAs, IndexType, KeyOrIndexDisplay, ProcedureParam, ReferentialAction,
-    UserDefinedTypeCompositeAttributeDef, UserDefinedTypeRepresentation,
+    TableConstraint, UserDefinedTypeCompositeAttributeDef, UserDefinedTypeRepresentation,
 };
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
@@ -1580,6 +1580,15 @@ pub enum Statement {
         params: CreateFunctionBody,
     },
     /// ```sql
+    /// CREATE PROCEDURE
+    /// ```
+    CreateProcedure {
+        or_alter: bool,
+        name: ObjectName,
+        params: Option<Vec<ProcedureParam>>,
+        body: Vec<Statement>,
+    },
+    /// ```sql
     /// CREATE MACRO
     /// ```
     ///
@@ -2111,6 +2120,30 @@ impl fmt::Display for Statement {
                 write!(f, "{params}")?;
                 Ok(())
             }
+            Statement::CreateProcedure {
+                name,
+                or_alter,
+                params,
+                body,
+            } => {
+                write!(
+                    f,
+                    "CREATE {or_alter}PROCEDURE {name}",
+                    or_alter = if *or_alter { "OR ALTER " } else { "" },
+                    name = name
+                )?;
+
+                if let Some(p) = params {
+                    if !p.is_empty() {
+                        write!(f, " ({})", display_comma_separated(p))?;
+                    }
+                }
+                write!(
+                    f,
+                    " AS BEGIN {body} END",
+                    body = display_separated(body, "; ")
+                )
+            }
             Statement::CreateMacro {
                 or_replace,
                 temporary,

src/keywords.rs

@@ -690,6 +690,7 @@ pub const RESERVED_FOR_TABLE_ALIAS: &[Keyword] = &[
     Keyword::SET,
     Keyword::QUALIFY,
     Keyword::WINDOW,
+    Keyword::END,
 ];
 
 /// Can't be used as a column alias, so that `SELECT <expr> alias`
@@ -719,4 +720,5 @@ pub const RESERVED_FOR_COLUMN_ALIAS: &[Keyword] = &[
     // Reserved only as a column alias in the `SELECT` clause
     Keyword::FROM,
     Keyword::INTO,
+    Keyword::END,
 ];

src/parser.rs

@@ -346,9 +346,14 @@ impl<'a> Parser<'a> {
                 expecting_statement_delimiter = false;
             }
 
-            if self.peek_token() == Token::EOF {
+            match self.peek_token().token {
-                break;
+                Token::EOF => break,
+
+                // end of statement
+                Token::Word(word) if word.keyword == Keyword::END => break,
+                _ => {}
             }
+
             if expecting_statement_delimiter {
                 return self.expected("end of statement", self.peek_token());
             }
@@ -2324,6 +2329,7 @@ impl<'a> Parser<'a> {
     /// Parse a SQL CREATE statement
     pub fn parse_create(&mut self) -> Result<Statement, ParserError> {
         let or_replace = self.parse_keywords(&[Keyword::OR, Keyword::REPLACE]);
+        let or_alter = self.parse_keywords(&[Keyword::OR, Keyword::ALTER]);
         let local = self.parse_one_of_keywords(&[Keyword::LOCAL]).is_some();
         let global = self.parse_one_of_keywords(&[Keyword::GLOBAL]).is_some();
         let transient = self.parse_one_of_keywords(&[Keyword::TRANSIENT]).is_some();
@@ -2369,6 +2375,8 @@ impl<'a> Parser<'a> {
             self.parse_create_sequence(temporary)
         } else if self.parse_keyword(Keyword::TYPE) {
             self.parse_create_type()
+        } else if self.parse_keyword(Keyword::PROCEDURE) {
+            self.parse_create_procedure(or_alter)
         } else {
             self.expected("an object type after CREATE", self.peek_token())
         }
@@ -3503,6 +3511,28 @@ impl<'a> Parser<'a> {
             .build())
     }
 
+    pub fn parse_optional_procedure_parameters(
+        &mut self,
+    ) -> Result<Option<Vec<ProcedureParam>>, ParserError> {
+        let mut params = vec![];
+        if !self.consume_token(&Token::LParen) || self.consume_token(&Token::RParen) {
+            return Ok(Some(params));
+        }
+        loop {
+            if let Token::Word(_) = self.peek_token().token {
+                params.push(self.parse_procedure_param()?)
+            }
+            let comma = self.consume_token(&Token::Comma);
+            if self.consume_token(&Token::RParen) {
+                // allow a trailing comma, even though it's not in standard
+                break;
+            } else if !comma {
+                return self.expected("',' or ')' after parameter definition", self.peek_token());
+            }
+        }
+        Ok(Some(params))
+    }
+
     pub fn parse_columns(&mut self) -> Result<(Vec<ColumnDef>, Vec<TableConstraint>), ParserError> {
         let mut columns = vec![];
         let mut constraints = vec![];
@@ -3530,6 +3560,12 @@ impl<'a> Parser<'a> {
         Ok((columns, constraints))
     }
 
+    pub fn parse_procedure_param(&mut self) -> Result<ProcedureParam, ParserError> {
+        let name = self.parse_identifier()?;
+        let data_type = self.parse_data_type()?;
+        Ok(ProcedureParam { name, data_type })
+    }
+
     pub fn parse_column_def(&mut self) -> Result<ColumnDef, ParserError> {
         let name = self.parse_identifier()?;
         let data_type = self.parse_data_type()?;
@@ -7082,6 +7118,21 @@ impl<'a> Parser<'a> {
         Ok(NamedWindowDefinition(ident, window_spec))
     }
 
+    pub fn parse_create_procedure(&mut self, or_alter: bool) -> Result<Statement, ParserError> {
+        let name = self.parse_object_name()?;
+        let params = self.parse_optional_procedure_parameters()?;
+        self.expect_keyword(Keyword::AS)?;
+        self.expect_keyword(Keyword::BEGIN)?;
+        let statements = self.parse_statements()?;
+        self.expect_keyword(Keyword::END)?;
+        Ok(Statement::CreateProcedure {
+            name,
+            or_alter,
+            params,
+            body: statements,
+        })
+    }
+
     pub fn parse_window_spec(&mut self) -> Result<WindowSpec, ParserError> {
         let partition_by = if self.parse_keywords(&[Keyword::PARTITION, Keyword::BY]) {
             self.parse_comma_separated(Parser::parse_expr)?

tests/sqlparser_mssql.rs

@@ -55,6 +55,88 @@ fn parse_mssql_delimited_identifiers() {
     );
 }
 
+#[test]
+fn parse_create_procedure() {
+    let sql = "CREATE OR ALTER PROCEDURE test (@foo INT, @bar VARCHAR(256)) AS BEGIN SELECT 1 END";
+
+    #[cfg(feature = "bigdecimal")]
+    let one = Value::Number(bigdecimal::BigDecimal::from(1), false);
+
+    #[cfg(not(feature = "bigdecimal"))]
+    let one = Value::Number("1".to_string(), false);
+
+    assert_eq!(
+        ms().verified_stmt(sql),
+        Statement::CreateProcedure {
+            or_alter: true,
+            body: vec![Statement::Query(Box::new(Query {
+                with: None,
+                limit: None,
+                offset: None,
+                fetch: None,
+                locks: vec![],
+                order_by: vec![],
+                body: Box::new(SetExpr::Select(Box::new(Select {
+                    distinct: None,
+                    top: None,
+                    projection: vec![SelectItem::UnnamedExpr(Expr::Value(one))],
+                    into: None,
+                    from: vec![],
+                    lateral_views: vec![],
+                    selection: None,
+                    group_by: vec![],
+                    cluster_by: vec![],
+                    distribute_by: vec![],
+                    sort_by: vec![],
+                    having: None,
+                    named_window: vec![],
+                    qualify: None
+                })))
+            }))],
+            params: Some(vec![
+                ProcedureParam {
+                    name: Ident {
+                        value: "@foo".into(),
+                        quote_style: None
+                    },
+                    data_type: DataType::Int(None)
+                },
+                ProcedureParam {
+                    name: Ident {
+                        value: "@bar".into(),
+                        quote_style: None
+                    },
+                    data_type: DataType::Varchar(Some(CharacterLength {
+                        length: 256,
+                        unit: None
+                    }))
+                }
+            ]),
+            name: ObjectName(vec![Ident {
+                value: "test".into(),
+                quote_style: None
+            }])
+        }
+    )
+}
+
+#[test]
+fn parse_mssql_create_procedure() {
+    let _ = ms_and_generic().verified_stmt("CREATE OR ALTER PROCEDURE foo AS BEGIN SELECT 1 END");
+    let _ = ms_and_generic().verified_stmt("CREATE PROCEDURE foo AS BEGIN SELECT 1 END");
+    let _ = ms().verified_stmt(
+        "CREATE PROCEDURE foo AS BEGIN SELECT [myColumn] FROM [myschema].[mytable] END",
+    );
+    let _ = ms_and_generic().verified_stmt(
+        "CREATE PROCEDURE foo (@CustomerName NVARCHAR(50)) AS BEGIN SELECT * FROM DEV END",
+    );
+    let _ = ms().verified_stmt("CREATE PROCEDURE [foo] AS BEGIN UPDATE bar SET col = 'test' END");
+    // Test a statement with END in it
+    let _ = ms().verified_stmt("CREATE PROCEDURE [foo] AS BEGIN SELECT [foo], CASE WHEN [foo] IS NULL THEN 'empty' ELSE 'notempty' END AS [foo] END");
+    // Multiple statements
+    let _ = ms().verified_stmt("CREATE PROCEDURE [foo] AS BEGIN UPDATE bar SET col = 'test'; SELECT [foo] FROM BAR WHERE [FOO] > 10 END");
+}
+
 #[test]
 fn parse_mssql_apply_join() {
     let _ = ms_and_generic().verified_only_select(