enhance SchemaStatVisitor

This commit is contained in:
温少
2012-02-23 20:39:21 +08:00
parent 94c5f6e420
commit a08c9e8688
16 changed files with 195 additions and 10 deletions
@@ -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);
}
@@ -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");
}
}
@@ -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<TableStat.Name, TableStat> tableStats = new HashMap<TableStat.Name, TableStat>();
protected Set<Column> fields = new HashSet<Column>();
protected final HashMap<TableStat.Name, TableStat> tableStats = new HashMap<TableStat.Name, TableStat>();
protected final Set<Column> fields = new HashSet<Column>();
protected final Set<Condition> conditions = new HashSet<Condition>();
protected final static ThreadLocal<Map<String, String>> aliasLocal = new ThreadLocal<Map<String, String>>();
protected final static ThreadLocal<String> currentTableLocal = new ThreadLocal<String>();
protected final static ThreadLocal<Mode> modeLocal = new ThreadLocal<Mode>();
public Set<Condition> 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<String, String> 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<String, String> 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<String, String> 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;
}
@@ -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;
@@ -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")));
@@ -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")));
@@ -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")));
@@ -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);
@@ -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")));
@@ -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")));
@@ -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")));
@@ -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")));
@@ -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")));
@@ -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")));
@@ -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")));
@@ -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();