From 0c0bdc5f5543cde04fc172d4db6265a5ec54883e Mon Sep 17 00:00:00 2001 From: wenshao Date: Sun, 25 Nov 2012 19:51:43 +0800 Subject: [PATCH] sharding support --- .../sharding/sql/MySqlShardingVisitor.java | 103 +++++++++++++----- .../mysql/visitor/MySqlOutputVisitor.java | 4 + .../bvt/sharding/MySqlShardingDeleteTest.java | 81 ++++++++++++++ .../bvt/sharding/MySqlShardingUpdateTest.java | 83 ++++++++++++++ .../bvt/sql/cobar/DMLDeleteParserTest.java | 4 +- .../MySqlParameterizedOutputVisitorTest.java | 2 +- 6 files changed, 246 insertions(+), 31 deletions(-) create mode 100644 src/test/java/com/alibaba/druid/bvt/sharding/MySqlShardingDeleteTest.java create mode 100644 src/test/java/com/alibaba/druid/bvt/sharding/MySqlShardingUpdateTest.java 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 0b026fd526..ff7e00e4dc 100644 --- a/src/main/java/com/alibaba/druid/sharding/sql/MySqlShardingVisitor.java +++ b/src/main/java/com/alibaba/druid/sharding/sql/MySqlShardingVisitor.java @@ -33,13 +33,17 @@ 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.SQLDeleteStatement; 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.ast.statement.SQLUpdateStatement; +import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlDeleteStatement; 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.ast.statement.MySqlUpdateStatement; import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitorAdapter; import com.alibaba.druid.sql.visitor.SQLEvalVisitorUtils; import com.alibaba.druid.util.JdbcConstants; @@ -76,6 +80,47 @@ public class MySqlShardingVisitor extends MySqlASTVisitorAdapter implements Shar return true; } + public boolean visit(MySqlDeleteStatement x) { + input = x; + + if (x.getFrom() != null) { + x.getFrom().setParent(x); + x.getFrom().accept(this); + } + + if (x.getTableSource() != null) { + x.getTableSource().setParent(x); + x.getTableSource().accept(this); + } + + if (x.getWhere() != null) { + x.getWhere().setParent(x); + x.getWhere().accept(this); + } + + result.add(x); + + return false; + } + + @Override + public boolean visit(MySqlUpdateStatement x) { + input = x; + + if (x.getTableSource() != null) { + x.getTableSource().setParent(x); + x.getTableSource().accept(this); + } + + if (x.getWhere() != null) { + x.getWhere().setParent(x); + x.getWhere().accept(this); + } + + result.add(x); + return false; + } + public boolean visit(MySqlInsertStatement x) { input = x; String table = x.getTableName().getSimleName(); @@ -160,12 +205,7 @@ public class MySqlShardingVisitor extends MySqlASTVisitorAdapter implements Shar @Override public boolean visit(SQLExprTableSource x) { - SQLSelectQueryBlock select = getSQLSelectQueryBlock(x); - - Map aliasMap = null; - if (select != null) { - aliasMap = getAliasMap(select); - } + Map aliasMap = getAliasMap(x); if (aliasMap != null) { if (x.getAlias() != null) { @@ -266,12 +306,12 @@ public class MySqlShardingVisitor extends MySqlASTVisitorAdapter implements Shar return tableSource; } - SQLSelectQueryBlock select = getSQLSelectQueryBlock(x.getParent()); - if (select != null && select.getFrom() instanceof SQLExprTableSource) { - SQLExpr expr = ((SQLExprTableSource) select.getFrom()).getExpr(); + SQLTableSource defaltTableSource = getDefaultTableSource(x.getParent()); + if (defaltTableSource instanceof SQLExprTableSource) { + SQLExpr expr = ((SQLExprTableSource) defaltTableSource).getExpr(); if (expr instanceof SQLIdentifierExpr) { - x.putAttribute(ATTR_TABLE_SOURCE, select.getFrom()); - return select.getFrom(); + x.putAttribute(ATTR_TABLE_SOURCE, defaltTableSource); + return defaltTableSource; } } @@ -298,13 +338,7 @@ public class MySqlShardingVisitor extends MySqlASTVisitorAdapter implements Shar } public static SQLTableSource getTableSource(String name, SQLObject parent) { - SQLSelectQueryBlock select = getSQLSelectQueryBlock(parent); - - if (select == null) { - return null; - } - - Map aliasMap = getAliasMap(select); + Map aliasMap = getAliasMap(parent); if (aliasMap == null) { return null; @@ -327,27 +361,40 @@ public class MySqlShardingVisitor extends MySqlASTVisitorAdapter implements Shar @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); + if (x == null) { + return null; } - return map; + + if (x instanceof SQLSelectQueryBlock || x instanceof SQLDeleteStatement) { + Map map = (Map) x.getAttribute(ATTR_ALIAS); + if (map == null) { + map = new HashMap(); + x.putAttribute(ATTR_ALIAS, map); + } + return map; + } + + return getAliasMap(x.getParent()); } - public static SQLSelectQueryBlock getSQLSelectQueryBlock(SQLObject x) { + public static SQLTableSource getDefaultTableSource(SQLObject x) { if (x == null) { return null; } if (x instanceof SQLSelectQueryBlock) { - return (SQLSelectQueryBlock) x; + return ((SQLSelectQueryBlock) x).getFrom(); } - if (x.getParent() != null) { - return getSQLSelectQueryBlock(x.getParent()); + if (x instanceof SQLDeleteStatement) { + return ((SQLDeleteStatement) x).getTableSource(); } - return null; + if (x instanceof SQLUpdateStatement) { + return ((SQLUpdateStatement) x).getTableSource(); + } + + return getDefaultTableSource(x.getParent()); } + } diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java index 58227c65e6..b6db0ca7cd 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java @@ -743,9 +743,11 @@ public class MySqlOutputVisitor extends SQLASTOutputVisitor implements MySqlASTV if (x.getWhere() != null) { println(); + incrementIndent(); print("WHERE "); x.getWhere().setParent(x); x.getWhere().accept(this); + decrementIndent(); } if (x.getOrderBy() != null) { @@ -1505,9 +1507,11 @@ public class MySqlOutputVisitor extends SQLASTOutputVisitor implements MySqlASTV if (x.getWhere() != null) { println(); + incrementIndent(); print("WHERE "); x.getWhere().setParent(x); x.getWhere().accept(this); + decrementIndent(); } if (x.getOrderBy() != null) { diff --git a/src/test/java/com/alibaba/druid/bvt/sharding/MySqlShardingDeleteTest.java b/src/test/java/com/alibaba/druid/bvt/sharding/MySqlShardingDeleteTest.java new file mode 100644 index 0000000000..8bb4e02754 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/sharding/MySqlShardingDeleteTest.java @@ -0,0 +1,81 @@ +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 MySqlShardingDeleteTest 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_delete_m() throws Exception { + String sql = "delete 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("DELETE FROM person_001" + // + "\nWHERE name = 'xx'" + // + "\n\tAND gender = 'M'", result); + Assert.assertEquals("db0", stmt.getAttribute(ShardingVisitor.ATTR_DB)); + } + + public void test_delete_f_params() throws Exception { + String sql = "delete 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("DELETE FROM 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/MySqlShardingUpdateTest.java b/src/test/java/com/alibaba/druid/bvt/sharding/MySqlShardingUpdateTest.java new file mode 100644 index 0000000000..c536628cc8 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/sharding/MySqlShardingUpdateTest.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 MySqlShardingUpdateTest 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_update_m() throws Exception { + String sql = "update person set X = null 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("UPDATE person_001" + // + "\nSET X = NULL" + // + "\nWHERE name = 'xx'" + // + "\n\tAND gender = 'M'", result); + Assert.assertEquals("db0", stmt.getAttribute(ShardingVisitor.ATTR_DB)); + } + + public void test_delete_f_params() throws Exception { + String sql = "update person set X = null 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("UPDATE person_002" + // + "\nSET X = NULL" + // + "\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/sql/cobar/DMLDeleteParserTest.java b/src/test/java/com/alibaba/druid/bvt/sql/cobar/DMLDeleteParserTest.java index 2899c74da7..9710b5e6eb 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/cobar/DMLDeleteParserTest.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/cobar/DMLDeleteParserTest.java @@ -56,7 +56,7 @@ public class DMLDeleteParserTest extends TestCase { Assert.assertEquals("DELETE FROM offer.*, wp_image.*\n" + // "USING offer a, wp_image b\n" + // "WHERE a.member_id = b.member_id\n" + // - "AND a.member_id = 'abc'", output); + "\tAND a.member_id = 'abc'", output); } public void testDelete_3() throws Exception { @@ -91,7 +91,7 @@ public class DMLDeleteParserTest extends TestCase { Assert.assertEquals("DELETE id.*\n" + // "FROM t1, t2\n" + // "WHERE col1 = 'adf'\n" + // - "AND col2 = 1", output); + "\tAND col2 = 1", output); } public void testDelete_6() throws Exception { diff --git a/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlParameterizedOutputVisitorTest.java b/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlParameterizedOutputVisitorTest.java index 7e530125c0..b5f299780b 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlParameterizedOutputVisitorTest.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/mysql/MySqlParameterizedOutputVisitorTest.java @@ -41,7 +41,7 @@ public class MySqlParameterizedOutputVisitorTest extends TestCase { validateOracle("SELECT * FROM T WHERE A.ID = ? OR A.ID = ?", "SELECT *\nFROM T\nWHERE A.ID = ?"); validate("INSERT INTO T (F1, F2) VALUES(?, ?), (?, ?), (?, ?)", "INSERT INTO T (F1, F2) VALUES (?, ?)"); validate("update net_device d, sys_user u set d.resp_user_id=u.id where d.resp_user_login_name=u.username and d.id in (42354)", // - "UPDATE net_device d, sys_user u\nSET d.resp_user_id = u.id\nWHERE d.resp_user_login_name = u.username\nAND d.id IN (?)"); + "UPDATE net_device d, sys_user u\nSET d.resp_user_id = u.id\nWHERE d.resp_user_login_name = u.username\n\tAND d.id IN (?)"); // System.out.println("SELECT * FORM A WHERE ID = (##)".replaceAll("##", "?")); }