Refactor ShadowDataSourceMappingsFinderFactory

features/shadow/core/src/main/java/org/apache/shardingsphere/shadow/route/ShadowSQLRouter.java

@@ -22,24 +22,48 @@
 import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
 import org.apache.shardingsphere.infra.route.DecorateSQLRouter;
 import org.apache.shardingsphere.infra.route.context.RouteContext;
+import org.apache.shardingsphere.infra.route.context.RouteMapper;
+import org.apache.shardingsphere.infra.route.context.RouteUnit;
 import org.apache.shardingsphere.infra.session.query.QueryContext;
 import org.apache.shardingsphere.shadow.constant.ShadowOrder;
+import org.apache.shardingsphere.shadow.route.engine.finder.ShadowDataSourceMappingsFinder;
 import org.apache.shardingsphere.shadow.route.engine.finder.ShadowDataSourceMappingsFinderFactory;
-import org.apache.shardingsphere.shadow.route.engine.ShadowRouteEngine;
 import org.apache.shardingsphere.shadow.rule.ShadowRule;
 
+import java.util.Collection;
+import java.util.LinkedList;
+import java.util.Map;
+import java.util.Optional;
+
 /**
  * Shadow SQL router.
  */
 @HighFrequencyInvocation
 public final class ShadowSQLRouter implements DecorateSQLRouter<ShadowRule> {
 
-    private final ShadowRouteEngine routeEngine = new ShadowRouteEngine();
-
     @Override
     public void decorateRouteContext(final RouteContext routeContext, final QueryContext queryContext, final ShardingSphereDatabase database,
                                      final ShadowRule rule, final ConfigurationProperties props) {
-        routeEngine.route(routeContext, rule, ShadowDataSourceMappingsFinderFactory.newInstance(queryContext));
+        route(routeContext, rule, ShadowDataSourceMappingsFinderFactory.newInstance(queryContext));
+    }
+
+    private void route(final RouteContext routeContext, final ShadowRule rule, final ShadowDataSourceMappingsFinder finder) {
+        Collection<RouteUnit> toBeRemovedRouteUnit = new LinkedList<>();
+        Collection<RouteUnit> toBeAddedRouteUnit = new LinkedList<>();
+        Map<String, String> shadowDataSourceMappings = finder.find(rule);
+        for (RouteUnit each : routeContext.getRouteUnits()) {
+            String logicName = each.getDataSourceMapper().getLogicName();
+            String actualName = each.getDataSourceMapper().getActualName();
+            Optional<String> productionDataSourceName = rule.findProductionDataSourceName(actualName);
+            if (productionDataSourceName.isPresent()) {
+                String shadowDataSourceName = shadowDataSourceMappings.get(productionDataSourceName.get());
+                toBeRemovedRouteUnit.add(each);
+                String dataSourceName = null == shadowDataSourceName ? productionDataSourceName.get() : shadowDataSourceName;
+                toBeAddedRouteUnit.add(new RouteUnit(new RouteMapper(logicName, dataSourceName), each.getTableMappers()));
+            }
+        }
+        routeContext.getRouteUnits().removeAll(toBeRemovedRouteUnit);
+        routeContext.getRouteUnits().addAll(toBeAddedRouteUnit);
     }
 
     @Override