Add SETOF support for PostgreSQL function return types

Author: fmguerreiro

src/ast/data_type.rs

@@ -61,6 +61,11 @@ pub enum DataType {
         /// Table columns.
         columns: Vec<ColumnDef>,
     },
+    /// 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<DataType>),
     /// Fixed-length character type, e.g. CHARACTER(10).
     Character(Option<CharacterLength>),
     /// 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"),

src/keywords.rs

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

src/parser/mod.rs

@@ -5488,7 +5488,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 is_setof = self.parse_keyword(Keyword::SETOF);
+            let base_type = self.parse_data_type()?;
+            Some(if is_setof {
+                DataType::SetOf(Box::new(base_type))
+            } else {
+                base_type
+            })
         } else {
             None
         };

tests/sqlparser_postgres.rs

@@ -4564,6 +4564,27 @@ 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(
+        r#"CREATE FUNCTION get_all() RETURNS SETOF my_schema."MyType" 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 =