Merge pull request #203 from wenshao/master

sharding support

src/main/java/com/alibaba/druid/sharding/sql/MySqlShardingVisitor.java

@@ -17,18 +17,32 @@ package com.alibaba.druid.sharding.sql;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import com.alibaba.druid.sharding.ShardingRuntimeException;
 import com.alibaba.druid.sharding.config.MappingRule;
 import com.alibaba.druid.sharding.config.RouteConfig;
 import com.alibaba.druid.sharding.config.TablePartition;
 import com.alibaba.druid.sql.ast.SQLExpr;
+import com.alibaba.druid.sql.ast.SQLObject;
 import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
 import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
+import com.alibaba.druid.sql.ast.expr.SQLLiteralExpr;
+import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
+import com.alibaba.druid.sql.ast.expr.SQLVariantRefExpr;
+import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
+import com.alibaba.druid.sql.ast.statement.SQLJoinTableSource;
+import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock;
+import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
+import com.alibaba.druid.sql.ast.statement.SQLTableSource;
 import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlInsertStatement;
+import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
 import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitorAdapter;
 import com.alibaba.druid.sql.visitor.SQLEvalVisitorUtils;
+import com.alibaba.druid.util.JdbcConstants;
 
 public class MySqlShardingVisitor extends MySqlASTVisitorAdapter implements ShardingVisitor {
 
@@ -38,6 +52,8 @@ public class MySqlShardingVisitor extends MySqlASTVisitorAdapter implements Shar
 
     private List<SQLStatement> result = new ArrayList<SQLStatement>(2);
 
+    private SQLStatement       input  = null;
+
     public MySqlShardingVisitor(RouteConfig routeConfig, List<Object> parameters){
         this.routeConfig = routeConfig;
         this.parameters = parameters;
@@ -55,7 +71,13 @@ public class MySqlShardingVisitor extends MySqlASTVisitorAdapter implements Shar
         return parameters;
     }
 
+    public boolean visit(SQLSelectStatement x) {
+        input = x;
+        return true;
+    }
+
     public boolean visit(MySqlInsertStatement x) {
+        input = x;
         String table = x.getTableName().getSimleName();
 
         MappingRule mappingRule = routeConfig.getMappingRule(table);
@@ -99,13 +121,233 @@ public class MySqlShardingVisitor extends MySqlASTVisitorAdapter implements Shar
         TablePartition tablePartition = routeConfig.getPartition(table, partition);
 
         x.setTableName(new SQLIdentifierExpr(tablePartition.getTable()));
-        
+
         if (tablePartition.getDatabase() != null) {
             x.putAttribute(ATTR_DB, tablePartition.getDatabase());
         }
 
         result.add(x);
-        
+
         return false;
     }
+
+    @Override
+    public boolean visit(MySqlSelectQueryBlock x) {
+        if (x.getFrom() != null) {
+            x.getFrom().setParent(x);
+            x.getFrom().accept(this);
+        }
+
+        if (x.getWhere() != null) {
+            x.getWhere().setParent(x);
+            x.getWhere().accept(this);
+        }
+
+        return false;
+    }
+
+    @Override
+    public boolean visit(SQLJoinTableSource x) {
+        x.getLeft().setParent(x);
+        x.getRight().setParent(x);
+
+        if (x.getCondition() != null) {
+            x.getCondition().setParent(x);
+        }
+
+        return true;
+    }
+
+    @Override
+    public boolean visit(SQLExprTableSource x) {
+        SQLSelectQueryBlock select = getSQLSelectQueryBlock(x);
+
+        Map<String, SQLTableSource> aliasMap = null;
+        if (select != null) {
+            aliasMap = getAliasMap(select);
+        }
+
+        if (aliasMap != null) {
+            if (x.getAlias() != null) {
+                aliasMap.put(x.getAlias(), x);
+            }
+
+            if (x.getExpr() instanceof SQLIdentifierExpr) {
+                String tableName = ((SQLIdentifierExpr) x.getExpr()).getName();
+                aliasMap.put(tableName, x);
+            }
+        }
+
+        return false;
+    }
+
+    @Override
+    public boolean visit(SQLBinaryOpExpr x) {
+        x.getLeft().setParent(x);
+        x.getRight().setParent(x);
+
+        x.getLeft().accept(this);
+        x.getRight().accept(this);
+
+        String column = null;
+        if ((column = getColumn(x.getLeft())) != null && isValue(x.getRight())) {
+            boolean isMappingColumn = false;
+            SQLTableSource tableSource = getBinaryOpExprLeftOrRightTableSource(x.getLeft());
+            MappingRule mappingRule = getMappingRule(tableSource);
+            if (mappingRule != null) {
+                if (mappingRule.getColumn().equalsIgnoreCase(column)) {
+                    isMappingColumn = true;
+                }
+            }
+            if (isMappingColumn) {
+                Object value = SQLEvalVisitorUtils.eval(JdbcConstants.MYSQL, x.getRight(), parameters);
+                String partitionName = null;
+                switch (x.getOperator()) {
+                    case Equality:
+                        partitionName = mappingRule.getPartition(value);
+                        break;
+                    default:
+                        throw new ShardingRuntimeException("not support operator " + x.getOperator());
+                }
+
+                if (partitionName == null) {
+                    throw new ShardingRuntimeException("sharding rule violation, partition not match, value : " + value);
+                }
+
+                TablePartition tablePartition = routeConfig.getPartition(mappingRule.getTable(), partitionName);
+
+                if (tableSource.getAttribute(ATTR_PARTITION) == null) {
+                    ((SQLExprTableSource) tableSource).setExpr(new SQLIdentifierExpr(tablePartition.getTable()));
+
+                    tableSource.putAttribute(ATTR_PARTITION, tablePartition);
+
+                    if (tablePartition.getDatabase() != null && input != null) {
+                        input.putAttribute(ATTR_DB, tablePartition.getDatabase());
+                    }
+                } else if (tableSource.getAttribute(ATTR_PARTITION) != tablePartition) {
+                    throw new ShardingRuntimeException("sharding rule violation, multi-partition matched, value : "
+                                                       + value);
+                }
+            }
+        }
+
+        return false;
+    }
+
+    MappingRule getMappingRule(SQLTableSource tableSource) {
+        if (tableSource instanceof SQLExprTableSource) {
+            SQLExpr expr = ((SQLExprTableSource) tableSource).getExpr();
+            if (expr instanceof SQLIdentifierExpr) {
+                String table = ((SQLIdentifierExpr) expr).getName();
+                return routeConfig.getMappingRule(table);
+            }
+        }
+
+        return null;
+    }
+
+    static boolean isValue(SQLExpr x) {
+        return x instanceof SQLLiteralExpr || x instanceof SQLVariantRefExpr;
+    }
+
+    static String getColumn(SQLExpr x) {
+        if (x instanceof SQLPropertyExpr) {
+            return ((SQLPropertyExpr) x).getName();
+        }
+        if (x instanceof SQLIdentifierExpr) {
+            return ((SQLIdentifierExpr) x).getName();
+        }
+        return null;
+    }
+
+    static SQLTableSource getBinaryOpExprLeftOrRightTableSource(SQLExpr x) {
+        SQLTableSource tableSource = (SQLTableSource) x.getAttribute(ATTR_TABLE_SOURCE);
+        if (tableSource != null) {
+            return tableSource;
+        }
+
+        SQLSelectQueryBlock select = getSQLSelectQueryBlock(x.getParent());
+        if (select != null && select.getFrom() instanceof SQLExprTableSource) {
+            SQLExpr expr = ((SQLExprTableSource) select.getFrom()).getExpr();
+            if (expr instanceof SQLIdentifierExpr) {
+                x.putAttribute(ATTR_TABLE_SOURCE, select.getFrom());
+                return select.getFrom();
+            }
+        }
+
+        return null;
+    }
+
+    @Override
+    public boolean visit(SQLPropertyExpr x) {
+        x.getOwner().setParent(x);
+
+        x.getOwner().accept(this);
+
+        return false;
+    }
+
+    @Override
+    public boolean visit(SQLIdentifierExpr x) {
+        SQLTableSource tableSource = getTableSource(x.getName(), x.getParent());
+        if (tableSource != null) {
+            x.putAttribute(ATTR_TABLE_SOURCE, tableSource);
+        }
+
+        return false;
+    }
+
+    public static SQLTableSource getTableSource(String name, SQLObject parent) {
+        SQLSelectQueryBlock select = getSQLSelectQueryBlock(parent);
+
+        if (select == null) {
+            return null;
+        }
+
+        Map<String, SQLTableSource> aliasMap = getAliasMap(select);
+
+        if (aliasMap == null) {
+            return null;
+        }
+
+        SQLTableSource tableSource = aliasMap.get(name);
+
+        if (tableSource != null) {
+            return tableSource;
+        }
+
+        for (Map.Entry<String, SQLTableSource> entry : aliasMap.entrySet()) {
+            if (name.equalsIgnoreCase(entry.getKey())) {
+                return tableSource;
+            }
+        }
+
+        return null;
+    }
+
+    @SuppressWarnings("unchecked")
+    public static Map<String, SQLTableSource> getAliasMap(SQLObject x) {
+        Map<String, SQLTableSource> map = (Map<String, SQLTableSource>) x.getAttribute(ATTR_ALIAS);
+        if (map == null) {
+            map = new HashMap<String, SQLTableSource>();
+            x.putAttribute(ATTR_ALIAS, map);
+        }
+        return map;
+    }
+
+    public static SQLSelectQueryBlock getSQLSelectQueryBlock(SQLObject x) {
+        if (x == null) {
+            return null;
+        }
+
+        if (x instanceof SQLSelectQueryBlock) {
+            return (SQLSelectQueryBlock) x;
+        }
+
+        if (x.getParent() != null) {
+            return getSQLSelectQueryBlock(x.getParent());
+        }
+
+        return null;
+    }
 }

