From 7e9b6797eefc113dfb131a1e8fae1f56c1ff7231 Mon Sep 17 00:00:00 2001 From: Jefffrey <22608443+Jefffrey@users.noreply.github.com> Date: Mon, 2 Jan 2023 20:43:08 +1100 Subject: [PATCH] Support RENAME for wildcard SELECTs --- src/ast/mod.rs | 8 +- src/ast/query.rs | 74 ++++++++++++++++- src/parser.rs | 39 ++++++++- tests/sqlparser_bigquery.rs | 63 +++++---------- tests/sqlparser_snowflake.rs | 149 ++++++++++++++++++++--------------- 5 files changed, 215 insertions(+), 118 deletions(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 45f1c8162..143778fec 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -34,10 +34,10 @@ pub use self::ddl::{ }; pub use self::operator::{BinaryOperator, UnaryOperator}; pub use self::query::{ - Cte, ExceptSelectItem, ExcludeSelectItem, Fetch, Join, JoinConstraint, JoinOperator, - LateralView, LockClause, LockType, NonBlock, Offset, OffsetRows, OrderByExpr, Query, Select, - SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Table, TableAlias, TableFactor, - TableWithJoins, Top, Values, WildcardAdditionalOptions, With, + Cte, ExceptSelectItem, ExcludeSelectItem, Fetch, IdentWithAlias, Join, JoinConstraint, + JoinOperator, LateralView, LockClause, LockType, NonBlock, Offset, OffsetRows, OrderByExpr, + Query, RenameSelectItem, Select, SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, + Table, TableAlias, TableFactor, TableWithJoins, Top, Values, WildcardAdditionalOptions, With, }; pub use self::value::{ escape_quoted_string, DateTimeField, DollarQuotedString, TrimWhereField, Value, diff --git a/src/ast/query.rs b/src/ast/query.rs index 0a01cf33c..9f9aca6b1 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -365,7 +365,27 @@ pub enum SelectItem { Wildcard(WildcardAdditionalOptions), } -/// Additional options for wildcards, e.g. Snowflake `EXCLUDE` and Bigquery `EXCEPT`. +/// Single aliased identifier +/// +/// # Syntax +/// ```plaintext +/// AS +/// ``` +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit))] +pub struct IdentWithAlias { + pub ident: Ident, + pub alias: Ident, +} + +impl fmt::Display for IdentWithAlias { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{} AS {}", self.ident, self.alias) + } +} + +/// Additional options for wildcards, e.g. Snowflake `EXCLUDE`/`RENAME` and Bigquery `EXCEPT`. #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash, Default)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr(feature = "visitor", derive(Visit))] @@ -374,6 +394,8 @@ pub struct WildcardAdditionalOptions { pub opt_exclude: Option, /// `[EXCEPT...]`. pub opt_except: Option, + /// `[RENAME ...]`. + pub opt_rename: Option, } impl fmt::Display for WildcardAdditionalOptions { @@ -384,6 +406,9 @@ impl fmt::Display for WildcardAdditionalOptions { if let Some(except) = &self.opt_except { write!(f, " {except}")?; } + if let Some(rename) = &self.opt_rename { + write!(f, " {rename}")?; + } Ok(()) } } @@ -429,6 +454,47 @@ impl fmt::Display for ExcludeSelectItem { } } +/// Snowflake `RENAME` information. +/// +/// # Syntax +/// ```plaintext +/// AS +/// | ( AS , AS , ...) +/// ``` +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit))] +pub enum RenameSelectItem { + /// Single column name with alias without parenthesis. + /// + /// # Syntax + /// ```plaintext + /// AS + /// ``` + Single(IdentWithAlias), + /// Multiple column names with aliases inside parenthesis. + /// # Syntax + /// ```plaintext + /// ( AS , AS , ...) + /// ``` + Multiple(Vec), +} + +impl fmt::Display for RenameSelectItem { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "RENAME")?; + match self { + Self::Single(column) => { + write!(f, " {column}")?; + } + Self::Multiple(columns) => { + write!(f, " ({})", display_comma_separated(columns))?; + } + } + Ok(()) + } +} + /// Bigquery `EXCEPT` information, with at least one column. /// /// # Syntax @@ -440,7 +506,7 @@ impl fmt::Display for ExcludeSelectItem { #[cfg_attr(feature = "visitor", derive(Visit))] pub struct ExceptSelectItem { /// First guaranteed column. - pub fist_elemnt: Ident, + pub first_element: Ident, /// Additional columns. This list can be empty. pub additional_elements: Vec, } @@ -449,12 +515,12 @@ impl fmt::Display for ExceptSelectItem { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "EXCEPT ")?; if self.additional_elements.is_empty() { - write!(f, "({})", self.fist_elemnt)?; + write!(f, "({})", self.first_element)?; } else { write!( f, "({}, {})", - self.fist_elemnt, + self.first_element, display_comma_separated(&self.additional_elements) )?; } diff --git a/src/parser.rs b/src/parser.rs index 19914253f..cd1a1ac35 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -4401,6 +4401,14 @@ impl<'a> Parser<'a> { Ok(values) } + /// Strictly parse `identifier AS identifier` + pub fn parse_identifier_with_alias(&mut self) -> Result { + let ident = self.parse_identifier()?; + self.expect_keyword(Keyword::AS)?; + let alias = self.parse_identifier()?; + Ok(IdentWithAlias { ident, alias }) + } + /// Parse `AS identifier` (or simply `identifier` if it's not a reserved keyword) /// Some examples with aliases: `SELECT 1 foo`, `SELECT COUNT(*) AS cnt`, /// `SELECT ... FROM t1 foo, t2 bar`, `SELECT ... FROM (...) AS bar` @@ -4432,7 +4440,7 @@ impl<'a> Parser<'a> { // ignore the and treat the multiple strings as // a single ." Token::SingleQuotedString(s) => Ok(Some(Ident::with_quote('\'', s))), - // Support for MySql dialect double qouted string, `AS "HOUR"` for example + // Support for MySql dialect double quoted string, `AS "HOUR"` for example Token::DoubleQuotedString(s) => Ok(Some(Ident::with_quote('\"', s))), _ => { if after_as { @@ -6055,10 +6063,16 @@ impl<'a> Parser<'a> { } else { None }; + let opt_rename = if dialect_of!(self is GenericDialect | SnowflakeDialect) { + self.parse_optional_select_item_rename()? + } else { + None + }; Ok(WildcardAdditionalOptions { opt_exclude, opt_except, + opt_rename, }) } @@ -6100,7 +6114,7 @@ impl<'a> Parser<'a> { )?; } [first, idents @ ..] => Some(ExceptSelectItem { - fist_elemnt: first.clone(), + first_element: first.clone(), additional_elements: idents.to_vec(), }), } @@ -6111,6 +6125,27 @@ impl<'a> Parser<'a> { Ok(opt_except) } + /// Parse a [`Rename`](RenameSelectItem) information for wildcard select items. + pub fn parse_optional_select_item_rename( + &mut self, + ) -> Result, ParserError> { + let opt_rename = if self.parse_keyword(Keyword::RENAME) { + if self.consume_token(&Token::LParen) { + let idents = + self.parse_comma_separated(|parser| parser.parse_identifier_with_alias())?; + self.expect_token(&Token::RParen)?; + Some(RenameSelectItem::Multiple(idents)) + } else { + let ident = self.parse_identifier_with_alias()?; + Some(RenameSelectItem::Single(ident)) + } + } else { + None + }; + + Ok(opt_rename) + } + /// Parse an expression, optionally followed by ASC or DESC (used in ORDER BY) pub fn parse_order_by_expr(&mut self) -> Result { let expr = self.parse_expr()?; diff --git a/tests/sqlparser_bigquery.rs b/tests/sqlparser_bigquery.rs index 4bf26ba81..ada5fec5e 100644 --- a/tests/sqlparser_bigquery.rs +++ b/tests/sqlparser_bigquery.rs @@ -265,51 +265,26 @@ fn parse_array_agg_func() { #[test] fn test_select_wildcard_with_except() { - match bigquery_and_generic().verified_stmt("SELECT * EXCEPT (col_a) FROM data") { - Statement::Query(query) => match *query.body { - SetExpr::Select(select) => match &select.projection[0] { - SelectItem::Wildcard(WildcardAdditionalOptions { - opt_except: Some(except), - .. - }) => { - assert_eq!( - *except, - ExceptSelectItem { - fist_elemnt: Ident::new("col_a"), - additional_elements: vec![] - } - ) - } - _ => unreachable!(), - }, - _ => unreachable!(), - }, - _ => unreachable!(), - }; + let select = bigquery_and_generic().verified_only_select("SELECT * EXCEPT (col_a) FROM data"); + let expected = SelectItem::Wildcard(WildcardAdditionalOptions { + opt_except: Some(ExceptSelectItem { + first_element: Ident::new("col_a"), + additional_elements: vec![], + }), + ..Default::default() + }); + assert_eq!(expected, select.projection[0]); - match bigquery_and_generic() - .verified_stmt("SELECT * EXCEPT (department_id, employee_id) FROM employee_table") - { - Statement::Query(query) => match *query.body { - SetExpr::Select(select) => match &select.projection[0] { - SelectItem::Wildcard(WildcardAdditionalOptions { - opt_except: Some(except), - .. - }) => { - assert_eq!( - *except, - ExceptSelectItem { - fist_elemnt: Ident::new("department_id"), - additional_elements: vec![Ident::new("employee_id")] - } - ) - } - _ => unreachable!(), - }, - _ => unreachable!(), - }, - _ => unreachable!(), - }; + let select = bigquery_and_generic() + .verified_only_select("SELECT * EXCEPT (department_id, employee_id) FROM employee_table"); + let expected = SelectItem::Wildcard(WildcardAdditionalOptions { + opt_except: Some(ExceptSelectItem { + first_element: Ident::new("department_id"), + additional_elements: vec![Ident::new("employee_id")], + }), + ..Default::default() + }); + assert_eq!(expected, select.projection[0]); assert_eq!( bigquery_and_generic() diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index e10141545..78a31e383 100644 --- a/tests/sqlparser_snowflake.rs +++ b/tests/sqlparser_snowflake.rs @@ -389,70 +389,91 @@ fn snowflake_and_generic() -> TestedDialects { #[test] fn test_select_wildcard_with_exclude() { - match snowflake_and_generic().verified_stmt("SELECT * EXCLUDE (col_a) FROM data") { - Statement::Query(query) => match *query.body { - SetExpr::Select(select) => match &select.projection[0] { - SelectItem::Wildcard(WildcardAdditionalOptions { - opt_exclude: Some(exclude), - .. - }) => { - assert_eq!( - *exclude, - ExcludeSelectItem::Multiple(vec![Ident::new("col_a")]) - ) - } - _ => unreachable!(), - }, - _ => unreachable!(), + let select = snowflake_and_generic().verified_only_select("SELECT * EXCLUDE (col_a) FROM data"); + let expected = SelectItem::Wildcard(WildcardAdditionalOptions { + opt_exclude: Some(ExcludeSelectItem::Multiple(vec![Ident::new("col_a")])), + ..Default::default() + }); + assert_eq!(expected, select.projection[0]); + + let select = snowflake_and_generic() + .verified_only_select("SELECT name.* EXCLUDE department_id FROM employee_table"); + let expected = SelectItem::QualifiedWildcard( + ObjectName(vec![Ident::new("name")]), + WildcardAdditionalOptions { + opt_exclude: Some(ExcludeSelectItem::Single(Ident::new("department_id"))), + ..Default::default() }, - _ => unreachable!(), - }; - - match snowflake_and_generic() - .verified_stmt("SELECT name.* EXCLUDE department_id FROM employee_table") - { - Statement::Query(query) => match *query.body { - SetExpr::Select(select) => match &select.projection[0] { - SelectItem::QualifiedWildcard( - _, - WildcardAdditionalOptions { - opt_exclude: Some(exclude), - .. - }, - ) => { - assert_eq!( - *exclude, - ExcludeSelectItem::Single(Ident::new("department_id")) - ) - } - _ => unreachable!(), - }, - _ => unreachable!(), - }, - _ => unreachable!(), - }; - - match snowflake_and_generic() - .verified_stmt("SELECT * EXCLUDE (department_id, employee_id) FROM employee_table") - { - Statement::Query(query) => match *query.body { - SetExpr::Select(select) => match &select.projection[0] { - SelectItem::Wildcard(WildcardAdditionalOptions { - opt_exclude: Some(exclude), - .. - }) => { - assert_eq!( - *exclude, - ExcludeSelectItem::Multiple(vec![ - Ident::new("department_id"), - Ident::new("employee_id") - ]) - ) - } - _ => unreachable!(), - }, - _ => unreachable!(), + ); + assert_eq!(expected, select.projection[0]); + + let select = snowflake_and_generic() + .verified_only_select("SELECT * EXCLUDE (department_id, employee_id) FROM employee_table"); + let expected = SelectItem::Wildcard(WildcardAdditionalOptions { + opt_exclude: Some(ExcludeSelectItem::Multiple(vec![ + Ident::new("department_id"), + Ident::new("employee_id"), + ])), + ..Default::default() + }); + assert_eq!(expected, select.projection[0]); +} + +#[test] +fn test_select_wildcard_with_rename() { + let select = + snowflake_and_generic().verified_only_select("SELECT * RENAME col_a AS col_b FROM data"); + let expected = SelectItem::Wildcard(WildcardAdditionalOptions { + opt_rename: Some(RenameSelectItem::Single(IdentWithAlias { + ident: Ident::new("col_a"), + alias: Ident::new("col_b"), + })), + ..Default::default() + }); + assert_eq!(expected, select.projection[0]); + + let select = snowflake_and_generic().verified_only_select( + "SELECT name.* RENAME (department_id AS new_dep, employee_id AS new_emp) FROM employee_table", + ); + let expected = SelectItem::QualifiedWildcard( + ObjectName(vec![Ident::new("name")]), + WildcardAdditionalOptions { + opt_rename: Some(RenameSelectItem::Multiple(vec![ + IdentWithAlias { + ident: Ident::new("department_id"), + alias: Ident::new("new_dep"), + }, + IdentWithAlias { + ident: Ident::new("employee_id"), + alias: Ident::new("new_emp"), + }, + ])), + ..Default::default() }, - _ => unreachable!(), - }; + ); + assert_eq!(expected, select.projection[0]); +} + +#[test] +fn test_select_wildcard_with_exclude_and_rename() { + let select = snowflake_and_generic() + .verified_only_select("SELECT * EXCLUDE col_z RENAME col_a AS col_b FROM data"); + let expected = SelectItem::Wildcard(WildcardAdditionalOptions { + opt_exclude: Some(ExcludeSelectItem::Single(Ident::new("col_z"))), + opt_rename: Some(RenameSelectItem::Single(IdentWithAlias { + ident: Ident::new("col_a"), + alias: Ident::new("col_b"), + })), + ..Default::default() + }); + assert_eq!(expected, select.projection[0]); + + // rename cannot precede exclude + assert_eq!( + snowflake_and_generic() + .parse_sql_statements("SELECT * RENAME col_a AS col_b EXCLUDE col_z FROM data") + .unwrap_err() + .to_string(), + "sql parser error: Expected end of statement, found: EXCLUDE" + ); }