diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlUpdateTableSource.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlUpdateTableSource.java new file mode 100644 index 0000000000..404079863f --- /dev/null +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/ast/statement/MySqlUpdateTableSource.java @@ -0,0 +1,55 @@ +/* + * 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.sql.dialect.mysql.ast.statement; + +import com.alibaba.druid.sql.ast.statement.SQLTableSourceImpl; +import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitor; +import com.alibaba.druid.sql.visitor.SQLASTVisitor; + +// just for alibaba mysql +public class MySqlUpdateTableSource extends SQLTableSourceImpl { + + private MySqlUpdateStatement update; + + public MySqlUpdateTableSource(MySqlUpdateStatement update){ + this.update = update; + } + + @Override + protected void accept0(SQLASTVisitor visitor) { + if (visitor instanceof MySqlASTVisitor) { + accept0((MySqlASTVisitor) visitor); + } else { + throw new IllegalArgumentException("not support visitor type : " + visitor.getClass().getName()); + } + } + + public void accept0(MySqlASTVisitor visitor) { + if (visitor.visit(this)) { + acceptChild(visitor, update); + } + visitor.endVisit(this); + } + + public MySqlUpdateStatement getUpdate() { + return update; + } + + public void setUpdate(MySqlUpdateStatement update) { + this.update = update; + } + +} 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 accba3cd56..e6b003ae99 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 @@ -19,11 +19,16 @@ 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; +import com.alibaba.druid.sql.ast.statement.SQLExprTableSource; +import com.alibaba.druid.sql.ast.statement.SQLSelect; import com.alibaba.druid.sql.ast.statement.SQLSelectGroupByClause; import com.alibaba.druid.sql.ast.statement.SQLSelectQuery; import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock; +import com.alibaba.druid.sql.ast.statement.SQLSubqueryTableSource; import com.alibaba.druid.sql.ast.statement.SQLTableSource; import com.alibaba.druid.sql.ast.statement.SQLUnionQuery; +import com.alibaba.druid.sql.ast.statement.SQLUnionQueryTableSource; +import com.alibaba.druid.sql.ast.statement.SQLUpdateSetItem; import com.alibaba.druid.sql.dialect.mysql.ast.MySqlForceIndexHint; import com.alibaba.druid.sql.dialect.mysql.ast.MySqlIgnoreIndexHint; import com.alibaba.druid.sql.dialect.mysql.ast.MySqlIndexHint; @@ -32,8 +37,10 @@ import com.alibaba.druid.sql.dialect.mysql.ast.MySqlUseIndexHint; import com.alibaba.druid.sql.dialect.mysql.ast.expr.MySqlOutFileExpr; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectGroupBy; 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.ast.statement.MySqlSelectQueryBlock.Limit; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUnionQuery; +import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUpdateTableSource; import com.alibaba.druid.sql.parser.ParserException; import com.alibaba.druid.sql.parser.SQLExprParser; import com.alibaba.druid.sql.parser.SQLSelectParser; @@ -166,6 +173,94 @@ public class MySqlSelectParser extends SQLSelectParser { return queryRest(queryBlock); } + public SQLTableSource parseTableSource() { + if (lexer.token() == Token.LPAREN) { + lexer.nextToken(); + SQLTableSource tableSource; + if (lexer.token() == Token.SELECT || lexer.token() == Token.WITH) { + SQLSelect select = select(); + accept(Token.RPAREN); + SQLSelectQuery query = queryRest(select.getQuery()); + if (query instanceof SQLUnionQuery) { + tableSource = new SQLUnionQueryTableSource((SQLUnionQuery) query); + } else { + tableSource = new SQLSubqueryTableSource(select); + } + } else if (lexer.token() == Token.LPAREN) { + tableSource = parseTableSource(); + accept(Token.RPAREN); + } else { + tableSource = parseTableSource(); + accept(Token.RPAREN); + } + + return parseTableSourceRest(tableSource); + } + + if(lexer.token() == Token.UPDATE) { + SQLTableSource tableSource = new MySqlUpdateTableSource(parseUpdateStatment()); + return parseTableSourceRest(tableSource); + } + + if (lexer.token() == Token.SELECT) { + throw new ParserException("TODO"); + } + + SQLExprTableSource tableReference = new SQLExprTableSource(); + + parseTableSourceQueryTableExpr(tableReference); + + SQLTableSource tableSrc = parseTableSourceRest(tableReference); + + if (lexer.hasComment() && lexer.isKeepComments()) { + tableSrc.addAfterComment(lexer.readAndResetComments()); + } + + return tableSrc; + } + + private MySqlUpdateStatement parseUpdateStatment() { + MySqlUpdateStatement update = new MySqlUpdateStatement(); + + lexer.nextToken(); + + if (identifierEquals("LOW_PRIORITY")) { + lexer.nextToken(); + update.setLowPriority(true); + } + + if (identifierEquals("IGNORE")) { + lexer.nextToken(); + update.setIgnore(true); + } + + SQLTableSource updateTableSource = this.exprParser.createSelectParser().parseTableSource(); + update.setTableSource(updateTableSource); + + accept(Token.SET); + + for (;;) { + SQLUpdateSetItem item = this.exprParser.parseUpdateSetItem(); + update.addItem(item); + + if (lexer.token() != Token.COMMA) { + break; + } + + lexer.nextToken(); + } + + if (lexer.token() == (Token.WHERE)) { + lexer.nextToken(); + update.setWhere(this.exprParser.expr()); + } + + update.setOrderBy(this.exprParser.parseOrderBy()); + update.setLimit(parseLimit()); + + return update; + } + protected void parseInto(SQLSelectQueryBlock queryBlock) { if (lexer.token() == (Token.INTO)) { lexer.nextToken(); diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlASTVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlASTVisitor.java index 74e7918cdb..fc2c334ba5 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlASTVisitor.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlASTVisitor.java @@ -136,6 +136,7 @@ import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlTableIndex; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUnionQuery; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUnlockTablesStatement; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUpdateStatement; +import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUpdateTableSource; import com.alibaba.druid.sql.visitor.SQLASTVisitor; public interface MySqlASTVisitor extends SQLASTVisitor { @@ -650,5 +651,9 @@ public interface MySqlASTVisitor extends SQLASTVisitor { boolean visit(MySqlCursorDeclareStatement x); void endVisit(MySqlCursorDeclareStatement x); + + boolean visit(MySqlUpdateTableSource x); + + void endVisit(MySqlUpdateTableSource x); } // diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlASTVisitorAdapter.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlASTVisitorAdapter.java index 6fd4e3964d..3c129826cc 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlASTVisitorAdapter.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlASTVisitorAdapter.java @@ -141,6 +141,7 @@ import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlTableIndex; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUnionQuery; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUnlockTablesStatement; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUpdateStatement; +import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUpdateTableSource; import com.alibaba.druid.sql.visitor.SQLASTVisitorAdapter; public class MySqlASTVisitorAdapter extends SQLASTVisitorAdapter implements MySqlASTVisitor { @@ -1466,5 +1467,15 @@ public class MySqlASTVisitorAdapter extends SQLASTVisitorAdapter implements MySq // TODO Auto-generated method stub } + + @Override + public boolean visit(MySqlUpdateTableSource x) { + return true; + } + + @Override + public void endVisit(MySqlUpdateTableSource x) { + + } } // 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 12e0b106ab..c59ed17cd7 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 @@ -168,6 +168,7 @@ import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlTableIndex; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUnionQuery; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUnlockTablesStatement; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUpdateStatement; +import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUpdateTableSource; import com.alibaba.druid.sql.visitor.SQLASTOutputVisitor; public class MySqlOutputVisitor extends SQLASTOutputVisitor implements MySqlASTVisitor { @@ -3585,4 +3586,17 @@ public class MySqlOutputVisitor extends SQLASTOutputVisitor implements MySqlASTV } + @Override + public boolean visit(MySqlUpdateTableSource x) { + MySqlUpdateStatement update = x.getUpdate(); + if (update != null) { + update.accept0(this); + } + return false; + } + + @Override + public void endVisit(MySqlUpdateTableSource x) { + + } } // 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 5c261e8f21..444234d994 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 @@ -156,6 +156,7 @@ import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlTableIndex; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUnionQuery; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUnlockTablesStatement; import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUpdateStatement; +import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlUpdateTableSource; import com.alibaba.druid.sql.visitor.SchemaStatVisitor; import com.alibaba.druid.stat.TableStat; import com.alibaba.druid.stat.TableStat.Mode; @@ -1570,5 +1571,18 @@ public class MySqlSchemaStatVisitor extends SchemaStatVisitor implements MySqlAS // TODO Auto-generated method stub } + + @Override + public boolean visit(MySqlUpdateTableSource x) { + if (x.getUpdate() != null) { + return this.visit(x.getUpdate()); + } + return false; + } + + @Override + public void endVisit(MySqlUpdateTableSource x) { + + } } 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 b1be24a49e..e450cd9237 100644 --- a/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java +++ b/src/main/java/com/alibaba/druid/sql/parser/SQLSelectParser.java @@ -355,7 +355,7 @@ public class SQLSelectParser extends SQLParser { return tableSrc; } - private void parseTableSourceQueryTableExpr(SQLExprTableSource tableReference) { + protected void parseTableSourceQueryTableExpr(SQLExprTableSource tableReference) { if (lexer.token() == Token.LITERAL_ALIAS || lexer.token() == Token.IDENTIFIED || lexer.token() == Token.LITERAL_CHARS) { tableReference.setExpr(this.exprParser.name()); 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 new file mode 100644 index 0000000000..833078c922 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlSelectTest_update.java @@ -0,0 +1,81 @@ +/* + * 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.ast.statement.SQLSelect; +import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; +import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock; +import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser; +import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlSchemaStatVisitor; + +public class MySqlSelectTest_update extends MysqlTest { + + public void test_0() throws Exception { + String sql = "select current_no from " + + "update wlb_waybill_branch_rule set current_no = current_no + ? ,gmt_modified = now() " + + "where id = ? and status = ? and end_no > current_no " + + "order by current_no desc " + + "limit 10"; + + MySqlStatementParser parser = new MySqlStatementParser(sql); + List statementList = parser.parseStatementList(); + SQLStatement stmt = statementList.get(0); + + SQLSelectStatement selectStmt = (SQLSelectStatement) stmt; + + SQLSelect select = selectStmt.getSelect(); + Assert.assertNotNull(select.getQuery()); + MySqlSelectQueryBlock queryBlock = (MySqlSelectQueryBlock) select.getQuery(); + Assert.assertNull(queryBlock.getOrderBy()); + + print(statementList); + + Assert.assertEquals(1, statementList.size()); + + String output = SQLUtils.toMySqlString(stmt); + + String expected = "SELECT current_no" + + "\nFROM UPDATE wlb_waybill_branch_rule" + + "\nSET current_no = current_no + ?, gmt_modified = now()" + + "\nWHERE id = ?" + + "\n\tAND status = ?" + + "\n\tAND end_no > current_no" + + "\nORDER BY current_no DESC" + + "\nLIMIT 10"; + + Assert.assertEquals(expected, output); + + 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(5, visitor.getColumns().size()); + Assert.assertEquals(4, visitor.getConditions().size()); + Assert.assertEquals(0, visitor.getOrderByColumns().size()); + } +}