mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-08-04 06:18:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Support RENAME for wildcard SELECTs (#784)

Browse source
This commit is contained in:
Jeffrey 2023-01-03 02:29:06 +11:00 • committed by GitHub
parent 98403c07b1
commit 17f604f757
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 214 additions and 117 deletions
Download patch file Download diff file Expand all files Collapse all files

8
src/ast/mod.rs
View file

@ -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,

74
src/ast/query.rs
View file

@ -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
/// <ident> AS <alias>
/// ```
#[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, VisitMut))]
@ -374,6 +394,8 @@ pub struct WildcardAdditionalOptions {
pub opt_exclude: Option<ExcludeSelectItem>,
/// `[EXCEPT...]`.
pub opt_except: Option<ExceptSelectItem>,
/// `[RENAME ...]`.
pub opt_rename: Option<RenameSelectItem>,
}
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
/// <col_name> AS <col_alias>
/// | (<col_name> AS <col_alias>, <col_name> AS <col_alias>, ...)
/// ```
#[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
/// <col_name> AS <col_alias>
/// ```
Single(IdentWithAlias),
/// Multiple column names with aliases inside parenthesis.
/// # Syntax
/// ```plaintext
/// (<col_name> AS <col_alias>, <col_name> AS <col_alias>, ...)
/// ```
Multiple(Vec<IdentWithAlias>),
}
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, VisitMut))]
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<Ident>,
}
@ -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)
)?;
}

39
src/parser.rs
View file

@ -4401,6 +4401,14 @@ impl<'a> Parser<'a> {
Ok(values)
}
/// Strictly parse `identifier AS identifier`
pub fn parse_identifier_with_alias(&mut self) -> Result<IdentWithAlias, ParserError> {
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 <separator> and treat the multiple strings as
// a single <literal>."
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 {
@ -6063,10 +6071,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,
})
}
@ -6108,7 +6122,7 @@ impl<'a> Parser<'a> {
)?;
}
[first, idents @ ..] => Some(ExceptSelectItem {
fist_elemnt: first.clone(),
first_element: first.clone(),
additional_elements: idents.to_vec(),
}),
}
@ -6119,6 +6133,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<Option<RenameSelectItem>, 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<OrderByExpr, ParserError> {
let expr = self.parse_expr()?;

63
tests/sqlparser_bigquery.rs
View file

@ -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()

147
tests/sqlparser_snowflake.rs
View file

@ -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!(),
},
_ => 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]);
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!(),
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!(),
};
);
assert_eq!(expected, select.projection[0]);
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!(),
},
_ => unreachable!(),
};
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()
},
);
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"
);
}