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 e94afbee0e..9a036e85ce 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,10 +27,10 @@ public class SQLUpdateStatement extends SQLStatementImpl { private static final long serialVersionUID = 1L; - private final List items = new ArrayList(); - private SQLExpr where; + protected final List items = new ArrayList(); + protected SQLExpr where; - private SQLTableSource tableSource; + protected SQLTableSource tableSource; public SQLUpdateStatement(){ diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlExprParser.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlExprParser.java index 2e7825a89e..64fe05a29f 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlExprParser.java +++ b/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); } diff --git a/src/main/java/com/alibaba/druid/sql/dialect/oracle/parser/OracleExprParser.java b/src/main/java/com/alibaba/druid/sql/dialect/oracle/parser/OracleExprParser.java index 75bf0d90c2..03ff6f7c7f 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/oracle/parser/OracleExprParser.java +++ b/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); } 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 new file mode 100644 index 0000000000..5e3e18b3d9 --- /dev/null +++ b/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 returning = new ArrayList(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 getReturning() { + return returning; + } + + public void setReturning(List 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); + } + +} 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 85a7522b4a..f27dbbf17f 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 @@ -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); diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGASTVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGASTVisitor.java index d9fe19dad1..4505a2b12e 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGASTVisitor.java +++ b/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); } diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGASTVisitorAdapter.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGASTVisitorAdapter.java index 45339130e5..7f75acbe8d 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGASTVisitorAdapter.java +++ b/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; + } } diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java index b8bbb91553..b95a30931a 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java +++ b/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; + } } 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 fc5958da26..360519b1be 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 @@ -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 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()); - + 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()); + 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 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 oldAliasMap = aliasLocal.get(); + + aliasLocal.set(new HashMap()); + + 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 aliasMap = aliasLocal.get(); + aliasMap.put(ident, ident); + + accept(x.getItems()); + accept(x.getWhere()); + + aliasLocal.set(oldAliasMap); + + return false; + } } diff --git a/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java b/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java index a04e47adeb..d661f45d4d 100644 --- a/src/main/java/com/alibaba/druid/sql/parser/SQLExprParser.java +++ b/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); } diff --git a/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java b/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java index 36a725e446..b14b43c2c0 100644 --- a/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java +++ b/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 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 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 aliasMap = aliasLocal.get(); if (aliasMap != null) { if (x.getAlias() != null) { aliasMap.put(x.getAlias(), ident);