Skip to content

Commit 3017265

Browse files
aharpervcaharpervc
authored
Add support for table valued functions for SQL Server (#1839)
1 parent bf2b72f commit 3017265

File tree

5 files changed

+201
-21
lines changed

5 files changed

+201
-21
lines changed

‎src/ast/data_type.rs

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,17 @@ pub enum DataType {
4848
/// Table type in [PostgreSQL], e.g. CREATE FUNCTION RETURNS TABLE(...).
4949
///
5050
/// [PostgreSQL]: https://www.postgresql.org/docs/15/sql-createfunction.html
51-
Table(Vec<ColumnDef>),
51+
/// [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
52+
Table(Option<Vec<ColumnDef>>),
53+
/// Table type with a name, e.g. CREATE FUNCTION RETURNS @result TABLE(...).
54+
///
55+
/// [MsSQl]: https://learn.microsoft.com/en-us/sql/t-sql/statements/create-function-transact-sql?view=sql-server-ver16#table
56+
NamedTable {
57+
/// Table name.
58+
name: ObjectName,
59+
/// Table columns.
60+
columns: Vec<ColumnDef>,
61+
},
5262
/// Fixed-length character type, e.g. CHARACTER(10).
5363
Character(Option<CharacterLength>),
5464
/// Fixed-length char type, e.g. CHAR(10).
@@ -716,7 +726,17 @@ impl fmt::Display for DataType {
716726
DataType::Unspecified => Ok(()),
717727
DataType::Trigger => write!(f, "TRIGGER"),
718728
DataType::AnyType => write!(f, "ANY TYPE"),
719-
DataType::Table(fields) => write!(f, "TABLE({})", display_comma_separated(fields)),
729+
DataType::Table(fields) => match fields {
730+
Some(fields) => {
731+
write!(f, "TABLE({})", display_comma_separated(fields))
732+
}
733+
None => {
734+
write!(f, "TABLE")
735+
}
736+
},
737+
DataType::NamedTable { name, columns } => {
738+
write!(f, "{} TABLE ({})", name, display_comma_separated(columns))
739+
}
720740
DataType::GeometricType(kind) => write!(f, "{}", kind),
721741
}
722742
}

‎src/ast/ddl.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2356,6 +2356,12 @@ impl fmt::Display for CreateFunction {
23562356
if let Some(CreateFunctionBody::Return(function_body)) = &self.function_body {
23572357
write!(f, " RETURN {function_body}")?;
23582358
}
2359+
if let Some(CreateFunctionBody::AsReturnExpr(function_body)) = &self.function_body {
2360+
write!(f, " AS RETURN {function_body}")?;
2361+
}
2362+
if let Some(CreateFunctionBody::AsReturnSelect(function_body)) = &self.function_body {
2363+
write!(f, " AS RETURN {function_body}")?;
2364+
}
23592365
if let Some(using) = &self.using {
23602366
write!(f, " {using}")?;
23612367
}

‎src/ast/mod.rs

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8780,6 +8780,30 @@ pub enum CreateFunctionBody {
87808780
///
87818781
/// [PostgreSQL]: https://www.postgresql.org/docs/current/sql-createfunction.html
87828782
Return(Expr),
8783+
8784+
/// Function body expression using the 'AS RETURN' keywords
8785+
///
8786+
/// Example:
8787+
/// ```sql
8788+
/// CREATE FUNCTION myfunc(a INT, b INT)
8789+
/// RETURNS TABLE
8790+
/// AS RETURN (SELECT a + b AS sum);
8791+
/// ```
8792+
///
8793+
/// [MsSql]: https://learn.microsoft.com/en-us/sql/t-sql/statements/create-function-transact-sql
8794+
AsReturnExpr(Expr),
8795+
8796+
/// Function body expression using the 'AS RETURN' keywords, with an un-parenthesized SELECT query
8797+
///
8798+
/// Example:
8799+
/// ```sql
8800+
/// CREATE FUNCTION myfunc(a INT, b INT)
8801+
/// RETURNS TABLE
8802+
/// AS RETURN SELECT a + b AS sum;
8803+
/// ```
8804+
///
8805+
/// [MsSql]: https://learn.microsoft.com/en-us/sql/t-sql/statements/create-function-transact-sql?view=sql-server-ver16#select_stmt
8806+
AsReturnSelect(Select),
87838807
}
87848808

87858809
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]

‎src/parser/mod.rs

Lines changed: 62 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -5204,19 +5204,62 @@ impl<'a> Parser<'a> {
52045204
let (name, args) = self.parse_create_function_name_and_params()?;
52055205

52065206
self.expect_keyword(Keyword::RETURNS)?;
5207-
let return_type = Some(self.parse_data_type()?);
52085207

5209-
self.expect_keyword_is(Keyword::AS)?;
5208+
let return_table = self.maybe_parse(|p| {
5209+
let return_table_name = p.parse_identifier()?;
5210+
5211+
p.expect_keyword_is(Keyword::TABLE)?;
5212+
p.prev_token();
5213+
5214+
let table_column_defs = match p.parse_data_type()? {
5215+
DataType::Table(Some(table_column_defs)) if !table_column_defs.is_empty() => {
5216+
table_column_defs
5217+
}
5218+
_ => parser_err!(
5219+
"Expected table column definitions after TABLE keyword",
5220+
p.peek_token().span.start
5221+
)?,
5222+
};
5223+
5224+
Ok(DataType::NamedTable {
5225+
name: ObjectName(vec![ObjectNamePart::Identifier(return_table_name)]),
5226+
columns: table_column_defs,
5227+
})
5228+
})?;
52105229

5211-
let begin_token = self.expect_keyword(Keyword::BEGIN)?;
5212-
let statements = self.parse_statement_list(&[Keyword::END])?;
5213-
let end_token = self.expect_keyword(Keyword::END)?;
5230+
let return_type = if return_table.is_some() {
5231+
return_table
5232+
} else {
5233+
Some(self.parse_data_type()?)
5234+
};
52145235

5215-
let function_body = Some(CreateFunctionBody::AsBeginEnd(BeginEndStatements {
5216-
begin_token: AttachedToken(begin_token),
5217-
statements,
5218-
end_token: AttachedToken(end_token),
5219-
}));
5236+
let _ = self.parse_keyword(Keyword::AS);
5237+
5238+
let function_body = if self.peek_keyword(Keyword::BEGIN) {
5239+
let begin_token = self.expect_keyword(Keyword::BEGIN)?;
5240+
let statements = self.parse_statement_list(&[Keyword::END])?;
5241+
let end_token = self.expect_keyword(Keyword::END)?;
5242+
5243+
Some(CreateFunctionBody::AsBeginEnd(BeginEndStatements {
5244+
begin_token: AttachedToken(begin_token),
5245+
statements,
5246+
end_token: AttachedToken(end_token),
5247+
}))
5248+
} else if self.parse_keyword(Keyword::RETURN) {
5249+
if self.peek_token() == Token::LParen {
5250+
Some(CreateFunctionBody::AsReturnExpr(self.parse_expr()?))
5251+
} else if self.peek_keyword(Keyword::SELECT) {
5252+
let select = self.parse_select()?;
5253+
Some(CreateFunctionBody::AsReturnSelect(select))
5254+
} else {
5255+
parser_err!(
5256+
"Expected a subquery (or bare SELECT statement) after RETURN",
5257+
self.peek_token().span.start
5258+
)?
5259+
}
5260+
} else {
5261+
parser_err!("Unparsable function body", self.peek_token().span.start)?
5262+
};
52205263

52215264
Ok(Statement::CreateFunction(CreateFunction {
52225265
or_alter,
@@ -9797,8 +9840,14 @@ impl<'a> Parser<'a> {
97979840
Ok(DataType::AnyType)
97989841
}
97999842
Keyword::TABLE => {
9800-
let columns = self.parse_returns_table_columns()?;
9801-
Ok(DataType::Table(columns))
9843+
// an LParen after the TABLE keyword indicates that table columns are being defined
9844+
// whereas no LParen indicates an anonymous table expression will be returned
9845+
if self.peek_token() == Token::LParen {
9846+
let columns = self.parse_returns_table_columns()?;
9847+
Ok(DataType::Table(Some(columns)))
9848+
} else {
9849+
Ok(DataType::Table(None))
9850+
}
98029851
}
98039852
Keyword::SIGNED => {
98049853
if self.parse_keyword(Keyword::INTEGER) {
@@ -9839,13 +9888,7 @@ impl<'a> Parser<'a> {
98399888
}
98409889

98419890
fn parse_returns_table_column(&mut self) -> Result<ColumnDef, ParserError> {
9842-
let name = self.parse_identifier()?;
9843-
let data_type = self.parse_data_type()?;
9844-
Ok(ColumnDef {
9845-
name,
9846-
data_type,
9847-
options: Vec::new(), // No constraints expected here
9848-
})
9891+
self.parse_column_def()
98499892
}
98509893

98519894
fn parse_returns_table_columns(&mut self) -> Result<Vec<ColumnDef>, ParserError> {

‎tests/sqlparser_mssql.rs

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,12 @@ fn parse_create_function() {
254254
";
255255
let _ = ms().verified_stmt(multi_statement_function);
256256

257+
let multi_statement_function_without_as = multi_statement_function.replace(" AS", "");
258+
let _ = ms().one_statement_parses_to(
259+
&multi_statement_function_without_as,
260+
multi_statement_function,
261+
);
262+
257263
let create_function_with_conditional = "\
258264
CREATE FUNCTION some_scalar_udf() \
259265
RETURNS INT \
@@ -288,6 +294,87 @@ fn parse_create_function() {
288294
END\
289295
";
290296
let _ = ms().verified_stmt(create_function_with_return_expression);
297+
298+
let create_inline_table_value_function = "\
299+
CREATE FUNCTION some_inline_tvf(@foo INT, @bar VARCHAR(256)) \
300+
RETURNS TABLE \
301+
AS \
302+
RETURN (SELECT 1 AS col_1)\
303+
";
304+
let _ = ms().verified_stmt(create_inline_table_value_function);
305+
306+
let create_inline_table_value_function_without_parentheses = "\
307+
CREATE FUNCTION some_inline_tvf(@foo INT, @bar VARCHAR(256)) \
308+
RETURNS TABLE \
309+
AS \
310+
RETURN SELECT 1 AS col_1\
311+
";
312+
let _ = ms().verified_stmt(create_inline_table_value_function_without_parentheses);
313+
314+
let create_inline_table_value_function_without_as =
315+
create_inline_table_value_function.replace(" AS", "");
316+
let _ = ms().one_statement_parses_to(
317+
&create_inline_table_value_function_without_as,
318+
create_inline_table_value_function,
319+
);
320+
321+
let create_multi_statement_table_value_function = "\
322+
CREATE FUNCTION some_multi_statement_tvf(@foo INT, @bar VARCHAR(256)) \
323+
RETURNS @t TABLE (col_1 INT) \
324+
AS \
325+
BEGIN \
326+
INSERT INTO @t SELECT 1; \
327+
RETURN; \
328+
END\
329+
";
330+
let _ = ms().verified_stmt(create_multi_statement_table_value_function);
331+
332+
let create_multi_statement_table_value_function_without_as =
333+
create_multi_statement_table_value_function.replace(" AS", "");
334+
let _ = ms().one_statement_parses_to(
335+
&create_multi_statement_table_value_function_without_as,
336+
create_multi_statement_table_value_function,
337+
);
338+
339+
let create_multi_statement_table_value_function_with_constraints = "\
340+
CREATE FUNCTION some_multi_statement_tvf(@foo INT, @bar VARCHAR(256)) \
341+
RETURNS @t TABLE (col_1 INT NOT NULL) \
342+
AS \
343+
BEGIN \
344+
INSERT INTO @t SELECT 1; \
345+
RETURN @t; \
346+
END\
347+
";
348+
let _ = ms().verified_stmt(create_multi_statement_table_value_function_with_constraints);
349+
350+
let create_multi_statement_tvf_without_table_definition = "\
351+
CREATE FUNCTION incorrect_tvf(@foo INT, @bar VARCHAR(256)) \
352+
RETURNS @t TABLE ()
353+
AS \
354+
BEGIN \
355+
INSERT INTO @t SELECT 1; \
356+
RETURN @t; \
357+
END\
358+
";
359+
assert_eq!(
360+
ParserError::ParserError("Unparsable function body".to_owned()),
361+
ms().parse_sql_statements(create_multi_statement_tvf_without_table_definition)
362+
.unwrap_err()
363+
);
364+
365+
let create_inline_tvf_without_subquery_or_bare_select = "\
366+
CREATE FUNCTION incorrect_tvf(@foo INT, @bar VARCHAR(256)) \
367+
RETURNS TABLE
368+
AS \
369+
RETURN 'hi'\
370+
";
371+
assert_eq!(
372+
ParserError::ParserError(
373+
"Expected a subquery (or bare SELECT statement) after RETURN".to_owned()
374+
),
375+
ms().parse_sql_statements(create_inline_tvf_without_subquery_or_bare_select)
376+
.unwrap_err()
377+
);
291378
}
292379

293380
#[test]

0 commit comments

Comments
 (0)