refactor(ast): move SETOF from DataType into FunctionReturnType enum

SETOF is a PostgreSQL-specific modifier for function return types, not a
standalone data type. Introduce FunctionReturnType enum with DataType and
SetOf variants, replacing DataType::SetOf.

Addresses review feedback on #2217.

Author: fmguerreiro

src/ast/data_type.rs

@@ -61,11 +61,6 @@ 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).
@@ -813,7 +808,6 @@ 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/ddl.rs

@@ -3466,6 +3466,31 @@ impl fmt::Display for CreateDomain {
     }
 }
 
+/// The return type of a `CREATE FUNCTION` statement.
+///
+/// Separates `SETOF` from plain data types because `SETOF` is a
+/// [PostgreSQL]-specific modifier, not a standalone data type.
+///
+/// [PostgreSQL]: https://www.postgresql.org/docs/current/sql-createfunction.html
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum FunctionReturnType {
+    /// A normal data type, e.g. `RETURNS TEXT`.
+    DataType(DataType),
+    /// `RETURNS SETOF <type>`, e.g. `RETURNS SETOF TEXT`.
+    SetOf(DataType),
+}
+
+impl fmt::Display for FunctionReturnType {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            FunctionReturnType::DataType(data_type) => write!(f, "{data_type}"),
+            FunctionReturnType::SetOf(data_type) => write!(f, "SETOF {data_type}"),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
@@ -3486,7 +3511,7 @@ pub struct CreateFunction {
     /// List of arguments for the function.
     pub args: Option<Vec<OperateFunctionArg>>,
     /// The return type of the function.
-    pub return_type: Option<DataType>,
+    pub return_type: Option<FunctionReturnType>,
     /// The expression that defines the function.
     ///
     /// Examples:

src/ast/mod.rs

@@ -71,15 +71,16 @@ pub use self::ddl::{
     CreateIndex, CreateOperator, CreateOperatorClass, CreateOperatorFamily, CreatePolicy,
     CreatePolicyCommand, CreatePolicyType, CreateTable, CreateTrigger, CreateView, Deduplicate,
     DeferrableInitial, DropBehavior, DropExtension, DropFunction, DropOperator, DropOperatorClass,
-    DropOperatorFamily, DropOperatorSignature, DropPolicy, DropTrigger, ForValues, GeneratedAs,
-    GeneratedExpressionMode, IdentityParameters, IdentityProperty, IdentityPropertyFormatKind,
-    IdentityPropertyKind, IdentityPropertyOrder, IndexColumn, IndexOption, IndexType,
-    KeyOrIndexDisplay, Msck, NullsDistinctOption, OperatorArgTypes, OperatorClassItem,
-    OperatorFamilyDropItem, OperatorFamilyItem, OperatorOption, OperatorPurpose, Owner, Partition,
-    PartitionBoundValue, ProcedureParam, ReferentialAction, RenameTableNameKind, ReplicaIdentity,
-    TagsColumnOption, TriggerObjectKind, Truncate, UserDefinedTypeCompositeAttributeDef,
-    UserDefinedTypeInternalLength, UserDefinedTypeRangeOption, UserDefinedTypeRepresentation,
-    UserDefinedTypeSqlDefinitionOption, UserDefinedTypeStorage, ViewColumnDef,
+    DropOperatorFamily, DropOperatorSignature, DropPolicy, DropTrigger, ForValues,
+    FunctionReturnType, GeneratedAs, GeneratedExpressionMode, IdentityParameters, IdentityProperty,
+    IdentityPropertyFormatKind, IdentityPropertyKind, IdentityPropertyOrder, IndexColumn,
+    IndexOption, IndexType, KeyOrIndexDisplay, Msck, NullsDistinctOption, OperatorArgTypes,
+    OperatorClassItem, OperatorFamilyDropItem, OperatorFamilyItem, OperatorOption, OperatorPurpose,
+    Owner, Partition, PartitionBoundValue, ProcedureParam, ReferentialAction, RenameTableNameKind,
+    ReplicaIdentity, TagsColumnOption, TriggerObjectKind, Truncate,
+    UserDefinedTypeCompositeAttributeDef, UserDefinedTypeInternalLength,
+    UserDefinedTypeRangeOption, UserDefinedTypeRepresentation, UserDefinedTypeSqlDefinitionOption,
+    UserDefinedTypeStorage, ViewColumnDef,
 };
 pub use self::dml::{
     Delete, Insert, Merge, MergeAction, MergeClause, MergeClauseKind, MergeInsertExpr,

src/ast/spans.rs

@@ -597,6 +597,31 @@ impl Spanned for CreateTable {
     }
 }
 
+impl Spanned for PartitionBoundValue {
+    fn span(&self) -> Span {
+        match self {
+            PartitionBoundValue::Expr(expr) => expr.span(),
+            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())),
+            ),
+            ForValues::With { .. } => Span::empty(),
+            ForValues::Default => Span::empty(),
+        }
+    }
+}
+
 impl Spanned for ColumnDef {
     fn span(&self) -> Span {
         let ColumnDef {

src/parser/mod.rs

@@ -5545,7 +5545,7 @@ impl<'a> Parser<'a> {
         self.expect_token(&Token::RParen)?;
 
         let return_type = if self.parse_keyword(Keyword::RETURNS) {
-            Some(self.parse_data_type()?)
+            Some(self.parse_function_return_type()?)
         } else {
             None
         };
@@ -5724,7 +5724,7 @@ impl<'a> Parser<'a> {
         let (name, args) = self.parse_create_function_name_and_params()?;
 
         let return_type = if self.parse_keyword(Keyword::RETURNS) {
-            Some(self.parse_data_type()?)
+            Some(self.parse_function_return_type()?)
         } else {
             None
         };
@@ -5827,10 +5827,10 @@ impl<'a> Parser<'a> {
             })
         })?;
 
-        let return_type = if return_table.is_some() {
-            return_table
+        let return_type = if let Some(table_type) = return_table {
+            Some(FunctionReturnType::DataType(table_type))
         } else {
-            Some(self.parse_data_type()?)
+            Some(FunctionReturnType::DataType(self.parse_data_type()?))
         };
 
         let _ = self.parse_keyword(Keyword::AS);
@@ -5883,6 +5883,17 @@ impl<'a> Parser<'a> {
         })
     }
 
+    /// Parse a [FunctionReturnType] after the `RETURNS` keyword.
+    ///
+    /// Handles `RETURNS SETOF <type>` (PostgreSQL) and plain `RETURNS <type>`.
+    fn parse_function_return_type(&mut self) -> Result<FunctionReturnType, ParserError> {
+        if self.parse_keyword(Keyword::SETOF) {
+            Ok(FunctionReturnType::SetOf(self.parse_data_type()?))
+        } else {
+            Ok(FunctionReturnType::DataType(self.parse_data_type()?))
+        }
+    }
+
     fn parse_create_function_name_and_params(
         &mut self,
     ) -> Result<(ObjectName, Vec<OperateFunctionArg>), ParserError> {
@@ -12064,10 +12075,6 @@ 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_bigquery.rs

@@ -2279,7 +2279,7 @@ fn test_bigquery_create_function() {
                 Ident::new("myfunction"),
             ]),
             args: Some(vec![OperateFunctionArg::with_name("x", DataType::Float64),]),
-            return_type: Some(DataType::Float64),
+            return_type: Some(FunctionReturnType::DataType(DataType::Float64)),
             function_body: Some(CreateFunctionBody::AsAfterOptions(Expr::Value(
                 number("42").with_empty_span()
             ))),

tests/sqlparser_mssql.rs

@@ -255,7 +255,7 @@ fn parse_create_function() {
                     default_expr: None,
                 },
             ]),
