PostgreSQL RETURNS SETOF Support in CREATE FUNCTION

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<DataType>),
     /// 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) => {

src/keywords.rs

@@ -933,6 +933,7 @@ define_keywords!(
     SESSION_USER,
     SET,
     SETERROR,
+    SETOF,
     SETS,
     SETTINGS,
     SHARE,

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
         };

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 =