improve SchemaStatVisitor

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();

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<String, Object> attributes = new HashMap<String, Object>();
 
@@ -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;
         }

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) {