diff --git a/src/ast/data_type.rs b/src/ast/data_type.rs index 285eec505..51ed582ee 100644 --- a/src/ast/data_type.rs +++ b/src/ast/data_type.rs @@ -47,6 +47,10 @@ pub enum EnumMember { #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] pub enum DataType { + /// Set-returning function type in [PostgreSQL], e.g. CREATE FUNCTION RETURNS SETOF UUID. + /// + /// [PostgreSQL]: https://www.postgresql.org/docs/current/sql-createfunction.html + SetOf(Box), /// Table type in [PostgreSQL], e.g. CREATE FUNCTION RETURNS TABLE(...). /// /// [PostgreSQL]: https://www.postgresql.org/docs/15/sql-createfunction.html @@ -501,6 +505,7 @@ pub enum DataType { impl fmt::Display for DataType { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { + DataType::SetOf(data_type) => write!(f, "SETOF {data_type}"), DataType::Character(size) => format_character_string_type(f, "CHARACTER", size), DataType::Char(size) => format_character_string_type(f, "CHAR", size), DataType::CharacterVarying(size) => { diff --git a/src/keywords.rs b/src/keywords.rs index 80f679c07..9f171cb24 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -933,6 +933,7 @@ define_keywords!( SESSION_USER, SET, SETERROR, + SETOF, SETS, SETTINGS, SHARE, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 75450f75d..89143478b 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -5581,7 +5581,13 @@ impl<'a> Parser<'a> { self.expect_token(&Token::RParen)?; let return_type = if self.parse_keyword(Keyword::RETURNS) { - Some(self.parse_data_type()?) + let setof = self.parse_keyword(Keyword::SETOF); + let return_type = self.parse_data_type()?; + Some(if setof { + DataType::SetOf(Box::new(return_type)) + } else { + return_type + }) } else { None }; diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 434c5fd7b..8d9930aa8 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -4642,6 +4642,31 @@ fn parse_create_function_detailed() { ); } +#[test] +fn parse_create_function_returns_setof() { + let cases = [ + "CREATE FUNCTION any_ids() RETURNS SETOF UUID LANGUAGE sql STABLE AS 'SELECT ''00000000-0000-0000-0000-000000000000''::uuid'", + "CREATE FUNCTION ids_for_user(p_user_id UUID) RETURNS SETOF UUID LANGUAGE sql STABLE AS 'SELECT p_user_id'", + "CREATE FUNCTION ids_from_array() RETURNS SETOF UUID LANGUAGE sql STABLE AS 'SELECT unnest(ARRAY[''00000000-0000-0000-0000-000000000000''::uuid])'", + ]; + + for sql in cases { + match pg_and_generic().verified_stmt(sql) { + Statement::CreateFunction(CreateFunction { + return_type, + language, + behavior, + .. + }) => { + assert_eq!(return_type, Some(DataType::SetOf(Box::new(DataType::Uuid)))); + assert_eq!(language, Some(Ident::new("sql"))); + assert_eq!(behavior, Some(FunctionBehavior::Stable)); + } + _ => panic!("Expected CreateFunction"), + } + } +} + #[test] fn parse_create_function_with_security() { let sql =