feat: Support explicit DISTINCT option with UNION

Signed-off-by: Alex Qyoun-ae <[email protected]>

Author: MazterQyou

src/ast/mod.rs

@@ -37,8 +37,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, SetOperatorOption,
+    TableAlias, TableFactor, TableWithJoins, Top, Values, With,
 };
 pub use self::value::{escape_single_quote_string, DateTimeField, TrimWhereField, Value};
 

src/ast/query.rs

@@ -78,7 +78,7 @@ pub enum SetExpr {
     /// UNION/EXCEPT/INTERSECT of two queries
     SetOperation {
         op: SetOperator,
-        all: bool,
+        option: Option<SetOperatorOption>,
         left: Box<SetExpr>,
         right: Box<SetExpr>,
     },
@@ -98,10 +98,13 @@ impl fmt::Display for SetExpr {
                 left,
                 right,
                 op,
-                all,
+                option,
             } => {
-                let all_str = if *all { " ALL" } else { "" };
-                write!(f, "{} {}{} {}", left, op, all_str, right)
+                let option_str = option
+                    .as_ref()
+                    .map(|option| format!(" {}", option))
+                    .unwrap_or_else(|| "".to_string());
+                write!(f, "{} {}{} {}", left, op, option_str, right)
             }
         }
     }
@@ -125,6 +128,22 @@ impl fmt::Display for SetOperator {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum SetOperatorOption {
+    All,
+    Distinct,
+}
+
+impl fmt::Display for SetOperatorOption {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.write_str(match self {
+            SetOperatorOption::All => "ALL",
+            SetOperatorOption::Distinct => "DISTINCT",
+        })
+    }
+}
+
 /// A restricted variant of `SELECT` (without CTEs/`ORDER BY`), which may
 /// appear either as the only body item of a `Query`, or as an operand
 /// to a set operation like `UNION`.

src/parser.rs

@@ -3473,10 +3473,11 @@ impl<'a> Parser<'a> {
                 break;
             }
             self.next_token(); // skip past the set operator
+            let option = self.parse_set_operator_option();
             expr = SetExpr::SetOperation {
                 left: Box::new(expr),
                 op: op.unwrap(),
-                all: self.parse_keyword(Keyword::ALL),
+                option,
                 right: Box::new(self.parse_query_body(next_precedence)?),
             };
         }
@@ -3493,6 +3494,16 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn parse_set_operator_option(&mut self) -> Option<SetOperatorOption> {
+        if self.parse_keyword(Keyword::ALL) {
+            return Some(SetOperatorOption::All);
+        }
+        if self.parse_keyword(Keyword::DISTINCT) {
+            return Some(SetOperatorOption::Distinct);
+        }
+        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> {

tests/sqlparser_common.rs

@@ -3658,6 +3658,7 @@ fn parse_union() {
     // 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 2");
     verified_stmt("SELECT 1 EXCEPT SELECT 2");
     verified_stmt("SELECT 1 EXCEPT ALL SELECT 2");
     verified_stmt("SELECT 1 INTERSECT SELECT 2");

tests/sqlparser_postgres.rs

@@ -1272,7 +1272,7 @@ fn parse_array_subquery_expr() {
             with: None,
             body: SetExpr::SetOperation {
                 op: SetOperator::Union,
-                all: false,
+                option: None,
                 left: Box::new(SetExpr::Select(Box::new(Select {
                     distinct: false,
                     top: None,