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/ast/spans.rs

@@ -632,33 +632,6 @@ impl Spanned for TableConstraint {
     }
 }
 
-impl Spanned for PartitionBoundValue {
-    fn span(&self) -> Span {
-        match self {
-            PartitionBoundValue::Expr(expr) => expr.span(),
-            // MINVALUE and MAXVALUE are keywords without tracked spans
-            PartitionBoundValue::MinValue => Span::empty(),
-            PartitionBoundValue::MaxValue => Span::empty(),
-        }
-    }
-}
-
-impl Spanned for ForValues {
-    fn span(&self) -> Span {
-        match self {
-            ForValues::In(exprs) => union_spans(exprs.iter().map(|e| e.span())),
-            ForValues::From { from, to } => union_spans(
-                from.iter()
-                    .map(|v| v.span())
-                    .chain(to.iter().map(|v| v.span())),
-            ),
-            // WITH (MODULUS n, REMAINDER r) - u64 values have no spans
-            ForValues::With { .. } => Span::empty(),
-            ForValues::Default => Span::empty(),
-        }
-    }
-}
-
 impl Spanned for CreateIndex {
     fn span(&self) -> Span {
         let CreateIndex {

src/keywords.rs

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

src/parser/mod.rs

@@ -12064,6 +12064,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);

tests/sqlparser_postgres.rs

@@ -4603,6 +4603,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 =