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

feat: add support for except clause on wildcards (#745)

Browse source
This commit is contained in:
Augusto Fotino 2022-12-05 16:03:59 -03:00 • committed by GitHub
parent 5b53df97c4
commit b3688513eb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 232 additions and 40 deletions
Download patch file Download diff file Expand all files Collapse all files

7
src/ast/mod.rs
View file

@ -31,9 +31,10 @@ pub use self::ddl::{
};
pub use self::operator::{BinaryOperator, UnaryOperator};
pub use self::query::{
Cte, ExcludeSelectItem, Fetch, Join, JoinConstraint, JoinOperator, LateralView, LockType,
Offset, OffsetRows, OrderByExpr, Query, Select, SelectInto, SelectItem, SetExpr, SetOperator,
SetQuantifier, Table, TableAlias, TableFactor, TableWithJoins, Top, Values, With,
Cte, ExceptSelectItem, ExcludeSelectItem, Fetch, Join, JoinConstraint, JoinOperator,
LateralView, LockType, Offset, OffsetRows, OrderByExpr, Query, Select, SelectInto, SelectItem,
SetExpr, SetOperator, SetQuantifier, Table, TableAlias, TableFactor, TableWithJoins, Top,
Values, WildcardAdditionalOptions, With,
};
pub use self::value::{escape_quoted_string, DateTimeField, TrimWhereField, Value};

70
src/ast/query.rs
View file

@ -347,9 +347,31 @@ pub enum SelectItem {
/// An expression, followed by `[ AS ] alias`
ExprWithAlias { expr: Expr, alias: Ident },
/// `alias.*` or even `schema.table.*`
QualifiedWildcard(ObjectName, Option<ExcludeSelectItem>),
QualifiedWildcard(ObjectName, WildcardAdditionalOptions),
/// An unqualified `*`
Wildcard(Option<ExcludeSelectItem>),
Wildcard(WildcardAdditionalOptions),
}
/// Additional options for wildcards, e.g. Snowflake `EXCLUDE` and Bigquery `EXCEPT`.
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash, Default)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct WildcardAdditionalOptions {
/// `[EXCLUDE...]`.
pub opt_exclude: Option<ExcludeSelectItem>,
/// `[EXCEPT...]`.
pub opt_except: Option<ExceptSelectItem>,
}
impl fmt::Display for WildcardAdditionalOptions {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
if let Some(exclude) = &self.opt_exclude {
write!(f, " {exclude}")?;
}
if let Some(except) = &self.opt_except {
write!(f, " {except}")?;
}
Ok(())
}
}
/// Snowflake `EXCLUDE` information.
@ -392,23 +414,51 @@ impl fmt::Display for ExcludeSelectItem {
}
}
/// Bigquery `EXCEPT` information, with at least one column.
///
/// # Syntax
/// ```plaintext
/// EXCEPT (<col_name> [, ...])
/// ```
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct ExceptSelectItem {
/// First guaranteed column.
pub fist_elemnt: Ident,
/// Additional columns. This list can be empty.
pub additional_elements: Vec<Ident>,
}
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)?;
} else {
write!(
f,
"({}, {})",
self.fist_elemnt,
display_comma_separated(&self.additional_elements)
)?;
}
Ok(())
}
}
impl fmt::Display for SelectItem {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match &self {
SelectItem::UnnamedExpr(expr) => write!(f, "{}", expr),
SelectItem::ExprWithAlias { expr, alias } => write!(f, "{} AS {}", expr, alias),
SelectItem::QualifiedWildcard(prefix, opt_exclude) => {
SelectItem::QualifiedWildcard(prefix, additional_options) => {
write!(f, "{}.*", prefix)?;
if let Some(exclude) = opt_exclude {
write!(f, " {exclude}")?;
}
write!(f, "{additional_options}")?;
Ok(())
}
SelectItem::Wildcard(opt_exclude) => {
SelectItem::Wildcard(additional_options) => {
write!(f, "*")?;
if let Some(exclude) = opt_exclude {
write!(f, " {exclude}")?;
}
write!(f, "{additional_options}")?;
Ok(())
}
}

75
src/parser.rs
View file

@ -5643,27 +5643,39 @@ impl<'a> Parser<'a> {
None => SelectItem::UnnamedExpr(expr),
})
}
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))
}
WildcardExpr::QualifiedWildcard(prefix) => Ok(SelectItem::QualifiedWildcard(
prefix,
self.parse_wildcard_additional_options()?,
)),
WildcardExpr::Wildcard => Ok(SelectItem::Wildcard(
self.parse_wildcard_additional_options()?,
)),
}
}
/// Parse an [`WildcardAdditionalOptions`](WildcardAdditionalOptions) information for wildcard select items.
///
/// If it is not possible to parse it, will return an option.
pub fn parse_wildcard_additional_options(
&mut self,
) -> Result<WildcardAdditionalOptions, ParserError> {
let opt_exclude = if dialect_of!(self is GenericDialect | SnowflakeDialect) {
self.parse_optional_select_item_exclude()?
} else {
None
};
let opt_except = if dialect_of!(self is GenericDialect | BigQueryDialect) {
self.parse_optional_select_item_except()?
} else {
None
};
Ok(WildcardAdditionalOptions {
opt_exclude,
opt_except,
})
}
/// Parse an [`Exclude`](ExcludeSelectItem) information for wildcard select items.
///
/// If it is not possible to parse it, will return an option.
@ -5686,6 +5698,33 @@ impl<'a> Parser<'a> {
Ok(opt_exclude)
}
/// Parse an [`Except`](ExceptSelectItem) information for wildcard select items.
///
/// If it is not possible to parse it, will return an option.
pub fn parse_optional_select_item_except(
&mut self,
) -> Result<Option<ExceptSelectItem>, ParserError> {
let opt_except = if self.parse_keyword(Keyword::EXCEPT) {
let idents = self.parse_parenthesized_column_list(Mandatory)?;
match &idents[..] {
[] => {
return self.expected(
"at least one column should be parsed by the expect clause",
self.peek_token(),
)?;
}
[first, idents @ ..] => Some(ExceptSelectItem {
fist_elemnt: first.clone(),
additional_elements: idents.to_vec(),
}),
}
} else {
None
};
Ok(opt_except)
}
/// 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()?;

