improved sql parser, support pg udpate_from_clause. issue #1430

src/main/java/com/alibaba/druid/sql/ast/statement/SQLUpdateStatement.java

@@ -27,6 +27,7 @@ public class SQLUpdateStatement extends SQLStatementImpl {
 
     protected final List<SQLUpdateSetItem> items = new ArrayList<SQLUpdateSetItem>();
     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 ");

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<SQLExpr>  returning = new ArrayList<SQLExpr>(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;
     }

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());

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<String, String> aliasMap = getAliasMap();
         aliasMap.put(ident, ident);
 

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;
     }

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<SQLStatement> 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());
+    }
+}