From 1be060fcc103f209252aa0dfb48f88a565dc9f75 Mon Sep 17 00:00:00 2001 From: lovasoa Date: Fri, 15 Nov 2024 22:51:36 +0100 Subject: [PATCH 1/4] This adds support for ```sql SELECT * FROM some_fun() AS x (a TEXT, b INT) ``` fixes https://github.com/apache/datafusion-sqlparser-rs/issues/1524 --- src/ast/mod.rs | 4 +-- src/ast/query.rs | 36 +++++++++++++++++++++++- src/parser/mod.rs | 21 ++++++++++++-- tests/sqlparser_common.rs | 59 +++++++++++++++++++++++++++++++++------ 4 files changed, 106 insertions(+), 14 deletions(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 39c742153..dd81b0a5c 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -60,8 +60,8 @@ pub use self::query::{ OrderBy, OrderByExpr, PivotValueSource, ProjectionSelect, Query, RenameSelectItem, RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select, SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Setting, SymbolDefinition, Table, - TableAlias, TableFactor, TableFunctionArgs, TableVersion, TableWithJoins, Top, TopQuantity, - ValueTableMode, Values, WildcardAdditionalOptions, With, WithFill, + TableAlias, TableAliasColumnDef, TableFactor, TableFunctionArgs, TableVersion, TableWithJoins, + Top, TopQuantity, ValueTableMode, Values, WildcardAdditionalOptions, With, WithFill, }; pub use self::trigger::{ diff --git a/src/ast/query.rs b/src/ast/query.rs index 60ebe3765..1fd923fbc 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -1597,7 +1597,7 @@ impl fmt::Display for TableFactor { #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] pub struct TableAlias { pub name: Ident, - pub columns: Vec, + pub columns: Vec, } impl fmt::Display for TableAlias { @@ -1610,6 +1610,40 @@ impl fmt::Display for TableAlias { } } +/// SQL column definition in a table expression alias. +/// Most of the time, the data type is not specified. +/// But some table-valued functions do require specifying the data type. +/// +/// See https://www.postgresql.org/docs/17/queries-table-expressions.html#QUERIES-TABLEFUNCTIONS +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +pub struct TableAliasColumnDef { + /// Column name alias + pub name: Ident, + /// Some table-valued functions require specifying the data type in the alias. + pub data_type: Option, +} + +impl TableAliasColumnDef { + pub fn from_column_name(name: &str) -> Self { + TableAliasColumnDef { + name: Ident::new(name), + data_type: None, + } + } +} + +impl fmt::Display for TableAliasColumnDef { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{}", self.name)?; + if let Some(ref data_type) = self.data_type { + write!(f, " {}", data_type)?; + } + Ok(()) + } +} + #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] diff --git a/src/parser/mod.rs b/src/parser/mod.rs index a583112a7..a0cc0f542 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -8270,7 +8270,7 @@ impl<'a> Parser<'a> { ) -> Result, ParserError> { match self.parse_optional_alias(reserved_kwds)? { Some(name) => { - let columns = self.parse_parenthesized_column_list(Optional, false)?; + let columns = self.parse_table_alias_column_defs()?; Ok(Some(TableAlias { name, columns })) } None => Ok(None), @@ -8607,6 +8607,23 @@ impl<'a> Parser<'a> { } } + /// Parse a parenthesized comma-separated list of unqualified, possibly quoted identifiers + pub fn parse_table_alias_column_defs( + &mut self, + ) -> Result, ParserError> { + if self.consume_token(&Token::LParen) { + let cols = self.parse_comma_separated(|p| { + let name = p.parse_identifier(false)?; + let data_type = p.maybe_parse(|p| p.parse_data_type())?; + Ok(TableAliasColumnDef { name, data_type }) + })?; + self.expect_token(&Token::RParen)?; + Ok(cols) + } else { + Ok(vec![]) + } + } + pub fn parse_precision(&mut self) -> Result { self.expect_token(&Token::LParen)?; let n = self.parse_literal_uint()?; @@ -9174,7 +9191,7 @@ impl<'a> Parser<'a> { materialized: is_materialized, } } else { - let columns = self.parse_parenthesized_column_list(Optional, false)?; + let columns = self.parse_table_alias_column_defs()?; self.expect_keyword(Keyword::AS)?; let mut is_materialized = None; if dialect_of!(self is PostgreSqlDialect) { diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index 2ffb5f44b..44314ee47 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -532,7 +532,11 @@ fn parse_select_with_table_alias() { name: ObjectName(vec![Ident::new("lineitem")]), alias: Some(TableAlias { name: Ident::new("l"), - columns: vec![Ident::new("A"), Ident::new("B"), Ident::new("C"),], + columns: vec![ + TableAliasColumnDef::from_column_name("A"), + TableAliasColumnDef::from_column_name("B"), + TableAliasColumnDef::from_column_name("C"), + ], }), args: None, with_hints: vec![], @@ -5576,6 +5580,40 @@ fn parse_table_function() { ); } +#[test] +fn parse_table_valued_function_with_alias_and_column_defs() { + let sql = r#"SELECT * FROM jsonb_to_record('{"a": "x", "b": 2}'::JSONB) AS x (a TEXT, b INT)"#; + let select = verified_only_select(sql); + + match only(&select.from) { + TableWithJoins { + relation: TableFactor::Table { + alias: Some(alias), .. + }, + .. + } => { + assert_eq!(alias.name.value, "x"); + assert_eq!( + alias.columns, + vec![ + TableAliasColumnDef { + name: Ident::new("a"), + data_type: Some(DataType::Text), + }, + TableAliasColumnDef { + name: Ident::new("b"), + data_type: Some(DataType::Int(None)), + }, + ] + ); + } + _ => unreachable!( + "Expecting only TableWithJoins with TableFactor::Table, got {:#?}", + select.from + ), + } +} + #[test] fn parse_unnest() { let sql = "SELECT UNNEST(make_array(1, 2, 3))"; @@ -6335,7 +6373,10 @@ fn parse_cte_renamed_columns() { let sql = "WITH cte (col1, col2) AS (SELECT foo, bar FROM baz) SELECT * FROM cte"; let query = all_dialects().verified_query(sql); assert_eq!( - vec![Ident::new("col1"), Ident::new("col2")], + vec![ + TableAliasColumnDef::from_column_name("col1"), + TableAliasColumnDef::from_column_name("col2") + ], query .with .unwrap() @@ -6364,10 +6405,7 @@ fn parse_recursive_cte() { value: "nums".to_string(), quote_style: None, }, - columns: vec![Ident { - value: "val".to_string(), - quote_style: None, - }], + columns: vec![TableAliasColumnDef::from_column_name("val")], }, query: Box::new(cte_query), from: None, @@ -9310,7 +9348,10 @@ fn parse_pivot_table() { value: "p".to_string(), quote_style: None }, - columns: vec![Ident::new("c"), Ident::new("d")], + columns: vec![ + TableAliasColumnDef::from_column_name("c"), + TableAliasColumnDef::from_column_name("d"), + ], }), } ); @@ -9371,8 +9412,8 @@ fn parse_unpivot_table() { name: Ident::new("u"), columns: ["product", "quarter", "quantity"] .into_iter() - .map(Ident::new) - .collect() + .map(TableAliasColumnDef::from_column_name) + .collect(), }), } ); From 8494ce86d23072ac3889ac41e5fc513cb5d6b469 Mon Sep 17 00:00:00 2001 From: lovasoa Date: Fri, 15 Nov 2024 23:10:31 +0100 Subject: [PATCH 2/4] fix cargo doc --- src/ast/query.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ast/query.rs b/src/ast/query.rs index 1fd923fbc..e99864556 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -1614,7 +1614,7 @@ impl fmt::Display for TableAlias { /// Most of the time, the data type is not specified. /// But some table-valued functions do require specifying the data type. /// -/// See https://www.postgresql.org/docs/17/queries-table-expressions.html#QUERIES-TABLEFUNCTIONS +/// See #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] From 9f6f8e33258404fa4e1026357c3aa62de5c527aa Mon Sep 17 00:00:00 2001 From: lovasoa Date: Sat, 16 Nov 2024 23:22:13 +0100 Subject: [PATCH 3/4] apply suggestions --- src/ast/query.rs | 1 + src/parser/mod.rs | 6 ++---- tests/sqlparser_common.rs | 2 +- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/ast/query.rs b/src/ast/query.rs index e99864556..b9aaadb99 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -1626,6 +1626,7 @@ pub struct TableAliasColumnDef { } impl TableAliasColumnDef { + /// Create a new table alias column definition with only a name and no type pub fn from_column_name(name: &str) -> Self { TableAliasColumnDef { name: Ident::new(name), diff --git a/src/parser/mod.rs b/src/parser/mod.rs index a0cc0f542..98363f9ad 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -8607,10 +8607,8 @@ impl<'a> Parser<'a> { } } - /// Parse a parenthesized comma-separated list of unqualified, possibly quoted identifiers - pub fn parse_table_alias_column_defs( - &mut self, - ) -> Result, ParserError> { + /// Parse a parenthesized comma-separated list of table alias column definitions. + fn parse_table_alias_column_defs(&mut self) -> Result, ParserError> { if self.consume_token(&Token::LParen) { let cols = self.parse_comma_separated(|p| { let name = p.parse_identifier(false)?; diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index 44314ee47..f8c158510 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -5581,7 +5581,7 @@ fn parse_table_function() { } #[test] -fn parse_table_valued_function_with_alias_and_column_defs() { +fn parse_select_with_alias_and_column_defs() { let sql = r#"SELECT * FROM jsonb_to_record('{"a": "x", "b": 2}'::JSONB) AS x (a TEXT, b INT)"#; let select = verified_only_select(sql); From 85d8b4f60f865de112ea6ab1b95cfb5ae94a530e Mon Sep 17 00:00:00 2001 From: lovasoa Date: Sun, 17 Nov 2024 10:39:33 +0100 Subject: [PATCH 4/4] fix https://github.com/apache/datafusion-sqlparser-rs/pull/1526#discussion_r1845295217 --- src/ast/query.rs | 2 +- tests/sqlparser_common.rs | 18 +++++++++--------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/ast/query.rs b/src/ast/query.rs index b9aaadb99..b9849358f 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -1627,7 +1627,7 @@ pub struct TableAliasColumnDef { impl TableAliasColumnDef { /// Create a new table alias column definition with only a name and no type - pub fn from_column_name(name: &str) -> Self { + pub fn from_name>(name: S) -> Self { TableAliasColumnDef { name: Ident::new(name), data_type: None, diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index f8c158510..f4be74224 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -533,9 +533,9 @@ fn parse_select_with_table_alias() { alias: Some(TableAlias { name: Ident::new("l"), columns: vec![ - TableAliasColumnDef::from_column_name("A"), - TableAliasColumnDef::from_column_name("B"), - TableAliasColumnDef::from_column_name("C"), + TableAliasColumnDef::from_name("A"), + TableAliasColumnDef::from_name("B"), + TableAliasColumnDef::from_name("C"), ], }), args: None, @@ -6374,8 +6374,8 @@ fn parse_cte_renamed_columns() { let query = all_dialects().verified_query(sql); assert_eq!( vec![ - TableAliasColumnDef::from_column_name("col1"), - TableAliasColumnDef::from_column_name("col2") + TableAliasColumnDef::from_name("col1"), + TableAliasColumnDef::from_name("col2") ], query .with @@ -6405,7 +6405,7 @@ fn parse_recursive_cte() { value: "nums".to_string(), quote_style: None, }, - columns: vec![TableAliasColumnDef::from_column_name("val")], + columns: vec![TableAliasColumnDef::from_name("val")], }, query: Box::new(cte_query), from: None, @@ -9349,8 +9349,8 @@ fn parse_pivot_table() { quote_style: None }, columns: vec![ - TableAliasColumnDef::from_column_name("c"), - TableAliasColumnDef::from_column_name("d"), + TableAliasColumnDef::from_name("c"), + TableAliasColumnDef::from_name("d"), ], }), } @@ -9412,7 +9412,7 @@ fn parse_unpivot_table() { name: Ident::new("u"), columns: ["product", "quarter", "quantity"] .into_iter() - .map(TableAliasColumnDef::from_column_name) + .map(TableAliasColumnDef::from_name) .collect(), }), }