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 14fb53b9ea..2c53d6c0f0 100644 --- a/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java +++ b/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java @@ -56,6 +56,7 @@ import com.alibaba.druid.sql.ast.statement.SQLInsertStatement; import com.alibaba.druid.sql.ast.statement.SQLJoinTableSource; import com.alibaba.druid.sql.ast.statement.SQLRollbackStatement; import com.alibaba.druid.sql.ast.statement.SQLSelect; +import com.alibaba.druid.sql.ast.statement.SQLSelectGroupByClause; import com.alibaba.druid.sql.ast.statement.SQLSelectItem; import com.alibaba.druid.sql.ast.statement.SQLSelectOrderByItem; import com.alibaba.druid.sql.ast.statement.SQLSelectQuery; @@ -71,6 +72,7 @@ import com.alibaba.druid.stat.TableStat.Column; import com.alibaba.druid.stat.TableStat.Condition; import com.alibaba.druid.stat.TableStat.Mode; import com.alibaba.druid.stat.TableStat.Relationship; +import com.alibaba.druid.util.StringUtils; public class SchemaStatVisitor extends SQLASTVisitorAdapter { @@ -120,8 +122,11 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter { tableName = handleName(tableName); columnName = handleName(columnName); - Column column = new Column(tableName, columnName); - columns.add(column); + Column column = this.getColumn(tableName, columnName); + if (column == null) { + column = new Column(tableName, columnName); + columns.add(column); + } return column; } @@ -170,7 +175,7 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter { if (flag2) { ident = ident.replaceAll(" ", ""); } - + if (flag3) { ident = ident.replaceAll("'", ""); } @@ -576,7 +581,7 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter { if (x.getFrom() == null) { return false; } - + setMode(x, Mode.Select); if (x.getFrom() instanceof SQLSubqueryTableSource) { @@ -660,6 +665,11 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter { if (owner != null) { Column column = addColumn(owner, x.getName()); x.putAttribute(ATTR_COLUMN, column); + if (column != null) { + if (column != null) { + setColumn(x, column); + } + } } } return false; @@ -730,18 +740,66 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter { return false; } + Column column; if (currentTable != null) { - Column column = addColumn(currentTable, ident); + column = addColumn(currentTable, ident); x.putAttribute(ATTR_COLUMN, column); } else { - Column column = handleUnkownColumn(ident); + column = handleUnkownColumn(ident); if (column != null) { x.putAttribute(ATTR_COLUMN, column); } } + if (column != null) { + setColumn(x, column); + } return false; } + private void setColumn(SQLExpr x, Column column) { + SQLObject current = x; + for (;;) { + SQLObject parent = current.getParent(); + + if (parent == null) { + break; + } + + if (parent instanceof SQLSelectQueryBlock) { + SQLSelectQueryBlock query = (SQLSelectQueryBlock) parent; + if (query.getWhere() == current) { + column.setWhere(true); + } + break; + } + + if (parent instanceof SQLSelectGroupByClause) { + SQLSelectGroupByClause groupBy = (SQLSelectGroupByClause) parent; + if (current == groupBy.getHaving()) { + column.setHaving(true); + } else if (groupBy.getItems().contains(current)) { + column.setGroupBy(true); + } + break; + } + + if (parent instanceof SQLSelectItem) { + column.setSelec(true); + break; + } + + if (parent instanceof SQLJoinTableSource) { + SQLJoinTableSource join = (SQLJoinTableSource) parent; + if (join.getCondition() == current) { + column.setJoin(true); + } + break; + } + + current = parent; + } + } + protected Column handleUnkownColumn(String columnName) { return addColumn("UNKNOWN", columnName); } @@ -771,6 +829,16 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter { return columns; } + public Column getColumn(String tableName, String columnName) { + for (Column column : this.columns) { + if (StringUtils.equalsIgnoreCase(tableName, column.getTable()) + && StringUtils.equalsIgnoreCase(columnName, column.getName())) { + return column; + } + } + return null; + } + public boolean visit(SQLSelectStatement x) { setAliasMap(); return true; @@ -1051,22 +1119,22 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter { public boolean visit(SQLAlterTableDropForeinKey x) { return false; } - + @Override public boolean visit(SQLUseStatement x) { return false; } - + @Override public boolean visit(SQLAlterTableDisableConstraint x) { return false; } - + @Override public boolean visit(SQLAlterTableEnableConstraint x) { return false; } - + @Override public boolean visit(SQLAlterTableStatement x) { String tableName = x.getName().toString(); diff --git a/src/main/java/com/alibaba/druid/stat/TableStat.java b/src/main/java/com/alibaba/druid/stat/TableStat.java index f52d0f06ef..5bbbd65cf7 100644 --- a/src/main/java/com/alibaba/druid/stat/TableStat.java +++ b/src/main/java/com/alibaba/druid/stat/TableStat.java @@ -366,6 +366,11 @@ public class TableStat { private String table; private String name; + private boolean where; + private boolean select; + private boolean groupBy; + private boolean having; + private boolean join; private Map attributes = new HashMap(); @@ -386,6 +391,46 @@ public class TableStat { this.table = table; } + public boolean isWhere() { + return where; + } + + public void setWhere(boolean where) { + this.where = where; + } + + public boolean isSelect() { + return select; + } + + public void setSelec(boolean select) { + this.select = select; + } + + public boolean isGroupBy() { + return groupBy; + } + + public void setGroupBy(boolean groupBy) { + this.groupBy = groupBy; + } + + public boolean isHaving() { + return having; + } + + public boolean isJoin() { + return join; + } + + public void setJoin(boolean join) { + this.join = join; + } + + public void setHaving(boolean having) { + this.having = having; + } + public String getName() { return name; } diff --git a/src/main/java/com/alibaba/druid/util/StringUtils.java b/src/main/java/com/alibaba/druid/util/StringUtils.java index e80081f37c..f53e006e1f 100644 --- a/src/main/java/com/alibaba/druid/util/StringUtils.java +++ b/src/main/java/com/alibaba/druid/util/StringUtils.java @@ -83,6 +83,13 @@ public class StringUtils { } return a.equals(b); } + + public static boolean equalsIgnoreCase(String a, String b) { + if (a == null) { + return b == null; + } + return a.equalsIgnoreCase(b); + } public static boolean isEmpty(String value) { if (value == null || value.length() == 0) { diff --git a/src/test/java/com/alibaba/druid/bvt/sql/mysql/visitor/MySqlSchemaStatVisitorTest1.java b/src/test/java/com/alibaba/druid/bvt/sql/mysql/visitor/MySqlSchemaStatVisitorTest1.java index 3789035b62..c71d8787c9 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/mysql/visitor/MySqlSchemaStatVisitorTest1.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/mysql/visitor/MySqlSchemaStatVisitorTest1.java @@ -17,9 +17,10 @@ package com.alibaba.druid.bvt.sql.mysql.visitor; import java.util.List; -import junit.framework.Assert; import junit.framework.TestCase; +import org.junit.Assert; + import com.alibaba.druid.sql.ast.SQLStatement; import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser; import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlSchemaStatVisitor; @@ -81,5 +82,13 @@ public class MySqlSchemaStatVisitorTest1 extends TestCase { Assert.assertEquals(true, visitor.getColumns().contains(new Column("usergroups", "id"))); Assert.assertEquals(true, visitor.getColumns().contains(new Column("usergroups", "name"))); + Assert.assertTrue(visitor.getColumn("users", "name").isSelect()); + Assert.assertTrue(visitor.getColumn("usergroups", "name").isSelect()); + + Assert.assertTrue(visitor.getColumn("users", "groupId").isJoin()); + Assert.assertTrue(visitor.getColumn("usergroups", "id").isJoin()); + + Assert.assertTrue(visitor.getColumn("users", "groupId").isWhere()); + } } diff --git a/src/test/java/com/alibaba/druid/bvt/sql/mysql/visitor/MySqlSchemaStatVisitorTest2.java b/src/test/java/com/alibaba/druid/bvt/sql/mysql/visitor/MySqlSchemaStatVisitorTest2.java index 4d55a0b49a..c7911c7502 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/mysql/visitor/MySqlSchemaStatVisitorTest2.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/mysql/visitor/MySqlSchemaStatVisitorTest2.java @@ -17,9 +17,10 @@ package com.alibaba.druid.bvt.sql.mysql.visitor; import java.util.List; -import junit.framework.Assert; import junit.framework.TestCase; +import org.junit.Assert; + import com.alibaba.druid.sql.ast.SQLStatement; import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser; import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlSchemaStatVisitor;