Add basic CREATE FUNCTION support for SQL Server

aharpervc · aharpervc · commit f3af75ffdb75 · 2025-04-21T14:46:04.000-04:00
- in this dialect, functions can have statement(s) bodies like stored procedures (including `BEGIN`..`END`)
- functions must end with `RETURN`, so a corresponding statement type is also introduced
diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs
@@ -2272,6 +2272,12 @@ impl fmt::Display for CreateFunction {
         if let Some(CreateFunctionBody::AsAfterOptions(function_body)) = &self.function_body {
             write!(f, " AS {function_body}")?;
         }
+        if let Some(CreateFunctionBody::MultiStatement(statements)) = &self.function_body {
+            write!(f, " AS")?;
+            write!(f, " BEGIN")?;
+            write!(f, " {}", display_separated(statements, "; "))?;
+            write!(f, " END")?;
+        }
         Ok(())
     }
 }
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -3614,6 +3614,7 @@ pub enum Statement {
     /// 1. [Hive](https://cwiki.apache.org/confluence/display/hive/languagemanual+ddl#LanguageManualDDL-Create/Drop/ReloadFunction)
     /// 2. [PostgreSQL](https://www.postgresql.org/docs/15/sql-createfunction.html)
     /// 3. [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#create_function_statement)
+    /// 4. [MsSql](https://learn.microsoft.com/en-us/sql/t-sql/statements/create-function-transact-sql)
     CreateFunction(CreateFunction),
     /// CREATE TRIGGER
     ///
@@ -4060,6 +4061,12 @@ pub enum Statement {
     ///
     /// See: <https://learn.microsoft.com/en-us/sql/t-sql/statements/print-transact-sql>
     Print(PrintStatement),
+    /// ```sql
+    /// RETURN [ expression ]
+    /// ```
+    ///
+    /// See [ReturnStatement]
+    Return(ReturnStatement),
 }
 
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
@@ -5752,6 +5759,7 @@ impl fmt::Display for Statement {
                 Ok(())
             }
             Statement::Print(s) => write!(f, "{s}"),
+            Statement::Return(r) => write!(f, "{r}"),
             Statement::List(command) => write!(f, "LIST {command}"),
             Statement::Remove(command) => write!(f, "REMOVE {command}"),
         }
