Support UNION DISTINCT with subquery (#31)

Author: git-hulk

parser/ast.go

@@ -3008,6 +3008,7 @@ type SelectQuery struct {
 	Limit         *LimitExpr
 	Settings      *SettingsExprList
 	UnionAll      *SelectQuery
+	UnionDistinct *SelectQuery
 	Except        *SelectQuery
 }
 
@@ -3090,6 +3091,10 @@ func (s *SelectQuery) String(level int) string { // nolint: funlen
 		builder.WriteString(NewLine(level))
 		builder.WriteString(" UNION ALL ")
 		builder.WriteString(s.UnionAll.String(level))
+	} else if s.UnionDistinct != nil {
+		builder.WriteString(NewLine(level))
+		builder.WriteString(" UNION DISTINCT ")
+		builder.WriteString(s.UnionDistinct.String(level))
 	} else if s.Except != nil {
 		builder.WriteString(NewLine(level))
 		builder.WriteString(" EXCEPT ")

parser/parser_query.go

@@ -671,20 +671,28 @@ func (p *Parser) parseSelectQuery(_ Pos) (*SelectQuery, error) {
 	}
 	switch {
 	case p.tryConsumeKeyword(KeywordUnion) != nil:
-		if err := p.consumeKeyword(KeywordAll); err != nil {
-			return nil, err
-		}
-		unionAllExpr, err := p.parseSelectStatement(p.Pos())
-		if err != nil {
-			return nil, err
+		switch {
+		case p.tryConsumeKeyword(KeywordAll) != nil:
+			unionAllExpr, err := p.parseSelectStatement(p.Pos())
+			if err != nil {
+				return nil, err
+			}
+			selectExpr.UnionAll = unionAllExpr
+		case p.tryConsumeKeyword(KeywordDistinct) != nil:
+			unionDistinctExpr, err := p.parseSelectStatement(p.Pos())
+			if err != nil {
+				return nil, err
+			}
+			selectExpr.UnionDistinct = unionDistinctExpr
+		default:
+			return nil, fmt.Errorf("expected ALL or DISTINCT, got %s", p.lastTokenKind())
 		}
-		selectExpr.UnionAll = unionAllExpr
 	case p.tryConsumeKeyword(KeywordExcept) != nil:
 		exceptExpr, err := p.parseSelectStatement(p.Pos())
 		if err != nil {
 			return nil, err
 		}
-		selectExpr.UnionAll = exceptExpr
+		selectExpr.Except = exceptExpr
 	}
 	if hasParen {
 		if _, err := p.consumeTokenKind(")"); err != nil {

parser/testdata/ddl/output/create_live_view_basic.sql.golden.json

@@ -118,6 +118,7 @@
         "Limit": null,
         "Settings": null,
         "UnionAll": null,
+        "UnionDistinct": null,
         "Except": null
       }
     }

parser/testdata/ddl/output/create_materialized_view_with_empty_table_schema.sql.golden.json

@@ -423,6 +423,7 @@
                 "Limit": null,
                 "Settings": null,
                 "UnionAll": null,
+                "UnionDistinct": null,
                 "Except": null
               },
               "AliasPos": 441,
@@ -466,6 +467,7 @@
         "Limit": null,
         "Settings": null,
         "UnionAll": null,
+        "UnionDistinct": null,
         "Except": null
       }
     },

parser/testdata/ddl/output/create_view_basic.sql.golden.json

@@ -127,6 +127,7 @@
         "Limit": null,
         "Settings": null,
         "UnionAll": null,
+        "UnionDistinct": null,
         "Except": null
       }
     }

parser/testdata/ddl/output/create_view_on_cluster_with_uuid.sql.golden.json

@@ -88,6 +88,7 @@
         "Limit": null,
         "Settings": null,
         "UnionAll": null,
+        "UnionDistinct": null,
         "Except": null
       }
     }

parser/testdata/dml/output/insert_with_select.sql.golden.json

@@ -88,6 +88,7 @@
       "Limit": null,
       "Settings": null,
       "UnionAll": null,
+      "UnionDistinct": null,
       "Except": null
     }
   }

parser/testdata/query/format/select_with_union_distinct.sql

@@ -0,0 +1,14 @@
+-- Origin SQL:
+SELECT replica_name FROM system.ha_replicas UNION DISTINCT SELECT replica_name FROM system.ha_unique_replicas format JSON
+
+-- Format SQL:
+
+SELECT 
+  replica_name
+FROM
+  system.ha_replicas
+ UNION DISTINCT 
+SELECT 
+  replica_name
+FROM
+  system.ha_unique_replicas;

parser/testdata/query/output/select_simple.sql.golden.json

@@ -375,6 +375,7 @@
     "Limit": null,
     "Settings": null,
     "UnionAll": null,
+    "UnionDistinct": null,
     "Except": null
   }
 ]

parser/testdata/query/output/select_simple_with_bracket.sql.golden.json

@@ -155,6 +155,7 @@
     "Limit": null,
     "Settings": null,
     "UnionAll": null,
+    "UnionDistinct": null,
     "Except": null
   }
 ]