diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/stmt/PGSelectStatement.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/stmt/PGSelectStatement.java index cd570b372d..1c760e401e 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/stmt/PGSelectStatement.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/stmt/PGSelectStatement.java @@ -1,5 +1,6 @@ package com.alibaba.druid.sql.dialect.postgresql.ast.stmt; +import com.alibaba.druid.sql.ast.statement.SQLSelect; import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; import com.alibaba.druid.sql.dialect.postgresql.ast.PGWithClause; import com.alibaba.druid.sql.dialect.postgresql.visitor.PGASTVisitor; @@ -10,6 +11,13 @@ public class PGSelectStatement extends SQLSelectStatement implements PGSQLStatem private static final long serialVersionUID = 1L; private PGWithClause with; + public PGSelectStatement(){ + } + + public PGSelectStatement(SQLSelect select){ + super(select); + } + public PGWithClause getWith() { return with; } @@ -17,7 +25,7 @@ public class PGSelectStatement extends SQLSelectStatement implements PGSQLStatem public void setWith(PGWithClause with) { this.with = with; } - + protected void accept0(SQLASTVisitor visitor) { accept0((PGASTVisitor) visitor); } 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 f27dbbf17f..8c9e059d18 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 @@ -15,6 +15,7 @@ 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.PGUpdateStatement; import com.alibaba.druid.sql.parser.Lexer; @@ -311,6 +312,10 @@ public class PGSQLStatementParser extends SQLStatementParser { return withQuery; } + public PGSelectStatement parseSelect() throws ParserException { + return new PGSelectStatement(createSQLSelectParser().select()); + } + public SQLStatement parseWith() { PGWithClause with = this.parseWithClause(); if (lexer.token() == Token.INSERT) { @@ -318,7 +323,12 @@ public class PGSQLStatementParser extends SQLStatementParser { stmt.setWith(with); return stmt; } - + + if (lexer.token() == Token.SELECT) { + PGSelectStatement stmt = this.parseSelect(); + stmt.setWith(with); + return stmt; + } throw new ParserException("TODO"); } } 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 2c460b5233..5e472701fc 100644 --- a/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java +++ b/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java @@ -6,11 +6,14 @@ import java.util.List; import java.util.Map; import java.util.Set; +import com.alibaba.druid.sql.ast.SQLExpr; import com.alibaba.druid.sql.ast.SQLName; import com.alibaba.druid.sql.ast.SQLObject; import com.alibaba.druid.sql.ast.expr.SQLAggregateExpr; import com.alibaba.druid.sql.ast.expr.SQLAllColumnExpr; +import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr; import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr; +import com.alibaba.druid.sql.ast.expr.SQLInListExpr; import com.alibaba.druid.sql.ast.expr.SQLInSubQueryExpr; import com.alibaba.druid.sql.ast.expr.SQLMethodInvokeExpr; import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr; @@ -27,17 +30,90 @@ 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.Condition; import com.alibaba.druid.stat.TableStat.Mode; public class SchemaStatVisitor extends SQLASTVisitorAdapter { - protected HashMap tableStats = new HashMap(); - protected Set fields = new HashSet(); + protected final HashMap tableStats = new HashMap(); + protected final Set fields = new HashSet(); + protected final Set conditions = new HashSet(); protected final static ThreadLocal> aliasLocal = new ThreadLocal>(); protected final static ThreadLocal currentTableLocal = new ThreadLocal(); protected final static ThreadLocal modeLocal = new ThreadLocal(); + public Set getConditions() { + return conditions; + } + + public boolean visit(SQLBinaryOpExpr x) { + switch (x.getOperator()) { + case Equality: + case NotEqual: + case GreaterThan: + case GreaterThanOrEqual: + case LessThan: + case LessThanOrEqual: + case LessThanOrEqualOrGreaterThan: + case Like: + case NotLike: + case Is: + case IsNot: + handleCondition(x.getLeft(), x.getOperator().name); + handleCondition(x.getRight(), x.getOperator().name); + break; + default: + break; + } + return true; + } + + protected void handleCondition(SQLExpr expr, String operator) { + Map aliasMap = aliasLocal.get(); + if (aliasMap == null) { + return; + } + + if (expr instanceof SQLPropertyExpr) { + SQLExpr owner = ((SQLPropertyExpr) expr).getOwner(); + String column = ((SQLPropertyExpr) expr).getName(); + + if (owner instanceof SQLIdentifierExpr) { + String table = ((SQLIdentifierExpr) owner).getName(); + if (aliasMap.containsKey(table)) { + table = aliasMap.get(table); + } + + if (table != null) { + Condition condition = new Condition(); + condition.setColumn(new Column(table, column)); + condition.setOperator(operator); + this.conditions.add(condition); + } + } + + return; + } + + if (expr instanceof SQLIdentifierExpr) { + String column = ((SQLIdentifierExpr) expr).getName(); + String table = currentTableLocal.get(); + if (table != null) { + if (aliasMap.containsKey(table)) { + table = aliasMap.get(table); + } + } + + if (table != null) { + Condition condition = new Condition(); + condition.setColumn(new Column(table, column)); + condition.setOperator(operator); + this.conditions.add(condition); + } + } + } + @Override public boolean visit(SQLTruncateStatement x) { x.putAttribute("_original_use_mode", modeLocal.get()); @@ -67,7 +143,7 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter { return false; } - + @Override public boolean visit(SQLDropTableStatement x) { x.putAttribute("_original_use_mode", modeLocal.get()); @@ -162,7 +238,7 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter { SQLExprTableSource tableSource = (SQLExprTableSource) x.getFrom(); if (tableSource.getExpr() instanceof SQLName) { String ident = tableSource.getExpr().toString(); - + Map aliasMap = aliasLocal.get(); if (aliasMap.containsKey(ident) && aliasMap.get(ident) == null) { currentTableLocal.set(null); @@ -266,13 +342,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(); @@ -378,8 +454,23 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter { return false; } + public boolean visit(SQLInListExpr x) { + if (x.isNot()) { + handleCondition(x.getExpr(), "NOT IN"); + } else { + handleCondition(x.getExpr(), "IN"); + } + + return true; + } + @Override public boolean visit(SQLInSubQueryExpr x) { + if (x.isNot()) { + handleCondition(x.getExpr(), "NOT IN"); + } else { + handleCondition(x.getExpr(), "IN"); + } return true; } diff --git a/src/main/java/com/alibaba/druid/stat/TableStat.java b/src/main/java/com/alibaba/druid/stat/TableStat.java index bbc1a49525..c8aeb099b8 100644 --- a/src/main/java/com/alibaba/druid/stat/TableStat.java +++ b/src/main/java/com/alibaba/druid/stat/TableStat.java @@ -26,7 +26,7 @@ public class TableStat { public int getDropCount() { return dropCount; } - + public void incrementDropCount() { dropCount++; } @@ -135,6 +135,70 @@ public class TableStat { } } + public static class Condition { + + private Column column; + private String operator; + + public Column getColumn() { + return column; + } + + public void setColumn(Column column) { + this.column = column; + } + + public String getOperator() { + return operator; + } + + public void setOperator(String operator) { + this.operator = operator; + } + + @Override + public int hashCode() { + final int prime = 31; + int result = 1; + result = prime * result + ((column == null) ? 0 : column.hashCode()); + result = prime * result + ((operator == null) ? 0 : operator.hashCode()); + return result; + } + + @Override + public boolean equals(Object obj) { + if (this == obj) { + return true; + } + if (obj == null) { + return false; + } + if (getClass() != obj.getClass()) { + return false; + } + Condition other = (Condition) obj; + if (column == null) { + if (other.column != null) { + return false; + } + } else if (!column.equals(other.column)) { + return false; + } + if (operator == null) { + if (other.operator != null) { + return false; + } + } else if (!operator.equals(other.operator)) { + return false; + } + return true; + } + + public String toString() { + return this.column.toString() + " " + this.operator; + } + } + public static class Column { private String table; diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest.java index 22dccc427b..4ddbedea5d 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest.java @@ -28,6 +28,7 @@ public class PGDeleteTest extends PGTest { System.out.println("Tables : " + visitor.getTables()); System.out.println("fields : " + visitor.getFields()); + System.out.println("coditions : " + visitor.getConditions()); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("films"))); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("producers"))); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest2.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest2.java index a5ace197be..6908d90dee 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest2.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest2.java @@ -27,6 +27,7 @@ public class PGDeleteTest2 extends PGTest { System.out.println("Tables : " + visitor.getTables()); System.out.println("fields : " + visitor.getFields()); + System.out.println("coditions : " + visitor.getConditions()); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("tasks"))); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest3.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest3.java index d49397e8f0..1588797bae 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest3.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest3.java @@ -27,6 +27,7 @@ public class PGDeleteTest3 extends PGTest { System.out.println("Tables : " + visitor.getTables()); System.out.println("fields : " + visitor.getFields()); + System.out.println("coditions : " + visitor.getConditions()); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("tasks"))); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest4.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest4.java index efa14b7516..ccb25a3928 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest4.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest4.java @@ -30,6 +30,7 @@ public class PGDeleteTest4 extends PGTest { Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("films"))); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("producers"))); + System.out.println("coditions : " + visitor.getConditions()); Assert.assertTrue(visitor.getFields().size() == 3); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest5.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest5.java index 1734388cee..46d47c8fce 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest5.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest5.java @@ -27,6 +27,7 @@ public class PGDeleteTest5 extends PGTest { System.out.println("Tables : " + visitor.getTables()); System.out.println("fields : " + visitor.getFields()); + System.out.println("coditions : " + visitor.getConditions()); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("films"))); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("producers"))); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest0.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest0.java index bed33fac79..e99712dc23 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest0.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest0.java @@ -27,6 +27,7 @@ public class PGInsertTest0 extends PGTest { System.out.println("Tables : " + visitor.getTables()); System.out.println("fields : " + visitor.getFields()); + System.out.println("coditions : " + visitor.getConditions()); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("films"))); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest1.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest1.java index 805a594863..78efe5323f 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest1.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest1.java @@ -27,6 +27,7 @@ public class PGInsertTest1 extends PGTest { System.out.println("Tables : " + visitor.getTables()); System.out.println("fields : " + visitor.getFields()); + System.out.println("coditions : " + visitor.getConditions()); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("films"))); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest3.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest3.java index 45883abd4d..497ff30bfc 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest3.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest3.java @@ -27,6 +27,7 @@ public class PGInsertTest3 extends PGTest { System.out.println("Tables : " + visitor.getTables()); System.out.println("fields : " + visitor.getFields()); + System.out.println("coditions : " + visitor.getConditions()); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("films"))); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("tmp_films"))); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest4.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest4.java index 6ec42097da..b3b4da579e 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest4.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest4.java @@ -27,6 +27,7 @@ public class PGInsertTest4 extends PGTest { System.out.println("Tables : " + visitor.getTables()); System.out.println("fields : " + visitor.getFields()); + System.out.println("coditions : " + visitor.getConditions()); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("distributors"))); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest5.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest5.java index 911aa17f62..8573628669 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest5.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGInsertTest5.java @@ -32,6 +32,7 @@ public class PGInsertTest5 extends PGTest { System.out.println("Tables : " + visitor.getTables()); System.out.println("fields : " + visitor.getFields()); + System.out.println("coditions : " + visitor.getConditions()); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("employees_log"))); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("employees"))); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGSelectIntoTest.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGSelectIntoTest.java index fc7817122c..ca75e61c51 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGSelectIntoTest.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGSelectIntoTest.java @@ -27,6 +27,7 @@ public class PGSelectIntoTest extends PGTest { System.out.println("Tables : " + visitor.getTables()); System.out.println("fields : " + visitor.getFields()); + System.out.println("coditions : " + visitor.getConditions()); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("films_recent"))); Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("films"))); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PostgresqlResourceTest.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PostgresqlResourceTest.java index 6c540323f8..bdd427d23f 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PostgresqlResourceTest.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PostgresqlResourceTest.java @@ -49,6 +49,7 @@ public class PostgresqlResourceTest extends PGTest { System.out.println(sql); System.out.println("Tables : " + visitor.getTables()); System.out.println("fields : " + visitor.getFields()); + System.out.println("coditions : " + visitor.getConditions()); System.out.println(); System.out.println();