-            return_type: Some(DataType::Int(None)),
+            return_type: Some(FunctionReturnType::DataType(DataType::Int(None))),
             function_body: Some(CreateFunctionBody::AsBeginEnd(BeginEndStatements {
                 begin_token: AttachedToken::empty(),
                 statements: vec![Statement::Return(ReturnStatement {
@@ -430,7 +430,7 @@ fn parse_create_function_parameter_default_values() {
                 data_type: DataType::Int(None),
                 default_expr: Some(Expr::Value((number("42")).with_empty_span())),
             },]),
-            return_type: Some(DataType::Int(None)),
+            return_type: Some(FunctionReturnType::DataType(DataType::Int(None))),
             function_body: Some(CreateFunctionBody::AsBeginEnd(BeginEndStatements {
                 begin_token: AttachedToken::empty(),
                 statements: vec![Statement::Return(ReturnStatement {

tests/sqlparser_postgres.rs

@@ -4346,7 +4346,7 @@ $$"#;
                     DataType::Varchar(None),
                 ),
             ]),
-            return_type: Some(DataType::Boolean),
+            return_type: Some(FunctionReturnType::DataType(DataType::Boolean)),
             language: Some("plpgsql".into()),
             behavior: None,
             called_on_null: None,
@@ -4389,7 +4389,7 @@ $$"#;
                     DataType::Int(None)
                 )
             ]),
-            return_type: Some(DataType::Boolean),
+            return_type: Some(FunctionReturnType::DataType(DataType::Boolean)),
             language: Some("plpgsql".into()),
             behavior: None,
             called_on_null: None,
@@ -4436,7 +4436,7 @@ $$"#;
                     DataType::Int(None)
                 ),
             ]),
-            return_type: Some(DataType::Boolean),
+            return_type: Some(FunctionReturnType::DataType(DataType::Boolean)),
             language: Some("plpgsql".into()),
             behavior: None,
             called_on_null: None,
@@ -4483,7 +4483,7 @@ $$"#;
                     DataType::Int(None)
                 ),
             ]),