@@ -8354,6 +8362,7 @@ impl fmt::Display for FunctionDeterminismSpecifier {
 ///
 /// [BigQuery]: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#syntax_11
 /// [PostgreSQL]: https://www.postgresql.org/docs/15/sql-createfunction.html
+/// [MsSql]: https://learn.microsoft.com/en-us/sql/t-sql/statements/create-function-transact-sql
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
@@ -8382,6 +8391,22 @@ pub enum CreateFunctionBody {
     ///
     /// [BigQuery]: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#syntax_11
     AsAfterOptions(Expr),
+    /// Function body with statements before the `RETURN` keyword.
+    ///
+    /// Example:
+    /// ```sql
+    /// CREATE FUNCTION my_scalar_udf(a INT, b INT)
+    /// RETURNS INT
+    /// AS
+    /// BEGIN
+    ///     DECLARE c INT;
+    ///     SET c = a + b;
+    ///     RETURN c;
+    /// END
+    /// ```
+    ///
+    /// [MsSql]: https://learn.microsoft.com/en-us/sql/t-sql/statements/create-function-transact-sql
+    MultiStatement(Vec<Statement>),
     /// Function body expression using the 'RETURN' keyword.
     ///
     /// Example:
@@ -9230,6 +9255,41 @@ impl fmt::Display for PrintStatement {
     }
 }
 
+/// Return (MsSql)
+///
+/// for Functions:
+/// RETURN scalar_expression
+///
+/// See <https://learn.microsoft.com/en-us/sql/t-sql/statements/create-function-transact-sql>
+///
+/// for Triggers:
+/// RETURN
+///
+/// See <https://learn.microsoft.com/en-us/sql/t-sql/statements/create-trigger-transact-sql>
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct ReturnStatement {
+    pub value: Option<ReturnStatementValue>,
+}
+
+impl fmt::Display for ReturnStatement {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match &self.value {
+            Some(ReturnStatementValue::Expr(expr)) => write!(f, "RETURN {}", expr),
+            None => write!(f, "RETURN"),
+        }
+    }
+}
+
+/// Variants of a `RETURN` statement
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum ReturnStatementValue {
+    Expr(Expr),
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/ast/spans.rs b/src/ast/spans.rs
@@ -520,6 +520,7 @@ impl Spanned for Statement {
             Statement::RenameTable { .. } => Span::empty(),
             Statement::RaisError { .. } => Span::empty(),
             Statement::Print { .. } => Span::empty(),
+            Statement::Return { .. } => Span::empty(),
             Statement::List(..) | Statement::Remove(..) => Span::empty(),
         }
     }
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -577,13 +577,7 @@ impl<'a> Parser<'a> {
                 Keyword::GRANT => self.parse_grant(),
                 Keyword::REVOKE => self.parse_revoke(),
                 Keyword::START => self.parse_start_transaction(),
-                // `BEGIN` is a nonstandard but common alias for the
-                // standard `START TRANSACTION` statement. It is supported
-                // by at least PostgreSQL and MySQL.
                 Keyword::BEGIN => self.parse_begin(),
-                // `END` is a nonstandard but common alias for the
-                // standard `COMMIT TRANSACTION` statement. It is supported
-                // by PostgreSQL.
                 Keyword::END => self.parse_end(),
                 Keyword::SAVEPOINT => self.parse_savepoint(),
                 Keyword::RELEASE => self.parse_release(),
@@ -618,6 +612,7 @@ impl<'a> Parser<'a> {
                 // `COMMENT` is snowflake specific https://docs.snowflake.com/en/sql-reference/sql/comment
                 Keyword::COMMENT if self.dialect.supports_comment_on() => self.parse_comment(),
                 Keyword::PRINT => self.parse_print(),
+                Keyword::RETURN => self.parse_return(),
                 _ => self.expected("an SQL statement", next_token),
             },
             Token::LParen => {
@@ -4880,6 +4875,8 @@ impl<'a> Parser<'a> {
             self.parse_create_macro(or_replace, temporary)
         } else if dialect_of!(self is BigQueryDialect) {
             self.parse_bigquery_create_function(or_replace, temporary)
+        } else if dialect_of!(self is MsSqlDialect) {
+            self.parse_mssql_create_function(or_replace, temporary)
         } else {
             self.prev_token();
             self.expected("an object type after CREATE", self.peek_token())
@@ -5134,6 +5131,72 @@ impl<'a> Parser<'a> {
         }))
     }
 
+    /// Parse `CREATE FUNCTION` for [MsSql]
+    ///
+    /// [MsSql]: https://learn.microsoft.com/en-us/sql/t-sql/statements/create-function-transact-sql
+    fn parse_mssql_create_function(
+        &mut self,
+        or_replace: bool,
+        temporary: bool,
+    ) -> Result<Statement, ParserError> {
+        let name = self.parse_object_name(false)?;
+
+        let parse_function_param =
+            |parser: &mut Parser| -> Result<OperateFunctionArg, ParserError> {
+                let name = parser.parse_identifier()?;
+                let data_type = parser.parse_data_type()?;
+                Ok(OperateFunctionArg {
+                    mode: None,
+                    name: Some(name),
+                    data_type,
+                    default_expr: None,
+                })
+            };
+        self.expect_token(&Token::LParen)?;
+        let args = self.parse_comma_separated0(parse_function_param, Token::RParen)?;
+        self.expect_token(&Token::RParen)?;
+
+        let return_type = if self.parse_keyword(Keyword::RETURNS) {
+            Some(self.parse_data_type()?)
+        } else {
+            return parser_err!("Expected RETURNS keyword", self.peek_token().span.start);
+        };
+
+        self.expect_keyword_is(Keyword::AS)?;
+        self.expect_keyword_is(Keyword::BEGIN)?;
+        let mut result = self.parse_statements()?;
+        // note: `parse_statements` will consume the `END` token & produce a Commit statement...
+        if let Some(Statement::Commit {
+            chain,
+            end,
+            modifier,
+        }) = result.last()
+        {
+            if *chain == false && *end == true && *modifier == None {
+                result = result[..result.len() - 1].to_vec();
+            }
+        }
+        let function_body = Some(CreateFunctionBody::MultiStatement(result));
+
+        Ok(Statement::CreateFunction(CreateFunction {
+            or_replace,
+            temporary,
+            if_not_exists: false,
+            name,
+            args: Some(args),
+            return_type,
+            function_body,
+            language: None,
+            determinism_specifier: None,
+            options: None,
+            remote_connection: None,
+            using: None,
+            behavior: None,
+            called_on_null: None,
+            parallel: None,
+        }))
+    }
+
     fn parse_function_arg(&mut self) -> Result<OperateFunctionArg, ParserError> {
         let mode = if self.parse_keyword(Keyword::IN) {
             Some(ArgMode::In)
@@ -15064,6 +15127,14 @@ impl<'a> Parser<'a> {
         }))
     }
 
