feat: DataGrip compatibility improvements

.github/workflows/rust.yml

@@ -48,7 +48,7 @@ jobs:
   test:
     strategy:
       matrix:
-        rust: [stable, beta, nightly, nightly-2022-03-22]
+        rust: [stable, beta, nightly]
     runs-on: ubuntu-22.04
     steps:
     - name: Setup Rust

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

@@ -496,7 +496,17 @@ impl<'a> Parser<'a> {
                             }
                         }
 
-                        if self.consume_token(&Token::LParen) {
+                        #[cfg(feature = "std")]
+                        let is_array_agg = dialect_of!(self is PostgreSqlDialect | GenericDialect)
+                            && id_parts.len() == 2
+                            && id_parts[0].value.eq_ignore_ascii_case("pg_catalog")
+                            && id_parts[1].value.eq_ignore_ascii_case("array_agg");
+                        #[cfg(not(feature = "std"))]
+                        let is_array_agg = false;
+
+                        if is_array_agg {
+                            self.parse_array_agg_expr()
+                        } else if self.consume_token(&Token::LParen) {
                             self.prev_token();
                             self.parse_function(ObjectName(id_parts))
                         } else {
@@ -3473,10 +3483,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 +3504,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,