diff --git a/src/ast/data_type.rs b/src/ast/data_type.rs index 285eec505..4182bc620 100644 --- a/src/ast/data_type.rs +++ b/src/ast/data_type.rs @@ -61,6 +61,11 @@ pub enum DataType { /// Table columns. columns: Vec, }, + /// SETOF type modifier for [PostgreSQL] function return types, + /// e.g. `CREATE FUNCTION ... RETURNS SETOF text`. + /// + /// [PostgreSQL]: https://www.postgresql.org/docs/current/sql-createfunction.html + SetOf(Box), /// Fixed-length character type, e.g. CHARACTER(10). Character(Option), /// Fixed-length char type, e.g. CHAR(10). @@ -808,6 +813,7 @@ impl fmt::Display for DataType { DataType::NamedTable { name, columns } => { write!(f, "{} TABLE ({})", name, display_comma_separated(columns)) } + DataType::SetOf(inner) => write!(f, "SETOF {inner}"), DataType::GeometricType(kind) => write!(f, "{kind}"), DataType::TsVector => write!(f, "TSVECTOR"), DataType::TsQuery => write!(f, "TSQUERY"), diff --git a/src/keywords.rs b/src/keywords.rs index 7950b1918..f738276ff 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -932,6 +932,7 @@ define_keywords!( SESSION_USER, SET, SETERROR, + SETOF, SETS, SETTINGS, SHARE, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 7a2bda8ac..000d156d6 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -11976,6 +11976,10 @@ impl<'a> Parser<'a> { let field_defs = self.parse_click_house_tuple_def()?; Ok(DataType::Tuple(field_defs)) } + Keyword::SETOF => { + let inner = self.parse_data_type()?; + Ok(DataType::SetOf(Box::new(inner))) + } Keyword::TRIGGER => Ok(DataType::Trigger), Keyword::ANY if self.peek_keyword(Keyword::TYPE) => { let _ = self.parse_keyword(Keyword::TYPE); diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index f8c738136..5d6706e5f 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -4564,6 +4564,30 @@ fn parse_create_function_detailed() { ); } +#[test] +fn parse_create_function_returns_setof() { + pg_and_generic().verified_stmt( + "CREATE FUNCTION get_users() RETURNS SETOF TEXT LANGUAGE sql AS 'SELECT name FROM users'", + ); + pg_and_generic().verified_stmt( + "CREATE FUNCTION get_ids() RETURNS SETOF INTEGER LANGUAGE sql AS 'SELECT id FROM users'", + ); + pg_and_generic().verified_stmt( + r#"CREATE FUNCTION get_all() RETURNS SETOF my_schema."MyType" LANGUAGE sql AS 'SELECT * FROM t'"#, + ); + pg_and_generic().verified_stmt( + "CREATE FUNCTION get_rows() RETURNS SETOF RECORD LANGUAGE sql AS 'SELECT * FROM t'", + ); + + let sql = "CREATE FUNCTION get_names() RETURNS SETOF TEXT LANGUAGE sql AS 'SELECT name FROM t'"; + match pg_and_generic().verified_stmt(sql) { + Statement::CreateFunction(CreateFunction { return_type, .. }) => { + assert_eq!(return_type, Some(DataType::SetOf(Box::new(DataType::Text)))); + } + _ => panic!("Expected CreateFunction"), + } +} + #[test] fn parse_create_function_with_security() { let sql =