Skip to content

Commit

Permalink
Add support for the TABLE returns datatype in create function for pos…
Browse files Browse the repository at this point in the history
…tgresql.
  • Loading branch information
remysaissy committed Jan 29, 2025
1 parent 269967a commit f8637f2
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 1 deletion.
5 changes: 5 additions & 0 deletions src/ast/data_type.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,10 @@ pub enum EnumMember {
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
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>),
/// Fixed-length character type e.g. CHARACTER(10)
Character(Option<CharacterLength>),
/// Fixed-length char type e.g. CHAR(10)
Expand Down Expand Up @@ -630,6 +634,7 @@ 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)),
}
}
}
Expand Down
27 changes: 26 additions & 1 deletion src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4535,7 +4535,14 @@ impl<'a> Parser<'a> {
self.expect_token(&Token::RParen)?;

let return_type = if self.parse_keyword(Keyword::RETURNS) {
Some(self.parse_data_type()?)
if dialect_of!(self is PostgreSqlDialect | GenericDialect)
&& self.parse_keyword(Keyword::TABLE)
{
let columns = self.parse_parenthesized_columns()?;
Some(DataType::Table(columns))
} else {
Some(self.parse_data_type()?)
}
} else {
None
};
Expand Down Expand Up @@ -8866,6 +8873,24 @@ impl<'a> Parser<'a> {
Ok((data, trailing_bracket))
}

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,
collation: None,
options: Vec::new(), // No constraints expected here
})
}

fn parse_parenthesized_columns(&mut self) -> Result<Vec<ColumnDef>, ParserError> {
self.expect_token(&Token::LParen)?;
let columns = self.parse_comma_separated(Parser::parse_returns_table_column)?;
self.expect_token(&Token::RParen)?;
Ok(columns)
}

pub fn parse_string_values(&mut self) -> Result<Vec<String>, ParserError> {
self.expect_token(&Token::LParen)?;
let mut values = Vec::new();
Expand Down
1 change: 1 addition & 0 deletions tests/sqlparser_postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3802,6 +3802,7 @@ fn parse_create_function_detailed() {
pg_and_generic().verified_stmt("CREATE OR REPLACE FUNCTION add(a INTEGER, IN b INTEGER = 1) RETURNS INTEGER LANGUAGE SQL STABLE PARALLEL UNSAFE RETURN a + b");
pg_and_generic().verified_stmt("CREATE OR REPLACE FUNCTION add(a INTEGER, IN b INTEGER = 1) RETURNS INTEGER LANGUAGE SQL STABLE CALLED ON NULL INPUT PARALLEL UNSAFE RETURN a + b");
pg_and_generic().verified_stmt(r#"CREATE OR REPLACE FUNCTION increment(i INTEGER) RETURNS INTEGER LANGUAGE plpgsql AS $$ BEGIN RETURN i + 1; END; $$"#);
pg_and_generic().verified_stmt(r#"CREATE OR REPLACE FUNCTION return_table(i INTEGER) RETURNS TABLE(id UUID, is_active BOOLEAN) LANGUAGE plpgsql AS $$ BEGIN RETURN QUERY SELECT NULL::UUID, NULL::BOOLEAN; END; $$"#);
}
#[test]
fn parse_incorrect_create_function_parallel() {
Expand Down

0 comments on commit f8637f2

Please sign in to comment.