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 716a243f3c..c07b59b416 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 @@ -27,6 +27,7 @@ public class SQLUpdateStatement extends SQLStatementImpl { protected final List items = new ArrayList(); protected SQLExpr where; + protected SQLTableSource from; protected SQLTableSource tableSource; @@ -81,6 +82,17 @@ public class SQLUpdateStatement extends SQLStatementImpl { item.setParent(this); } + public SQLTableSource getFrom() { + return from; + } + + public void setFrom(SQLTableSource from) { + if (from != null) { + from.setParent(this); + } + this.from = from; + } + @Override public void output(StringBuffer buf) { buf.append("UPDATE "); diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/stmt/PGUpdateStatement.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/stmt/PGUpdateStatement.java index f4d1aa8ed5..d46b86abc8 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/stmt/PGUpdateStatement.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/stmt/PGUpdateStatement.java @@ -34,20 +34,12 @@ public class PGUpdateStatement extends SQLUpdateStatement implements PGSQLStatem private List returning = new ArrayList(2); - private SQLTableSource from; + public PGUpdateStatement(){ super (JdbcConstants.POSTGRESQL); } - public SQLTableSource getFrom() { - return from; - } - - public void setFrom(SQLTableSource from) { - this.from = from; - } - public boolean isOnly() { return only; } diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSQLStatementParser.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSQLStatementParser.java index 7f13a269fa..d2070ec03e 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSQLStatementParser.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSQLStatementParser.java @@ -36,10 +36,7 @@ import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGInsertStatement; import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectStatement; import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGShowStatement; import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGUpdateStatement; -import com.alibaba.druid.sql.parser.Lexer; -import com.alibaba.druid.sql.parser.ParserException; -import com.alibaba.druid.sql.parser.SQLStatementParser; -import com.alibaba.druid.sql.parser.Token; +import com.alibaba.druid.sql.parser.*; public class PGSQLStatementParser extends SQLStatementParser { @@ -60,11 +57,18 @@ public class PGSQLStatementParser extends SQLStatementParser { PGUpdateStatement udpateStatement = new PGUpdateStatement(); - SQLTableSource tableSource = this.exprParser.createSelectParser().parseTableSource(); + SQLSelectParser selectParser = this.exprParser.createSelectParser(); + SQLTableSource tableSource = selectParser.parseTableSource(); udpateStatement.setTableSource(tableSource); parseUpdateSet(udpateStatement); + if (lexer.token() == Token.FROM) { + lexer.nextToken(); + SQLTableSource from = selectParser.parseTableSource(); + udpateStatement.setFrom(from); + } + if (lexer.token() == (Token.WHERE)) { lexer.nextToken(); udpateStatement.setWhere(this.exprParser.expr()); diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGSchemaStatVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGSchemaStatVisitor.java index db265b434e..78bde57eb7 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGSchemaStatVisitor.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGSchemaStatVisitor.java @@ -239,6 +239,8 @@ public class PGSchemaStatVisitor extends SchemaStatVisitor implements PGASTVisit TableStat stat = getTableStat(ident); stat.incrementUpdateCount(); + accept(x.getFrom()); + Map aliasMap = getAliasMap(); aliasMap.put(ident, ident); diff --git a/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/ast/stmt/SQLServerUpdateStatement.java b/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/ast/stmt/SQLServerUpdateStatement.java index dff14515fc..bc6e2492e0 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/ast/stmt/SQLServerUpdateStatement.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/sqlserver/ast/stmt/SQLServerUpdateStatement.java @@ -27,7 +27,6 @@ import com.alibaba.druid.util.JdbcConstants; public class SQLServerUpdateStatement extends SQLUpdateStatement implements SQLServerStatement { private SQLServerTop top; - private SQLTableSource from; private SQLServerOutput output; public SQLServerUpdateStatement(){ @@ -42,14 +41,6 @@ public class SQLServerUpdateStatement extends SQLUpdateStatement implements SQLS this.top = top; } - public SQLTableSource getFrom() { - return from; - } - - public void setFrom(SQLTableSource from) { - this.from = from; - } - public SQLServerOutput getOutput() { return output; } diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/select/PGSelectTest42.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/select/PGSelectTest42.java new file mode 100644 index 0000000000..0b76c21ff9 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/select/PGSelectTest42.java @@ -0,0 +1,60 @@ +/* + * 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.postgresql.select; + +import com.alibaba.druid.sql.PGTest; +import com.alibaba.druid.sql.SQLUtils; +import com.alibaba.druid.sql.ast.SQLStatement; +import com.alibaba.druid.sql.dialect.postgresql.parser.PGSQLStatementParser; +import com.alibaba.druid.sql.dialect.postgresql.visitor.PGSchemaStatVisitor; +import org.junit.Assert; + +import java.util.List; + +public class PGSelectTest42 extends PGTest { + + public void test_0() throws Exception { + String sql = "UPDATE sys_account AS a SET online = 2 FROM auto_handler_online o WHERE a.id = o.account_id AND a.online != 2"; + + PGSQLStatementParser parser = new PGSQLStatementParser(sql); + List statementList = parser.parseStatementList(); + SQLStatement stmt = statementList.get(0); + + Assert.assertEquals("UPDATE sys_account\n" + + "SET online = 2\n" + + "FROM auto_handler_online o\n" + + "WHERE a.id = o.account_id\n" + + "\tAND a.online != 2", SQLUtils.toPGString(stmt)); + + Assert.assertEquals("update sys_account\n" + + "set online = 2\n" + + "from auto_handler_online o\n" + + "where a.id = o.account_id\n" + + "\tand a.online != 2", SQLUtils.toPGString(stmt, SQLUtils.DEFAULT_LCASE_FORMAT_OPTION)); + + Assert.assertEquals(1, statementList.size()); + + PGSchemaStatVisitor visitor = new PGSchemaStatVisitor(); + stmt.accept(visitor); + +// System.out.println("Tables : " + visitor.getTables()); +// System.out.println("fields : " + visitor.getColumns()); +// System.out.println("coditions : " + visitor.getConditions()); + + Assert.assertEquals(4, visitor.getColumns().size()); + Assert.assertEquals(2, visitor.getTables().size()); + } +}