improved big query sql parser

alibaba · Nov 6, 2024 · 003b1bd · 003b1bd
1 parent 4467a00
commit 003b1bd
Show file tree

Hide file tree

Showing 10 changed files with 87 additions and 27 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/ast/BigQueryCharExpr.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/ast/BigQueryCharExpr.java
@@ -3,13 +3,13 @@
 import com.alibaba.druid.sql.ast.SQLExpr;
 import com.alibaba.druid.sql.ast.SQLObject;
 import com.alibaba.druid.sql.ast.expr.SQLCharExpr;
-import com.alibaba.druid.sql.visitor.SQLASTOutputVisitor;
+import com.alibaba.druid.sql.dialect.bigquery.visitor.BigQueryVisitor;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 import java.util.Collections;
 import java.util.List;
 
-public class BigQueryCharExpr extends SQLCharExpr implements SQLExpr {
+public class BigQueryCharExpr extends SQLCharExpr implements SQLExpr, BigQueryObject {
     private String prefix;
     private boolean space;
     private boolean isAlias;
@@ -25,6 +25,10 @@ public void setPrefix(String prefix) {
         this.prefix = prefix;
     }
 
+    public String getPrefix() {
+        return prefix;
+    }
+
     public boolean isSpace() {
         return space;
     }
@@ -52,24 +56,19 @@ public BigQueryCharExpr(String text, String prefix, boolean space, boolean isAli
         this.isAlias = isAlias;
     }
 
+    public void accept0(BigQueryVisitor v) {
+        v.visit(this);
+        v.endVisit(this);
+    }
+
     @Override
-    protected void accept0(SQLASTVisitor v) {
-        if (v instanceof SQLASTOutputVisitor) {
-            SQLASTOutputVisitor visitor = (SQLASTOutputVisitor) v;
-            if (hasPrefix()) {
-                visitor.print(prefix);
-            }
-            if (isSpace()) {
-                visitor.print(" ");
-            }
-            if (!isAlias) {
-                visitor.print("'");
-            }
-            visitor.print(text);
-            if (!isAlias) {
-                visitor.print("'");
-            }
+    public void accept0(SQLASTVisitor v) {
+        if (v instanceof BigQueryVisitor) {
+            this.accept0((BigQueryVisitor) v);
+            return;
         }
+        v.visit(this);
+        v.endVisit(this);
     }
 
     @Override
@@ -86,5 +85,4 @@ public BigQueryCharExpr clone() {
         clone.setAlias(this.isAlias);
         return clone;
     }
-
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryLexer.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryLexer.java
@@ -104,6 +104,9 @@ protected Keywords loadKeywords() {
         map.put("WITH", Token.WITH);
         map.put("WHILE", Token.WHILE);
         map.put("VIEW", Token.VIEW);
+        map.put("TRUNCATE", Token.TRUNCATE);
+        map.put("BEGIN", Token.BEGIN);
+        map.put("END", Token.END);
 
         return new Keywords(map);
     }
@@ -145,7 +148,7 @@ protected void scanAlias() {
     @Override
     protected void initDialectFeature() {
         super.initDialectFeature();
-        this.dialectFeature.configFeature(SQLDateExpr, ScanSubAsIdentifier, GroupByAll, InRestSpecificOperation);
+        this.dialectFeature.configFeature(SQLDateExpr, GroupByAll, ScanSubAsIdentifier, InRestSpecificOperation);
     }
 
     @Override

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
@@ -21,6 +21,9 @@ protected BigQuerySelectQueryBlock createSelectQueryBlock() {
     }
 
     protected void querySelectListBefore(SQLSelectQueryBlock x) {
+        if (lexer.nextIf(Token.DISTINCT)) {
+            x.setDistinct();
+        }
         if (lexer.nextIf(Token.AS)) {
             acceptIdentifier("STRUCT");
             ((BigQuerySelectQueryBlock) x).setAsStruct(true);

diff --git a/.../src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryStatementParser.java b/.../src/main/java/com/alibaba/druid/sql/dialect/bigquery/parser/BigQueryStatementParser.java
@@ -117,6 +117,10 @@ public boolean parseStatementListDialect(List<SQLStatement> statementList) {
             statementList.add(parseAssert());
             return true;
         }
+        if (lexer.token() == Token.BEGIN) {
+            statementList.add(parseBlock());
+            return true;
+        }
         return false;
     }
 

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
@@ -363,4 +363,21 @@ public boolean visit(SQLAggregateExpr x) {
         this.parameterized = parameterized;
         return false;
     }
+
+    public boolean visit(BigQueryCharExpr x) {
+        if (x.hasPrefix()) {
+            print0(x.getPrefix());
+        }
+        if (x.isSpace()) {
+            print0(" ");
+        }
+        if (!x.isAlias()) {
+            print("'");
+        }
+        print(x.getText());
+        if (!x.isAlias()) {
+            print("'");
+        }
+        return false;
+    }
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/visitor/BigQueryVisitor.java b/core/src/main/java/com/alibaba/druid/sql/dialect/bigquery/visitor/BigQueryVisitor.java
@@ -38,4 +38,11 @@ default boolean visit(BigQueryDateTimeExpr x) {
 
     default void endVisit(BigQueryDateTimeExpr x) {
     }
+
+    default boolean visit(BigQueryCharExpr x) {
+        return true;
+    }
+
+    default void endVisit(BigQueryCharExpr x) {
+    }
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/parser/Lexer.java b/core/src/main/java/com/alibaba/druid/sql/parser/Lexer.java
@@ -2679,7 +2679,9 @@ public void scanIdentifier() {
                     hasLeftBrace = false;
                     bufPos++;
                     continue;
-                } else if (ch == '-' && dialectFeatureEnabled(ScanSubAsIdentifier)) {
+                } else if (ch == '-'
+                        && isIdentifierChar(charAt(pos + 1))
+                        && dialectFeatureEnabled(ScanSubAsIdentifier)) {
                     bufPos++;
                     continue;
                 }

diff --git a/core/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java b/core/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java
@@ -3550,18 +3550,16 @@ protected SQLDropProcedureStatement parseDropProcedure(boolean acceptDrop) {
 
     public SQLStatement parseTruncate() {
         accept(Token.TRUNCATE);
-        if (lexer.token == Token.TABLE) {
-            lexer.nextToken();
+        if (!lexer.nextIf(Token.TABLE)) {
+            lexer.nextIfIdentifier("TABLE");
         }
         SQLTruncateStatement stmt = new SQLTruncateStatement(getDbType());
 
-        if (lexer.token == Token.ONLY) {
-            lexer.nextToken();
+        if (lexer.nextIf(Token.ONLY)) {
             stmt.setOnly(true);
         }
 
-        if (lexer.token == Token.IF) {
-            lexer.nextToken();
+        if (lexer.nextIf(Token.IF)) {
             accept(Token.EXISTS);
             stmt.setIfExists(true);
         }

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
@@ -4256,6 +4256,7 @@ public boolean visit(SQLUnaryExpr x) {
             case BINARY:
             case Prior:
             case ConnectByRoot:
+            case NOT:
                 print(' ');
                 if (operator != SQLUnaryOperator.Prior && expr instanceof SQLBinaryOpExpr && !((SQLBinaryOpExpr) expr).isParenthesized()) {
                     print('(');

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,30 @@
+/* a
+   b
+   c
+*/
+select 1
+--------------------
+/* a
+   b
+   c
+*/
+SELECT 1
+------------------------------------------------------------------------------------------------------------------------
+
+select * FROM
+##`a-b.c.e`
+`a-b-c-e`
+--------------------
+SELECT *
+FROM  ##`a-b.c.e`
+`a-b-c-e`
+------------------------------------------------------------------------------------------------------------------------
+SELECT *
+FROM p-- [NEW!]
+--------------------
+SELECT *
+FROM p -- [NEW!]
+------------------------------------------------------------------------------------------------------------------------
 SELECT a.* EXCEPT (f1, f2) from t
 --------------------
 SELECT a.* EXCEPT(f1, f2)