From 2d5c96320179fffb747adf1f8d0d706201126401 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B8=A9=E5=B0=91?= Date: Thu, 23 Feb 2012 20:22:11 +0800 Subject: [PATCH] enhance SchemaStatVisitor --- .../druid/sql/visitor/SchemaStatVisitor.java | 105 ++++++++++++++++-- .../com/alibaba/druid/stat/TableStat.java | 66 ++++++++++- .../bvt/sql/postgresql/PGDeleteTest.java | 1 + .../bvt/sql/postgresql/PGDeleteTest2.java | 1 + .../bvt/sql/postgresql/PGDeleteTest3.java | 1 + .../bvt/sql/postgresql/PGDeleteTest4.java | 1 + 6 files changed, 167 insertions(+), 8 deletions(-) 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);