针对mysql增加with子查询的解析支持 alibaba#5761

lizongbo · May 5, 2024 · 0e4a3d8 · 0e4a3d8
1 parent 15b887e
commit 0e4a3d8
Show file tree

Hide file tree

Showing 9 changed files with 439 additions and 7 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLSelectExpr.java b/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLSelectExpr.java
@@ -0,0 +1,114 @@
+/*
+ * Copyright 1999-2024 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.druid.sql.ast.expr;
+
+import com.alibaba.druid.sql.ast.*;
+import com.alibaba.druid.sql.ast.statement.SQLSelect;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
+
+import java.io.Serializable;
+
+/**
+ * @author lizongbo
+ */
+public class SQLSelectExpr extends SQLExprImpl implements SQLReplaceable, Serializable {
+    private static final long serialVersionUID = 1L;
+    protected SQLSelect sqlSelect;
+    protected SQLExpr as;
+
+    public SQLSelectExpr() {
+    }
+
+    public SQLSelectExpr(SQLSelect sqlSelect) {
+        this.sqlSelect = sqlSelect;
+    }
+
+    public SQLSelect getSqlSelect() {
+        return sqlSelect;
+    }
+
+    public void setSqlSelect(SQLSelect sqlSelect) {
+        this.sqlSelect = sqlSelect;
+    }
+
+    public SQLExpr getAs() {
+        return as;
+    }
+
+    public void setAs(SQLExpr as) {
+        this.as = as;
+    }
+
+    public SQLSelectExpr clone() {
+        SQLSelectExpr x = new SQLSelectExpr();
+        if (sqlSelect != null) {
+            x.setSqlSelect(sqlSelect.clone());
+        }
+        if (as != null) {
+            x.setAs(as.clone());
+        }
+        x.setParenthesized(parenthesized);
+        return x;
+    }
+
+    protected void accept0(SQLASTVisitor visitor) {
+        if (visitor.visit(this)) {
+            if (this.sqlSelect != null) {
+                this.sqlSelect.accept(visitor);
+            }
+        }
+        visitor.endVisit(this);
+    }
+    @Override
+    public int hashCode() {
+        final int prime = 31;
+        int result = 1;
+        result = prime * result + ((sqlSelect == null) ? 0 : sqlSelect.hashCode());
+        return result;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj) {
+            return true;
+        }
+        if (obj == null) {
+            return false;
+        }
+        if (getClass() != obj.getClass()) {
+            return false;
+        }
+        SQLSelectExpr other = (SQLSelectExpr) obj;
+        if (sqlSelect == null) {
+            if (other.sqlSelect != null) {
+                return false;
+            }
+        } else if (!sqlSelect.equals(other.sqlSelect)) {
+            return false;
+        }
+        return true;
+    }
+
+    @Override
+    public SQLDataType computeDataType() {
+        return SQLBooleanExpr.DATA_TYPE;
+    }
+
+    @Override
+    public boolean replace(SQLExpr expr, SQLExpr target) {
+        return false;
+    }
+}
diff --git a/core/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java b/core/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java
@@ -395,7 +395,6 @@ public SQLExpr primary() {
                     sqlExpr = new SQLMethodInvokeExpr();
                     break;
                 }
-
                 sqlExpr = expr();
 
                 if (lexer.token == Token.COMMA) {
@@ -5912,7 +5911,6 @@ public SQLSelectItem parseSelectItem() {
         boolean connectByRoot = false;
         Token token = lexer.token;
         int startPos = lexer.startPos;
-
         if (token == Token.IDENTIFIER
                 && !(lexer.hashLCase() == -5808529385363204345L && lexer.charAt(lexer.pos) == '\'' && dbType == DbType.mysql) // x'123' X'123'
         ) {
@@ -6474,7 +6472,6 @@ public SQLSelectItem parseSelectItem() {
             while (lexer.token == Token.HINT) {
                 lexer.nextToken();
             }
-
             expr = expr();
         }
 

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
@@ -24,6 +24,7 @@
 import com.alibaba.druid.sql.dialect.hive.parser.HiveCreateTableParser;
 import com.alibaba.druid.sql.dialect.hive.stmt.HiveCreateTableStatement;
 import com.alibaba.druid.sql.dialect.mysql.ast.expr.MySqlOrderingExpr;
+import com.alibaba.druid.sql.parser.Lexer.SavePoint;
 import com.alibaba.druid.util.FnvHash;
 import com.alibaba.druid.util.StringUtils;
 
@@ -1074,6 +1075,32 @@ protected SQLExpr parseGroupByItem() {
     protected void parseSelectList(SQLSelectQueryBlock queryBlock) {
         final List<SQLSelectItem> selectList = queryBlock.getSelectList();
         for (; ; ) {
+            SavePoint savePoint = lexer.markOut();
+            if (lexer.token() == Token.LPAREN) {
+                lexer.nextToken();
+                if (lexer.token() == Token.WITH) {
+                    String alias = null;
+                    SQLSelect select = select();
+                    SQLSelectExpr sqlSelectExpr = new SQLSelectExpr(select);
+                    sqlSelectExpr.setParenthesized(true);
+                    SQLSelectItem selectItem = new SQLSelectItem(sqlSelectExpr, alias, false);
+                    selectList.add(selectItem);
+                    selectItem.setParent(queryBlock);
+                    accept(Token.RPAREN);
+                    if (lexer.token() == Token.AS) {
+                        accept(Token.AS);
+                        sqlSelectExpr.setAs(new SQLIdentifierExpr(lexer.stringVal()));
+                        lexer.nextToken();
+                        if (lexer.token != Token.COMMA) {
+                            break;
+                        }
+                        lexer.nextToken();
+                    }
+                } else {
+                    lexer.reset(savePoint);
+                }
+            }
+
             final SQLSelectItem selectItem = this.exprParser.parseSelectItem();
             selectList.add(selectItem);
             selectItem.setParent(queryBlock);

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
@@ -1152,11 +1152,26 @@ protected final void printExpr(SQLExpr x, boolean parameterized) {
             visit((SQLInListExpr) x);
         } else if (clazz == SQLNotExpr.class) {
             visit((SQLNotExpr) x);
+        } else if (clazz == SQLSelectExpr.class) {
+            visit((SQLSelectExpr) x);
         } else {
             x.accept(this);
         }
     }
-
+    public boolean visit(SQLSelectExpr x) {
+        if (x.isParenthesized()) {
+            print('(');
+        }
+        visit(x.getSqlSelect());
+        if (x.isParenthesized()) {
+            print(')');
+        }
+        if (x.getAs() != null) {
+            print0(ucase ? " AS " : " as ");
+            printExpr(x.getAs());
+        }
+        return false;
+    }
     public boolean visit(SQLCaseExpr x) {
         if (x.isParenthesized()) {
             print('(');

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
@@ -2564,4 +2564,11 @@ default boolean visit(SQLCostStatement x) {
     default void endVisit(SQLCostStatement x) {
     }
 
+    default boolean visit(SQLSelectExpr x) {
+        return true;
+    }
+
+    default void endVisit(SQLSelectExpr x) {
+    }
+
 }