From 1a28541e6394799ae1dcee5bec86df5be3ba28bb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=AB=98=E9=93=81?= Date: Thu, 5 May 2016 11:56:19 +0800 Subject: [PATCH] support taobao urgly update returning sql. for issue 1196. https://github.com/alibaba/druid/issues/1196 --- .../ast/statement/MySqlUpdateStatement.java | 46 ++++++++- .../mysql/parser/MySqlSelectParser.java | 22 +++++ .../mysql/visitor/MySqlOutputVisitor.java | 18 ++++ .../mysql/visitor/MySqlSchemaStatVisitor.java | 9 +- .../bvt/sql/mysql/MySqlUpdateTest_6.java | 94 +++++++++++++++++++ 5 files changed, 185 insertions(+), 4 deletions(-) create mode 100644 src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlUpdateTest_6.java diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlUpdateStatement.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlUpdateStatement.java index 685bf60df5..6109cfb435 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlUpdateStatement.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlUpdateStatement.java @@ -18,6 +18,7 @@ package com.alibaba.druid.sql.dialect.mysql.ast.statement; import java.util.ArrayList; import java.util.List; +import com.alibaba.druid.sql.ast.SQLExpr; import com.alibaba.druid.sql.ast.SQLOrderBy; import com.alibaba.druid.sql.ast.statement.SQLSelectItem; import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement; @@ -31,10 +32,14 @@ public class MySqlUpdateStatement extends SQLUpdateStatement implements MySqlSta private SQLOrderBy orderBy; private Limit limit; - private boolean lowPriority = false; - private boolean ignore = false; + private boolean lowPriority = false; + private boolean ignore = false; + private boolean commitOnSuccess = false; + private boolean rollBackOnFail = false; + private boolean queryOnPk = false; + private SQLExpr targetAffectRow; - private List returning = new ArrayList(); + private List returning = new ArrayList(); public MySqlUpdateStatement(){ super(JdbcConstants.MYSQL); @@ -98,6 +103,41 @@ public class MySqlUpdateStatement extends SQLUpdateStatement implements MySqlSta this.ignore = ignore; } + public boolean isCommitOnSuccess() { + return commitOnSuccess; + } + + public void setCommitOnSuccess(boolean commitOnSuccess) { + this.commitOnSuccess = commitOnSuccess; + } + + public boolean isRollBackOnFail() { + return rollBackOnFail; + } + + public void setRollBackOnFail(boolean rollBackOnFail) { + this.rollBackOnFail = rollBackOnFail; + } + + public boolean isQueryOnPk() { + return queryOnPk; + } + + public void setQueryOnPk(boolean queryOnPk) { + this.queryOnPk = queryOnPk; + } + + public SQLExpr getTargetAffectRow() { + return targetAffectRow; + } + + public void setTargetAffectRow(SQLExpr targetAffectRow) { + if (targetAffectRow != null) { + targetAffectRow.setParent(this); + } + this.targetAffectRow = targetAffectRow; + } + public SQLOrderBy getOrderBy() { return orderBy; } diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlSelectParser.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlSelectParser.java index e0f94f8dc4..3ff115c917 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlSelectParser.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlSelectParser.java @@ -15,6 +15,7 @@ */ package com.alibaba.druid.sql.dialect.mysql.parser; +import com.alibaba.druid.sql.ast.SQLExpr; import com.alibaba.druid.sql.ast.SQLSetQuantifier; import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr; import com.alibaba.druid.sql.ast.expr.SQLLiteralExpr; @@ -251,6 +252,27 @@ public class MySqlSelectParser extends SQLSelectParser { lexer.nextToken(); update.setIgnore(true); } + + if (identifierEquals("COMMIT_ON_SUCCESS")) { + lexer.nextToken(); + update.setCommitOnSuccess(true); + } + + if (identifierEquals("ROLLBACK_ON_FAIL")) { + lexer.nextToken(); + update.setRollBackOnFail(true); + } + + if (identifierEquals("QUEUE_ON_PK")) { + lexer.nextToken(); + update.setQueryOnPk(true); + } + + if (identifierEquals("TARGET_AFFECT_ROW")) { + lexer.nextToken(); + SQLExpr targetAffectRow = this.exprParser.expr(); + update.setTargetAffectRow(targetAffectRow); + } SQLTableSource updateTableSource = this.exprParser.createSelectParser().parseTableSource(); update.setTableSource(updateTableSource); diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java index 9546b7142b..ba6ae99b8c 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java @@ -1645,6 +1645,24 @@ public class MySqlOutputVisitor extends SQLASTOutputVisitor implements MySqlASTV if (x.isIgnore()) { print0(ucase ? "IGNORE " : "ignore "); } + + if (x.isCommitOnSuccess()) { + print0(ucase ? "COMMIT_ON_SUCCESS " : "commit_on_success "); + } + + if (x.isRollBackOnFail()) { + print0(ucase ? "ROLLBACK_ON_FAIL " : "rollback_on_fail "); + } + + if (x.isQueryOnPk()) { + print0(ucase ? "QUEUE_ON_PK " : "queue_on_pk "); + } + + if (x.getTargetAffectRow() != null) { + print0(ucase ? "TARGET_AFFECT_ROW " : "target_affect_row "); + x.getTargetAffectRow().accept(this); + print(' '); + } x.getTableSource().accept(this); diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlSchemaStatVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlSchemaStatVisitor.java index b81b2300b0..108eade106 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlSchemaStatVisitor.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlSchemaStatVisitor.java @@ -24,6 +24,7 @@ import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr; import com.alibaba.druid.sql.ast.statement.SQLAlterTableStatement; import com.alibaba.druid.sql.ast.statement.SQLCreateTableStatement; import com.alibaba.druid.sql.ast.statement.SQLExprTableSource; +import com.alibaba.druid.sql.ast.statement.SQLSelectItem; import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock; import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; import com.alibaba.druid.sql.ast.statement.SQLTableSource; @@ -538,7 +539,13 @@ public class MySqlSchemaStatVisitor extends SchemaStatVisitor implements MySqlAS @Override public boolean visit(MySqlUpdateStatement x) { - return visit((SQLUpdateStatement) x); + + visit((SQLUpdateStatement) x); + for (SQLSelectItem item : x.getReturning()) { + item.accept(this); + } + + return false; } @Override diff --git a/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlUpdateTest_6.java b/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlUpdateTest_6.java new file mode 100644 index 0000000000..8dc9e63916 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlUpdateTest_6.java @@ -0,0 +1,94 @@ +/* + * Copyright 1999-2101 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.bvt.sql.mysql; + +import java.util.List; + +import org.junit.Assert; + +import com.alibaba.druid.sql.MysqlTest; +import com.alibaba.druid.sql.SQLUtils; +import com.alibaba.druid.sql.ast.SQLStatement; +import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser; +import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlSchemaStatVisitor; +import com.alibaba.druid.stat.TableStat; +import com.alibaba.druid.stat.TableStat.Column; + +public class MySqlUpdateTest_6 extends MysqlTest { + + public void test_0() throws Exception { + String sql = "/* 6445c36514623747244871009d62f9/0//sub_order_id:1462374729595/order_id:1462374729594/ */" + + "SELECT id, user_id, item_id, sku_id, flag" + + " , sellable_quantity, fff03, status, feature, feature_lock" + + " , version, gmt_create, gmt_modified, root_id, parent_id, dist_type, dist_id" + + " , occupy_quantity, user_type " + + "FROM " + + "UPDATE COMMIT_ON_SUCCESS ROLLBACK_ON_FAIL TARGET_AFFECT_ROW ? " + + "`table_3966` " + + "SET `fff03` = `fff03` + ?" + + " , `flag` = `flag` & ~ (1 << 10) & ~ (1 << 11)" + + " , `version` = `version` + 3, `gmt_modified` = NOW() " + + "WHERE `root_id` = ? AND `status` = 1 AND `id` IN (?, ?) AND `fff03` + ? >= 0"; + + MySqlStatementParser parser = new MySqlStatementParser(sql); + List statementList = parser.parseStatementList(); + SQLStatement stmt = statementList.get(0); + print(statementList); + + Assert.assertEquals(1, statementList.size()); + + MySqlSchemaStatVisitor visitor = new MySqlSchemaStatVisitor(); + stmt.accept(visitor); + +// System.out.println("Tables : " + visitor.getTables()); +// System.out.println("fields : " + visitor.getColumns()); +// System.out.println("coditions : " + visitor.getConditions()); +// System.out.println("orderBy : " + visitor.getOrderByColumns()); + + Assert.assertEquals(1, visitor.getTables().size()); + Assert.assertEquals(19, visitor.getColumns().size()); + Assert.assertEquals(3, visitor.getConditions().size()); + + Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("table_3966"))); + + Assert.assertTrue(visitor.getColumns().contains(new Column("table_3966", "user_id"))); + Assert.assertTrue(visitor.getColumns().contains(new Column("table_3966", "id"))); + Assert.assertTrue(visitor.getColumns().contains(new Column("table_3966", "item_id"))); + + { + String output = SQLUtils.toMySqlString(stmt); + Assert.assertEquals("SELECT id, user_id, item_id, sku_id, flag, sellable_quantity, fff03, status, feature, feature_lock, version, gmt_create, gmt_modified, root_id, parent_id, dist_type, dist_id, occupy_quantity, user_type" + + "\nFROM UPDATE COMMIT_ON_SUCCESS ROLLBACK_ON_FAIL TARGET_AFFECT_ROW ? `table_3966`" + + "\nSET `fff03` = `fff03` + ?, `flag` = `flag` & ~(1 << 10 & ~(1 << 11)), `version` = `version` + 3, `gmt_modified` = NOW()" + + "\nWHERE `root_id` = ?" + + "\n\tAND `status` = 1" + + "\n\tAND `id` IN (?, ?)" + + "\n\tAND `fff03` + ? >= 0", // + output); + } + { + String output = SQLUtils.toMySqlString(stmt, SQLUtils.DEFAULT_LCASE_FORMAT_OPTION); + Assert.assertEquals("select id, user_id, item_id, sku_id, flag, sellable_quantity, fff03, status, feature, feature_lock, version, gmt_create, gmt_modified, root_id, parent_id, dist_type, dist_id, occupy_quantity, user_type" + + "\nfrom update commit_on_success rollback_on_fail target_affect_row ? `table_3966`" + + "\nset `fff03` = `fff03` + ?, `flag` = `flag` & ~(1 << 10 & ~(1 << 11)), `version` = `version` + 3, `gmt_modified` = NOW()" + + "\nwhere `root_id` = ?" + + "\n\tand `status` = 1" + + "\n\tand `id` in (?, ?)" + + "\n\tand `fff03` + ? >= 0", // + output); + } + } +}