bigquery sql parser support more control statement

alibaba · Dec 12, 2024 · 8c9aee5 · 8c9aee5
1 parent 54875e5
commit 8c9aee5
Show file tree

Hide file tree

Showing 11 changed files with 132 additions and 34 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLContinueStatement.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLContinueStatement.java
@@ -0,0 +1,12 @@
+package com.alibaba.druid.sql.ast.statement;
+
+import com.alibaba.druid.sql.ast.SQLStatementImpl;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
+
+public class SQLContinueStatement extends SQLStatementImpl {
+    @Override
+    public void accept0(SQLASTVisitor visitor) {
+        visitor.visit(this);
+        visitor.endVisit(this);
+    }
+}
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLLeaveStatement.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLLeaveStatement.java
@@ -0,0 +1,12 @@
+package com.alibaba.druid.sql.ast.statement;
+
+import com.alibaba.druid.sql.ast.SQLStatementImpl;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
+
+public class SQLLeaveStatement extends SQLStatementImpl {
+    @Override
+    public void accept0(SQLASTVisitor visitor) {
+        visitor.visit(this);
+        visitor.endVisit(this);
+    }
+}
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
@@ -102,10 +102,12 @@ protected Keywords loadKeywords() {
         map.put("WINDOW", Token.WINDOW);
         map.put("WITH", Token.WITH);
         map.put("WHILE", Token.WHILE);
+        map.put("LOOP", Token.LOOP);
+        map.put("LEAVE", Token.LEAVE);
+        map.put("CONTINUE", Token.CONTINUE);
         map.put("VIEW", Token.VIEW);
         map.put("TRUNCATE", Token.TRUNCATE);
         map.put("BEGIN", Token.BEGIN);
-        map.put("END", Token.END);
         map.put("TABLE", Token.TABLE);
         map.put("EXCEPTION", Token.EXCEPTION);
         map.put("RAISE", Token.RAISE);

diff --git a/...rc/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlExecuteStatement.java b/...rc/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlExecuteStatement.java
@@ -17,43 +17,42 @@
 
 import com.alibaba.druid.sql.ast.SQLExpr;
 import com.alibaba.druid.sql.ast.SQLName;
-import com.alibaba.druid.sql.ast.SQLObject;
+import com.alibaba.druid.sql.ast.statement.SQLExecuteImmediateStatement;
 import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitor;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 import java.util.ArrayList;
 import java.util.List;
 
-public class MySqlExecuteStatement extends MySqlStatementImpl {
-    private SQLName statementName;
+public class MySqlExecuteStatement extends SQLExecuteImmediateStatement implements MySqlStatement {
     private final List<SQLExpr> parameters = new ArrayList<SQLExpr>();
 
     public SQLName getStatementName() {
-        return statementName;
+        return (SQLName) dynamicSql;
     }
 
-    public void setStatementName(SQLName statementName) {
-        this.statementName = statementName;
+    public void setStatementName(SQLName x) {
+        super.setDynamicSql(x);
     }
 
     public List<SQLExpr> getParameters() {
         return parameters;
     }
 
-    public void accept0(MySqlASTVisitor visitor) {
-        if (visitor.visit(this)) {
-            acceptChild(visitor, statementName);
-            acceptChild(visitor, parameters);
+    public void accept0(MySqlASTVisitor v) {
+        if (v.visit(this)) {
+            super.acceptChild(v);
+            acceptChild(v, parameters);
         }
-        visitor.endVisit(this);
+        v.endVisit(this);
     }
 
     @Override
-    public List<SQLObject> getChildren() {
-        List<SQLObject> children = new ArrayList<SQLObject>();
-        if (statementName != null) {
-            children.add(statementName);
+    public void accept0(SQLASTVisitor v) {
+        if (v instanceof MySqlASTVisitor) {
+            ((MySqlASTVisitor) v).visit(this);
+        } else {
+            super.accept0(v);
         }
-        children.addAll(this.parameters);
-        return children;
     }
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlStatementParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlStatementParser.java
@@ -1077,19 +1077,6 @@ public boolean parseStatementListDialect(List<SQLStatement> statementList) {
             return true;
         }
 
-        if (lexer.identifierEquals("EXECUTE")) {
-            acceptIdentifier("EXECUTE");
-
-            if (lexer.identifierEquals("RESTART") || lexer.identifierEquals("UPDATE")) {
-                MySqlExecuteForAdsStatement stmt = parseExecuteForAds();
-                statementList.add(stmt);
-            } else {
-                MySqlExecuteStatement stmt = parseExecute();
-                statementList.add(stmt);
-            }
-            return true;
-        }
-
         if (lexer.identifierEquals("DEALLOCATE")) {
             MysqlDeallocatePrepareStatement stmt = parseDeallocatePrepare();
             statementList.add(stmt);
@@ -4575,7 +4562,13 @@ public MySqlPrepareStatement parsePrepare() {
         return new MySqlPrepareStatement(name, from);
     }
 
-    public MySqlExecuteStatement parseExecute() {
+    @Override
+    public SQLStatement parseExecute() {
+        acceptIdentifier("EXECUTE");
+        if (lexer.identifierEquals("RESTART") || lexer.identifierEquals("UPDATE")) {
+            return parseExecuteForAds();
+        }
+
         MySqlExecuteStatement stmt = new MySqlExecuteStatement();
 
         SQLName statementName = exprParser.name();
@@ -8900,6 +8893,7 @@ public MySqlSelectIntoStatement parseSelectInto() {
     /**
      * parse loop statement
      */
+    @Override
     public SQLLoopStatement parseLoop() {
         SQLLoopStatement loopStmt = new SQLLoopStatement();
         accept(Token.LOOP);

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/oracle/parser/OracleStatementParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/oracle/parser/OracleStatementParser.java
@@ -1237,6 +1237,7 @@ public OracleForStatement parseFor() {
         return stmt;
     }
 
+    @Override
     public SQLLoopStatement parseLoop() {
         accept(Token.LOOP);
 

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
@@ -654,6 +654,27 @@ && dialectFeatureEnabled(ParseStatementListCommitReturn)) {
                 continue;
             }
 
+            if (lexer.token == LOOP) {
+                SQLStatement stmt = parseLoop();
+                statementList.add(stmt);
+                stmt.setParent(parent);
+                continue;
+            }
+
+            if (lexer.token == CONTINUE) {
+                SQLStatement stmt = parseContinue();
+                statementList.add(stmt);
+                stmt.setParent(parent);
+                continue;
+            }
+
+            if (lexer.token == LEAVE) {
+                SQLStatement stmt = parseLeave();
+                statementList.add(stmt);
+                stmt.setParent(parent);
+                continue;
+            }
+
             if (lexer.identifierEquals("EXECUTE")) {
                 SQLStatement stmt = parseExecute();
                 statementList.add(stmt);
@@ -1322,16 +1343,31 @@ public SQLWhileStatement parseWhile() {
         return stmt;
     }
 
+    public SQLStatement parseLoop() {
+        accept(Token.LOOP);
+        SQLLoopStatement stmt = new SQLLoopStatement();
+        this.parseStatementList(stmt.getStatements(), -1, stmt);
+        accept(Token.END);
+        accept(Token.LOOP);
+        return stmt;
+    }
+
     public SQLStatement parseDeclare() {
         throw new ParserException("not supported. " + lexer.info());
     }
 
+    public SQLStatement parseContinue() {
+        accept(Token.CONTINUE);
+        return new SQLContinueStatement();
+    }
+
     public SQLStatement parseRepeat() {
         throw new ParserException("not supported. " + lexer.info());
     }
 
     public SQLStatement parseLeave() {
-        throw new ParserException("not supported. " + lexer.info());
+        accept(LEAVE);
+        return new SQLLeaveStatement();
     }
 
     public SQLStatement parseCache() {

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
@@ -11353,6 +11353,16 @@ public boolean visit(SQLExecuteImmediateStatement x) {
         return false;
     }
 
+    public boolean visit(SQLContinueStatement x) {
+        print(ucase ? "CONTINUE" : "continue");
+        return false;
+    }
+
+    public boolean visit(SQLLeaveStatement x) {
+        print(ucase ? "LEAVE" : "leave");
+        return false;
+    }
+
     protected void printCreateTableLike(SQLCreateTableStatement x) {
         SQLExprTableSource like = x.getLike();
         if (like == null) {

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
@@ -2694,6 +2694,19 @@ default boolean visit(SQLExceptionStatement.Item x) {
     }
 
     default void endVisit(SQLExceptionStatement.Item x) {}
+
+    default boolean visit(SQLContinueStatement x) {
+        return true;
+    }
+
+    default void endVisit(SQLContinueStatement x) {}
+
+    default boolean visit(SQLLeaveStatement x) {
+        return true;
+    }
+
+    default void endVisit(SQLLeaveStatement x) {}
+
     static SQLASTVisitor ofMethodInvoke(Consumer<SQLMethodInvokeExpr> p) {
         return new SQLASTVisitor() {
             public boolean visit(SQLMethodInvokeExpr x) {

diff --git a/core/src/test/java/com/alibaba/druid/bvt/sql/SQLResourceTest.java b/core/src/test/java/com/alibaba/druid/bvt/sql/SQLResourceTest.java
@@ -125,7 +125,7 @@ protected void fileTest(File file) throws IOException {
 
             SQLStatementParser parser = SQLParserUtils.createSQLStatementParser(sql, dbType);
             SQLStatement stmt = parser.parseStatement();
-            assertEquals(Token.EOF, parser.getLexer().token());
+            assertEquals(parser.getLexer().info(), Token.EOF, parser.getLexer().token());
             String result = SQLUtils.toSQLString(stmt, dbType).trim();
             assertEquals(expected, result);
 

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,22 @@
+LOOP
+  SET x = x + 1;
+  IF x >= 10 THEN
+    LEAVE;
+  END IF;
+END LOOP;
+--------------------
+LOOP
+	SET x = x + 1;
+	IF x >= 10
+		THEN
+		LEAVE;
+	END IF;
+END LOOP;
+------------------------------------------------------------------------------------------------------------------------
+CONTINUE;
+--------------------
+CONTINUE;
+------------------------------------------------------------------------------------------------------------------------
 EXECUTE IMMEDIATE "SELECT @a * (@b + 2)" INTO y USING 1 as a, 3 as b;
 --------------------
 EXECUTE IMMEDIATE 'SELECT @a * (@b + 2)' INTO y USING 1 AS a, 3 AS b;