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 dd70ef9432..685bf60df5 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 @@ -15,7 +15,11 @@ */ package com.alibaba.druid.sql.dialect.mysql.ast.statement; +import java.util.ArrayList; +import java.util.List; + import com.alibaba.druid.sql.ast.SQLOrderBy; +import com.alibaba.druid.sql.ast.statement.SQLSelectItem; import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock.Limit; import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitor; @@ -24,14 +28,16 @@ import com.alibaba.druid.util.JdbcConstants; public class MySqlUpdateStatement extends SQLUpdateStatement implements MySqlStatement { - private SQLOrderBy orderBy; - private Limit limit; + private SQLOrderBy orderBy; + private Limit limit; - private boolean lowPriority = false; - private boolean ignore = false; - - public MySqlUpdateStatement() { - super (JdbcConstants.MYSQL); + private boolean lowPriority = false; + private boolean ignore = false; + + private List returning = new ArrayList(); + + public MySqlUpdateStatement(){ + super(JdbcConstants.MYSQL); } public Limit getLimit() { @@ -45,6 +51,17 @@ public class MySqlUpdateStatement extends SQLUpdateStatement implements MySqlSta this.limit = limit; } + public List getReturning() { + return returning; + } + + public void addReturning(List returning) { + for (SQLSelectItem item : returning) { + item.setParent(this); + this.returning.add(item); + } + } + @Override protected void accept0(SQLASTVisitor visitor) { if (visitor instanceof MySqlASTVisitor) { 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 51e40dfb84..e0f94f8dc4 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 @@ -45,6 +45,9 @@ import com.alibaba.druid.sql.parser.Token; public class MySqlSelectParser extends SQLSelectParser { + protected boolean returningFlag = false; + protected MySqlUpdateStatement updateStmt; + public MySqlSelectParser(SQLExprParser exprParser){ super(exprParser); } @@ -52,7 +55,25 @@ public class MySqlSelectParser extends SQLSelectParser { public MySqlSelectParser(String sql){ this(new MySqlExprParser(sql)); } + + public void parseFrom(SQLSelectQueryBlock queryBlock) { + if (lexer.token() != Token.FROM) { + return; + } + + lexer.nextToken(); + + if (lexer.token() == Token.UPDATE) { // taobao returning to urgly syntax + updateStmt = this.parseUpdateStatment(); + updateStmt.addReturning(queryBlock.getSelectList()); + returningFlag = true; + return; + } + + queryBlock.setFrom(parseTableSource()); + } + @Override public SQLSelectQuery query() { if (lexer.token() == (Token.LPAREN)) { diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlStatementParser.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlStatementParser.java index 506ba8a937..6f7cfcac31 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlStatementParser.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlStatementParser.java @@ -18,7 +18,6 @@ package com.alibaba.druid.sql.dialect.mysql.parser; import java.util.ArrayList; import java.util.List; - import com.alibaba.druid.sql.ast.SQLCommentHint; import com.alibaba.druid.sql.ast.SQLDataTypeImpl; import com.alibaba.druid.sql.ast.SQLDeclareItem; @@ -230,9 +229,16 @@ public class MySqlStatementParser extends SQLStatementParser { return parser.parseCrateTable(); } - public SQLSelectStatement parseSelect() { + public SQLStatement parseSelect() { MySqlSelectParser selectParser = new MySqlSelectParser(this.exprParser); - return new SQLSelectStatement(selectParser.select(), JdbcConstants.MYSQL); + + SQLSelect select = selectParser.select(); + + if (selectParser.returningFlag) { + return selectParser.updateStmt; + } + + return new SQLSelectStatement(select, JdbcConstants.MYSQL); } public SQLUpdateStatement parseUpdateStatement() { @@ -3664,7 +3670,8 @@ public class MySqlStatementParser extends SQLStatementParser { accept(Token.FOR); - stmt.setSelect(parseSelect()); + SQLSelectStatement selelctStmt = (SQLSelectStatement)parseSelect(); + stmt.setSelect(selelctStmt); accept(Token.SEMI); 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 15df4346b9..9546b7142b 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 @@ -1629,6 +1629,13 @@ public class MySqlOutputVisitor extends SQLASTOutputVisitor implements MySqlASTV @Override public boolean visit(MySqlUpdateStatement x) { + if (x.getReturning() != null && x.getReturning().size() > 0) { + print0(ucase ? "SELECT " : "select "); + printAndAccept(x.getReturning(), ", "); + println(); + print0(ucase ? "FROM " : "from "); + } + print0(ucase ? "UPDATE " : "update "); if (x.isLowPriority()) { diff --git a/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java b/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java index 7a9031e8bd..fce98f8cb9 100644 --- a/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java +++ b/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java @@ -372,9 +372,9 @@ public class SQLSelectParser extends SQLParser { if (lexer.token() != Token.FROM) { return; } - + lexer.nextToken(); - + queryBlock.setFrom(parseTableSource()); } diff --git a/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java b/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java index 89de536de4..48b6e67b73 100644 --- a/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java +++ b/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java @@ -1766,7 +1766,7 @@ public class SQLStatementParser extends SQLParser { return new SQLCreateTableParser(this.exprParser); } - public SQLSelectStatement parseSelect() { + public SQLStatement parseSelect() { SQLSelectParser selectParser = createSQLSelectParser(); SQLSelect select = selectParser.select(); return new SQLSelectStatement(select); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlSelectTest_update.java b/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlSelectTest_update.java index 84512ec83c..059df58729 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlSelectTest_update.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlSelectTest_update.java @@ -23,8 +23,9 @@ import com.alibaba.druid.sql.MysqlTest; import com.alibaba.druid.sql.SQLUtils; import com.alibaba.druid.sql.ast.SQLStatement; import com.alibaba.druid.sql.ast.statement.SQLSelect; -import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; +import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock; +import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUpdateStatement; import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser; import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlSchemaStatVisitor; @@ -41,14 +42,9 @@ public class MySqlSelectTest_update extends MysqlTest { List statementList = parser.parseStatementList(); SQLStatement stmt = statementList.get(0); - SQLSelectStatement selectStmt = (SQLSelectStatement) stmt; + MySqlUpdateStatement updateStmt = (MySqlUpdateStatement) stmt; - SQLSelect select = selectStmt.getSelect(); - Assert.assertNotNull(select.getQuery()); - MySqlSelectQueryBlock queryBlock = (MySqlSelectQueryBlock) select.getQuery(); - Assert.assertNull(queryBlock.getOrderBy()); - -// print(statementList); + print(statementList); Assert.assertEquals(1, statementList.size()); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlUpdateTest_5.java b/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlUpdateTest_5.java new file mode 100644 index 0000000000..40fe83acd4 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlUpdateTest_5.java @@ -0,0 +1,76 @@ +/* + * 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_5 extends MysqlTest { + + public void test_0() throws Exception { + String sql = "select num from update ttt set num=num+1 where id = 1"; + + 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(2, visitor.getColumns().size()); + Assert.assertEquals(1, visitor.getConditions().size()); + + Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("ttt"))); + + Assert.assertTrue(visitor.getColumns().contains(new Column("ttt", "num"))); + Assert.assertTrue(visitor.getColumns().contains(new Column("ttt", "id"))); + + { + String output = SQLUtils.toMySqlString(stmt); + Assert.assertEquals("SELECT num" + + "\nFROM UPDATE ttt" + + "\nSET num = num + 1" + + "\nWHERE id = 1", // + output); + } + { + String output = SQLUtils.toMySqlString(stmt, SQLUtils.DEFAULT_LCASE_FORMAT_OPTION); + Assert.assertEquals("select num" + + "\nfrom update ttt" + + "\nset num = num + 1" + + "\nwhere id = 1", // + output); + } + } +}