From 4efdf7a34a37107a341d9c33f4693c6447afe181 Mon Sep 17 00:00:00 2001 From: wenshao Date: Thu, 19 Dec 2024 09:37:38 +0800 Subject: [PATCH] improve sql parser bigquery support --- .../sql/ast/statement/SQLUnionQuery.java | 34 +++++++++++++++---- .../druid/sql/parser/SQLSelectParser.java | 4 +++ .../sql/visitor/SQLASTOutputVisitor.java | 5 +++ .../test/resources/bvt/parser/bigquery/9.txt | 28 +++++++++++++++ pom.xml | 4 +-- 5 files changed, 66 insertions(+), 9 deletions(-) diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUnionQuery.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUnionQuery.java index 1b2de383bb..b6e2f1a9e5 100644 --- a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUnionQuery.java +++ b/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 relations = new ArrayList(); 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 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); diff --git a/core/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java b/core/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java index 0e166a8788..47e4a27b7d 100644 --- a/core/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java +++ b/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)); } diff --git a/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java b/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java index 97dccba5bc..9675077009 100644 --- a/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java +++ b/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 relations = x.getRelations(); diff --git a/core/src/test/resources/bvt/parser/bigquery/9.txt b/core/src/test/resources/bvt/parser/bigquery/9.txt index cf33b7371f..09df262324 100644 --- a/core/src/test/resources/bvt/parser/bigquery/9.txt +++ b/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 diff --git a/pom.xml b/pom.xml index eb7139d2dd..d9e28c92e3 100644 --- a/pom.xml +++ b/pom.xml @@ -6,8 +6,8 @@ com.alibaba druid-parent - 1.2.25-SNAPSHOT - ${project.artifactId} + 1.2.25_preview39 + druid-parent A JDBC datasource implementation. pom https://github.com/alibaba/druid