-            return_type: Some(DataType::Boolean),
+            return_type: Some(FunctionReturnType::DataType(DataType::Boolean)),
             language: Some("plpgsql".into()),
             behavior: None,
             called_on_null: None,
@@ -4523,7 +4523,7 @@ $$"#;
                 ),
                 OperateFunctionArg::with_name("b", DataType::Varchar(None)),
             ]),
-            return_type: Some(DataType::Boolean),
+            return_type: Some(FunctionReturnType::DataType(DataType::Boolean)),
             language: Some("plpgsql".into()),
             behavior: None,
             called_on_null: None,
@@ -4566,7 +4566,7 @@ fn parse_create_function() {
                 OperateFunctionArg::unnamed(DataType::Integer(None)),
                 OperateFunctionArg::unnamed(DataType::Integer(None)),
             ]),
-            return_type: Some(DataType::Integer(None)),
+            return_type: Some(FunctionReturnType::DataType(DataType::Integer(None))),
             language: Some("SQL".into()),
             behavior: Some(FunctionBehavior::Immutable),
             called_on_null: Some(FunctionCalledOnNull::Strict),
@@ -4621,7 +4621,7 @@ fn parse_create_function_returns_setof() {
     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))));
+            assert_eq!(return_type, Some(FunctionReturnType::SetOf(DataType::Text)));
         }
         _ => panic!("Expected CreateFunction"),
     }
@@ -4702,10 +4702,10 @@ fn parse_create_function_c_with_module_pathname() {
                 "input",
                 DataType::Custom(ObjectName::from(vec![Ident::new("cstring")]), vec![]),
             ),]),
-            return_type: Some(DataType::Custom(
+            return_type: Some(FunctionReturnType::DataType(DataType::Custom(
                 ObjectName::from(vec![Ident::new("cas")]),
                 vec![]
-            )),
+            ))),
             language: Some("c".into()),
             behavior: Some(FunctionBehavior::Immutable),
             called_on_null: None,
@@ -6399,7 +6399,7 @@ fn parse_trigger_related_functions() {
             if_not_exists: false,
             name: ObjectName::from(vec![Ident::new("emp_stamp")]),
             args: Some(vec![]),
-            return_type: Some(DataType::Trigger),
+            return_type: Some(FunctionReturnType::DataType(DataType::Trigger)),
             function_body: Some(
                 CreateFunctionBody::AsBeforeOptions {
                     body: Expr::Value((