Support visit pivot/unpivot node.

alibaba · Jan 16, 2025 · 75ea4f7 · 75ea4f7
1 parent a5a2980
commit 75ea4f7
Show file tree

Hide file tree

Showing 13 changed files with 42 additions and 43 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/SQLAdhocTableSource.java b/core/src/main/java/com/alibaba/druid/sql/ast/SQLAdhocTableSource.java
@@ -14,9 +14,8 @@ public SQLAdhocTableSource(SQLCreateTableStatement definition) {
     @Override
     protected void accept0(SQLASTVisitor v) {
         if (v.visit(this)) {
-            if (definition != null) {
-                definition.accept(v);
-            }
+            acceptChild(v, definition);
+            super.accept0(v);
         }
         v.endVisit(this);
     }

diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLExprTableSource.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLExprTableSource.java
@@ -258,13 +258,8 @@ public void addPartition(SQLName partition) {
     @Override
     protected void accept0(SQLASTVisitor visitor) {
         if (visitor.visit(this)) {
-            if (expr != null) {
-                expr.accept(visitor);
-            }
-
-            if (sampling != null) {
-                sampling.accept(visitor);
-            }
+            acceptChild(visitor, expr);
+            acceptChild(visitor, sampling);
         }
         visitor.endVisit(this);
     }

diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLGeneratedTableSource.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLGeneratedTableSource.java
@@ -29,6 +29,7 @@ protected void accept0(SQLASTVisitor v) {
             acceptChild(v, methodName);
             acceptChild(v, columns);
             acceptChild(v, items);
+            super.accept0(v);
         }
         v.endVisit(this);
     }

diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLJoinTableSource.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLJoinTableSource.java
@@ -55,28 +55,12 @@ public SQLJoinTableSource(SQLTableSource left, JoinType joinType, SQLTableSource
 
     protected void accept0(SQLASTVisitor visitor) {
         if (visitor.visit(this)) {
-            if (left != null) {
-                left.accept(visitor);
-            }
-
-            if (right != null) {
-                right.accept(visitor);
-            }
-
-            if (condition != null) {
-                condition.accept(visitor);
-            }
-
-            for (int i = 0; i < using.size(); i++) {
-                SQLExpr item = using.get(i);
-                if (item != null) {
-                    item.accept(visitor);
-                }
-            }
-
-            if (udj != null) {
-                udj.accept(visitor);
-            }
+            acceptChild(visitor, left);
+            acceptChild(visitor, right);
+            acceptChild(visitor, condition);
+            acceptChild(visitor, using);
+            acceptChild(visitor, udj);
+            super.accept0(visitor);
         }
 
         visitor.endVisit(this);

diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLLateralViewTableSource.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLLateralViewTableSource.java
@@ -43,6 +43,7 @@ protected void accept0(SQLASTVisitor visitor) {
             acceptChild(visitor, tableSource);
             acceptChild(visitor, method);
             acceptChild(visitor, columns);
+            super.accept0(visitor);
         }
         visitor.endVisit(this);
     }

diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSubqueryTableSource.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSubqueryTableSource.java
@@ -64,12 +64,8 @@ public void setSelect(SQLSelect x) {
     @Override
     protected void accept0(SQLASTVisitor visitor) {
         if (visitor.visit(this)) {
-            if (select != null) {
-                select.accept(visitor);
-            }
-            if (pivot != null) {
-                pivot.accept(visitor);
-            }
+            acceptChild(visitor, this.select);
+            super.accept0(visitor);
         }
         visitor.endVisit(this);
     }

diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLTableSourceImpl.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLTableSourceImpl.java
@@ -19,6 +19,7 @@
 import com.alibaba.druid.sql.ast.*;
 import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
 import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 import com.alibaba.druid.util.FnvHash;
 
 import java.util.ArrayList;
@@ -248,4 +249,14 @@ public void setUnpivot(SQLUnpivot x) {
         }
         this.unpivot = x;
     }
+
+    @Override
+    protected void accept0(SQLASTVisitor visitor) {
+        if (visitor.visit(this)) {
+            acceptChild(visitor, this.flashback);
+            acceptChild(visitor, this.pivot);
+            acceptChild(visitor, this.unpivot);
+        }
+        visitor.endVisit(this);
+    }
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUnionQueryTableSource.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUnionQueryTableSource.java
@@ -45,9 +45,9 @@ public SQLUnionQueryTableSource(SQLUnionQuery union) {
     @Override
     protected void accept0(SQLASTVisitor visitor) {
         if (visitor.visit(this)) {
-            if (union != null) {
-                union.accept(visitor);
-            }
+            acceptChild(visitor, union);
+            acceptChild(visitor, columns);
+            super.accept0(visitor);
         }
         visitor.endVisit(this);
     }

diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUnnestTableSource.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUnnestTableSource.java
@@ -24,6 +24,7 @@ protected void accept0(SQLASTVisitor v) {
             acceptChild(v, items);
             acceptChild(v, columns);
             acceptChild(v, offset);
+            super.accept0(v);
         }
         v.endVisit(this);
     }

diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLValuesTableSource.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLValuesTableSource.java
@@ -68,6 +68,7 @@ public void accept0(SQLASTVisitor visitor) {
         if (visitor.visit(this)) {
             acceptChild(visitor, values);
             acceptChild(visitor, columns);
+            super.accept0(visitor);
         }
         visitor.endVisit(this);
     }

diff --git a/...c/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlUpdateTableSource.java b/...c/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlUpdateTableSource.java
@@ -32,9 +32,8 @@ public MySqlUpdateTableSource(MySqlUpdateStatement update) {
     protected void accept0(SQLASTVisitor visitor) {
         if (visitor instanceof MySqlASTVisitor) {
             accept0((MySqlASTVisitor) visitor);
-        } else {
-            throw new IllegalArgumentException("not support visitor type : " + visitor.getClass().getName());
         }
+        super.accept0(visitor);
     }
 
     public void accept0(MySqlASTVisitor visitor) {

diff --git a/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitor.java b/core/src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitor.java
@@ -120,6 +120,11 @@ default boolean visit(SQLAllColumnExpr x) {
         return true;
     }
 
+    default boolean visit(SQLTableSourceImpl x) {
+        return true;
+    }
+
+    default void endVisit(SQLTableSourceImpl x) {}
     default boolean visit(SQLBetweenExpr x) {
         return true;
     }

diff --git a/core/src/test/resources/bvt/parser/bigquery/0.txt b/core/src/test/resources/bvt/parser/bigquery/0.txt
@@ -1,3 +1,9 @@
+select * from a unpivot(ct for type in ('1'))
+--------------------
+SELECT *
+FROM a
+UNPIVOT (ct FOR type IN ('1'))
+------------------------------------------------------------------------------------------------------------------------
 select timestamp_add(a, interval 1 millisecond) from b
 --------------------
 SELECT timestamp_add(a, INTERVAL 1 MILLISECOND)