From fa6bd01b19fabc9087d3b306e7c92f12cb8d368d Mon Sep 17 00:00:00 2001 From: Augusto Fotino <75763288+AugustoFKL@users.noreply.github.com> Date: Wed, 30 Nov 2022 14:29:43 -0300 Subject: [PATCH] Support EXCLUDE support for snowflake and generic dialect (#721) The exclude clause can be used after a possibly qualified on SELECT --- src/ast/mod.rs | 6 ++-- src/ast/query.rs | 60 +++++++++++++++++++++++++++++--- src/keywords.rs | 1 + src/parser.rs | 42 +++++++++++++++++++++-- tests/sqlparser_common.rs | 14 ++++---- tests/sqlparser_postgres.rs | 2 +- tests/sqlparser_snowflake.rs | 66 +++++++++++++++++++++++++++++++++--- 7 files changed, 170 insertions(+), 21 deletions(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 226dc8b43..fbf318921 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -31,9 +31,9 @@ pub use self::ddl::{ }; pub use self::operator::{BinaryOperator, UnaryOperator}; pub use self::query::{ - Cte, Fetch, Join, JoinConstraint, JoinOperator, LateralView, LockType, Offset, OffsetRows, - OrderByExpr, Query, Select, SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, - TableAlias, TableFactor, TableWithJoins, Top, Values, With, + Cte, ExcludeSelectItem, Fetch, Join, JoinConstraint, JoinOperator, LateralView, LockType, + Offset, OffsetRows, OrderByExpr, Query, Select, SelectInto, SelectItem, SetExpr, SetOperator, + SetQuantifier, TableAlias, TableFactor, TableWithJoins, Top, Values, With, }; pub use self::value::{escape_quoted_string, DateTimeField, TrimWhereField, Value}; diff --git a/src/ast/query.rs b/src/ast/query.rs index 34b4ed7f5..0c07ce640 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -321,9 +321,49 @@ pub enum SelectItem { /// An expression, followed by `[ AS ] alias` ExprWithAlias { expr: Expr, alias: Ident }, /// `alias.*` or even `schema.table.*` - QualifiedWildcard(ObjectName), + QualifiedWildcard(ObjectName, Option), /// An unqualified `*` - Wildcard, + Wildcard(Option), +} + +/// Snowflake `EXCLUDE` information. +/// +/// # Syntax +/// ```plaintext +/// +/// | (, , ...) +/// ``` +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub enum ExcludeSelectItem { + /// Single column name without parenthesis. + /// + /// # Syntax + /// ```plaintext + /// + /// ``` + Single(Ident), + /// Multiple column names inside parenthesis. + /// # Syntax + /// ```plaintext + /// (, , ...) + /// ``` + Multiple(Vec), +} + +impl fmt::Display for ExcludeSelectItem { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "EXCLUDE")?; + match self { + Self::Single(column) => { + write!(f, " {column}")?; + } + Self::Multiple(columns) => { + write!(f, " ({})", display_comma_separated(columns))?; + } + } + Ok(()) + } } impl fmt::Display for SelectItem { @@ -331,8 +371,20 @@ impl fmt::Display for SelectItem { match &self { SelectItem::UnnamedExpr(expr) => write!(f, "{}", expr), SelectItem::ExprWithAlias { expr, alias } => write!(f, "{} AS {}", expr, alias), - SelectItem::QualifiedWildcard(prefix) => write!(f, "{}.*", prefix), - SelectItem::Wildcard => write!(f, "*"), + SelectItem::QualifiedWildcard(prefix, opt_exclude) => { + write!(f, "{}.*", prefix)?; + if let Some(exclude) = opt_exclude { + write!(f, " {exclude}")?; + } + Ok(()) + } + SelectItem::Wildcard(opt_exclude) => { + write!(f, "*")?; + if let Some(exclude) = opt_exclude { + write!(f, " {exclude}")?; + } + Ok(()) + } } } } diff --git a/src/keywords.rs b/src/keywords.rs index 79d456f15..1cbe9f289 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -227,6 +227,7 @@ define_keywords!( EVENT, EVERY, EXCEPT, + EXCLUDE, EXEC, EXECUTE, EXISTS, diff --git a/src/parser.rs b/src/parser.rs index deb6d6bf9..741135a6e 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -5423,11 +5423,49 @@ impl<'a> Parser<'a> { None => SelectItem::UnnamedExpr(expr), }) } - WildcardExpr::QualifiedWildcard(prefix) => Ok(SelectItem::QualifiedWildcard(prefix)), - WildcardExpr::Wildcard => Ok(SelectItem::Wildcard), + WildcardExpr::QualifiedWildcard(prefix) => { + let opt_exclude = if dialect_of!(self is GenericDialect | SnowflakeDialect) { + self.parse_optional_select_item_exclude()? + } else { + None + }; + + Ok(SelectItem::QualifiedWildcard(prefix, opt_exclude)) + } + WildcardExpr::Wildcard => { + let opt_exclude = if dialect_of!(self is GenericDialect | SnowflakeDialect) { + self.parse_optional_select_item_exclude()? + } else { + None + }; + + Ok(SelectItem::Wildcard(opt_exclude)) + } } } + /// Parse an [`Exclude`](ExcludeSelectItem) information for wildcard select items. + /// + /// If it is not possible to parse it, will return an option. + pub fn parse_optional_select_item_exclude( + &mut self, + ) -> Result, ParserError> { + let opt_exclude = if self.parse_keyword(Keyword::EXCLUDE) { + if self.consume_token(&Token::LParen) { + let columns = self.parse_comma_separated(|parser| parser.parse_identifier())?; + self.expect_token(&Token::RParen)?; + Some(ExcludeSelectItem::Multiple(columns)) + } else { + let column = self.parse_identifier()?; + Some(ExcludeSelectItem::Single(column)) + } + } else { + None + }; + + Ok(opt_exclude) + } + /// 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_common.rs b/tests/sqlparser_common.rs index 6790bbff6..7041a0609 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -578,22 +578,22 @@ fn parse_select_into() { fn parse_select_wildcard() { let sql = "SELECT * FROM foo"; let select = verified_only_select(sql); - assert_eq!(&SelectItem::Wildcard, only(&select.projection)); + assert_eq!(&SelectItem::Wildcard(None), only(&select.projection)); let sql = "SELECT foo.* FROM foo"; let select = verified_only_select(sql); assert_eq!( - &SelectItem::QualifiedWildcard(ObjectName(vec![Ident::new("foo")])), + &SelectItem::QualifiedWildcard(ObjectName(vec![Ident::new("foo")]), None), only(&select.projection) ); let sql = "SELECT myschema.mytable.* FROM myschema.mytable"; let select = verified_only_select(sql); assert_eq!( - &SelectItem::QualifiedWildcard(ObjectName(vec![ - Ident::new("myschema"), - Ident::new("mytable"), - ])), + &SelectItem::QualifiedWildcard( + ObjectName(vec![Ident::new("myschema"), Ident::new("mytable"),]), + None + ), only(&select.projection) ); @@ -5432,7 +5432,7 @@ fn parse_merge() { body: Box::new(SetExpr::Select(Box::new(Select { distinct: false, top: None, - projection: vec![SelectItem::Wildcard], + projection: vec![SelectItem::Wildcard(None)], into: None, from: vec![TableWithJoins { relation: TableFactor::Table { diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index d7e07a762..2a871a320 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -1229,7 +1229,7 @@ fn parse_pg_returning() { pg_and_generic().verified_stmt("DELETE FROM tasks WHERE status = 'DONE' RETURNING *"); match stmt { Statement::Delete { returning, .. } => { - assert_eq!(Some(vec![SelectItem::Wildcard,]), returning); + assert_eq!(Some(vec![SelectItem::Wildcard(None),]), returning); } _ => unreachable!(), }; diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index a201c5db7..127528b82 100644 --- a/tests/sqlparser_snowflake.rs +++ b/tests/sqlparser_snowflake.rs @@ -14,14 +14,14 @@ //! Test SQL syntax specific to Snowflake. The parser based on the //! generic dialect is also tested (on the inputs it can handle). -#[macro_use] -mod test_utils; -use test_utils::*; - use sqlparser::ast::*; use sqlparser::dialect::{GenericDialect, SnowflakeDialect}; use sqlparser::parser::ParserError; use sqlparser::tokenizer::*; +use test_utils::*; + +#[macro_use] +mod test_utils; #[test] fn test_snowflake_create_table() { @@ -364,3 +364,61 @@ fn snowflake_and_generic() -> TestedDialects { dialects: vec![Box::new(SnowflakeDialect {}), Box::new(GenericDialect {})], } } + +#[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(Some(exclude)) => { + assert_eq!( + *exclude, + ExcludeSelectItem::Multiple(vec![Ident::new("col_a")]) + ) + } + _ => unreachable!(), + }, + _ => unreachable!(), + }, + _ => 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(_, 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(Some(exclude)) => { + assert_eq!( + *exclude, + ExcludeSelectItem::Multiple(vec![ + Ident::new("department_id"), + Ident::new("employee_id") + ]) + ) + } + _ => unreachable!(), + }, + _ => unreachable!(), + }, + _ => unreachable!(), + }; +}