From 021126367cc9f8745a22027d927b216d82763127 Mon Sep 17 00:00:00 2001 From: wenshao Date: Sun, 25 Nov 2012 18:18:36 +0800 Subject: [PATCH] add sharding support --- .../sharding/sql/MySqlShardingVisitor.java | 246 +++++++++++++++++- .../druid/sharding/sql/ShardingVisitor.java | 21 +- .../bvt/sharding/MySqlShardingSelectTest.java | 83 ++++++ ...Test.java => MysqlShardingInsertTest.java} | 2 +- 4 files changed, 348 insertions(+), 4 deletions(-) create mode 100644 src/test/java/com/alibaba/druid/bvt/sharding/MySqlShardingSelectTest.java rename src/test/java/com/alibaba/druid/bvt/sharding/{InsertTest.java => MysqlShardingInsertTest.java} (98%) diff --git a/src/main/java/com/alibaba/druid/sharding/sql/MySqlShardingVisitor.java b/src/main/java/com/alibaba/druid/sharding/sql/MySqlShardingVisitor.java index d336c1e876..0b026fd526 100644 --- a/src/main/java/com/alibaba/druid/sharding/sql/MySqlShardingVisitor.java +++ b/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 result = new ArrayList(2); + private SQLStatement input = null; + public MySqlShardingVisitor(RouteConfig routeConfig, List 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 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 aliasMap = getAliasMap(select); + + if (aliasMap == null) { + return null; + } + + SQLTableSource tableSource = aliasMap.get(name); + + if (tableSource != null) { + return tableSource; + } + + for (Map.Entry entry : aliasMap.entrySet()) { + if (name.equalsIgnoreCase(entry.getKey())) { + return tableSource; + } + } + + return null; + } + + @SuppressWarnings("unchecked") + public static Map getAliasMap(SQLObject x) { + Map map = (Map) x.getAttribute(ATTR_ALIAS); + if (map == null) { + map = new HashMap(); + 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; + } } diff --git a/src/main/java/com/alibaba/druid/sharding/sql/ShardingVisitor.java b/src/main/java/com/alibaba/druid/sharding/sql/ShardingVisitor.java index a72fdabb09..a0f4f7a008 100644 --- a/src/main/java/com/alibaba/druid/sharding/sql/ShardingVisitor.java +++ b/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"; } diff --git a/src/test/java/com/alibaba/druid/bvt/sharding/MySqlShardingSelectTest.java b/src/test/java/com/alibaba/druid/bvt/sharding/MySqlShardingSelectTest.java new file mode 100644 index 0000000000..bc00b3ef7d --- /dev/null +++ b/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 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 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)); + } +} diff --git a/src/test/java/com/alibaba/druid/bvt/sharding/InsertTest.java b/src/test/java/com/alibaba/druid/bvt/sharding/MysqlShardingInsertTest.java similarity index 98% rename from src/test/java/com/alibaba/druid/bvt/sharding/InsertTest.java rename to src/test/java/com/alibaba/druid/bvt/sharding/MysqlShardingInsertTest.java index 9c7728de3d..7b2f295dd9 100644 --- a/src/test/java/com/alibaba/druid/bvt/sharding/InsertTest.java +++ b/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();