From 07ae91eed4bc475fbaa60ef43bf4fb19cb636b38 Mon Sep 17 00:00:00 2001
From: jiangML <1060319118@qq.com>
Date: Wed, 1 Nov 2023 11:23:30 +0800
Subject: [PATCH] Fix getInUsedStorageUnitNameAndRulesMap error for
 RuleMetaData

---
 .../metadata/database/rule/RuleMetaData.java  | 20 +++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/infra/common/src/main/java/org/apache/shardingsphere/infra/metadata/database/rule/RuleMetaData.java b/infra/common/src/main/java/org/apache/shardingsphere/infra/metadata/database/rule/RuleMetaData.java
index 97a3254e7d4b6..0cb3627a14714 100644
--- a/infra/common/src/main/java/org/apache/shardingsphere/infra/metadata/database/rule/RuleMetaData.java
+++ b/infra/common/src/main/java/org/apache/shardingsphere/infra/metadata/database/rule/RuleMetaData.java
@@ -30,6 +30,7 @@
 import java.util.LinkedHashSet;
 import java.util.LinkedList;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.stream.Collectors;
@@ -106,9 +107,9 @@ public Map<String, Collection<Class<? extends ShardingSphereRule>>> getInUsedSto
         Map<String, Collection<Class<? extends ShardingSphereRule>>> result = new LinkedHashMap<>();
         for (ShardingSphereRule each : rules) {
             if (each instanceof DataSourceContainedRule) {
-                result.putAll(getInUsedStorageUnitNameAndRulesMap(each, getInUsedStorageUnitNames((DataSourceContainedRule) each)));
+                mergeInUsedStorageUnitNameAndRules(result, getInUsedStorageUnitNameAndRulesMap(each, getInUsedStorageUnitNames((DataSourceContainedRule) each)));
             } else if (each instanceof DataNodeContainedRule) {
-                result.putAll(getInUsedStorageUnitNameAndRulesMap(each, getInUsedStorageUnitNames((DataNodeContainedRule) each)));
+                mergeInUsedStorageUnitNameAndRules(result, getInUsedStorageUnitNameAndRulesMap(each, getInUsedStorageUnitNames((DataNodeContainedRule) each)));
             }
         }
         return result;
@@ -132,4 +133,19 @@ private Collection<String> getInUsedStorageUnitNames(final DataSourceContainedRu
     private Collection<String> getInUsedStorageUnitNames(final DataNodeContainedRule rule) {
         return rule.getAllDataNodes().values().stream().flatMap(each -> each.stream().map(DataNode::getDataSourceName).collect(Collectors.toSet()).stream()).collect(Collectors.toSet());
     }
+    
+    private void mergeInUsedStorageUnitNameAndRules(final Map<String, Collection<Class<? extends ShardingSphereRule>>> storageUnitNameAndRulesMap,
+                                                    final Map<String, Collection<Class<? extends ShardingSphereRule>>> toBeMergeStorageUnitNameAndRulesMap) {
+        for (Entry<String, Collection<Class<? extends ShardingSphereRule>>> entry : toBeMergeStorageUnitNameAndRulesMap.entrySet()) {
+            if (storageUnitNameAndRulesMap.containsKey(entry.getKey())) {
+                for (Class<? extends ShardingSphereRule> each : entry.getValue()) {
+                    if (!storageUnitNameAndRulesMap.get(entry.getKey()).contains(each)) {
+                        storageUnitNameAndRulesMap.get(entry.getKey()).add(each);
+                    }
+                }
+            } else {
+                storageUnitNameAndRulesMap.put(entry.getKey(), entry.getValue());
+            }
+        }
+    }
 }