Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support DISTINCT for SetOperator #689

Merged
merged 11 commits into from Nov 7, 2022
4 changes: 2 additions & 2 deletions src/ast/mod.rs
Expand Up @@ -32,8 +32,8 @@ 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, TableAlias,
TableFactor, TableWithJoins, Top, Values, With,
OrderByExpr, Query, Select, SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier,
TableAlias, TableFactor, TableWithJoins, Top, Values, With,
};
pub use self::value::{DateTimeField, TrimWhereField, Value};

Expand Down
34 changes: 30 additions & 4 deletions src/ast/query.rs
Expand Up @@ -64,6 +64,25 @@ impl fmt::Display for Query {
}
}

/// Options for SetOperator
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would be nice having more description about the structure. Maybe a reference of which big dialects use it...

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum SetQuantifier {
All,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In general it should be possible to round trip structures in sqlparser -- so it should be possible to distinguish between ALL, DISTINCT or nothing was specified in the query

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll add the None to the enum SetOperatorOption, and remove Option<T> for SetOperatorOption.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I fixed this at ac41a8e

Distinct,
None,
}

impl fmt::Display for SetQuantifier {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
SetQuantifier::All => write!(f, "ALL"),
SetQuantifier::Distinct => write!(f, "DISTINCT"),
SetQuantifier::None => write!(f, ""),
}
}
}

/// A node in a tree, representing a "query body" expression, roughly:
/// `SELECT ... [ {UNION|EXCEPT|INTERSECT} SELECT ...]`
#[allow(clippy::large_enum_variant)]
Expand All @@ -78,7 +97,7 @@ pub enum SetExpr {
/// UNION/EXCEPT/INTERSECT of two queries
SetOperation {
op: SetOperator,
all: bool,
set_quantifier: SetQuantifier,
alamb marked this conversation as resolved.
Show resolved Hide resolved
left: Box<SetExpr>,
right: Box<SetExpr>,
},
Expand All @@ -98,10 +117,17 @@ impl fmt::Display for SetExpr {
left,
right,
op,
all,
set_quantifier,
} => {
let all_str = if *all { " ALL" } else { "" };
write!(f, "{} {}{} {}", left, op, all_str, right)
write!(f, "{} {}", left, op)?;
match set_quantifier {
SetQuantifier::All | SetQuantifier::Distinct => {
write!(f, " {}", set_quantifier)?
}
SetQuantifier::None => write!(f, "{}", set_quantifier)?,
}
write!(f, " {}", right)?;
Ok(())
}
}
}
Expand Down
27 changes: 26 additions & 1 deletion src/parser.rs
Expand Up @@ -4173,10 +4173,11 @@ impl<'a> Parser<'a> {
break;
}
self.next_token(); // skip past the set operator
let set_quantifier = self.parse_set_quantifier(&op);
expr = SetExpr::SetOperation {
left: Box::new(expr),
op: op.unwrap(),
all: self.parse_keyword(Keyword::ALL),
set_quantifier,
right: Box::new(self.parse_query_body(next_precedence)?),
};
}
Expand All @@ -4193,6 +4194,30 @@ impl<'a> Parser<'a> {
}
}

pub fn parse_set_quantifier(&mut self, op: &Option<SetOperator>) -> SetQuantifier {
match op {
Some(SetOperator::Union) => {
if self.parse_keyword(Keyword::ALL) {
SetQuantifier::All
} else if self.parse_keyword(Keyword::DISTINCT) {
SetQuantifier::Distinct
} else {
SetQuantifier::None
}
}
Some(SetOperator::Except) | Some(SetOperator::Intersect) => {
if self.parse_keyword(Keyword::ALL) {
SetQuantifier::All
} else if self.parse_keyword(Keyword::DISTINCT) {
SetQuantifier::Distinct
} else {
SetQuantifier::None
}
}
_ => SetQuantifier::None,
}
}

/// Parse a restricted `SELECT` statement (no CTEs / `UNION` / `ORDER BY`),
/// assuming the initial `SELECT` was already consumed
pub fn parse_select(&mut self) -> Result<Select, ParserError> {
Expand Down
6 changes: 5 additions & 1 deletion tests/sqlparser_common.rs
Expand Up @@ -4078,14 +4078,17 @@ fn parse_derived_tables() {
}

#[test]
fn parse_union() {
fn parse_union_except_intersect() {
// TODO: add assertions
verified_stmt("SELECT 1 UNION SELECT 2");
verified_stmt("SELECT 1 UNION ALL SELECT 2");
verified_stmt("SELECT 1 UNION DISTINCT SELECT 1");
verified_stmt("SELECT 1 EXCEPT SELECT 2");
verified_stmt("SELECT 1 EXCEPT ALL SELECT 2");
verified_stmt("SELECT 1 EXCEPT DISTINCT SELECT 1");
verified_stmt("SELECT 1 INTERSECT SELECT 2");
verified_stmt("SELECT 1 INTERSECT ALL SELECT 2");
verified_stmt("SELECT 1 INTERSECT DISTINCT SELECT 1");
verified_stmt("SELECT 1 UNION SELECT 2 UNION SELECT 3");
verified_stmt("SELECT 1 EXCEPT SELECT 2 UNION SELECT 3"); // Union[Except[1,2], 3]
verified_stmt("SELECT 1 INTERSECT (SELECT 2 EXCEPT SELECT 3)");
Expand All @@ -4094,6 +4097,7 @@ fn parse_union() {
verified_stmt("SELECT 1 UNION SELECT 2 INTERSECT SELECT 3"); // Union[1, Intersect[2,3]]
verified_stmt("SELECT foo FROM tab UNION SELECT bar FROM TAB");
verified_stmt("(SELECT * FROM new EXCEPT SELECT * FROM old) UNION ALL (SELECT * FROM old EXCEPT SELECT * FROM new) ORDER BY 1");
verified_stmt("(SELECT * FROM new EXCEPT DISTINCT SELECT * FROM old) UNION DISTINCT (SELECT * FROM old EXCEPT DISTINCT SELECT * FROM new) ORDER BY 1");
}

#[test]
Expand Down
2 changes: 1 addition & 1 deletion tests/sqlparser_postgres.rs
Expand Up @@ -1310,7 +1310,7 @@ fn parse_array_subquery_expr() {
with: None,
body: Box::new(SetExpr::SetOperation {
op: SetOperator::Union,
all: false,
set_quantifier: SetQuantifier::None,
left: Box::new(SetExpr::Select(Box::new(Select {
distinct: false,
top: None,
Expand Down