improve sql parser bigquery support

alibaba · Dec 15, 2024 · 96e9edf · 96e9edf
1 parent b0af2a0
commit 96e9edf
Show file tree

Hide file tree

Showing 18 changed files with 432 additions and 60 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLBinaryOpExpr.java b/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLBinaryOpExpr.java
@@ -625,6 +625,11 @@ public SQLDataType computeDataType() {
                         )) {
                             return new SQLDataTypeImpl("BIGING");
                         }
+
+                        if ((leftDataType.isString() && !rightDataType.isString())
+                                || (!leftDataType.isString() && rightDataType.isString())) {
+                            return null;
+                        }
                     }
 
                     return leftDataType;

diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLStructExpr.java b/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLStructExpr.java
@@ -2,13 +2,14 @@
 
 import com.alibaba.druid.sql.ast.SQLExpr;
 import com.alibaba.druid.sql.ast.SQLExprImpl;
+import com.alibaba.druid.sql.ast.SQLReplaceable;
 import com.alibaba.druid.sql.ast.SQLStructDataType;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 import java.util.ArrayList;
 import java.util.List;
 
-public class SQLStructExpr extends SQLExprImpl {
+public class SQLStructExpr extends SQLExprImpl implements SQLReplaceable {
     private SQLStructDataType dataType;
     private final List<SQLAliasedExpr> items = new ArrayList<>();
 
@@ -65,4 +66,16 @@ protected void cloneTo(SQLStructExpr x) {
             x.addItem(item.clone());
         }
     }
+
+    @Override
+    public boolean replace(SQLExpr expr, SQLExpr target) {
+        for (int i = 0; i < items.size(); i++) {
+            if (items.get(i) == expr) {
+                target.setParent(this);
+                items.set(i, (SQLAliasedExpr) target);
+                return true;
+            }
+        }
+        return false;
+    }
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelectQueryBlock.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelectQueryBlock.java
@@ -30,6 +30,9 @@
 import java.util.TreeSet;
 
 public class SQLSelectQueryBlock extends SQLSelectQueryBase implements SQLReplaceable, SQLDbTypedObject {
+    // for bigquery
+    protected SQLWithSubqueryClause with;
+
     protected int distionOption;
     protected final List<SQLSelectItem> selectList = new ArrayList<SQLSelectItem>();
 
@@ -732,73 +735,77 @@ protected void accept0(SQLASTVisitor visitor) {
         visitor.endVisit(this);
     }
 
-    protected void acceptChild(SQLASTVisitor visitor) {
+    protected void acceptChild(SQLASTVisitor v) {
+        if (with != null) {
+            with.accept(v);
+        }
+
         for (int i = 0; i < this.selectList.size(); i++) {
             SQLSelectItem item = this.selectList.get(i);
             if (item != null) {
-                item.accept(visitor);
+                item.accept(v);
             }
         }
 
         if (this.from != null) {
-            this.from.accept(visitor);
+            this.from.accept(v);
         }
 
         if (this.windows != null) {
             for (int i = 0; i < windows.size(); i++) {
                 SQLWindow item = windows.get(i);
-                item.accept(visitor);
+                item.accept(v);
             }
         }
 
         if (this.into != null) {
-            this.into.accept(visitor);
+            this.into.accept(v);
         }
 
         if (this.where != null) {
-            this.where.accept(visitor);
+            this.where.accept(v);
         }
 
         if (this.startWith != null) {
-            this.startWith.accept(visitor);
+            this.startWith.accept(v);
         }
 
         if (this.connectBy != null) {
-            this.connectBy.accept(visitor);
+            this.connectBy.accept(v);
         }
 
         if (this.groupBy != null) {
-            this.groupBy.accept(visitor);
+            this.groupBy.accept(v);
         }
 
         if (this.qualify != null) {
-            this.qualify.accept(visitor);
+            this.qualify.accept(v);
         }
 
         if (this.orderBy != null) {
-            this.orderBy.accept(visitor);
+            this.orderBy.accept(v);
         }
 
         if (this.distributeBy != null) {
             for (int i = 0; i < distributeBy.size(); i++) {
                 SQLSelectOrderByItem item = distributeBy.get(i);
-                item.accept(visitor);
+                item.accept(v);
             }
         }
 
         if (this.sortBy != null) {
             for (int i = 0; i < sortBy.size(); i++) {
                 SQLSelectOrderByItem item = sortBy.get(i);
-                item.accept(visitor);
+                item.accept(v);
             }
         }
 
         if (this.waitTime != null) {
-            this.waitTime.accept(visitor);
+            this.waitTime.accept(v);
         }
 
         if (this.limit != null) {
-            this.limit.accept(visitor);
+            this.limit.accept(v);
         }
     }
 
@@ -1027,6 +1034,10 @@ public void cloneTo(SQLSelectQueryBlock x) {
         x.parenthesized = parenthesized;
         x.distionOption = distionOption;
 
+        if (with != null) {
+            x.setWith(with.clone());
+        }
+
         if (x.selectList.size() > 0) {
             x.selectList.clear();
         }
@@ -1641,7 +1652,17 @@ public boolean clearMapJoinHint() {
                 }
             }
         }
-
         return removeCount > 0;
     }
+
+    public SQLWithSubqueryClause getWith() {
+        return with;
+    }
+
+    public void setWith(SQLWithSubqueryClause x) {
+        if (x != null) {
+            x.setParent(this);
+        }
+        this.with = x;
+    }
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/ast/BigQueryModelExpr.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/ast/BigQueryModelExpr.java
@@ -0,0 +1,56 @@
+package com.alibaba.druid.sql.dialect.bigquery.ast;
+
+import com.alibaba.druid.sql.ast.SQLExpr;
+import com.alibaba.druid.sql.ast.SQLExprImpl;
+import com.alibaba.druid.sql.ast.SQLName;
+import com.alibaba.druid.sql.dialect.bigquery.visitor.BigQueryVisitor;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
+
+public class BigQueryModelExpr extends SQLExprImpl implements BigQueryObject {
+    private SQLName name;
+
+    public SQLName getName() {
+        return name;
+    }
+
+    public void setName(SQLName x) {
+        if (x != null) {
+            x.setParent(this);
+        }
+        this.name = x;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        return false;
+    }
+
+    @Override
+    public int hashCode() {
+        return 0;
+    }
+
+    @Override
+    public void accept0(SQLASTVisitor v) {
+        if (v instanceof BigQueryVisitor) {
+            accept0((BigQueryVisitor) v);
+        }
+    }
+
+    @Override
+    public void accept0(BigQueryVisitor v) {
+        if (v.visit(this)) {
+            acceptChild(v);
+        }
+        v.endVisit(this);
+    }
+
+    @Override
+    public SQLExpr clone() {
+        return null;
+    }
+
+    public void acceptChild(SQLASTVisitor v) {
+        acceptChild(v, name);
+    }
+}
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/ast/BigQueryObject.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/ast/BigQueryObject.java
@@ -11,5 +11,5 @@ default void accept0(SQLASTVisitor v) {
         }
     }
 
