bug fixed for schema stat

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

@@ -27,10 +27,10 @@ public class SQLUpdateStatement extends SQLStatementImpl {
 
     private static final long            serialVersionUID = 1L;
 
-    private final List<SQLUpdateSetItem> items            = new ArrayList<SQLUpdateSetItem>();
-    private SQLExpr                      where;
+    protected final List<SQLUpdateSetItem> items            = new ArrayList<SQLUpdateSetItem>();
+    protected SQLExpr                      where;
 
-    private SQLTableSource               tableSource;
+    protected SQLTableSource               tableSource;
 
     public SQLUpdateStatement(){
 

src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlExprParser.java

@@ -385,7 +385,7 @@ public class MySqlExprParser extends SQLExprParser {
         return super.primaryRest(expr);
     }
 
-    protected SQLSelectParser createSelectParser() {
+    public SQLSelectParser createSelectParser() {
         return new MySqlSelectParser(this.lexer);
     }
 

src/main/java/com/alibaba/druid/sql/dialect/oracle/parser/OracleExprParser.java

@@ -478,7 +478,7 @@ public class OracleExprParser extends SQLExprParser {
     }
 
     @Override
-    protected OracleSelectParser createSelectParser() {
+    public OracleSelectParser createSelectParser() {
         return new OracleSelectParser(this.lexer);
     }
 

src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/stmt/PGUpdateStatement.java

@@ -0,0 +1,72 @@
+package com.alibaba.druid.sql.dialect.postgresql.ast.stmt;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import com.alibaba.druid.sql.ast.SQLExpr;
+import com.alibaba.druid.sql.ast.statement.SQLTableSource;
+import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement;
+import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithClause;
+import com.alibaba.druid.sql.dialect.postgresql.visitor.PGASTVisitor;
+import com.alibaba.druid.sql.visitor.SQLASTVisitor;
+
+public class PGUpdateStatement extends SQLUpdateStatement implements PGSQLStatement {
+
+    private static final long serialVersionUID = 1L;
+
+    private boolean           only             = false;
+
+    private PGWithClause      with;
+
+    private List<SQLExpr>     returning        = new ArrayList<SQLExpr>(2);
+
+    private SQLTableSource    from;
+
+    public SQLTableSource getFrom() {
+        return from;
+    }
+
+    public void setFrom(SQLTableSource from) {
+        this.from = from;
+    }
+
+    public boolean isOnly() {
+        return only;
+    }
+
+    public void setOnly(boolean only) {
+        this.only = only;
+    }
+
+    public List<SQLExpr> getReturning() {
+        return returning;
+    }
+
+    public void setReturning(List<SQLExpr> returning) {
+        this.returning = returning;
+    }
+
+    public PGWithClause getWith() {
+        return with;
+    }
+
+    public void setWith(PGWithClause with) {
+        this.with = with;
+    }
+
+    protected void accept0(SQLASTVisitor visitor) {
+        accept0((PGASTVisitor) visitor);
+    }
+
+    @Override
+    public void accept0(PGASTVisitor visitor) {
+        if (visitor.visit(this)) {
+            acceptChild(visitor, with);
+            acceptChild(visitor, tableSource);
+            acceptChild(visitor, items);
+            acceptChild(visitor, where);
+        }
+        visitor.endVisit(this);
+    }
+
+}

src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSQLStatementParser.java

@@ -7,12 +7,16 @@ import com.alibaba.druid.sql.ast.SQLName;
 import com.alibaba.druid.sql.ast.SQLStatement;
 import com.alibaba.druid.sql.ast.expr.SQLQueryExpr;
 import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
+import com.alibaba.druid.sql.ast.statement.SQLTableSource;
+import com.alibaba.druid.sql.ast.statement.SQLUpdateSetItem;
+import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement;
 import com.alibaba.druid.sql.dialect.postgresql.ast.PGCurrentOfExpr;
 import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithClause;
 import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithQuery;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGDeleteStatement;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGInsertStatement;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGTruncateStatement;
+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;
@@ -32,6 +36,53 @@ public class PGSQLStatementParser extends SQLStatementParser {
     public PGSelectParser createSQLSelectParser() {
         return new PGSelectParser(this.lexer);
     }
+    
+    public SQLUpdateStatement parseUpdateStatement() throws ParserException {
+        accept(Token.UPDATE);
+
+        PGUpdateStatement udpateStatement = new PGUpdateStatement();
+
+        SQLTableSource tableSource = this.exprParser.createSelectParser().parseTableSource();
+        udpateStatement.setTableSource(tableSource);
+
+        accept(Token.SET);
+
+        for (;;) {
+            SQLUpdateSetItem item = new SQLUpdateSetItem();
+            item.setColumn(this.exprParser.name());
+            accept(Token.EQ);
+            item.setValue(this.exprParser.expr());
+
+            udpateStatement.getItems().add(item);
+
+            if (lexer.token() == (Token.COMMA)) {
+                lexer.nextToken();
+                continue;
+            }
+
+            break;
+        }
+
+        if (lexer.token() == (Token.WHERE)) {
+            lexer.nextToken();
+            udpateStatement.setWhere(this.exprParser.expr());
+        }
+        
+        if (lexer.token() == Token.RETURNING) {
+            lexer.nextToken();
+            
+            for (;;) {
+                udpateStatement.getReturning().add(this.exprParser.expr());
+                if (lexer.token() == Token.COMMA) {
+                    lexer.nextToken();
+                    continue;
+                }
+                break;
+            }
+        }
+
+        return udpateStatement;
+    }
 
     public PGInsertStatement parseInsert() {
         accept(Token.INSERT);

src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGASTVisitor.java

@@ -8,6 +8,7 @@ import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGInsertStatement;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectStatement;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGTruncateStatement;
+import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGUpdateStatement;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
 public interface PGASTVisitor extends SQLASTVisitor {
@@ -55,4 +56,8 @@ public interface PGASTVisitor extends SQLASTVisitor {
     void endVisit(PGSelectStatement x);
     
     boolean visit(PGSelectStatement x);
+    
+    void endVisit(PGUpdateStatement x);
+    
+    boolean visit(PGUpdateStatement x);
 }

src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGASTVisitorAdapter.java

@@ -1,16 +1,17 @@
 package com.alibaba.druid.sql.dialect.postgresql.visitor;
 
 import com.alibaba.druid.sql.dialect.postgresql.ast.PGCurrentOfExpr;
+import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithClause;
+import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithQuery;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGDeleteStatement;
 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.PGTruncateStatement;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.FetchClause;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.ForClause;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.IntoClause;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.WindowClause;
-import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithClause;
-import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithQuery;
+import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectStatement;
+import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGTruncateStatement;
+import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGUpdateStatement;
 import com.alibaba.druid.sql.visitor.SQLASTVisitorAdapter;
 
 public class PGASTVisitorAdapter extends SQLASTVisitorAdapter implements PGASTVisitor {
@@ -132,5 +133,15 @@ public class PGASTVisitorAdapter extends SQLASTVisitorAdapter implements PGASTVi
     public boolean visit(PGSelectStatement x) {
         return true;
     }
+    
+    @Override
+    public void endVisit(PGUpdateStatement x) {
+        
+    }
+    
+    @Override
+    public boolean visit(PGUpdateStatement x) {
+        return true;
+    }
 
 }

src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java

@@ -4,17 +4,18 @@ import com.alibaba.druid.sql.ast.SQLSetQuantifier;
 import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock;
 import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
 import com.alibaba.druid.sql.dialect.postgresql.ast.PGCurrentOfExpr;
+import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithClause;
+import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithQuery;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGDeleteStatement;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGInsertStatement;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock;
-import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectStatement;
-import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGTruncateStatement;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.FetchClause;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.ForClause;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.IntoClause;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.WindowClause;
-import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithClause;
-import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithQuery;
+import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectStatement;
+import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGTruncateStatement;
+import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGUpdateStatement;
 import com.alibaba.druid.sql.visitor.SQLASTOutputVisitor;
 
 public class PGOutputVisitor extends SQLASTOutputVisitor implements PGASTVisitor {
@@ -395,4 +396,54 @@ public class PGOutputVisitor extends SQLASTOutputVisitor implements PGASTVisitor
         
         return visit((SQLSelectStatement) x);
     }
+    
+    @Override
+    public void endVisit(PGUpdateStatement x) {
+        
+    }
+    
+    @Override
+    public boolean visit(PGUpdateStatement x) {
+        if (x.getWith() != null) {
+            x.getWith().accept(this);
+            println();
+        }
+        
+        print("UPDATE ");
+        
+        if (x.isOnly()) {
+            print("ONLY ");
+        }
+
+        x.getTableSource().accept(this);
+
+        println();
+        print("SET ");
+        for (int i = 0, size = x.getItems().size(); i < size; ++i) {
+            if (i != 0) {
+                print(", ");
+            }
+            x.getItems().get(i).accept(this);
+        }
+        
+        if (x.getFrom() != null) {
+            println();
+            print("FROM ");
+            x.getFrom().accept(this);
+        }
+
+        if (x.getWhere() != null) {
+            println();
+            print("WHERE ");
+            x.getWhere().accept(this);
+        }
+        
+        if (x.getReturning().size() > 0) {
+            println();
+            print("RETURNING ");
+            printAndAccept(x.getReturning(), ", ");
+        }
+
+        return false;
+    }
 }

src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGSchemaStatVisitor.java

@@ -5,20 +5,20 @@ import java.util.Map;
 
 import com.alibaba.druid.sql.ast.SQLName;
 import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
-import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
 import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
 import com.alibaba.druid.sql.ast.statement.SQLTruncateStatement;
 import com.alibaba.druid.sql.dialect.postgresql.ast.PGCurrentOfExpr;
+import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithClause;
+import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithQuery;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGDeleteStatement;
 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.PGTruncateStatement;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.FetchClause;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.ForClause;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.IntoClause;
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectQueryBlock.WindowClause;
-import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithClause;
-import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithQuery;
+import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGSelectStatement;
+import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGTruncateStatement;
+import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.PGUpdateStatement;
 import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
 import com.alibaba.druid.stat.TableStat;
 import com.alibaba.druid.stat.TableStat.Mode;
@@ -63,8 +63,19 @@ public class PGSchemaStatVisitor extends SchemaStatVisitor implements PGASTVisit
 
     @Override
     public boolean visit(PGWithQuery x) {
+        Map<String, String> aliasMap = aliasLocal.get();
+        if (aliasMap != null) {
+            String alias = null;
+            if (x.getName() != null) {
+                alias = x.getName().toString();
+            }
 
-        return true;
+            if (alias != null) {
+                aliasMap.put(alias, null);
+            }
+        }
+        x.getQuery().accept(this);
+        return false;
     }
 
     @Override
@@ -74,7 +85,6 @@ public class PGSchemaStatVisitor extends SchemaStatVisitor implements PGASTVisit
 
     @Override
     public boolean visit(PGWithClause x) {
-
         return true;
     }
 
@@ -117,9 +127,9 @@ public class PGSchemaStatVisitor extends SchemaStatVisitor implements PGASTVisit
         if (x.getWith() != null) {
             x.getWith().accept(this);
         }
-        
+
         aliasLocal.set(new HashMap<String, String>());
-        
+
         for (SQLName name : x.getUsing()) {
             String ident = name.toString();
 
@@ -135,7 +145,6 @@ public class PGSchemaStatVisitor extends SchemaStatVisitor implements PGASTVisit
                 aliasMap.put(ident, ident);
             }
         }
-        
 
         x.putAttribute("_original_use_mode", modeLocal.get());
         modeLocal.set(Mode.Delete);
@@ -160,7 +169,7 @@ public class PGSchemaStatVisitor extends SchemaStatVisitor implements PGASTVisit
 
     @Override
     public void endVisit(PGCurrentOfExpr x) {
-        
+
     }
 
     @Override
@@ -170,28 +179,96 @@ public class PGSchemaStatVisitor extends SchemaStatVisitor implements PGASTVisit
 
     @Override
     public void endVisit(PGInsertStatement x) {
-        
+
     }
 
     @Override
     public boolean visit(PGInsertStatement x) {
+        aliasLocal.set(new HashMap<String, String>());
+
         if (x.getWith() != null) {
             x.getWith().accept(this);
         }
-        
-        return visit((SQLInsertStatement) x);
+
+        x.putAttribute("_original_use_mode", modeLocal.get());
+        modeLocal.set(Mode.Insert);
+
+        String originalTable = currentTableLocal.get();
+
+        if (x.getTableName() instanceof SQLName) {
+            String ident = ((SQLName) x.getTableName()).toString();
+            currentTableLocal.set(ident);
+            x.putAttribute("_old_local_", originalTable);
+
+            TableStat stat = tableStats.get(ident);
+            if (stat == null) {
+                stat = new TableStat();
+                tableStats.put(new TableStat.Name(ident), stat);
+            }
+            stat.incrementInsertCount();
+
+            Map<String, String> aliasMap = aliasLocal.get();
+            if (aliasMap != null) {
+                if (x.getAlias() != null) {
+                    aliasMap.put(x.getAlias(), ident);
+                }
+                aliasMap.put(ident, ident);
+            }
+        }
+
+        accept(x.getColumns());
+        accept(x.getQuery());
+
+        return false;
     }
+
     @Override
     public void endVisit(PGSelectStatement x) {
-        
+
     }
-    
+
     @Override
     public boolean visit(PGSelectStatement x) {
         if (x.getWith() != null) {
             x.getWith().accept(this);
         }
-        
+
         return visit((SQLSelectStatement) x);
     }
+
+    @Override
+    public void endVisit(PGUpdateStatement x) {
+
+    }
+
+    @Override
+    public boolean visit(PGUpdateStatement x) {
+        Map<String, String> oldAliasMap = aliasLocal.get();
+        
+        aliasLocal.set(new HashMap<String, String>());
+
+        if (x.getWith() != null) {
+            x.getWith().accept(this);
+        }
+
+        String ident = x.getTableName().toString();
+        currentTableLocal.set(ident);
+
+        TableStat stat = tableStats.get(ident);
+        if (stat == null) {
+            stat = new TableStat();
+            tableStats.put(new TableStat.Name(ident), stat);
+        }
+        stat.incrementUpdateCount();
+
+        Map<String, String> aliasMap = aliasLocal.get();
+        aliasMap.put(ident, ident);
+
+        accept(x.getItems());
+        accept(x.getWhere());
+        
+        aliasLocal.set(oldAliasMap);
+
+        return false;
+    }
 }

src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java

@@ -418,7 +418,7 @@ public class SQLExprParser extends SQLParser {
         throw new ParserException("TODO");
     }
 
-    protected SQLSelectParser createSelectParser() {
+    public SQLSelectParser createSelectParser() {
         return new SQLSelectParser(lexer);
     }
 

src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java

@@ -162,8 +162,14 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter {
             SQLExprTableSource tableSource = (SQLExprTableSource) x.getFrom();
             if (tableSource.getExpr() instanceof SQLName) {
                 String ident = tableSource.getExpr().toString();
-                currentTableLocal.set(ident);
-                x.putAttribute("_old_local_", originalTable);
+                
+                Map<String, String> aliasMap = aliasLocal.get();
+                if (aliasMap.containsKey(ident) && aliasMap.get(ident) == null) {
+                    
+                } else {
+                    currentTableLocal.set(ident);
+                    x.putAttribute("_old_local_", originalTable);
+                }
             }
         }
 
@@ -260,6 +266,13 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter {
     public boolean visit(SQLExprTableSource x) {
         if (isSimpleExprTableSource(x)) {
             String ident = x.getExpr().toString();
+            
+            Map<String, String> aliasMap = aliasLocal.get();
+            
+            if (aliasMap.containsKey(ident) && aliasMap.get(ident) == null) {
+                return false;
+            }
+            
             TableStat stat = tableStats.get(ident);
             if (stat == null) {
                 stat = new TableStat();
@@ -284,7 +297,6 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter {
                     break;
             }
 
-            Map<String, String> aliasMap = aliasLocal.get();
             if (aliasMap != null) {
                 if (x.getAlias() != null) {
                     aliasMap.put(x.getAlias(), ident);