src/main/java/com/alibaba/druid/sharding/sql/ShardingVisitor.java

@@ -1,5 +1,24 @@
+/*
+ * Copyright 1999-2011 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
 package com.alibaba.druid.sharding.sql;
 
 public interface ShardingVisitor {
-    public static final String ATTR_DB = "sharding.db";
+
+    public static final String ATTR_DB           = "sharding.db";
+    public static final String ATTR_PARTITION    = "sharding.partition";
+    public static final String ATTR_TABLE_SOURCE = "sharding.tableSource";
+    public static final String ATTR_ALIAS        = "sharding.alias";
 }

src/test/java/com/alibaba/druid/bvt/sharding/MySqlShardingSelectTest.java

@@ -0,0 +1,83 @@
+package com.alibaba.druid.bvt.sharding;
+
+import java.util.Arrays;
+import java.util.List;
+
+import junit.framework.Assert;
+import junit.framework.TestCase;
+
+import com.alibaba.druid.sharding.config.LogicTable;
+import com.alibaba.druid.sharding.config.MappingRuleAdapter;
+import com.alibaba.druid.sharding.config.MappingRuleListEntry;
+import com.alibaba.druid.sharding.config.RouteConfig;
+import com.alibaba.druid.sharding.sql.MySqlShardingVisitor;
+import com.alibaba.druid.sharding.sql.ShardingVisitor;
+import com.alibaba.druid.sql.SQLUtils;
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser;
+import com.alibaba.druid.sql.parser.SQLStatementParser;
+
+public class MySqlShardingSelectTest extends TestCase {
+
+    private RouteConfig routeConfig = new RouteConfig();
+
+    protected void setUp() throws Exception {
+        MappingRuleAdapter rule = new MappingRuleAdapter();
+
+        rule.setColumn("gender");
+        rule.setTable("person");
+        rule.getEntries().add(new MappingRuleListEntry("person_m", "M"));
+        rule.getEntries().add(new MappingRuleListEntry("person_f", "F"));
+
+        routeConfig.getMappingRules().put("person", rule);
+
+        LogicTable logicTable = new LogicTable();
+        logicTable.setLogicTable("person");
+
+        logicTable.addPartition("person_m", "db0", "person_001");
+        logicTable.addPartition("person_f", "db1", "person_002");
+
+        routeConfig.getLogicTables().put("person", logicTable);
+    }
+
+    public void test_insert_m() throws Exception {
+        String sql = "select * from person where name = 'xx' AND gender = 'M'";
+
+        SQLStatementParser parser = new MySqlStatementParser(sql);
+        List<SQLStatement> stmtList = parser.parseStatementList();
+        Assert.assertEquals(1, stmtList.size());
+
+        SQLStatement stmt = stmtList.get(0);
+        MySqlShardingVisitor visitor = new MySqlShardingVisitor(routeConfig);
+        stmt.accept(visitor);
+
+        String result = SQLUtils.toMySqlString(stmt);
+
+        Assert.assertEquals("SELECT *" + //
+                            "\nFROM person_001" + //
+                            "\nWHERE name = 'xx'" + //
+                            "\n\tAND gender = 'M'", result);
+        Assert.assertEquals("db0", stmt.getAttribute(ShardingVisitor.ATTR_DB));
+    }
+
+    public void test_insert_f_params() throws Exception {
+        String sql = "select * from person where name = 'xx' AND gender = ?";
+
+        SQLStatementParser parser = new MySqlStatementParser(sql);
+        List<SQLStatement> stmtList = parser.parseStatementList();
+        Assert.assertEquals(1, stmtList.size());
+
+        SQLStatement stmt = stmtList.get(0);
+        MySqlShardingVisitor visitor = new MySqlShardingVisitor(routeConfig, "F"); // 参数化
+
+        stmt.accept(visitor);
+
+        String result = SQLUtils.toMySqlString(stmt);
+
+        Assert.assertEquals("SELECT *" + //
+                            "\nFROM person_002" + //
+                            "\nWHERE name = 'xx'" + //
+                            "\n\tAND gender = ?", result);
+        Assert.assertEquals("db1", stmt.getAttribute(ShardingVisitor.ATTR_DB));
+    }
+}

src/test/java/com/alibaba/druid/bvt/sharding/MysqlShardingInsertTest.java

@@ -17,7 +17,7 @@ import com.alibaba.druid.sql.ast.SQLStatement;
 import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser;
 import com.alibaba.druid.sql.parser.SQLStatementParser;
 
-public class InsertTest extends TestCase {
+public class MysqlShardingInsertTest extends TestCase {
 
     private RouteConfig routeConfig = new RouteConfig();