Adjust DatabaseDiscoveryDataSourceRule getDataSourceMapper for primar…

…y datasource first (apache#23214)
Byeongcheol-Kim · Dec 31, 2022 · 1d2194b · 1d2194b
1 parent ea9cee4
commit 1d2194b
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 24 deletions.
diff --git a/...main/java/org/apache/shardingsphere/dbdiscovery/rule/DatabaseDiscoveryDataSourceRule.java b/...main/java/org/apache/shardingsphere/dbdiscovery/rule/DatabaseDiscoveryDataSourceRule.java
@@ -27,7 +27,7 @@
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.LinkedList;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -125,8 +125,14 @@ public Map<String, DataSource> getDataSourceGroup(final Map<String, DataSource>
      */
     public Map<String, Collection<String>> getDataSourceMapper() {
         Map<String, Collection<String>> result = new HashMap<>(1, 1);
-        Collection<String> actualDataSourceNames = new LinkedList<>(dataSourceNames);
-        result.put(groupName, actualDataSourceNames);
+        result.put(groupName, getActualDataSourceNames());
+        return result;
+    }
+
+    private Collection<String> getActualDataSourceNames() {
+        Collection<String> result = new LinkedHashSet<>();
+        result.add(primaryDataSourceName);
+        result.addAll(dataSourceNames);
         return result;
     }
 }
diff --git a/.../java/org/apache/shardingsphere/dbdiscovery/rule/DatabaseDiscoveryDataSourceRuleTest.java b/.../java/org/apache/shardingsphere/dbdiscovery/rule/DatabaseDiscoveryDataSourceRuleTest.java
@@ -22,10 +22,8 @@
 import org.junit.Test;
 
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.Collections;
-import java.util.LinkedHashMap;
-import java.util.Map;
+import java.util.HashSet;
 import java.util.Properties;
 
 import static org.hamcrest.CoreMatchers.is;
@@ -62,12 +60,7 @@ public void assertGetDataSourceNamesWithoutDisabledDataSourceNames() {
 
     @Test
     public void assertGetDataSourceMapper() {
-        assertThat(databaseDiscoveryDataSourceRule.getDataSourceMapper(), is(getExpectedDataSourceMapper()));
-    }
-
-    private Map<String, Collection<String>> getExpectedDataSourceMapper() {
-        Map<String, Collection<String>> result = new LinkedHashMap<>(2, 1);
-        result.put("test_pr", Arrays.asList("ds_0", "ds_1"));
-        return result;
+        databaseDiscoveryDataSourceRule.changePrimaryDataSourceName("ds_1");
+        assertThat(databaseDiscoveryDataSourceRule.getDataSourceMapper(), is(Collections.singletonMap("test_pr", new HashSet<>(Arrays.asList("ds_1", "ds_0")))));
     }
 }
diff --git a/...e/src/test/java/org/apache/shardingsphere/dbdiscovery/rule/DatabaseDiscoveryRuleTest.java b/...e/src/test/java/org/apache/shardingsphere/dbdiscovery/rule/DatabaseDiscoveryRuleTest.java
@@ -23,8 +23,8 @@
 import org.apache.shardingsphere.infra.config.algorithm.AlgorithmConfiguration;
 import org.apache.shardingsphere.infra.config.mode.ModeConfiguration;
 import org.apache.shardingsphere.infra.config.mode.PersistRepositoryConfiguration;
-import org.apache.shardingsphere.infra.rule.identifier.type.exportable.constant.ExportableConstants;
 import org.apache.shardingsphere.infra.instance.InstanceContext;
+import org.apache.shardingsphere.infra.rule.identifier.type.exportable.constant.ExportableConstants;
 import org.apache.shardingsphere.schedule.core.ScheduleContextFactory;
 import org.apache.shardingsphere.test.fixture.jdbc.MockedDataSource;
 import org.junit.BeforeClass;
@@ -34,7 +34,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
@@ -76,14 +76,7 @@ private void assertDataSourceRule(final DatabaseDiscoveryDataSourceRule actual)
     public void assertGetDataSourceMapper() {
         DatabaseDiscoveryRule databaseDiscoveryRule = createRule();
         Map<String, Collection<String>> actual = databaseDiscoveryRule.getDataSourceMapper();
-        Map<String, Collection<String>> expected = getDataSourceMapper();
-        assertThat(actual, is(expected));
-    }
-
-    private Map<String, Collection<String>> getDataSourceMapper() {
-        Map<String, Collection<String>> result = new HashMap<>(1, 1);
-        result.put("replica_ds", Arrays.asList("primary_ds", "replica_ds_0", "replica_ds_1"));
-        return result;
+        assertThat(actual, is(Collections.singletonMap("replica_ds", new HashSet<>(Arrays.asList("primary_ds", "replica_ds_0", "replica_ds_1")))));
     }
 
     @Test

diff --git a/.../java/org/apache/shardingsphere/mask/distsql/parser/core/MaskDistSQLStatementVisitor.java b/.../java/org/apache/shardingsphere/mask/distsql/parser/core/MaskDistSQLStatementVisitor.java
@@ -107,7 +107,7 @@ private Properties getProperties(final PropertiesDefinitionContext ctx) {
         }
         return result;
     }
-
+    
     @Override
     public ASTNode visitDatabaseName(final DatabaseNameContext ctx) {
         return new DatabaseSegment(ctx.getStart().getStartIndex(), ctx.getStop().getStopIndex(), new IdentifierValue(ctx.getText()));