+    /// Parse [Statement::Return]
+    fn parse_return(&mut self) -> Result<Statement, ParserError> {
+        let expr = self.parse_expr()?;
+        Ok(Statement::Return(ReturnStatement {
+            value: Some(ReturnStatementValue::Expr(expr)),
+        }))
+    }
+
     /// Consume the parser and return its underlying token buffer
     pub fn into_tokens(self) -> Vec<TokenWithSpan> {
         self.tokens
diff --git a/tests/sqlparser_hive.rs b/tests/sqlparser_hive.rs
@@ -25,7 +25,7 @@ use sqlparser::ast::{
     Expr, Function, FunctionArgumentList, FunctionArguments, Ident, ObjectName, OrderByExpr,
     OrderByOptions, SelectItem, Set, Statement, TableFactor, UnaryOperator, Use, Value,
 };
-use sqlparser::dialect::{GenericDialect, HiveDialect, MsSqlDialect};
+use sqlparser::dialect::{AnsiDialect, GenericDialect, HiveDialect};
 use sqlparser::parser::ParserError;
 use sqlparser::test_utils::*;
 
@@ -423,7 +423,7 @@ fn parse_create_function() {
     }
 
     // Test error in dialect that doesn't support parsing CREATE FUNCTION
-    let unsupported_dialects = TestedDialects::new(vec![Box::new(MsSqlDialect {})]);
+    let unsupported_dialects = TestedDialects::new(vec![Box::new(AnsiDialect {})]);
 
     assert_eq!(
         unsupported_dialects.parse_sql_statements(sql).unwrap_err(),
diff --git a/tests/sqlparser_mssql.rs b/tests/sqlparser_mssql.rs

Original file line number	Diff line number	Diff line change
`@@ -2272,6 +2272,12 @@ impl fmt::Display for CreateFunction {`
`2272`	`2272`	`if let Some(CreateFunctionBody::AsAfterOptions(function_body)) = &self.function_body {`
`2273`	`2273`	`write!(f, " AS {function_body}")?;`
`2274`	`2274`	`}`
	`2275`	`+ if let Some(CreateFunctionBody::MultiStatement(statements)) = &self.function_body {`
	`2276`	`+ write!(f, " AS")?;`
	`2277`	`+ write!(f, " BEGIN")?;`
	`2278`	`+ write!(f, " {}", display_separated(statements, "; "))?;`
	`2279`	`+ write!(f, " END")?;`
	`2280`	`+ }`
`2275`	`2281`	`Ok(())`
`2276`	`2282`	`}`
`2277`	`2283`	`}`
Original file line number	Diff line number	Diff line change
`@@ -520,6 +520,7 @@ impl Spanned for Statement {`
`520`	`520`	`Statement::RenameTable { .. } => Span::empty(),`
`521`	`521`	`Statement::RaisError { .. } => Span::empty(),`
`522`	`522`	`Statement::Print { .. } => Span::empty(),`
	`523`	`+ Statement::Return { .. } => Span::empty(),`
`523`	`524`	`Statement::List(..) \| Statement::Remove(..) => Span::empty(),`
`524`	`525`	`}`
`525`	`526`	`}`