improve sql parser bigquery support

core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUnionQuery.java

@@ -21,12 +21,10 @@
 import com.alibaba.druid.sql.ast.SQLOrderBy;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
+import java.util.*;
 
 public class SQLUnionQuery extends SQLSelectQueryBase implements SQLDbTypedObject {
+    protected SQLWithSubqueryClause with;
     private List<SQLSelectQuery> relations = new ArrayList<SQLSelectQuery>();
     private SQLUnionOperator operator = SQLUnionOperator.UNION;
     private SQLOrderBy orderBy;
@@ -135,6 +133,10 @@ public void setOrderBy(SQLOrderBy orderBy) {
     @Override
     protected void accept0(SQLASTVisitor visitor) {
         if (visitor.visit(this)) {
+            if (with != null) {
+                with.accept(visitor);
+            }
+
             for (SQLSelectQuery relation : relations) {
                 relation.accept(visitor);
             }
@@ -165,6 +167,9 @@ public SQLUnionQuery clone() {
         SQLUnionQuery x = new SQLUnionQuery();
 
         x.parenthesized = parenthesized;
+        if (with != null) {
+            x.setWith(with.clone());
+        }
 
         for (SQLSelectQuery relation : relations) {
             SQLSelectQuery r = relation.clone();
@@ -271,6 +276,17 @@ public List<SQLSelectQuery> getChildren() {
         return Arrays.asList(left, right);
     }
 
+    public SQLWithSubqueryClause getWith() {
+        return with;
+    }
+
+    public void setWith(SQLWithSubqueryClause x) {
+        if (x != null) {
+            x.setParent(this);
+        }
+        this.with = x;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {
@@ -285,16 +301,19 @@ public boolean equals(Object o) {
         if (parenthesized != that.parenthesized) {
             return false;
         }
-        if (relations != null ? !relations.equals(that.relations) : that.relations != null) {
+        if (!Objects.equals(with, that.with)) {
+            return false;
+        }
+        if (!Objects.equals(relations, that.relations)) {
             return false;
         }
         if (operator != that.operator) {
             return false;
         }
-        if (orderBy != null ? !orderBy.equals(that.orderBy) : that.orderBy != null) {
+        if (!Objects.equals(orderBy, that.orderBy)) {
             return false;
         }
-        if (limit != null ? !limit.equals(that.limit) : that.limit != null) {
+        if (!Objects.equals(limit, that.limit)) {
             return false;
         }
         return dbType == that.dbType;
@@ -303,6 +322,7 @@ public boolean equals(Object o) {
     @Override
     public int hashCode() {
         int result = (parenthesized ? 1 : 0);
+        result = 31 * result + (with != null ? with.hashCode() : 0);
         result = 31 * result + (relations != null ? relations.hashCode() : 0);
         result = 31 * result + (operator != null ? operator.hashCode() : 0);
         result = 31 * result + (orderBy != null ? orderBy.hashCode() : 0);

core/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java

@@ -1196,6 +1196,10 @@ public SQLTableSource parseTableSource(boolean forFrom) {
                 SQLSelectQuery query = queryRest(selectQuery, acceptUnion);
                 if (query instanceof SQLUnionQuery) {
                     tableSource = new SQLUnionQueryTableSource((SQLUnionQuery) query);
+                    SQLWithSubqueryClause with = select.getWithSubQuery();
+                    if (with != null) {
+                        ((SQLUnionQuery) query).setWith(with);
+                    }
                 } else {
                     tableSource = SQLSubqueryTableSource.fixParenthesized(new SQLSubqueryTableSource(select));
                 }

core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java

@@ -4174,6 +4174,11 @@ public boolean visit(SQLNullConstraint x) {
 
     @Override
     public boolean visit(SQLUnionQuery x) {
+        SQLWithSubqueryClause with = x.getWith();
+        if (with != null) {
+            with.accept(this);
+            println();
+        }
         SQLUnionOperator operator = x.getOperator();
 
         List<SQLSelectQuery> relations = x.getRelations();

core/src/test/resources/bvt/parser/bigquery/9.txt

@@ -44,6 +44,34 @@ SET tbl.x1 = oba.x2, tbl.x3 = oba.x4, tbl.x5 = oba.x6, tbl.x7 = oba.x8, tbl.x9 =
 FROM (
 		SELECT y1 AS z1, y2 AS z2
 		FROM (
+			WITH a1 AS (
+					SELECT t1.y3 AS y4, t1.y5 AS y6
+					FROM a.b.c t1
+						LEFT JOIN UNNEST(t1.y7) AS oba
+					WHERE LOWER(t1.y6) = 'completed'
+				),
+				a2 AS (
+					SELECT x1 AS x2, x3 AS x4, x5 AS x6
+					FROM tbl1
+					WHERE rn = 1
+						AND x5 = 'SUCCESS'
+				),
+				all_data AS (
+					SELECT t
+					FROM a.b.c tbl
+						JOIN (
+							SELECT _PARTITIONDATE AS dt, x1 AS y19, x2 AS y20, x3 AS y21, x4 AS y22
+								, x5 AS y23, x6 AS y24, x7 AS y25, x8 AS y26, x9 AS y27
+							FROM a.b.e
+							WHERE x9 = 'ID'
+						) temp3
+						ON tbl.y4 = temp3.y4
+				),
+				diff AS (
+					SELECT *
+					FROM all_data
+					WHERE y2 != y3
+				)
 			SELECT *
 			FROM all_data
 			UNION ALL

pom.xml

@@ -6,8 +6,8 @@
 
 	<groupId>com.alibaba</groupId>
 	<artifactId>druid-parent</artifactId>
-	<version>1.2.25-SNAPSHOT</version>
-	<name>${project.artifactId}</name>
+	<version>1.2.25_preview39</version>
+	<name>druid-parent</name>
 	<description>A JDBC datasource implementation.</description>
 	<packaging>pom</packaging>
 	<url>https://github.com/alibaba/druid</url>