diff --git a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLTruncateStatement.java b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLTruncateStatement.java new file mode 100644 index 0000000000..05d931b156 --- /dev/null +++ b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLTruncateStatement.java @@ -0,0 +1,29 @@ +package com.alibaba.druid.sql.ast.statement; + +import java.util.ArrayList; +import java.util.List; + +import com.alibaba.druid.sql.ast.SQLName; +import com.alibaba.druid.sql.ast.SQLStatementImpl; +import com.alibaba.druid.sql.visitor.SQLASTVisitor; + +public class SQLTruncateStatement extends SQLStatementImpl { + + private static final long serialVersionUID = 1L; + protected List tableNames = new ArrayList(2); + + public List getTableNames() { + return tableNames; + } + + public void setTableNames(List tableNames) { + this.tableNames = tableNames; + } + + @Override + protected void accept0(SQLASTVisitor visitor) { + if (visitor.visit(this)) { + acceptChild(visitor, tableNames); + } + } +} diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlLexer.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlLexer.java index 099681828d..9a8bb0b9fd 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlLexer.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlLexer.java @@ -103,6 +103,7 @@ public class MySqlLexer extends Lexer { map.put("LOCK", Token.LOCK); map.put("SOME", Token.SOME); map.put("ANY", Token.ANY); + map.put("TRUNCATE", Token.TRUNCATE); map.put("TRUE", Token.TRUE); map.put("FALSE", Token.FALSE); diff --git a/src/main/java/com/alibaba/druid/sql/dialect/oracle/parser/OracleLexer.java b/src/main/java/com/alibaba/druid/sql/dialect/oracle/parser/OracleLexer.java index 29136f2f5c..d32eb5e23c 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/oracle/parser/OracleLexer.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/oracle/parser/OracleLexer.java @@ -85,6 +85,7 @@ public class OracleLexer extends Lexer { map.put("LOCK", Token.LOCK); map.put("SOME", Token.SOME); map.put("ANY", Token.ANY); + map.put("TRUNCATE", Token.TRUNCATE); map.put("START", Token.START); map.put("CONNECT", Token.CONNECT); diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGSQLStatement.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGSQLStatement.java new file mode 100644 index 0000000000..621ba46a53 --- /dev/null +++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGSQLStatement.java @@ -0,0 +1,7 @@ +package com.alibaba.druid.sql.dialect.postgresql.ast; + +import com.alibaba.druid.sql.ast.SQLStatement; + +public interface PGSQLStatement extends SQLStatement, PGSQLObject { + +} diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGSQLStatementImpl.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGSQLStatementImpl.java new file mode 100644 index 0000000000..c6161c79ba --- /dev/null +++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGSQLStatementImpl.java @@ -0,0 +1,8 @@ +package com.alibaba.druid.sql.dialect.postgresql.ast; + + +public abstract class PGSQLStatementImpl extends PGSQLObjectImpl implements PGSQLStatement { + + private static final long serialVersionUID = 1L; + +} diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGTruncateStatement.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGTruncateStatement.java new file mode 100644 index 0000000000..97be5503db --- /dev/null +++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGTruncateStatement.java @@ -0,0 +1,43 @@ +package com.alibaba.druid.sql.dialect.postgresql.ast; + +import com.alibaba.druid.sql.ast.statement.SQLTruncateStatement; +import com.alibaba.druid.sql.dialect.postgresql.visitor.PGASTVisitor; + +public class PGTruncateStatement extends SQLTruncateStatement implements PGSQLStatement{ + + private static final long serialVersionUID = 1L; + + private boolean only; + private Boolean restartIdentity; + private Boolean cascade; + + public boolean isOnly() { + return only; + } + + public void setOnly(boolean only) { + this.only = only; + } + + public Boolean getRestartIdentity() { + return restartIdentity; + } + + public void setRestartIdentity(Boolean restartIdentity) { + this.restartIdentity = restartIdentity; + } + + public Boolean getCascade() { + return cascade; + } + + public void setCascade(Boolean cascade) { + this.cascade = cascade; + } + + public void accept0(PGASTVisitor visitor) { + if (visitor.visit(this)) { + acceptChild(visitor, tableNames); + } + } +} diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGLexer.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGLexer.java index fbe9b55cd4..39467e4477 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGLexer.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGLexer.java @@ -81,6 +81,7 @@ public class PGLexer extends Lexer { map.put("LOCK", Token.LOCK); map.put("SOME", Token.SOME); map.put("ANY", Token.ANY); + map.put("TRUNCATE", Token.TRUNCATE); map.put("LIMIT", Token.LIMIT); map.put("OFFSET", Token.OFFSET); 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 af27f91ca8..e276ccf6c3 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 @@ -2,7 +2,9 @@ package com.alibaba.druid.sql.dialect.postgresql.parser; import java.util.List; +import com.alibaba.druid.sql.ast.SQLName; import com.alibaba.druid.sql.ast.SQLStatement; +import com.alibaba.druid.sql.dialect.postgresql.ast.PGTruncateStatement; import com.alibaba.druid.sql.parser.Lexer; import com.alibaba.druid.sql.parser.ParserException; import com.alibaba.druid.sql.parser.SQLStatementParser; @@ -22,6 +24,53 @@ public class PGSQLStatementParser extends SQLStatementParser { return new PGSelectParser(this.lexer); } + public SQLStatement parseTruncate() { + accept(Token.TRUNCATE); + + PGTruncateStatement stmt = new PGTruncateStatement(); + + if (lexer.token() == Token.TABLE) { + lexer.nextToken(); + } + + if (lexer.token() == Token.ONLY) { + lexer.nextToken(); + stmt.setOnly(true); + } + + for (;;) { + SQLName name = this.exprParser.name(); + stmt.getTableNames().add(name); + + if (lexer.token() == Token.COMMA) { + lexer.nextToken(); + continue; + } + + break; + } + + if (lexer.token() == Token.RESTART) { + lexer.nextToken(); + accept(Token.IDENTITY); + stmt.setRestartIdentity(Boolean.TRUE); + } else if (lexer.token() == Token.SHARE) { + lexer.nextToken(); + accept(Token.IDENTITY); + stmt.setRestartIdentity(Boolean.FALSE); + } + + if (lexer.token() == Token.CASCADE) { + lexer.nextToken(); + stmt.setCascade(Boolean.TRUE); + } else if (lexer.token() == Token.RESTRICT) { + lexer.nextToken(); + stmt.setCascade(Boolean.FALSE); + } + + return stmt; + } + public boolean parseStatementListDialect(List statementList) { if (lexer.token() == Token.WITH) { SQLStatement stmt = parseSelect(); 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 8e03e23342..dc05dae3b5 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 @@ -1,30 +1,36 @@ package com.alibaba.druid.sql.dialect.postgresql.visitor; import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock; +import com.alibaba.druid.sql.dialect.postgresql.ast.PGTruncateStatement; import com.alibaba.druid.sql.visitor.SQLASTVisitor; public interface PGASTVisitor extends SQLASTVisitor { - void endVisit(PGSelectQueryBlock.WindowClause x); - boolean visit(PGSelectQueryBlock.WindowClause x); + void endVisit(PGSelectQueryBlock.WindowClause x); - void endVisit(PGSelectQueryBlock.FetchClause x); + boolean visit(PGSelectQueryBlock.WindowClause x); - boolean visit(PGSelectQueryBlock.FetchClause x); + void endVisit(PGSelectQueryBlock.FetchClause x); - void endVisit(PGSelectQueryBlock.ForClause x); + boolean visit(PGSelectQueryBlock.FetchClause x); - boolean visit(PGSelectQueryBlock.ForClause x); - - void endVisit(PGSelectQueryBlock.WithQuery x); - - boolean visit(PGSelectQueryBlock.WithQuery x); - - void endVisit(PGSelectQueryBlock.WithClause x); - - boolean visit(PGSelectQueryBlock.WithClause x); - - void endVisit(PGSelectQueryBlock.IntoClause x); - - boolean visit(PGSelectQueryBlock.IntoClause x); + void endVisit(PGSelectQueryBlock.ForClause x); + + boolean visit(PGSelectQueryBlock.ForClause x); + + void endVisit(PGSelectQueryBlock.WithQuery x); + + boolean visit(PGSelectQueryBlock.WithQuery x); + + void endVisit(PGSelectQueryBlock.WithClause x); + + boolean visit(PGSelectQueryBlock.WithClause x); + + void endVisit(PGSelectQueryBlock.IntoClause x); + + boolean visit(PGSelectQueryBlock.IntoClause x); + + void endVisit(PGTruncateStatement x); + + boolean visit(PGTruncateStatement 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 a73c754c80..83121db44f 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 @@ -6,75 +6,86 @@ import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.IntoClaus import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.WindowClause; import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.WithClause; import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.WithQuery; +import com.alibaba.druid.sql.dialect.postgresql.ast.PGTruncateStatement; import com.alibaba.druid.sql.visitor.SQLASTVisitorAdapter; -public class PGASTVisitorAdapter extends SQLASTVisitorAdapter implements - PGASTVisitor { +public class PGASTVisitorAdapter extends SQLASTVisitorAdapter implements PGASTVisitor { - @Override - public void endVisit(WindowClause x) { + @Override + public void endVisit(WindowClause x) { - } + } - @Override - public boolean visit(WindowClause x) { + @Override + public boolean visit(WindowClause x) { - return true; - } + return true; + } - @Override - public void endVisit(FetchClause x) { + @Override + public void endVisit(FetchClause x) { - } + } - @Override - public boolean visit(FetchClause x) { + @Override + public boolean visit(FetchClause x) { - return true; - } + return true; + } - @Override - public void endVisit(ForClause x) { + @Override + public void endVisit(ForClause x) { - } + } - @Override - public boolean visit(ForClause x) { + @Override + public boolean visit(ForClause x) { - return true; - } + return true; + } - @Override - public void endVisit(WithQuery x) { + @Override + public void endVisit(WithQuery x) { - } + } - @Override - public boolean visit(WithQuery x) { + @Override + public boolean visit(WithQuery x) { - return true; - } + return true; + } - @Override - public void endVisit(WithClause x) { + @Override + public void endVisit(WithClause x) { - } + } - @Override - public boolean visit(WithClause x) { + @Override + public boolean visit(WithClause x) { - return true; - } - - @Override - public void endVisit(IntoClause x) { - - } - - @Override - public boolean visit(IntoClause x) { - - return true; - } + return true; + } + + @Override + public void endVisit(IntoClause x) { + + } + + @Override + public boolean visit(IntoClause x) { + + return true; + } + + @Override + public void endVisit(PGTruncateStatement x) { + + } + + @Override + public boolean visit(PGTruncateStatement 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 9ed68008ca..3a4a538651 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 @@ -9,6 +9,7 @@ import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.IntoClaus import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.WindowClause; import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.WithClause; import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.WithQuery; +import com.alibaba.druid.sql.dialect.postgresql.ast.PGTruncateStatement; import com.alibaba.druid.sql.visitor.SQLASTOutputVisitor; public class PGOutputVisitor extends SQLASTOutputVisitor implements @@ -232,4 +233,36 @@ public class PGOutputVisitor extends SQLASTOutputVisitor implements x.getTable().accept(this); return false; } + + @Override + public void endVisit(PGTruncateStatement x) { + + } + + @Override + public boolean visit(PGTruncateStatement x) { + print("TRUNCATE TABLE "); + if (x.isOnly()) { + print("ONLY "); + } + + printlnAndAccept(x.getTableNames(), ", "); + + if (x.getRestartIdentity() != null) { + if (x.getRestartIdentity().booleanValue()) { + print(" RESTART IDENTITY"); + } else { + print(" CONTINUE IDENTITY"); + } + } + + if (x.getCascade() != null) { + if (x.getCascade().booleanValue()) { + print(" CASCADE"); + } else { + print(" RESTRICT"); + } + } + 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 00ecd4d419..2b29044df0 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 @@ -1,86 +1,97 @@ package com.alibaba.druid.sql.dialect.postgresql.visitor; +import com.alibaba.druid.sql.ast.statement.SQLTruncateStatement; import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.FetchClause; import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.ForClause; import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.IntoClause; import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.WindowClause; import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.WithClause; import com.alibaba.druid.sql.dialect.postgresql.ast.PGSelectQueryBlock.WithQuery; +import com.alibaba.druid.sql.dialect.postgresql.ast.PGTruncateStatement; import com.alibaba.druid.sql.visitor.SchemaStatVisitor; import com.alibaba.druid.stat.TableStat; -public class PGSchemaStatVisitor extends SchemaStatVisitor implements - PGASTVisitor { +public class PGSchemaStatVisitor extends SchemaStatVisitor implements PGASTVisitor { - @Override - public void endVisit(WindowClause x) { + @Override + public void endVisit(WindowClause x) { - } + } - @Override - public boolean visit(WindowClause x) { - return true; - } + @Override + public boolean visit(WindowClause x) { + return true; + } - @Override - public void endVisit(FetchClause x) { + @Override + public void endVisit(FetchClause x) { - } + } - @Override - public boolean visit(FetchClause x) { - return true; - } + @Override + public boolean visit(FetchClause x) { + return true; + } - @Override - public void endVisit(ForClause x) { + @Override + public void endVisit(ForClause x) { - } + } - @Override - public boolean visit(ForClause x) { + @Override + public boolean visit(ForClause x) { - return true; - } + return true; + } - @Override - public void endVisit(WithQuery x) { + @Override + public void endVisit(WithQuery x) { - } + } - @Override - public boolean visit(WithQuery x) { + @Override + public boolean visit(WithQuery x) { - return true; - } + return true; + } - @Override - public void endVisit(WithClause x) { + @Override + public void endVisit(WithClause x) { - } + } - @Override - public boolean visit(WithClause x) { + @Override + public boolean visit(WithClause x) { - return true; - } + return true; + } - @Override - public void endVisit(IntoClause x) { - - } + @Override + public void endVisit(IntoClause x) { - @Override - public boolean visit(IntoClause x) { - String ident = x.getTable().toString(); - - TableStat stat = tableStats.get(ident); - if (stat == null) { - stat = new TableStat(); - tableStats.put(new TableStat.Name(ident), stat); - } - stat.incrementInsertCount(); - return false; - } + } + @Override + public boolean visit(IntoClause x) { + String ident = x.getTable().toString(); + + TableStat stat = tableStats.get(ident); + if (stat == null) { + stat = new TableStat(); + tableStats.put(new TableStat.Name(ident), stat); + } + stat.incrementInsertCount(); + return false; + } + + @Override + public void endVisit(PGTruncateStatement x) { + + } + + @Override + public boolean visit(PGTruncateStatement x) { + this.visit((SQLTruncateStatement) x); + return false; + } } diff --git a/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java b/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java index fa8ce03473..8e7d04cdcd 100644 --- a/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java +++ b/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java @@ -116,6 +116,12 @@ public class SQLStatementParser extends SQLParser { continue; } } + + if (lexer.token() == Token.TRUNCATE) { + SQLStatement stmt = parseTruncate(); + statementList.add(stmt); + continue; + } if (identifierEquals("CALL")) { SQLCallStatement stmt = parseCall(); @@ -130,6 +136,10 @@ public class SQLStatementParser extends SQLParser { throw new ParserException("TODO " + lexer.token()); } } + + public SQLStatement parseTruncate() { + throw new ParserException("TODO " + lexer.token()); + } public SQLStatement parseInsert() { accept(Token.INSERT); diff --git a/src/main/java/com/alibaba/druid/sql/parser/Token.java b/src/main/java/com/alibaba/druid/sql/parser/Token.java index 4c4292544f..610c88e908 100644 --- a/src/main/java/com/alibaba/druid/sql/parser/Token.java +++ b/src/main/java/com/alibaba/druid/sql/parser/Token.java @@ -92,6 +92,7 @@ public enum Token { LOCK("LOCK"), SOME("SOME"), ANY("ANY"), + TRUNCATE("TRUNCATE"), // mysql TRUE("TRUE"), @@ -115,6 +116,11 @@ public enum Token { TEMPORARY("TEMPORARY"), TEMP("TEMP"), UNLOGGED("UNLOGGED"), + RESTART("RESTART"), + IDENTITY("IDENTITY"), + CONTINUE("CONTINUE"), + CASCADE("CASCADE"), + RESTRICT("RESTRICT"), // oracle START("START"), diff --git a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java index ad0c8402cf..cce75f0988 100644 --- a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java +++ b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java @@ -73,6 +73,7 @@ import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; import com.alibaba.druid.sql.ast.statement.SQLSetStatement; import com.alibaba.druid.sql.ast.statement.SQLSubqueryTableSource; import com.alibaba.druid.sql.ast.statement.SQLTableElement; +import com.alibaba.druid.sql.ast.statement.SQLTruncateStatement; import com.alibaba.druid.sql.ast.statement.SQLUnionQuery; import com.alibaba.druid.sql.ast.statement.SQLUniqueConstraint; import com.alibaba.druid.sql.ast.statement.SQLUpdateSetItem; @@ -864,4 +865,16 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter { return false; } + + @Override + public void endVisit(SQLTruncateStatement x) { + + } + + @Override + public boolean visit(SQLTruncateStatement x) { + print("TRUNCATE TABLE "); + printAndAccept(x.getTableNames(), ", "); + return false; + } } diff --git a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitor.java b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitor.java index 133a1c2c10..6bce389819 100644 --- a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitor.java +++ b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitor.java @@ -68,6 +68,7 @@ import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; import com.alibaba.druid.sql.ast.statement.SQLSetStatement; import com.alibaba.druid.sql.ast.statement.SQLSubqueryTableSource; import com.alibaba.druid.sql.ast.statement.SQLTableElement; +import com.alibaba.druid.sql.ast.statement.SQLTruncateStatement; import com.alibaba.druid.sql.ast.statement.SQLUnionQuery; import com.alibaba.druid.sql.ast.statement.SQLUniqueConstraint; import com.alibaba.druid.sql.ast.statement.SQLUpdateSetItem; @@ -314,4 +315,8 @@ public interface SQLASTVisitor { void endVisit(SQLSubqueryTableSource x); boolean visit(SQLSubqueryTableSource x); + + void endVisit(SQLTruncateStatement x); + + boolean visit(SQLTruncateStatement x); } diff --git a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitorAdapter.java b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitorAdapter.java index 9bfa7a7b90..57cc9b8cbd 100644 --- a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitorAdapter.java +++ b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTVisitorAdapter.java @@ -69,6 +69,7 @@ import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; import com.alibaba.druid.sql.ast.statement.SQLSetStatement; import com.alibaba.druid.sql.ast.statement.SQLSubqueryTableSource; import com.alibaba.druid.sql.ast.statement.SQLTableElement; +import com.alibaba.druid.sql.ast.statement.SQLTruncateStatement; import com.alibaba.druid.sql.ast.statement.SQLUnionQuery; import com.alibaba.druid.sql.ast.statement.SQLUniqueConstraint; import com.alibaba.druid.sql.ast.statement.SQLUpdateSetItem; @@ -539,4 +540,14 @@ public class SQLASTVisitorAdapter implements SQLASTVisitor { public boolean visit(SQLSubqueryTableSource x) { return true; } + + @Override + public void endVisit(SQLTruncateStatement x) { + + } + + @Override + public boolean visit(SQLTruncateStatement x) { + return true; + } } 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 00c214e62d..75db082ea9 100644 --- a/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java +++ b/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java @@ -21,266 +21,298 @@ import com.alibaba.druid.sql.ast.statement.SQLSelect; import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock; import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; import com.alibaba.druid.sql.ast.statement.SQLSubqueryTableSource; +import com.alibaba.druid.sql.ast.statement.SQLTruncateStatement; import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement; import com.alibaba.druid.stat.TableStat; import com.alibaba.druid.stat.TableStat.Column; import com.alibaba.druid.stat.TableStat.Mode; public class SchemaStatVisitor extends SQLASTVisitorAdapter { - protected HashMap tableStats = new HashMap(); - protected Set fields = new HashSet(); - protected final static ThreadLocal> aliasLocal = new ThreadLocal>(); - protected final static ThreadLocal currentTableLocal = new ThreadLocal(); - protected final static ThreadLocal modeLocal = new ThreadLocal(); + protected HashMap tableStats = new HashMap(); + protected Set fields = new HashSet(); - @Override - public boolean visit(SQLInsertStatement x) { - x.putAttribute("_original_use_mode", modeLocal.get()); - modeLocal.set(Mode.Insert); + protected final static ThreadLocal> aliasLocal = new ThreadLocal>(); + protected final static ThreadLocal currentTableLocal = new ThreadLocal(); + protected final static ThreadLocal modeLocal = new ThreadLocal(); - aliasLocal.set(new HashMap()); + @Override + public boolean visit(SQLTruncateStatement x) { + x.putAttribute("_original_use_mode", modeLocal.get()); + modeLocal.set(Mode.Insert); - String originalTable = currentTableLocal.get(); + aliasLocal.set(new HashMap()); - if (x.getTableName() instanceof SQLName) { - String ident = ((SQLName) x.getTableName()).toString(); - currentTableLocal.set(ident); - x.putAttribute("_old_local_", originalTable); + String originalTable = currentTableLocal.get(); - TableStat stat = tableStats.get(ident); - if (stat == null) { - stat = new TableStat(); - tableStats.put(new TableStat.Name(ident), stat); - } - stat.incrementInsertCount(); + for (SQLName name : x.getTableNames()) { + String ident = name.toString(); + currentTableLocal.set(ident); + x.putAttribute("_old_local_", originalTable); - Map aliasMap = aliasLocal.get(); - if (aliasMap != null) { - if (x.getAlias() != null) { - aliasMap.put(x.getAlias(), ident); - } - aliasMap.put(ident, ident); - } - } + TableStat stat = tableStats.get(ident); + if (stat == null) { + stat = new TableStat(); + tableStats.put(new TableStat.Name(ident), stat); + } + stat.incrementInsertCount(); - accept(x.getColumns()); - accept(x.getQuery()); + Map aliasMap = aliasLocal.get(); + if (aliasMap != null) { + aliasMap.put(ident, ident); + } + } - return false; - } + return false; + } - protected void accept(SQLObject x) { - if (x != null) { - x.accept(this); - } - } + @Override + public boolean visit(SQLInsertStatement x) { + x.putAttribute("_original_use_mode", modeLocal.get()); + modeLocal.set(Mode.Insert); - protected void accept(List nodes) { - for (int i = 0, size = nodes.size(); i < size; ++i) { - accept(nodes.get(i)); - } - } + aliasLocal.set(new HashMap()); - public boolean visit(SQLSelectQueryBlock x) { + String originalTable = currentTableLocal.get(); - if (x.getFrom() instanceof SQLSubqueryTableSource) { - x.getFrom().accept(this); - return false; - } + if (x.getTableName() instanceof SQLName) { + String ident = ((SQLName) x.getTableName()).toString(); + currentTableLocal.set(ident); + x.putAttribute("_old_local_", originalTable); - x.putAttribute("_original_use_mode", modeLocal.get()); - modeLocal.set(Mode.Select); + TableStat stat = tableStats.get(ident); + if (stat == null) { + stat = new TableStat(); + tableStats.put(new TableStat.Name(ident), stat); + } + stat.incrementInsertCount(); - String originalTable = currentTableLocal.get(); + Map aliasMap = aliasLocal.get(); + if (aliasMap != null) { + if (x.getAlias() != null) { + aliasMap.put(x.getAlias(), ident); + } + aliasMap.put(ident, ident); + } + } - if (x.getFrom() instanceof SQLExprTableSource) { - SQLExprTableSource tableSource = (SQLExprTableSource) x.getFrom(); - if (tableSource.getExpr() instanceof SQLName) { - String ident = tableSource.getExpr().toString(); - currentTableLocal.set(ident); - x.putAttribute("_old_local_", originalTable); - } - } + accept(x.getColumns()); + accept(x.getQuery()); - x.getFrom().accept(this); // 提前执行，获得aliasMap + return false; + } - return true; - } + protected void accept(SQLObject x) { + if (x != null) { + x.accept(this); + } + } - public void endVisit(SQLSelectQueryBlock x) { - String originalTable = (String) x.getAttribute("_old_local_"); - x.putAttribute("table", currentTableLocal.get()); - currentTableLocal.set(originalTable); + protected void accept(List nodes) { + for (int i = 0, size = nodes.size(); i < size; ++i) { + accept(nodes.get(i)); + } + } - Mode originalMode = (Mode) x.getAttribute("_original_use_mode"); - modeLocal.set(originalMode); - } + public boolean visit(SQLSelectQueryBlock x) { - public boolean visit(SQLJoinTableSource x) { - return true; - } + if (x.getFrom() instanceof SQLSubqueryTableSource) { + x.getFrom().accept(this); + return false; + } - public boolean visit(SQLPropertyExpr x) { - if (x.getOwner() instanceof SQLIdentifierExpr) { - String owner = ((SQLIdentifierExpr) x.getOwner()).getName(); + x.putAttribute("_original_use_mode", modeLocal.get()); + modeLocal.set(Mode.Select); - if (owner != null) { - Map aliasMap = aliasLocal.get(); - if (aliasMap != null) { - String table = aliasMap.get(owner); + String originalTable = currentTableLocal.get(); - // table == null时是SubQuery - if (table != null) { - fields.add(new Column(table, x.getName())); - } - } - } - } - return false; - } + if (x.getFrom() instanceof SQLExprTableSource) { + SQLExprTableSource tableSource = (SQLExprTableSource) x.getFrom(); + if (tableSource.getExpr() instanceof SQLName) { + String ident = tableSource.getExpr().toString(); + currentTableLocal.set(ident); + x.putAttribute("_old_local_", originalTable); + } + } - public boolean visit(SQLIdentifierExpr x) { - String currentTable = currentTableLocal.get(); + x.getFrom().accept(this); // 提前执行，获得aliasMap - if (currentTable != null) { - fields.add(new Column(currentTable, x.getName())); - } - return false; - } + return true; + } - public boolean visit(SQLAllColumnExpr x) { - String currentTable = currentTableLocal.get(); + public void endVisit(SQLSelectQueryBlock x) { + String originalTable = (String) x.getAttribute("_old_local_"); + x.putAttribute("table", currentTableLocal.get()); + currentTableLocal.set(originalTable); - if (currentTable != null) { - fields.add(new Column(currentTable, "*")); - } - return false; - } + Mode originalMode = (Mode) x.getAttribute("_original_use_mode"); + modeLocal.set(originalMode); + } - public Map getTables() { - return tableStats; - } + public boolean visit(SQLJoinTableSource x) { + return true; + } - public boolean containsTable(String tableName) { - return tableStats.containsKey(new TableStat.Name(tableName)); - } + public boolean visit(SQLPropertyExpr x) { + if (x.getOwner() instanceof SQLIdentifierExpr) { + String owner = ((SQLIdentifierExpr) x.getOwner()).getName(); - public Set getFields() { - return fields; - } + if (owner != null) { + Map aliasMap = aliasLocal.get(); + if (aliasMap != null) { + String table = aliasMap.get(owner); - public boolean visit(SQLSelectStatement x) { - aliasLocal.set(new HashMap()); - return true; - } + // table == null时是SubQuery + if (table != null) { + fields.add(new Column(table, x.getName())); + } + } + } + } + return false; + } - public void endVisit(SQLSelectStatement x) { - aliasLocal.set(null); - } + public boolean visit(SQLIdentifierExpr x) { + String currentTable = currentTableLocal.get(); - public boolean visit(SQLSubqueryTableSource x) { - Map aliasMap = aliasLocal.get(); - if (aliasMap != null) { - if (x.getAlias() != null) { - aliasMap.put(x.getAlias(), null); - } - } - return true; - } + if (currentTable != null) { + fields.add(new Column(currentTable, x.getName())); + } + return false; + } - protected boolean isSimpleExprTableSource(SQLExprTableSource x) { - return x.getExpr() instanceof SQLName; - } + public boolean visit(SQLAllColumnExpr x) { + String currentTable = currentTableLocal.get(); - public boolean visit(SQLExprTableSource x) { - if (isSimpleExprTableSource(x)) { - String ident = x.getExpr().toString(); - TableStat stat = tableStats.get(ident); - if (stat == null) { - stat = new TableStat(); - tableStats.put(new TableStat.Name(ident), stat); - } + if (currentTable != null) { + fields.add(new Column(currentTable, "*")); + } + return false; + } - Mode mode = modeLocal.get(); - switch (mode) { - case Delete: - stat.incrementDeleteCount(); - break; - case Insert: - stat.incrementInsertCount(); - break; - case Update: - stat.incrementUpdateCount(); - break; - case Select: - stat.incrementSelectCount(); - break; - default: - break; - } + public Map getTables() { + return tableStats; + } - Map aliasMap = aliasLocal.get(); - if (aliasMap != null) { - if (x.getAlias() != null) { - aliasMap.put(x.getAlias(), ident); - } - aliasMap.put(ident, ident); - } - } else { - accept(x.getExpr()); - } + public boolean containsTable(String tableName) { + return tableStats.containsKey(new TableStat.Name(tableName)); + } - return false; - } + public Set getFields() { + return fields; + } - public boolean visit(SQLSelect x) { - accept(x.getQuery()); + public boolean visit(SQLSelectStatement x) { + aliasLocal.set(new HashMap()); + return true; + } - String originalTable = currentTableLocal.get(); + public void endVisit(SQLSelectStatement x) { + aliasLocal.set(null); + } - currentTableLocal.set((String) x.getQuery().getAttribute("table")); - x.putAttribute("_old_local_", originalTable); + public boolean visit(SQLSubqueryTableSource x) { + Map aliasMap = aliasLocal.get(); + if (aliasMap != null) { + if (x.getAlias() != null) { + aliasMap.put(x.getAlias(), null); + } + } + return true; + } - accept(x.getOrderBy()); + protected boolean isSimpleExprTableSource(SQLExprTableSource x) { + return x.getExpr() instanceof SQLName; + } - currentTableLocal.set(originalTable); + public boolean visit(SQLExprTableSource x) { + if (isSimpleExprTableSource(x)) { + String ident = x.getExpr().toString(); + TableStat stat = tableStats.get(ident); + if (stat == null) { + stat = new TableStat(); + tableStats.put(new TableStat.Name(ident), stat); + } - return false; - } + Mode mode = modeLocal.get(); + switch (mode) { + case Delete: + stat.incrementDeleteCount(); + break; + case Insert: + stat.incrementInsertCount(); + break; + case Update: + stat.incrementUpdateCount(); + break; + case Select: + stat.incrementSelectCount(); + break; + default: + break; + } - public boolean visit(SQLAggregateExpr x) { - accept(x.getArguments()); - return false; - } + Map aliasMap = aliasLocal.get(); + if (aliasMap != null) { + if (x.getAlias() != null) { + aliasMap.put(x.getAlias(), ident); + } + aliasMap.put(ident, ident); + } + } else { + accept(x.getExpr()); + } - public boolean visit(SQLMethodInvokeExpr x) { - accept(x.getParameters()); - return false; - } + return false; + } - public boolean visit(SQLUpdateStatement x) { - aliasLocal.set(new HashMap()); + public boolean visit(SQLSelect x) { + accept(x.getQuery()); - String ident = x.getTableName().toString(); - currentTableLocal.set(ident); + String originalTable = currentTableLocal.get(); - TableStat stat = tableStats.get(ident); - if (stat == null) { - stat = new TableStat(); - tableStats.put(new TableStat.Name(ident), stat); - } - stat.incrementUpdateCount(); + currentTableLocal.set((String) x.getQuery().getAttribute("table")); + x.putAttribute("_old_local_", originalTable); - Map aliasMap = aliasLocal.get(); - aliasMap.put(ident, ident); + accept(x.getOrderBy()); - accept(x.getItems()); - accept(x.getWhere()); + currentTableLocal.set(originalTable); + + return false; + } + + public boolean visit(SQLAggregateExpr x) { + accept(x.getArguments()); + return false; + } + + public boolean visit(SQLMethodInvokeExpr x) { + accept(x.getParameters()); + return false; + } + + public boolean visit(SQLUpdateStatement x) { + aliasLocal.set(new HashMap()); + + 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()); + + return false; + } - return false; - } - public boolean visit(SQLDeleteStatement x) { aliasLocal.set(new HashMap()); @@ -308,7 +340,7 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter { aliasLocal.set(null); } - public void endVisit(SQLUpdateStatement x) { - aliasLocal.set(null); - } + public void endVisit(SQLUpdateStatement x) { + aliasLocal.set(null); + } } diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/TruncateTest.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/TruncateTest.java new file mode 100644 index 0000000000..8ca5384270 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/TruncateTest.java @@ -0,0 +1,36 @@ +package com.alibaba.druid.bvt.sql.postgresql; + +import java.util.List; + +import junit.framework.Assert; + +import com.alibaba.druid.sql.PGTest; +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 com.alibaba.druid.stat.TableStat; +import com.alibaba.druid.stat.TableStat.Column; + +public class TruncateTest extends PGTest { + public void test_0() throws Exception { + String sql = "TRUNCATE bigtable, fattable;"; + + PGSQLStatementParser parser = new PGSQLStatementParser(sql); + List statementList = parser.parseStatementList(); + SQLStatement statemen = statementList.get(0); + System.out.println(output(statementList)); + + Assert.assertEquals(1, statementList.size()); + + PGSchemaStatVisitor visitor = new PGSchemaStatVisitor(); + statemen.accept(visitor); + + System.out.println("Tables : " + visitor.getTables()); + System.out.println("fields : " + visitor.getFields()); + + Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("bigtable"))); + Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("fattable"))); + + Assert.assertTrue(visitor.getFields().size() == 0); + } +}