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

@@ -654,33 +654,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

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

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

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 =
@@ -8262,194 +8286,6 @@ fn parse_create_operator_class() {
         .is_err());
 }
 
-#[test]
-fn parse_create_table_partition_of_range() {
-    // RANGE partition with FROM ... TO
-    let sql = "CREATE TABLE measurement_y2006m02 PARTITION OF measurement FOR VALUES FROM ('2006-02-01') TO ('2006-03-01')";
-    match pg_and_generic().verified_stmt(sql) {
-        Statement::CreateTable(create_table) => {
-            assert_eq!("measurement_y2006m02", create_table.name.to_string());
-            assert_eq!(
-                Some(ObjectName::from(vec![Ident::new("measurement")])),
-                create_table.partition_of
-            );
-            match create_table.for_values {
-                Some(ForValues::From { from, to }) => {
-                    assert_eq!(1, from.len());
-                    assert_eq!(1, to.len());
-                    match &from[0] {
-                        PartitionBoundValue::Expr(Expr::Value(v)) => {
-                            assert_eq!("'2006-02-01'", v.to_string());
-                        }
-                        _ => panic!("Expected Expr value in from"),
-                    }
-                    match &to[0] {
-                        PartitionBoundValue::Expr(Expr::Value(v)) => {
-                            assert_eq!("'2006-03-01'", v.to_string());
-                        }
-                        _ => panic!("Expected Expr value in to"),
-                    }
-                }
-                _ => panic!("Expected ForValues::From"),
-            }
-        }
-        _ => panic!("Expected CreateTable"),
-    }
-}
-
-#[test]
-fn parse_create_table_partition_of_range_with_minvalue_maxvalue() {
-    // RANGE partition with MINVALUE/MAXVALUE
-    let sql =
-        "CREATE TABLE orders_old PARTITION OF orders FOR VALUES FROM (MINVALUE) TO ('2020-01-01')";
-    match pg_and_generic().verified_stmt(sql) {
-        Statement::CreateTable(create_table) => {
-            assert_eq!("orders_old", create_table.name.to_string());
-            assert_eq!(
-                Some(ObjectName::from(vec![Ident::new("orders")])),
-                create_table.partition_of
-            );
-            match create_table.for_values {
-                Some(ForValues::From { from, to }) => {
-                    assert_eq!(PartitionBoundValue::MinValue, from[0]);
-                    match &to[0] {
-                        PartitionBoundValue::Expr(Expr::Value(v)) => {
-                            assert_eq!("'2020-01-01'", v.to_string());
-                        }
-                        _ => panic!("Expected Expr value in to"),
-                    }
-                }
-                _ => panic!("Expected ForValues::From"),
-            }
-        }
-        _ => panic!("Expected CreateTable"),
-    }
-
-    // With MAXVALUE
-    let sql =
-        "CREATE TABLE orders_new PARTITION OF orders FOR VALUES FROM ('2024-01-01') TO (MAXVALUE)";
-    match pg_and_generic().verified_stmt(sql) {
-        Statement::CreateTable(create_table) => match create_table.for_values {
-            Some(ForValues::From { from, to }) => {
-                match &from[0] {
-                    PartitionBoundValue::Expr(Expr::Value(v)) => {
-                        assert_eq!("'2024-01-01'", v.to_string());
-                    }
-                    _ => panic!("Expected Expr value in from"),
-                }
-                assert_eq!(PartitionBoundValue::MaxValue, to[0]);
-            }
-            _ => panic!("Expected ForValues::From"),
-        },
-        _ => panic!("Expected CreateTable"),
-    }
-}
-
-#[test]
-fn parse_create_table_partition_of_list() {
-    // LIST partition
-    let sql = "CREATE TABLE orders_us PARTITION OF orders FOR VALUES IN ('US', 'CA', 'MX')";
-    match pg_and_generic().verified_stmt(sql) {
-        Statement::CreateTable(create_table) => {
-            assert_eq!("orders_us", create_table.name.to_string());
-            assert_eq!(
-                Some(ObjectName::from(vec![Ident::new("orders")])),
-                create_table.partition_of
-            );
-            match create_table.for_values {
-                Some(ForValues::In(values)) => {
-                    assert_eq!(3, values.len());
-                }
-                _ => panic!("Expected ForValues::In"),
-            }
-        }
-        _ => panic!("Expected CreateTable"),
-    }
-}
-
-#[test]
-fn parse_create_table_partition_of_hash() {
-    // HASH partition
-    let sql = "CREATE TABLE orders_p0 PARTITION OF orders FOR VALUES WITH (MODULUS 4, REMAINDER 0)";
-    match pg_and_generic().verified_stmt(sql) {
-        Statement::CreateTable(create_table) => {
-            assert_eq!("orders_p0", create_table.name.to_string());
-            assert_eq!(
-                Some(ObjectName::from(vec![Ident::new("orders")])),
-                create_table.partition_of
-            );
-            match create_table.for_values {
-                Some(ForValues::With { modulus, remainder }) => {
-                    assert_eq!(4, modulus);
-                    assert_eq!(0, remainder);
-                }
-                _ => panic!("Expected ForValues::With"),
-            }
-        }
-        _ => panic!("Expected CreateTable"),
-    }
-}
-
-#[test]
-fn parse_create_table_partition_of_default() {
-    // DEFAULT partition
-    let sql = "CREATE TABLE orders_default PARTITION OF orders DEFAULT";
-    match pg_and_generic().verified_stmt(sql) {
-        Statement::CreateTable(create_table) => {
-            assert_eq!("orders_default", create_table.name.to_string());
-            assert_eq!(
-                Some(ObjectName::from(vec![Ident::new("orders")])),
-                create_table.partition_of
-            );
-            assert_eq!(Some(ForValues::Default), create_table.for_values);
-        }
-        _ => panic!("Expected CreateTable"),
-    }
-}
-
-#[test]
-fn parse_create_table_partition_of_multicolumn_range() {
-    // Multi-column RANGE partition
-    let sql = "CREATE TABLE sales_2023_q1 PARTITION OF sales FOR VALUES FROM ('2023-01-01', 1) TO ('2023-04-01', 1)";
-    match pg_and_generic().verified_stmt(sql) {
-        Statement::CreateTable(create_table) => {
-            assert_eq!("sales_2023_q1", create_table.name.to_string());
-            match create_table.for_values {
-                Some(ForValues::From { from, to }) => {
-                    assert_eq!(2, from.len());
-                    assert_eq!(2, to.len());
-                }
-                _ => panic!("Expected ForValues::From"),
-            }
-        }
-        _ => panic!("Expected CreateTable"),
-    }
-}
-
-#[test]
-fn parse_create_table_partition_of_with_constraint() {
-    // With table constraint (not column constraint which has different syntax in PARTITION OF)
-    let sql = "CREATE TABLE orders_2023 PARTITION OF orders (\
-CONSTRAINT check_date CHECK (order_date >= '2023-01-01')\
-) FOR VALUES FROM ('2023-01-01') TO ('2024-01-01')";
-    match pg_and_generic().verified_stmt(sql) {
-        Statement::CreateTable(create_table) => {
-            assert_eq!("orders_2023", create_table.name.to_string());
-            assert_eq!(
-                Some(ObjectName::from(vec![Ident::new("orders")])),
-                create_table.partition_of
-            );
-            // Check that table constraint was parsed
-            assert_eq!(1, create_table.constraints.len());
-            match create_table.for_values {
-                Some(ForValues::From { .. }) => {}
-                _ => panic!("Expected ForValues::From"),
-            }
-        }
-        _ => panic!("Expected CreateTable"),
-    }
-}
-
 #[test]
 fn parse_identifiers_semicolon_handling() {
     let statement = "SHOW search_path; SELECT 1";