This is an automated email from the ASF dual-hosted git repository. github-bot pushed a commit to branch gh-readonly-queue/main/pr-2217-203ced42be28f4e418ca93a73fc75911be19f4b4 in repository https://gitbox.apache.org/repos/asf/datafusion-sqlparser-rs.git
commit b3e176daf302985feef46a3933428b196efb78e8 Author: Filipe Guerreiro <[email protected]> AuthorDate: Fri Mar 13 21:57:52 2026 +0900 Add SETOF support for PostgreSQL function return types (#2217) --- src/ast/ddl.rs | 24 +++++++++++++++++++++++- src/ast/mod.rs | 14 +++++++------- src/keywords.rs | 1 + src/parser/mod.rs | 22 +++++++++++++++------- tests/sqlparser_bigquery.rs | 2 +- tests/sqlparser_mssql.rs | 4 ++-- tests/sqlparser_postgres.rs | 42 +++++++++++++++++++++++++++++++++--------- 7 files changed, 82 insertions(+), 27 deletions(-) diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs index 157b209d..f0e79e73 100644 --- a/src/ast/ddl.rs +++ b/src/ast/ddl.rs @@ -3533,6 +3533,28 @@ impl fmt::Display for CreateDomain { } } +/// The return type of a `CREATE FUNCTION` statement. +#[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 { + /// `RETURNS <type>` + DataType(DataType), + /// `RETURNS SETOF <type>` + /// + /// [PostgreSQL](https://www.postgresql.org/docs/current/sql-createfunction.html) + 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))] @@ -3553,7 +3575,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: diff --git a/src/ast/mod.rs b/src/ast/mod.rs index cff089bc..c4d1b50c 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -72,13 +72,13 @@ pub use self::ddl::{ CreatePolicyCommand, CreatePolicyType, CreateTable, CreateTrigger, CreateView, Deduplicate, DeferrableInitial, DistStyle, 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, + 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, }; diff --git a/src/keywords.rs b/src/keywords.rs index d56d0484..94458ccb 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -938,6 +938,7 @@ define_keywords!( SESSION_USER, SET, SETERROR, + SETOF, SETS, SETTINGS, SHARE, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 3bc34bda..6adecb0c 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -5594,7 +5594,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 }; @@ -5774,7 +5774,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 }; @@ -5877,11 +5877,11 @@ impl<'a> Parser<'a> { }) })?; - let return_type = if return_table.is_some() { - return_table - } else { - Some(self.parse_data_type()?) + let data_type = match return_table { + Some(table_type) => table_type, + None => self.parse_data_type()?, }; + let return_type = Some(FunctionReturnType::DataType(data_type)); let _ = self.parse_keyword(Keyword::AS); @@ -5933,6 +5933,14 @@ impl<'a> Parser<'a> { }) } + 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> { @@ -8608,7 +8616,7 @@ impl<'a> Parser<'a> { } } - /// Parse a single [PartitionBoundValue]. + /// Parse a single partition bound value (MINVALUE, MAXVALUE, or expression). fn parse_partition_bound_value(&mut self) -> Result<PartitionBoundValue, ParserError> { if self.parse_keyword(Keyword::MINVALUE) { Ok(PartitionBoundValue::MinValue) diff --git a/tests/sqlparser_bigquery.rs b/tests/sqlparser_bigquery.rs index d3e47f99..79db34b0 100644 --- a/tests/sqlparser_bigquery.rs +++ b/tests/sqlparser_bigquery.rs @@ -2289,7 +2289,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() ))), diff --git a/tests/sqlparser_mssql.rs b/tests/sqlparser_mssql.rs index f2156e64..07dd0fcb 100644 --- a/tests/sqlparser_mssql.rs +++ b/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 { diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 9486af04..6b4f35d7 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -4441,7 +4441,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, @@ -4484,7 +4484,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, @@ -4531,7 +4531,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, @@ -4578,7 +4578,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, @@ -4618,7 +4618,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, @@ -4661,7 +4661,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), @@ -4698,6 +4698,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(FunctionReturnType::SetOf(DataType::Text))); + } + _ => panic!("Expected CreateFunction"), + } +} + #[test] fn parse_create_function_with_security() { let sql = @@ -4773,10 +4797,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, @@ -6493,7 +6517,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(( --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
