diff --git a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLDeleteStatement.java b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLDeleteStatement.java index fd1d1d5def..b79806ff8c 100644 --- a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLDeleteStatement.java +++ b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLDeleteStatement.java @@ -292,4 +292,13 @@ public class SQLDeleteStatement extends SQLStatementImpl implements SQLReplaceab return false; } + + public boolean addWhere(SQLExpr where) { + if (where == null) { + return false; + } + + this.addCondition(where); + return true; + } } diff --git a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelect.java b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelect.java index 3d95fba579..8f28cc3b32 100644 --- a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelect.java +++ b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelect.java @@ -20,6 +20,7 @@ import java.util.List; import com.alibaba.druid.sql.SQLUtils; import com.alibaba.druid.sql.ast.*; +import com.alibaba.druid.sql.ast.expr.SQLAllColumnExpr; import com.alibaba.druid.sql.dialect.oracle.ast.OracleSQLObject; import com.alibaba.druid.sql.visitor.SQLASTOutputVisitor; import com.alibaba.druid.sql.visitor.SQLASTVisitor; @@ -321,4 +322,26 @@ public class SQLSelect extends SQLObjectImpl { return null; } + + public boolean addWhere(SQLExpr where) { + if (where == null) { + return false; + } + + if (query instanceof SQLSelectQueryBlock) { + ((SQLSelectQueryBlock) query).addWhere(where); + return true; + } + + if (query instanceof SQLUnionQuery) { + SQLSelectQueryBlock queryBlock = new SQLSelectQueryBlock(); + queryBlock.setFrom(new SQLSelect(query), "u"); + queryBlock.addSelectItem(new SQLAllColumnExpr()); + queryBlock.setParent(queryBlock); + query = queryBlock; + return true; + } + + return false; + } } diff --git a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelectStatement.java b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelectStatement.java index 45bf022bf4..ffce0197e2 100644 --- a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelectStatement.java +++ b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLSelectStatement.java @@ -16,6 +16,7 @@ package com.alibaba.druid.sql.ast.statement; import com.alibaba.druid.sql.ast.SQLCommentHint; +import com.alibaba.druid.sql.ast.SQLExpr; import com.alibaba.druid.sql.ast.SQLObject; import com.alibaba.druid.sql.ast.SQLStatementImpl; import com.alibaba.druid.sql.visitor.SQLASTVisitor; @@ -85,4 +86,8 @@ public class SQLSelectStatement extends SQLStatementImpl { public List getChildren() { return Collections.singletonList(select); } + + public boolean addWhere(SQLExpr where) { + return select.addWhere(where); + } } diff --git a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUpdateStatement.java b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUpdateStatement.java index a2d0dfc817..5497074175 100644 --- a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUpdateStatement.java +++ b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLUpdateStatement.java @@ -312,4 +312,13 @@ public class SQLUpdateStatement extends SQLStatementImpl implements SQLReplaceab result = 31 * result + (hints != null ? hints.hashCode() : 0); return result; } + + public boolean addWhere(SQLExpr where) { + if (where == null) { + return false; + } + + this.addCondition(where); + return true; + } } diff --git a/src/main/java/com/alibaba/druid/util/MySqlUtils.java b/src/main/java/com/alibaba/druid/util/MySqlUtils.java index 75b9ad8d56..cd3925cd7d 100644 --- a/src/main/java/com/alibaba/druid/util/MySqlUtils.java +++ b/src/main/java/com/alibaba/druid/util/MySqlUtils.java @@ -104,7 +104,7 @@ public class MySqlUtils { e.printStackTrace(); } - } else if (major == 6) { + } else if (major == 6 || major == 8) { if (method_6_getValue == null && !method_6_getValue_error) { try { class_6_connection = Class.forName("com.mysql.cj.api.jdbc.JdbcConnection"); diff --git a/src/test/java/com/alibaba/druid/sql/refactor/AddWhere_0.java b/src/test/java/com/alibaba/druid/sql/refactor/AddWhere_0.java new file mode 100644 index 0000000000..61a96b3194 --- /dev/null +++ b/src/test/java/com/alibaba/druid/sql/refactor/AddWhere_0.java @@ -0,0 +1,76 @@ +package com.alibaba.druid.sql.refactor; + +import com.alibaba.druid.sql.SQLUtils; +import com.alibaba.druid.sql.ast.statement.SQLDeleteStatement; +import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; +import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement; +import com.alibaba.druid.util.JdbcConstants; +import junit.framework.TestCase; + +public class AddWhere_0 extends TestCase { + public void test_select_0() throws Exception { + String sql = "select * from t"; + SQLSelectStatement stmt = (SQLSelectStatement) SQLUtils.parseStatements(sql, JdbcConstants.MYSQL).get(0); + stmt.addWhere(SQLUtils.toSQLExpr("id = 1", JdbcConstants.MYSQL)); + + assertEquals("SELECT *\n" + + "FROM t\n" + + "WHERE id = 1", stmt.toString()); + } + + public void test_select_1() throws Exception { + String sql = "select * from t where name = 'xx'"; + SQLSelectStatement stmt = (SQLSelectStatement) SQLUtils.parseStatements(sql, JdbcConstants.MYSQL).get(0); + stmt.addWhere(SQLUtils.toSQLExpr("id = 1", JdbcConstants.MYSQL)); + + assertEquals("SELECT *\n" + + "FROM t\n" + + "WHERE name = 'xx'\n" + + "AND id = 1", stmt.toString()); + } + + public void test_select_1_union() throws Exception { + String sql = "select * from t1 union all select * from t2"; + SQLSelectStatement stmt = (SQLSelectStatement) SQLUtils.parseStatements(sql, JdbcConstants.MYSQL).get(0); + stmt.addWhere(SQLUtils.toSQLExpr("id = 1", JdbcConstants.MYSQL)); + + assertEquals("SELECT *\n" + + "FROM (\n" + + "\tSELECT *\n" + + "\tFROM t1\n" + + "\tUNION ALL\n" + + "\tSELECT *\n" + + "\tFROM t2\n" + + ") u", stmt.toString()); + } + + public void test_delete_0() throws Exception { + String sql = "delete from t"; + SQLDeleteStatement stmt = (SQLDeleteStatement) SQLUtils.parseStatements(sql, JdbcConstants.MYSQL).get(0); + stmt.addWhere(SQLUtils.toSQLExpr("id = 1", JdbcConstants.MYSQL)); + + assertEquals("DELETE FROM t\n" + + "WHERE id = 1", stmt.toString()); + } + + public void test_delete_1() throws Exception { + String sql = "delete from t where name = 'xx'"; + SQLDeleteStatement stmt = (SQLDeleteStatement) SQLUtils.parseStatements(sql, JdbcConstants.MYSQL).get(0); + stmt.addWhere(SQLUtils.toSQLExpr("id = 1", JdbcConstants.MYSQL)); + + assertEquals("DELETE FROM t\n" + + "WHERE name = 'xx'\n" + + "\tAND id = 1", stmt.toString()); + } + + public void test_update_0() throws Exception { + String sql = "update t set val = 'abc' where name = 'xx'"; + SQLUpdateStatement stmt = (SQLUpdateStatement) SQLUtils.parseStatements(sql, JdbcConstants.MYSQL).get(0); + stmt.addWhere(SQLUtils.toSQLExpr("id = 1", JdbcConstants.MYSQL)); + + assertEquals("UPDATE t\n" + + "SET val = 'abc'\n" + + "WHERE name = 'xx'\n" + + "\tAND id = 1", stmt.toString()); + } +}