Add multi-statement table valued function support for SQL Server

Author: aharpervc

src/ast/data_type.rs

@@ -50,6 +50,13 @@ pub enum DataType {
     /// [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?view=sql-server-ver16#c-create-a-multi-statement-table-valued-function
     Table(Vec<ColumnDef>),
+    /// Table type with a name, e.g. CREATE FUNCTION RETURNS @result TABLE(...).
+    NamedTable(
+        /// Table name.
+        ObjectName,
+        /// Table columns.
+        Vec<ColumnDef>,
+    ),
     /// Fixed-length character type, e.g. CHARACTER(10).
     Character(Option<CharacterLength>),
     /// Fixed-length char type, e.g. CHAR(10).
@@ -723,6 +730,9 @@ impl fmt::Display for DataType {
                 }
                 write!(f, "TABLE({})", display_comma_separated(fields))
             }
+            DataType::NamedTable(name, fields) => {
+                write!(f, "{} TABLE ({})", name, display_comma_separated(fields))
+            }
             DataType::GeometricType(kind) => write!(f, "{}", kind),
         }
     }

src/parser/mod.rs

@@ -5203,7 +5203,36 @@ 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()?);
+
+        let return_table = self.maybe_parse(|p| {
+            let return_table_name = p.parse_identifier()?;
+            let table_column_defs = if p.peek_keyword(Keyword::TABLE) {
+                match p.parse_data_type()? {
+                    DataType::Table(t) => t,
+                    _ => {
+                        return parser_err!(
+                            "Expected table data type after TABLE keyword, got {:?}",
+                            p.peek_token()
+                        )?
+                    }
+                }
+            } else {
+                return parser_err!(
+                    "Expected TABLE keyword after return type, got {:?}",
+                    p.peek_token()
+                );
+            };
+            Ok(DataType::NamedTable(
+                ObjectName(vec![ObjectNamePart::Identifier(return_table_name)]),
+                table_column_defs.clone(),
+            ))
+        })?;
+
+        let return_type = if return_table.is_some() {
+            return_table
+        } else {
+            Some(self.parse_data_type()?)
+        };
 
         self.expect_keyword_is(Keyword::AS)?;
 

tests/sqlparser_mssql.rs

@@ -296,6 +296,17 @@ fn parse_create_function() {
         RETURN (SELECT 1 AS col_1)\
     ";
     let _ = ms().verified_stmt(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);
 }
 
 #[test]