Fix select query block replace issue.

alibaba · Dec 19, 2024 · 5a1a5f0 · 5a1a5f0
1 parent 4efdf7a
commit 5a1a5f0
Show file tree

Hide file tree

Showing 5 changed files with 86 additions and 10 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLAliasedExpr.java b/core/src/main/java/com/alibaba/druid/sql/ast/expr/SQLAliasedExpr.java
@@ -3,11 +3,13 @@
 import com.alibaba.druid.sql.SQLUtils;
 import com.alibaba.druid.sql.ast.SQLDataType;
 import com.alibaba.druid.sql.ast.SQLExpr;
-import com.alibaba.druid.sql.ast.SQLObjectImpl;
+import com.alibaba.druid.sql.ast.SQLExprImpl;
 import com.alibaba.druid.sql.ast.SQLReplaceable;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
-public class SQLAliasedExpr extends SQLObjectImpl implements SQLReplaceable {
+import java.util.Objects;
+
+public class SQLAliasedExpr extends SQLExprImpl implements SQLReplaceable {
     protected SQLExpr expr;
     protected String alias;
 
@@ -119,4 +121,28 @@ protected void cloneTo(SQLAliasedExpr x) {
         }
         x.setAlias(alias);
     }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+
+        SQLAliasedExpr that = (SQLAliasedExpr) o;
+
+        if (!Objects.equals(expr, that.expr)) {
+            return false;
+        }
+        return Objects.equals(alias, that.alias);
+    }
+
+    @Override
+    public int hashCode() {
+        int result = expr != null ? expr.hashCode() : 0;
+        result = 31 * result + (alias != null ? alias.hashCode() : 0);
+        return result;
+    }
 }
diff --git a/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelectOrderByItem.java b/core/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelectOrderByItem.java
@@ -21,7 +21,7 @@
 
 import java.util.List;
 
-public final class SQLSelectOrderByItem extends SQLObjectImpl implements SQLReplaceable {
+public final class SQLSelectOrderByItem extends SQLExprImpl implements SQLReplaceable {
     protected SQLExpr expr;
     protected String collate;
     protected SQLExpr opclass;

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
@@ -1174,21 +1174,71 @@ public SQLTableSource findTableSourceWithColumn(long columnHash) {
 
     @Override
     public boolean replace(SQLExpr expr, SQLExpr target) {
+        boolean isReplaced = false;
         if (where == expr) {
             setWhere(target);
-            return true;
+            isReplaced = true;
         }
 
         if (startWith == expr) {
             setStartWith(target);
-            return true;
+            isReplaced = true;
         }
 
         if (connectBy == expr) {
             setConnectBy(target);
-            return true;
+            isReplaced = true;
         }
-        return false;
+
+        if (with == expr) {
+            setWith((SQLWithSubqueryClause) target);
+            isReplaced = true;
+        }
+
+        if (from == expr) {
+            setFrom((SQLTableSource) target);
+            isReplaced = true;
+        }
+
+        if (into == expr && target instanceof SQLExprTableSource) {
+            setInto((SQLExprTableSource) target);
+            isReplaced = true;
+        }
+
+        if (qualify == expr) {
+            setQualify(target);
+            isReplaced = true;
+        }
+
+        if (waitTime == expr) {
+            setWaitTime(target);
+            isReplaced = true;
+        }
+
+        if (expr instanceof SQLSelectItem && target instanceof SQLSelectItem) {
+            isReplaced = isReplaced || replaceList(selectList, (SQLSelectItem) expr, (SQLSelectItem) target);
+        }
+
+        if (expr instanceof SQLSelectOrderByItem && target instanceof SQLSelectOrderByItem) {
+            isReplaced = isReplaced || replaceList(distributeBy, (SQLSelectOrderByItem) expr, (SQLSelectOrderByItem) target)
+                || replaceList(sortBy, (SQLSelectOrderByItem) expr, (SQLSelectOrderByItem) target)
+                || replaceList(clusterBy, (SQLSelectOrderByItem) expr, (SQLSelectOrderByItem) target);
+        }
+
+        isReplaced = isReplaced || replaceList(forUpdateOf, expr, target);
+        return isReplaced;
+    }
+
+    protected <T extends SQLObject> boolean replaceList(List<T> exprList, T expr, T target) {
+        boolean isReplaced = false;
+        for (int i = 0; i < exprList.size(); i++) {
+            if (exprList.get(i) == expr) {
+                target.setParent(this);
+                exprList.set(i, target);
+                isReplaced = true;
+            }
+        }
+        return isReplaced;
     }
 
     public SQLSelectItem findSelectItem(String ident) {

diff --git a/core/src/test/java/com/alibaba/druid/bvt/filter/wall/mysql/MySqlWallTest91.java b/core/src/test/java/com/alibaba/druid/bvt/filter/wall/mysql/MySqlWallTest91.java
@@ -187,7 +187,7 @@ public void test_true4() {
         WallProvider provider = initWallProvider();
         {
             String sql = "SELECT 10006, @";
-            Assert.assertTrue(provider.checkValid(sql));
+            Assert.assertFalse(provider.checkValid(sql));
         }
     }
 

diff --git a/pom.xml b/pom.xml
@@ -6,8 +6,8 @@
 
 	<groupId>com.alibaba</groupId>
 	<artifactId>druid-parent</artifactId>
-	<version>1.2.25_preview39</version>
-	<name>druid-parent</name>
+	<version>1.2.25-SNAPSHOT</version>
+	<name>${project.artifactId}</name>
 	<description>A JDBC datasource implementation.</description>
 	<packaging>pom</packaging>
 	<url>https://github.com/alibaba/druid</url>