mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-07-07 17:04:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

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

Browse source
This commit is contained in:
Andrew Harper 2025-05-23 01:19:16 -04:00 • committed by GitHub
parent bf2b72fbe0
commit 301726541a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 201 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

24
src/ast/data_type.rs
View file

@ -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),
}
}

6
src/ast/ddl.rs
View file

@ -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}")?;
}

24
src/ast/mod.rs
View file

@ -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)]

81
src/parser/mod.rs
View file

@ -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)?;
let statements = self.parse_statement_list(&[Keyword::END])?;
let end_token = self.expect_keyword(Keyword::END)?;
p.expect_keyword_is(Keyword::TABLE)?;
p.prev_token();
let function_body = Some(CreateFunctionBody::AsBeginEnd(BeginEndStatements {
begin_token: AttachedToken(begin_token),
statements,
end_token: AttachedToken(end_token),
}));
let table_column_defs = match p.parse_data_type()? {
DataType::Table(Some(table_column_defs)) if !table_column_defs.is_empty() => {
table_column_defs
}
_ => 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()?;
Ok(DataType::Table(columns))
// an LParen after the TABLE keyword indicates that table columns are being defined
// 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()?;
let data_type = self.parse_data_type()?;
Ok(ColumnDef {
name,
data_type,
options: Vec::new(), // No constraints expected here
})
self.parse_column_def()
}
fn parse_returns_table_columns(&mut self) -> Result<Vec<ColumnDef>, ParserError> {

87
tests/sqlparser_mssql.rs
View file

@ -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]