enhance mapping for ORM sharding

src/main/java/com/alibaba/druid/mapping/Entity.java

@@ -71,4 +71,17 @@ public class Entity {
         this.properties = properties;
     }
 
+    public String toString() {
+        StringBuilder buf = new StringBuilder();
+        buf.append("{name=");
+        buf.append(name);
+        
+        buf.append(", tableName=");
+        buf.append(tableName);
+        
+        buf.append("}");
+
+        return buf.toString();
+    }
+
 }

src/main/java/com/alibaba/druid/mapping/MappingEngine.java

@@ -96,21 +96,29 @@ public class MappingEngine {
     public String explainToSelectSQL(String sql, List<Object> parameters) {
         SQLSelectQueryBlock query = explainToSelectSQLObject(sql);
 
-        query.accept(this.createMappingVisitor(parameters));
+        MappingVisitor visitor = this.createMappingVisitor(parameters);
+        query.accept(visitor);
+        afterResole(visitor);
 
         return toSQL(query);
     }
+    
+    public void afterResole(MappingVisitor visitor) {
+        
+    }
 
     public SQLDeleteStatement explainToDeleteSQLObject(String sql) {
         return provider.explainToDeleteSQLObject(this, sql);
     }
 
     public String explainToDeleteSQL(String sql) {
-        SQLDeleteStatement query = explainToDeleteSQLObject(sql);
+        SQLDeleteStatement stmt = explainToDeleteSQLObject(sql);
         
-        query.accept(this.createMappingVisitor());
+        MappingVisitor visitor = this.createMappingVisitor(Collections.emptyList());
+        stmt.accept(visitor);
+        afterResole(visitor);
 
-        return toSQL(query);
+        return toSQL(stmt);
     }
     
     public String resolveTableName(Entity entity, List<Object> parameters) {
@@ -126,11 +134,13 @@ public class MappingEngine {
     }
 
     public String explainToUpdateSQL(String sql) {
-        SQLUpdateStatement query = explainToUpdateSQLObject(sql);
+        SQLUpdateStatement stmt = explainToUpdateSQLObject(sql);
 
-        query.accept(this.createMappingVisitor());
+        MappingVisitor visitor = this.createMappingVisitor(Collections.emptyList());
+        stmt.accept(visitor);
+        afterResole(visitor);
 
-        return toSQL(query);
+        return toSQL(stmt);
     }
 
     public SQLInsertStatement explainToInsertSQLObject(String sql) {
@@ -138,11 +148,13 @@ public class MappingEngine {
     }
 
     public String explainToInsertSQL(String sql) {
-        SQLInsertStatement query = explainToInsertSQLObject(sql);
+        SQLInsertStatement stmt = explainToInsertSQLObject(sql);
 
-        query.accept(this.createMappingVisitor());
+        MappingVisitor visitor = this.createMappingVisitor(Collections.emptyList());
+        stmt.accept(visitor);
+        afterResole(visitor);
 
-        return toSQL(query);
+        return toSQL(stmt);
     }
 
     public List<Object> exportParameters(SQLObject sqlObject) {

src/main/java/com/alibaba/druid/mapping/Property.java

@@ -1,7 +1,6 @@
 package com.alibaba.druid.mapping;
 
 public class Property {
-
     private String name;
     private String desciption;
     private String dbColumnName;
@@ -39,4 +38,16 @@ public class Property {
         this.desciption = desciption;
     }
 
+    public String toString() {
+        StringBuilder buf = new StringBuilder();
+        buf.append("{name=");
+        buf.append(name);
+        
+        buf.append(", dbColumnName=");
+        buf.append(dbColumnName);
+        
+        buf.append("}");
+
+        return buf.toString();
+    }
 }

src/main/java/com/alibaba/druid/mapping/spi/MSSQLServerMappingVisitor.java

@@ -1,5 +1,6 @@
 package com.alibaba.druid.mapping.spi;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -7,8 +8,11 @@ import java.util.Map;
 
 import com.alibaba.druid.mapping.Entity;
 import com.alibaba.druid.mapping.MappingEngine;
+import com.alibaba.druid.mapping.Property;
+import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
 import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
 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.SQLSelectItem;
@@ -20,10 +24,12 @@ import com.alibaba.druid.sql.dialect.sqlserver.visitor.SQLServerASTVisitorAdapte
 
 public class MSSQLServerMappingVisitor extends SQLServerASTVisitorAdapter implements MappingVisitor {
 
-    private final MappingEngine                     engine;
-    private final Map<String, SQLTableSource> tableSources = new LinkedHashMap<String, SQLTableSource>();
+    private final MappingEngine               engine;
+    private final Map<String, SQLTableSource> tableSources   = new LinkedHashMap<String, SQLTableSource>();
 
     private final List<Object>                parameters;
+    private final List<PropertyValue>         propertyValues = new ArrayList<PropertyValue>();
+    private int                               variantIndex   = 0;
 
     public MSSQLServerMappingVisitor(MappingEngine engine){
         this(engine, Collections.emptyList());
@@ -42,6 +48,10 @@ public class MSSQLServerMappingVisitor extends SQLServerASTVisitorAdapter implem
         return parameters;
     }
 
+    public List<PropertyValue> getPropertyValues() {
+        return propertyValues;
+    }
+
     public LinkedHashMap<String, Entity> getEntities() {
         return engine.getEntities();
     }
@@ -50,6 +60,16 @@ public class MSSQLServerMappingVisitor extends SQLServerASTVisitorAdapter implem
         return tableSources;
     }
 
+    @Override
+    public String resolveTableName(Entity entity) {
+        return engine.resolveTableName(entity, parameters);
+    }
+
+    @Override
+    public String resovleColumnName(Entity entity, Property property) {
+        return engine.resovleColumnName(entity, property, parameters);
+    }
+
     @Override
     public Entity getFirstEntity() {
         return engine.getFirstEntity();
@@ -85,6 +105,16 @@ public class MSSQLServerMappingVisitor extends SQLServerASTVisitorAdapter implem
         return MappingVisitorUtils.visit(this, x);
     }
 
+    @Override
+    public boolean visit(SQLBinaryOpExpr x) {
+        return MappingVisitorUtils.visit(this, x);
+    }
+
+    @Override
+    public boolean visit(SQLVariantRefExpr x) {
+        return MappingVisitorUtils.visit(this, x);
+    }
+    
     @Override
     public boolean visit(SQLSubqueryTableSource x) {
         return MappingVisitorUtils.visit(this, x);
@@ -99,5 +129,13 @@ public class MSSQLServerMappingVisitor extends SQLServerASTVisitorAdapter implem
     public boolean visit(SQLExprTableSource x) {
         return MappingVisitorUtils.visit(this, x);
     }
-
+    
+    @Override
+    public int getAndIncrementVariantIndex() {
+        return variantIndex++;
+    }
+    
+    public int getVariantIndex() {
+        return variantIndex;
+    }
 }

src/main/java/com/alibaba/druid/mapping/spi/MappingVisitor.java

@@ -5,6 +5,7 @@ import java.util.Map;
 
 import com.alibaba.druid.mapping.Entity;
 import com.alibaba.druid.mapping.MappingEngine;
+import com.alibaba.druid.mapping.Property;
 import com.alibaba.druid.sql.ast.statement.SQLTableSource;
 import com.alibaba.druid.sql.visitor.SQLASTVisitor;
 
@@ -20,4 +21,12 @@ public interface MappingVisitor extends SQLASTVisitor {
     Map<String, Entity> getEntities();
 
     Map<String, SQLTableSource> getTableSources();
+    
+    String resolveTableName(Entity entity);
+    
+    String resovleColumnName(Entity entity, Property property);
+    
+    List<PropertyValue> getPropertyValues();
+    
+    int getAndIncrementVariantIndex();
 }

src/main/java/com/alibaba/druid/mapping/spi/MappingVisitorUtils.java

@@ -8,8 +8,13 @@ import com.alibaba.druid.mapping.MappingEngine;
 import com.alibaba.druid.mapping.Property;
 import com.alibaba.druid.sql.ast.SQLExpr;
 import com.alibaba.druid.sql.ast.expr.SQLAllColumnExpr;
+import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
+import com.alibaba.druid.sql.ast.expr.SQLBinaryOperator;
+import com.alibaba.druid.sql.ast.expr.SQLCharExpr;
 import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
+import com.alibaba.druid.sql.ast.expr.SQLNumericLiteralExpr;
 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.SQLInsertStatement;
@@ -34,7 +39,7 @@ public class MappingVisitorUtils {
             }
 
             x.putAttribute("mapping.entity", entity);
-            String tableName = visitor.getEngine().resolveTableName(entity, visitor.getParameters());
+            String tableName = visitor.resolveTableName(entity);
             tableExpr.setName(tableName);
         }
 
@@ -75,8 +80,10 @@ public class MappingVisitorUtils {
             throw new DruidMappingException("property not found : " + propertyName);
         }
 
-        String dbColumName = visitor.getEngine().resovleColumnName(entity, property, visitor.getParameters());
+        String dbColumName = visitor.resovleColumnName(entity, property);
         x.setName(dbColumName);
+        x.putAttribute("mapping.property", property);
+        x.putAttribute("mapping.entity", entity);
 
         if (x.getParent() instanceof SQLSelectItem) {
             SQLSelectItem selectItem = (SQLSelectItem) x.getParent();
@@ -87,7 +94,7 @@ public class MappingVisitorUtils {
 
         return false;
     }
-
+    
     public static boolean visit(MappingVisitor visitor, SQLIdentifierExpr x) {
         String propertyName = x.getName();
 
@@ -105,8 +112,11 @@ public class MappingVisitorUtils {
             throw new DruidMappingException("property not found : " + propertyName);
         }
 
-        String dbColumName = visitor.getEngine().resovleColumnName(propertyEntity, property, visitor.getParameters());
+        String dbColumName = visitor.resovleColumnName(propertyEntity, property);
         x.setName(dbColumName);
+        
+        x.putAttribute("mapping.property", property);
+        x.putAttribute("mapping.entity", propertyEntity);
 
         if (x.getParent() instanceof SQLSelectItem) {
             SQLSelectItem selectItem = (SQLSelectItem) x.getParent();
@@ -117,6 +127,76 @@ public class MappingVisitorUtils {
 
         return false;
     }
+    
+    public static boolean visit(MappingVisitor visitor, SQLBinaryOpExpr x) {
+        if (x.getOperator() == SQLBinaryOperator.Equality) {
+            if (x.getLeft() instanceof SQLIdentifierExpr && isSimpleValue(visitor, x.getRight())) {
+                visit(visitor, (SQLIdentifierExpr) x.getLeft());
+                x.getRight().accept(visitor);
+                
+                Entity entity = (Entity) x.getLeft().getAttribute("mapping.entity");
+                Property property = (Property) x.getLeft().getAttribute("mapping.property");
+                Object value = x.getRight().getAttribute("mapping.value");
+                
+                visitor.getPropertyValues().add(new PropertyValue(entity, property, value));
+                
+                return false;
+            }
+            
+            if (x.getLeft() instanceof SQLPropertyExpr && isSimpleValue(visitor, x.getRight())) {
+                visit(visitor, (SQLPropertyExpr) x.getLeft());
+                x.getRight().accept(visitor);
+                
+                Entity entity = (Entity) x.getLeft().getAttribute("mapping.entity");
+                Property property = (Property) x.getLeft().getAttribute("mapping.property");
+                Object value = x.getRight().getAttribute("mapping.value");
+                
+                visitor.getPropertyValues().add(new PropertyValue(entity, property, value));
+                
+                return false;
+            }
+        }
+        
+        return true;
+    }
+    
+    private static boolean isSimpleValue(MappingVisitor visitor, SQLExpr expr) {
+        if (expr instanceof SQLNumericLiteralExpr) {
+            expr.putAttribute("mapping.value", ((SQLNumericLiteralExpr) expr).getNumber());
+            return true;
+        }
+        
+        if (expr instanceof SQLCharExpr) {
+            expr.putAttribute("mapping.value", ((SQLCharExpr) expr).getText());
+            return true;
+        }
+        
+        if (expr instanceof SQLVariantRefExpr) {
+            Map<String, Object> attributes = expr.getAttributes();
+            Integer varIndex = (Integer) attributes.get("mapping.varIndex");
+            if (varIndex == null) {
+                varIndex = visitor.getAndIncrementVariantIndex();
+                expr.putAttribute("mapping.varIndex", varIndex);
+            }
+            
+            Object parameter = visitor.getParameters().get(varIndex);
+            
+            expr.putAttribute("mapping.value", parameter);
+            return true;
+        }
+        
+        return false;
+    }
+    
+    public static boolean visit(MappingVisitor visitor, SQLVariantRefExpr x) {
+        Map<String, Object> attributes = x.getAttributes();
+        Integer varIndex = (Integer) attributes.get("mapping.varIndex");
+        if (varIndex == null) {
+            varIndex = visitor.getAndIncrementVariantIndex();
+            x.putAttribute("mapping.varIndex", varIndex);
+        }
+        return false;
+    }
 
     public static boolean visit(MappingVisitor visitor, SQLSelectQueryBlock x) {
         if (x.getSelectList().size() == 0) {

src/main/java/com/alibaba/druid/mapping/spi/MySqlMappingVisitor.java

@@ -1,5 +1,6 @@
 package com.alibaba.druid.mapping.spi;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -7,8 +8,11 @@ import java.util.Map;
 
 import com.alibaba.druid.mapping.Entity;
 import com.alibaba.druid.mapping.MappingEngine;
+import com.alibaba.druid.mapping.Property;
+import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
 import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
 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.SQLSelectItem;
@@ -20,10 +24,12 @@ import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlASTVisitorAdapter;
 
 public class MySqlMappingVisitor extends MySqlASTVisitorAdapter implements MappingVisitor {
 
-    private final MappingEngine                     engine;
-    private final Map<String, SQLTableSource> tableSources = new LinkedHashMap<String, SQLTableSource>();
+    private final MappingEngine               engine;
+    private final Map<String, SQLTableSource> tableSources   = new LinkedHashMap<String, SQLTableSource>();
 
     private final List<Object>                parameters;
+    private final List<PropertyValue>         propertyValues = new ArrayList<PropertyValue>();
+    private int                               variantIndex   = 0;
 
     public MySqlMappingVisitor(MappingEngine engine){
         this(engine, Collections.emptyList());
@@ -42,6 +48,10 @@ public class MySqlMappingVisitor extends MySqlASTVisitorAdapter implements Mappi
         return parameters;
     }
 
+    public List<PropertyValue> getPropertyValues() {
+        return propertyValues;
+    }
+
     public Map<String, SQLTableSource> getTableSources() {
         return tableSources;
     }
@@ -50,6 +60,16 @@ public class MySqlMappingVisitor extends MySqlASTVisitorAdapter implements Mappi
         return engine.getEntities();
     }
 
+    @Override
+    public String resolveTableName(Entity entity) {
+        return engine.resolveTableName(entity, parameters);
+    }
+
+    @Override
+    public String resovleColumnName(Entity entity, Property property) {
+        return engine.resovleColumnName(entity, property, parameters);
+    }
+
     public Entity getFirstEntity() {
         return engine.getFirstEntity();
     }
@@ -84,6 +104,16 @@ public class MySqlMappingVisitor extends MySqlASTVisitorAdapter implements Mappi
         return MappingVisitorUtils.visit(this, x);
     }
 
+    @Override
+    public boolean visit(SQLBinaryOpExpr x) {
+        return MappingVisitorUtils.visit(this, x);
+    }
+    
+    @Override
+    public boolean visit(SQLVariantRefExpr x) {
+        return MappingVisitorUtils.visit(this, x);
+    }
+
     @Override
     public boolean visit(SQLSubqueryTableSource x) {
         return MappingVisitorUtils.visit(this, x);
@@ -98,4 +128,13 @@ public class MySqlMappingVisitor extends MySqlASTVisitorAdapter implements Mappi
     public boolean visit(SQLExprTableSource x) {
         return MappingVisitorUtils.visit(this, x);
     }
+    
+    @Override
+    public int getAndIncrementVariantIndex() {
+        return variantIndex++;
+    }
+    
+    public int getVariantIndex() {
+        return variantIndex;
+    }
 }

src/main/java/com/alibaba/druid/mapping/spi/OracleMappingVisitor.java

@@ -1,5 +1,6 @@
 package com.alibaba.druid.mapping.spi;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -7,8 +8,11 @@ import java.util.Map;
 
 import com.alibaba.druid.mapping.Entity;
 import com.alibaba.druid.mapping.MappingEngine;
+import com.alibaba.druid.mapping.Property;
+import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
 import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
 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.SQLSelectItem;
@@ -22,10 +26,12 @@ import com.alibaba.druid.sql.dialect.oracle.visitor.OracleASTVisitorAdapter;
 
 public class OracleMappingVisitor extends OracleASTVisitorAdapter implements MappingVisitor {
 
-    private final MappingEngine                     engine;
+    private final MappingEngine               engine;
     private final Map<String, SQLTableSource> tableSources = new LinkedHashMap<String, SQLTableSource>();
 
     private final List<Object>                parameters;
+    private final List<PropertyValue>         propertyValues = new ArrayList<PropertyValue>();
+    private int                               variantIndex   = 0;
 
     public OracleMappingVisitor(MappingEngine engine){
         this(engine, Collections.emptyList());
@@ -43,6 +49,10 @@ public class OracleMappingVisitor extends OracleASTVisitorAdapter implements Map
     public List<Object> getParameters() {
         return parameters;
     }
+    
+    public List<PropertyValue> getPropertyValues() {
+        return propertyValues;
+    }
 
     public LinkedHashMap<String, Entity> getEntities() {
         return engine.getEntities();
@@ -52,6 +62,16 @@ public class OracleMappingVisitor extends OracleASTVisitorAdapter implements Map
         return tableSources;
     }
 
+    @Override
+    public String resolveTableName(Entity entity) {
+        return engine.resolveTableName(entity, parameters);
+    }
+
+    @Override
+    public String resovleColumnName(Entity entity, Property property) {
+        return engine.resovleColumnName(entity, property, parameters);
+    }
+
     public Entity getFirstEntity() {
         return engine.getFirstEntity();
     }
@@ -86,6 +106,16 @@ public class OracleMappingVisitor extends OracleASTVisitorAdapter implements Map
         return MappingVisitorUtils.visit(this, x);
     }
 
+    @Override
+    public boolean visit(SQLBinaryOpExpr x) {
+        return MappingVisitorUtils.visit(this, x);
+    }
+    
+    @Override
+    public boolean visit(SQLVariantRefExpr x) {
+        return MappingVisitorUtils.visit(this, x);
+    }
+
     @Override
     public boolean visit(SQLSubqueryTableSource x) {
         return MappingVisitorUtils.visit(this, x);
@@ -105,10 +135,19 @@ public class OracleMappingVisitor extends OracleASTVisitorAdapter implements Map
     public boolean visit(OracleSelectTableReference x) {
         return MappingVisitorUtils.visit(this, x);
     }
-    
+
     @Override
     public boolean visit(OracleDeleteStatement x) {
 
         return true;
     }
+
+    @Override
+    public int getAndIncrementVariantIndex() {
+        return variantIndex++;
+    }
+    
+    public int getVariantIndex() {
+        return variantIndex;
+    }
 }

src/main/java/com/alibaba/druid/mapping/spi/PropertyValue.java

@@ -0,0 +1,45 @@
+package com.alibaba.druid.mapping.spi;
+
+import com.alibaba.druid.mapping.Entity;
+import com.alibaba.druid.mapping.Property;
+
+public class PropertyValue {
+
+    private Property property;
+    private Entity   entity;
+    private Object   value;
+
+    public PropertyValue(Entity entity, Property property, Object value){
+        this.property = property;
+        this.entity = entity;
+        this.value = value;
+    }
+
+    public PropertyValue(){
+    }
+
+    public Property getProperty() {
+        return property;
+    }
+
+    public void setProperty(Property property) {
+        this.property = property;
+    }
+
+    public Entity getEntity() {
+        return entity;
+    }
+
+    public void setEntity(Entity entity) {
+        this.entity = entity;
+    }
+
+    public Object getValue() {
+        return value;
+    }
+
+    public void setValue(Object value) {
+        this.value = value;
+    }
+
+}