mirror of
https://github.com/alibaba/druid.git
synced 2023-03-04 19:17:16 +00:00
enhance SchemaStatVisitor
This commit is contained in:
+9
-1
@@ -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);
|
||||
}
|
||||
|
||||
+11
-1
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user