sharding support

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<String, SQLTableSource> aliasMap = null;
-        if (select != null) {
-            aliasMap = getAliasMap(select);
-        }
+        Map<String, SQLTableSource> 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<String, SQLTableSource> aliasMap = getAliasMap(select);
+        Map<String, SQLTableSource> aliasMap = getAliasMap(parent);
 
         if (aliasMap == null) {
             return null;
@@ -327,27 +361,40 @@ public class MySqlShardingVisitor extends MySqlASTVisitorAdapter implements Shar
 
     @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);
+        if (x == null) {
+            return null;
         }
-        return map;
+
+        if (x instanceof SQLSelectQueryBlock || x instanceof SQLDeleteStatement) {
+            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;
+        }
+
+        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());
     }
+
 }

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) {

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<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("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<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("DELETE FROM 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/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<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("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<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("UPDATE person_002" + //
+                            "\nSET X = NULL" + //
+                            "\nWHERE name = 'xx'" + //
+                            "\n\tAND gender = ?", result);
+        Assert.assertEquals("db1", stmt.getAttribute(ShardingVisitor.ATTR_DB));
+    }
+}

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 {

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("##", "?"));
     }