79
tests/sqlparser_bigquery.rs
View file

@ -16,7 +16,7 @@ mod test_utils;
use test_utils::*;
use sqlparser::ast::*;
use sqlparser::dialect::BigQueryDialect;
use sqlparser::dialect::{BigQueryDialect, GenericDialect};
#[test]
fn parse_literal_string() {
@ -235,8 +235,85 @@ 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!(),
};
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!(),
};
match bigquery_and_generic().verified_stmt("SELECT * EXCEPT (col1, col2) FROM _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("col1"),
additional_elements: vec![Ident::new("col2")]
}
)
}
_ => unreachable!(),
},
_ => unreachable!(),
},
_ => unreachable!(),
};
}
fn bigquery() -> TestedDialects {
TestedDialects {
dialects: vec![Box::new(BigQueryDialect {})],
}
}
fn bigquery_and_generic() -> TestedDialects {
TestedDialects {
dialects: vec![Box::new(BigQueryDialect {}), Box::new(GenericDialect {})],
}
}

16
tests/sqlparser_common.rs
View file

@ -581,12 +581,18 @@ fn parse_select_into() {
fn parse_select_wildcard() {
let sql = "SELECT * FROM foo";
let select = verified_only_select(sql);
assert_eq!(&SelectItem::Wildcard(None), only(&select.projection));
assert_eq!(
&SelectItem::Wildcard(WildcardAdditionalOptions::default()),
only(&select.projection)
);
let sql = "SELECT foo.* FROM foo";
let select = verified_only_select(sql);
assert_eq!(
&SelectItem::QualifiedWildcard(ObjectName(vec![Ident::new("foo")]), None),
&SelectItem::QualifiedWildcard(
ObjectName(vec![Ident::new("foo")]),
WildcardAdditionalOptions::default()
),
only(&select.projection)
);
@ -595,7 +601,7 @@ fn parse_select_wildcard() {
assert_eq!(
&SelectItem::QualifiedWildcard(
ObjectName(vec![Ident::new("myschema"), Ident::new("mytable"),]),
None
WildcardAdditionalOptions::default(),
),
only(&select.projection)
);
@ -5588,7 +5594,9 @@ fn parse_merge() {
body: Box::new(SetExpr::Select(Box::new(Select {
distinct: false,
top: None,
projection: vec![SelectItem::Wildcard(None)],
projection: vec![SelectItem::Wildcard(
WildcardAdditionalOptions::default()
)],
into: None,
from: vec![TableWithJoins {
relation: TableFactor::Table {

7
tests/sqlparser_postgres.rs
View file

@ -1277,7 +1277,12 @@ 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(None),]), returning);
assert_eq!(
Some(vec![SelectItem::Wildcard(
WildcardAdditionalOptions::default()
),]),
returning
);
}
_ => unreachable!(),
};

18
tests/sqlparser_snowflake.rs
View file

@ -370,7 +370,10 @@ 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)) => {
SelectItem::Wildcard(WildcardAdditionalOptions {
opt_exclude: Some(exclude),
..
}) => {
assert_eq!(
*exclude,
ExcludeSelectItem::Multiple(vec![Ident::new("col_a")])
@ -388,7 +391,13 @@ fn test_select_wildcard_with_exclude() {
{
Statement::Query(query) => match *query.body {
SetExpr::Select(select) => match &select.projection[0] {
SelectItem::QualifiedWildcard(_, Some(exclude)) => {
SelectItem::QualifiedWildcard(
_,
WildcardAdditionalOptions {
opt_exclude: Some(exclude),
..
},
) => {
assert_eq!(
*exclude,
ExcludeSelectItem::Single(Ident::new("department_id"))
@ -406,7 +415,10 @@ fn test_select_wildcard_with_exclude() {
{
Statement::Query(query) => match *query.body {
SetExpr::Select(select) => match &select.projection[0] {
SelectItem::Wildcard(Some(exclude)) => {
SelectItem::Wildcard(WildcardAdditionalOptions {
opt_exclude: Some(exclude),
..
}) => {
assert_eq!(
*exclude,
ExcludeSelectItem::Multiple(vec![