-    void accept0(BigQueryVisitor visitor);
+    void accept0(BigQueryVisitor v);
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/ast/BigQuerySelectQueryBlock.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/ast/BigQuerySelectQueryBlock.java
@@ -44,13 +44,15 @@ public void accept0(SQLASTVisitor v) {
     }
 
     @Override
-    public void accept0(BigQueryVisitor visitor) {
-        if (visitor.visit(this)) {
-            if (differentialPrivacy != null) {
-                differentialPrivacy.accept(visitor);
-            }
-            acceptChild(visitor);
+    public void accept0(BigQueryVisitor v) {
+        super.accept0(v);
+    }
+
+    protected void acceptChild(SQLASTVisitor v) {
+        if (differentialPrivacy != null) {
+            differentialPrivacy.accept(v);
         }
+        super.acceptChild(v);
     }
 
     public static class DifferentialPrivacy

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/ast/BigQueryTableExpr.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/ast/BigQueryTableExpr.java
@@ -0,0 +1,56 @@
+package com.alibaba.druid.sql.dialect.bigquery.ast;
+
+import com.alibaba.druid.sql.ast.SQLExpr;
+import com.alibaba.druid.sql.ast.SQLExprImpl;
+import com.alibaba.druid.sql.ast.SQLName;
+import com.alibaba.druid.sql.dialect.bigquery.visitor.BigQueryVisitor;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
+
+public class BigQueryTableExpr extends SQLExprImpl implements BigQueryObject {
+    private SQLName name;
+
+    public SQLName getName() {
+        return name;
+    }
+
+    public void setName(SQLName x) {
+        if (x != null) {
+            x.setParent(this);
+        }
+        this.name = x;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        return false;
+    }
+
+    @Override
+    public int hashCode() {
+        return 0;
+    }
+
+    @Override
+    public void accept0(SQLASTVisitor v) {
+        if (v instanceof BigQueryVisitor) {
+            accept0((BigQueryVisitor) v);
+        }
+    }
+
+    @Override
+    public void accept0(BigQueryVisitor v) {
+        if (v.visit(this)) {
+            acceptChild(v);
+        }
+        v.endVisit(this);
+    }
+
+    @Override
+    public SQLExpr clone() {
+        return null;
+    }
+
+    public void acceptChild(SQLASTVisitor v) {
+        acceptChild(v, name);
+    }
+}
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryExprParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryExprParser.java
@@ -6,6 +6,8 @@
 import com.alibaba.druid.sql.ast.statement.*;
 import com.alibaba.druid.sql.dialect.bigquery.BQ;
 import com.alibaba.druid.sql.dialect.bigquery.ast.BigQueryCharExpr;
+import com.alibaba.druid.sql.dialect.bigquery.ast.BigQueryModelExpr;
+import com.alibaba.druid.sql.dialect.bigquery.ast.BigQueryTableExpr;
 import com.alibaba.druid.sql.parser.*;
 import com.alibaba.druid.util.FnvHash;
 
@@ -266,6 +268,33 @@ public SQLExpr primary() {
                     parseQueryExpr()
             );
         }
+
+        if (lexer.identifierEquals(FnvHash.Constants.MODEL)) {
+            Lexer.SavePoint mark = lexer.markOut();
+            lexer.nextToken();
+            if (lexer.token() == Token.IDENTIFIER) {
+                BigQueryModelExpr model = new BigQueryModelExpr();
+                model.setName(
+                        name()
+                );
+                return model;
+            } else {
+                lexer.reset(mark);
+            }
+        }
+        if (lexer.token() == Token.TABLE) {
+            Lexer.SavePoint mark = lexer.markOut();
+            lexer.nextToken();
+            if (lexer.token() == Token.IDENTIFIER) {
+                BigQueryTableExpr model = new BigQueryTableExpr();
+                model.setName(
+                        name()
+                );
+                return model;
+            } else {
+                lexer.reset(mark);
+            }
+        }
         return super.primary();
     }
 

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQuerySelectParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQuerySelectParser.java
@@ -50,4 +50,13 @@ protected String tableAlias(boolean must) {
         }
         return super.tableAlias(must);
     }
+
+    protected void queryBefore(SQLSelectQueryBlock x) {
+        if (lexer.token() == Token.WITH) {
+            BigQuerySelectQueryBlock queryBlock = (BigQuerySelectQueryBlock) x;
+            queryBlock.setWith(
+                    this.parseWith()
+            );
+        }
+    }
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/visitor/BigQueryOutputVisitor.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/visitor/BigQueryOutputVisitor.java
@@ -471,4 +471,16 @@ public boolean visit(BigQueryCreateModelStatement x) {
         print0(')');
         return false;
     }
+
+    public boolean visit(BigQueryModelExpr x) {
+        print0(ucase ? "MODEL " : "model ");
+        x.getName().accept(this);
+        return false;
+    }
+
+    public boolean visit(BigQueryTableExpr x) {
+        print0(ucase ? "TABLE " : "table ");
+        x.getName().accept(this);
+        return false;
+    }
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -11,5 +11,5 @@ default void accept0(SQLASTVisitor v) { @@
             }
         }
-        void accept0(BigQueryVisitor visitor);
+        void accept0(BigQueryVisitor v);
     }