Add support for table valued functions for SQL Server (#1839)

src/ast/data_type.rs

@@ -48,7 +48,17 @@ pub enum DataType {
     /// Table type in [PostgreSQL], e.g. CREATE FUNCTION RETURNS TABLE(...).
     ///
     /// [PostgreSQL]: https://www.postgresql.org/docs/15/sql-createfunction.html
-    Table(Vec<ColumnDef>),
+    /// [MsSQL]: https://learn.microsoft.com/en-us/sql/t-sql/statements/create-function-transact-sql?view=sql-server-ver16#c-create-a-multi-statement-table-valued-function
+    Table(Option<Vec<ColumnDef>>),
+    /// Table type with a name, e.g. CREATE FUNCTION RETURNS @result TABLE(...).
+    ///
+    /// [MsSQl]: https://learn.microsoft.com/en-us/sql/t-sql/statements/create-function-transact-sql?view=sql-server-ver16#table
+    NamedTable {
+        /// Table name.
+        name: ObjectName,
+        /// Table columns.
+        columns: Vec<ColumnDef>,
+    },
     /// Fixed-length character type, e.g. CHARACTER(10).
     Character(Option<CharacterLength>),
     /// Fixed-length char type, e.g. CHAR(10).
@@ -716,7 +726,17 @@ impl fmt::Display for DataType {
             DataType::Unspecified => Ok(()),
             DataType::Trigger => write!(f, "TRIGGER"),
             DataType::AnyType => write!(f, "ANY TYPE"),
-            DataType::Table(fields) => write!(f, "TABLE({})", display_comma_separated(fields)),
+            DataType::Table(fields) => match fields {
+                Some(fields) => {
+                    write!(f, "TABLE({})", display_comma_separated(fields))
+                }
+                None => {
+                    write!(f, "TABLE")
+                }
+            },
+            DataType::NamedTable { name, columns } => {
+                write!(f, "{} TABLE ({})", name, display_comma_separated(columns))
+            }
             DataType::GeometricType(kind) => write!(f, "{}", kind),
         }
     }

src/ast/ddl.rs

@@ -2356,6 +2356,12 @@ impl fmt::Display for CreateFunction {
         if let Some(CreateFunctionBody::Return(function_body)) = &self.function_body {
             write!(f, " RETURN {function_body}")?;
         }
+        if let Some(CreateFunctionBody::AsReturnExpr(function_body)) = &self.function_body {
+            write!(f, " AS RETURN {function_body}")?;
+        }
+        if let Some(CreateFunctionBody::AsReturnSelect(function_body)) = &self.function_body {
+            write!(f, " AS RETURN {function_body}")?;
+        }
         if let Some(using) = &self.using {
             write!(f, " {using}")?;
         }

src/ast/mod.rs

@@ -8780,6 +8780,30 @@ pub enum CreateFunctionBody {
     ///
     /// [PostgreSQL]: https://www.postgresql.org/docs/current/sql-createfunction.html
     Return(Expr),
+
+    /// Function body expression using the 'AS RETURN' keywords
+    ///
+    /// Example:
+    /// ```sql
+    /// CREATE FUNCTION myfunc(a INT, b INT)
+    /// RETURNS TABLE
+    /// AS RETURN (SELECT a + b AS sum);
+    /// ```
+    ///
+    /// [MsSql]: https://learn.microsoft.com/en-us/sql/t-sql/statements/create-function-transact-sql
+    AsReturnExpr(Expr),
+
+    /// Function body expression using the 'AS RETURN' keywords, with an un-parenthesized SELECT query
+    ///
+    /// Example:
+    /// ```sql
+    /// CREATE FUNCTION myfunc(a INT, b INT)
+    /// RETURNS TABLE
+    /// AS RETURN SELECT a + b AS sum;
+    /// ```
+    ///
+    /// [MsSql]: https://learn.microsoft.com/en-us/sql/t-sql/statements/create-function-transact-sql?view=sql-server-ver16#select_stmt
+    AsReturnSelect(Select),
 }
 
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]

src/parser/mod.rs

@@ -5204,19 +5204,62 @@ impl<'a> Parser<'a> {
         let (name, args) = self.parse_create_function_name_and_params()?;
 
         self.expect_keyword(Keyword::RETURNS)?;
-        let return_type = Some(self.parse_data_type()?);
 
-        self.expect_keyword_is(Keyword::AS)?;
+        let return_table = self.maybe_parse(|p| {
+            let return_table_name = p.parse_identifier()?;
 
-        let begin_token = self.expect_keyword(Keyword::BEGIN)?;
+            p.expect_keyword_is(Keyword::TABLE)?;
-        let statements = self.parse_statement_list(&[Keyword::END])?;
+            p.prev_token();
-        let end_token = self.expect_keyword(Keyword::END)?;
 
-        let function_body = Some(CreateFunctionBody::AsBeginEnd(BeginEndStatements {
+            let table_column_defs = match p.parse_data_type()? {
-            begin_token: AttachedToken(begin_token),
+                DataType::Table(Some(table_column_defs)) if !table_column_defs.is_empty() => {
-            statements,
+                    table_column_defs
-            end_token: AttachedToken(end_token),
+                }
-        }));
+                _ => parser_err!(
+                    "Expected table column definitions after TABLE keyword",
+                    p.peek_token().span.start
+                )?,
+            };
+
+            Ok(DataType::NamedTable {
+                name: ObjectName(vec![ObjectNamePart::Identifier(return_table_name)]),
+                columns: table_column_defs,
+            })
+        })?;
+
+        let return_type = if return_table.is_some() {
+            return_table
+        } else {
+            Some(self.parse_data_type()?)
+        };
+
+        let _ = self.parse_keyword(Keyword::AS);
+
+        let function_body = if self.peek_keyword(Keyword::BEGIN) {
+            let begin_token = self.expect_keyword(Keyword::BEGIN)?;
+            let statements = self.parse_statement_list(&[Keyword::END])?;
+            let end_token = self.expect_keyword(Keyword::END)?;
+
+            Some(CreateFunctionBody::AsBeginEnd(BeginEndStatements {
+                begin_token: AttachedToken(begin_token),
+                statements,
+                end_token: AttachedToken(end_token),
+            }))
+        } else if self.parse_keyword(Keyword::RETURN) {
+            if self.peek_token() == Token::LParen {
+                Some(CreateFunctionBody::AsReturnExpr(self.parse_expr()?))
+            } else if self.peek_keyword(Keyword::SELECT) {
+                let select = self.parse_select()?;
+                Some(CreateFunctionBody::AsReturnSelect(select))
+            } else {
+                parser_err!(
+                    "Expected a subquery (or bare SELECT statement) after RETURN",
+                    self.peek_token().span.start
+                )?
+            }
+        } else {
+            parser_err!("Unparsable function body", self.peek_token().span.start)?
+        };
 
         Ok(Statement::CreateFunction(CreateFunction {
             or_alter,
@@ -9797,8 +9840,14 @@ impl<'a> Parser<'a> {
                     Ok(DataType::AnyType)
                 }
                 Keyword::TABLE => {
-                    let columns = self.parse_returns_table_columns()?;
+                    // an LParen after the TABLE keyword indicates that table columns are being defined
-                    Ok(DataType::Table(columns))
+                    // whereas no LParen indicates an anonymous table expression will be returned
+                    if self.peek_token() == Token::LParen {
+                        let columns = self.parse_returns_table_columns()?;
+                        Ok(DataType::Table(Some(columns)))
+                    } else {
+                        Ok(DataType::Table(None))
+                    }
                 }
                 Keyword::SIGNED => {
                     if self.parse_keyword(Keyword::INTEGER) {
@@ -9839,13 +9888,7 @@ impl<'a> Parser<'a> {
     }
 
     fn parse_returns_table_column(&mut self) -> Result<ColumnDef, ParserError> {
-        let name = self.parse_identifier()?;
+        self.parse_column_def()
-        let data_type = self.parse_data_type()?;
-        Ok(ColumnDef {
-            name,
-            data_type,
-            options: Vec::new(), // No constraints expected here
-        })
     }
 
     fn parse_returns_table_columns(&mut self) -> Result<Vec<ColumnDef>, ParserError> {

tests/sqlparser_mssql.rs

@@ -254,6 +254,12 @@ fn parse_create_function() {
     ";
     let _ = ms().verified_stmt(multi_statement_function);
 
+    let multi_statement_function_without_as = multi_statement_function.replace(" AS", "");
+    let _ = ms().one_statement_parses_to(
+        &multi_statement_function_without_as,
+        multi_statement_function,
+    );
+
     let create_function_with_conditional = "\
         CREATE FUNCTION some_scalar_udf() \
         RETURNS INT \
@@ -288,6 +294,87 @@ fn parse_create_function() {
         END\
     ";
     let _ = ms().verified_stmt(create_function_with_return_expression);
+
+    let create_inline_table_value_function = "\
+        CREATE FUNCTION some_inline_tvf(@foo INT, @bar VARCHAR(256)) \
+        RETURNS TABLE \
+        AS \
+        RETURN (SELECT 1 AS col_1)\
+    ";
+    let _ = ms().verified_stmt(create_inline_table_value_function);
+
+    let create_inline_table_value_function_without_parentheses = "\
+        CREATE FUNCTION some_inline_tvf(@foo INT, @bar VARCHAR(256)) \
+        RETURNS TABLE \
+        AS \
+        RETURN SELECT 1 AS col_1\
+    ";
+    let _ = ms().verified_stmt(create_inline_table_value_function_without_parentheses);
+
+    let create_inline_table_value_function_without_as =
+        create_inline_table_value_function.replace(" AS", "");
+    let _ = ms().one_statement_parses_to(
+        &create_inline_table_value_function_without_as,
+        create_inline_table_value_function,
+    );
+
+    let create_multi_statement_table_value_function = "\
+        CREATE FUNCTION some_multi_statement_tvf(@foo INT, @bar VARCHAR(256)) \
+        RETURNS @t TABLE (col_1 INT) \
+        AS \
+        BEGIN \
+            INSERT INTO @t SELECT 1; \
+            RETURN; \
+        END\
+    ";
+    let _ = ms().verified_stmt(create_multi_statement_table_value_function);
+
+    let create_multi_statement_table_value_function_without_as =
+        create_multi_statement_table_value_function.replace(" AS", "");
+    let _ = ms().one_statement_parses_to(
+        &create_multi_statement_table_value_function_without_as,
+        create_multi_statement_table_value_function,
+    );
+
+    let create_multi_statement_table_value_function_with_constraints = "\
+        CREATE FUNCTION some_multi_statement_tvf(@foo INT, @bar VARCHAR(256)) \
+        RETURNS @t TABLE (col_1 INT NOT NULL) \
+        AS \
+        BEGIN \
+            INSERT INTO @t SELECT 1; \
+            RETURN @t; \
+        END\
+    ";
+    let _ = ms().verified_stmt(create_multi_statement_table_value_function_with_constraints);
+
+    let create_multi_statement_tvf_without_table_definition = "\
+        CREATE FUNCTION incorrect_tvf(@foo INT, @bar VARCHAR(256)) \
+        RETURNS @t TABLE ()
+        AS \
+        BEGIN \
+            INSERT INTO @t SELECT 1; \
+            RETURN @t; \
+        END\
+    ";
+    assert_eq!(
+        ParserError::ParserError("Unparsable function body".to_owned()),
+        ms().parse_sql_statements(create_multi_statement_tvf_without_table_definition)
+            .unwrap_err()
+    );
+
+    let create_inline_tvf_without_subquery_or_bare_select = "\
+        CREATE FUNCTION incorrect_tvf(@foo INT, @bar VARCHAR(256)) \
+        RETURNS TABLE
+        AS \
+        RETURN 'hi'\
+    ";
+    assert_eq!(
+        ParserError::ParserError(
+            "Expected a subquery (or bare SELECT statement) after RETURN".to_owned()
+        ),
+        ms().parse_sql_statements(create_inline_tvf_without_subquery_or_bare_select)
+            .unwrap_err()
+    );
 }
 
 #[test]