Move getAggregatedDataSourceMap from SingleTableLoadUtils to Physical…

…ResourceAggregator (#33202)

* Move getAggregatedDataSourceMap from SingleTableLoadUtils to PhysicalResourceAggregator

* fix unit test

features/sharding/core/src/main/java/org/apache/shardingsphere/sharding/rule/ShardingTable.java

@@ -214,15 +214,6 @@ public Map<String, List<DataNode>> getDataNodeGroups() {
         return DataNodeUtils.getDataNodeGroups(actualDataNodes);
     }
 
-    /**
-     * Get actual data source names.
-     *
-     * @return actual data source names
-     */
-    public Collection<String> getActualDataSourceNames() {
-        return actualDataSourceNames;
-    }
-
     /**
      * Get actual table names via target data source name.
      *

infra/common/src/main/java/org/apache/shardingsphere/infra/metadata/database/resource/PhysicalResourceAggregator.java

@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.shardingsphere.infra.metadata.database.resource;
+
+import com.cedarsoftware.util.CaseInsensitiveMap;
+import lombok.AccessLevel;
+import lombok.NoArgsConstructor;
+import org.apache.shardingsphere.infra.rule.ShardingSphereRule;
+import org.apache.shardingsphere.infra.rule.attribute.datasource.DataSourceMapperRuleAttribute;
+
+import javax.sql.DataSource;
+import java.util.Collection;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Optional;
+
+/**
+ * Physical resource aggregator.
+ */
+@NoArgsConstructor(access = AccessLevel.PRIVATE)
+public final class PhysicalResourceAggregator {
+
+    /**
+     * Get aggregated resources.
+     *
+     * @param dataSourceMap data source map
+     * @param builtRules built rules
+     * @return aggregated resources
+     */
+    public static Map<String, DataSource> getAggregatedResources(final Map<String, DataSource> dataSourceMap, final Collection<ShardingSphereRule> builtRules) {
+        Map<String, DataSource> result = new CaseInsensitiveMap<>(dataSourceMap);
+        for (ShardingSphereRule each : builtRules) {
+            Optional<DataSourceMapperRuleAttribute> ruleAttribute = each.getAttributes().findAttribute(DataSourceMapperRuleAttribute.class);
+            if (ruleAttribute.isPresent()) {
+                result = getAggregatedResources(result, ruleAttribute.get());
+            }
+        }
+        return result;
+    }
+
+    private static Map<String, DataSource> getAggregatedResources(final Map<String, DataSource> dataSourceMap, final DataSourceMapperRuleAttribute ruleAttribute) {
+        Map<String, DataSource> result = new CaseInsensitiveMap<>();
+        for (Entry<String, Collection<String>> entry : ruleAttribute.getDataSourceMapper().entrySet()) {
+            for (String each : entry.getValue()) {
+                if (dataSourceMap.containsKey(each)) {
+                    result.putIfAbsent(entry.getKey(), dataSourceMap.remove(each));
+                }
+            }
+        }
+        result.putAll(dataSourceMap);
+        return result;
+    }
+}

infra/database/core/src/main/java/org/apache/shardingsphere/infra/database/core/metadata/data/loader/type/SchemaMetaDataLoader.java

@@ -63,8 +63,8 @@ public final class SchemaMetaDataLoader {
      * @return loaded schema table names
      * @throws SQLException SQL exception
      */
-    public static Map<String, Collection<String>> loadSchemaTableNamesByExcludedTables(final String databaseName, final DatabaseType databaseType, final DataSource dataSource,
-                                                                                       final Collection<String> excludedTables) throws SQLException {
+    public static Map<String, Collection<String>> loadSchemaTableNames(final String databaseName, final DatabaseType databaseType, final DataSource dataSource,
+                                                                       final Collection<String> excludedTables) throws SQLException {
         try (MetaDataLoaderConnection connection = new MetaDataLoaderConnection(databaseType, dataSource.getConnection())) {
             Collection<String> schemaNames = loadSchemaNames(connection, databaseType);
             DialectDatabaseMetaData dialectDatabaseMetaData = new DatabaseTypeRegistry(databaseType).getDialectDatabaseMetaData();

infra/database/type/mysql/src/test/java/org/apache/shardingsphere/infra/database/mysql/metadata/data/loader/MySQLSchemaMetaDataLoaderTest.java

@@ -74,9 +74,9 @@ private ResultSet mockSchemaResultSet() throws SQLException {
     }
 
     @Test
-    void assertLoadSchemaTableNamesByExcludedTables() throws SQLException {
+    void assertLoadSchemaTableNames() throws SQLException {
         Map<String, Collection<String>> schemaTableNames = Collections.singletonMap(DefaultDatabase.LOGIC_NAME, Collections.singletonList("tbl"));
-        assertThat(SchemaMetaDataLoader.loadSchemaTableNamesByExcludedTables(DefaultDatabase.LOGIC_NAME,
+        assertThat(SchemaMetaDataLoader.loadSchemaTableNames(DefaultDatabase.LOGIC_NAME,
                 TypedSPILoader.getService(DatabaseType.class, "MySQL"), dataSource, Collections.emptyList()), is(schemaTableNames));
     }
 

infra/database/type/opengauss/src/test/java/org/apache/shardingsphere/infra/database/opengauss/metadata/data/loader/OpenGaussSchemaMetaDataLoaderTest.java

@@ -76,8 +76,8 @@ private ResultSet mockSchemaResultSet() throws SQLException {
     }
 
     @Test
-    void assertLoadSchemaTableNamesByExcludedTables() throws SQLException {
-        assertThat(SchemaMetaDataLoader.loadSchemaTableNamesByExcludedTables(DefaultDatabase.LOGIC_NAME,
+    void assertLoadSchemaTableNames() throws SQLException {
+        assertThat(SchemaMetaDataLoader.loadSchemaTableNames(DefaultDatabase.LOGIC_NAME,
                 TypedSPILoader.getService(DatabaseType.class, "openGauss"), dataSource, Collections.emptyList()), is(createSchemaTableNames()));
     }
 

infra/database/type/postgresql/src/test/java/org/apache/shardingsphere/infra/database/postgresql/metadata/data/loader/PostgreSQLSchemaMetaDataLoaderTest.java

@@ -76,8 +76,8 @@ private ResultSet mockSchemaResultSet() throws SQLException {
     }
 
     @Test
-    void assertLoadSchemaTableNamesByExcludedTables() throws SQLException {
-        assertThat(SchemaMetaDataLoader.loadSchemaTableNamesByExcludedTables(DefaultDatabase.LOGIC_NAME,
+    void assertLoadSchemaTableNames() throws SQLException {
+        assertThat(SchemaMetaDataLoader.loadSchemaTableNames(DefaultDatabase.LOGIC_NAME,
                 TypedSPILoader.getService(DatabaseType.class, "PostgreSQL"), dataSource, Collections.emptyList()), is(createSchemaTableNames()));
     }
 

kernel/single/core/src/main/java/org/apache/shardingsphere/single/datanode/SingleTableDataNodeLoader.java

@@ -182,7 +182,7 @@ private static Map<String, Map<String, Collection<String>>> getConfiguredTableMa
     public static Map<String, Collection<String>> loadSchemaTableNames(final String databaseName, final DatabaseType storageType,
                                                                        final DataSource dataSource, final String dataSourceName, final Collection<String> excludedTables) {
         try {
-            return SchemaMetaDataLoader.loadSchemaTableNamesByExcludedTables(databaseName, storageType, dataSource, excludedTables);
+            return SchemaMetaDataLoader.loadSchemaTableNames(databaseName, storageType, dataSource, excludedTables);
         } catch (final SQLException ex) {
             throw new SingleTablesLoadingException(databaseName, dataSourceName, ex);
         }

kernel/single/core/src/main/java/org/apache/shardingsphere/single/decorator/SingleRuleConfigurationDecorator.java

@@ -23,6 +23,7 @@
 import org.apache.shardingsphere.infra.database.core.type.DatabaseTypeRegistry;
 import org.apache.shardingsphere.infra.datanode.DataNode;
 import org.apache.shardingsphere.infra.exception.core.ShardingSpherePreconditions;
+import org.apache.shardingsphere.infra.metadata.database.resource.PhysicalResourceAggregator;
 import org.apache.shardingsphere.infra.rule.ShardingSphereRule;
 import org.apache.shardingsphere.single.config.SingleRuleConfiguration;
 import org.apache.shardingsphere.single.constant.SingleTableConstants;
@@ -62,7 +63,7 @@ private Collection<String> decorateTables(final String databaseName, final Map<S
         if (!isExpandRequired(splitTables)) {
             return splitTables;
         }
-        Map<String, DataSource> aggregatedDataSources = SingleTableLoadUtils.getAggregatedDataSourceMap(dataSources, builtRules);
+        Map<String, DataSource> aggregatedDataSources = PhysicalResourceAggregator.getAggregatedResources(dataSources, builtRules);
         DatabaseType databaseType = dataSources.isEmpty() ? DatabaseTypeEngine.getDefaultStorageType() : DatabaseTypeEngine.getStorageType(dataSources.values().iterator().next());
         Collection<String> excludedTables = SingleTableLoadUtils.getExcludedTables(builtRules);
         Map<String, Collection<DataNode>> actualDataNodes = SingleTableDataNodeLoader.load(databaseName, aggregatedDataSources, excludedTables);

kernel/single/core/src/main/java/org/apache/shardingsphere/single/rule/SingleRule.java

@@ -25,14 +25,14 @@
 import org.apache.shardingsphere.infra.database.core.type.DatabaseTypeRegistry;
 import org.apache.shardingsphere.infra.datanode.DataNode;
 import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
+import org.apache.shardingsphere.infra.metadata.database.resource.PhysicalResourceAggregator;
 import org.apache.shardingsphere.infra.metadata.database.schema.QualifiedTable;
 import org.apache.shardingsphere.infra.metadata.database.schema.util.IndexMetaDataUtils;
 import org.apache.shardingsphere.infra.rule.ShardingSphereRule;
 import org.apache.shardingsphere.infra.rule.attribute.RuleAttributes;
 import org.apache.shardingsphere.infra.rule.scope.DatabaseRule;
 import org.apache.shardingsphere.single.config.SingleRuleConfiguration;
 import org.apache.shardingsphere.single.datanode.SingleTableDataNodeLoader;
-import org.apache.shardingsphere.single.util.SingleTableLoadUtils;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.generic.table.SimpleTableSegment;
 
 import javax.sql.DataSource;
@@ -70,7 +70,7 @@ public SingleRule(final SingleRuleConfiguration ruleConfig, final String databas
                       final DatabaseType protocolType, final Map<String, DataSource> dataSourceMap, final Collection<ShardingSphereRule> builtRules) {
         configuration = ruleConfig;
         defaultDataSource = ruleConfig.getDefaultDataSource().orElse(null);
-        Map<String, DataSource> aggregateDataSourceMap = SingleTableLoadUtils.getAggregatedDataSourceMap(dataSourceMap, builtRules);
+        Map<String, DataSource> aggregateDataSourceMap = PhysicalResourceAggregator.getAggregatedResources(dataSourceMap, builtRules);
         dataSourceNames = aggregateDataSourceMap.keySet();
         this.protocolType = protocolType;
         singleTableDataNodes = SingleTableDataNodeLoader.load(databaseName, protocolType, aggregateDataSourceMap, builtRules, configuration.getTables());

kernel/single/core/src/main/java/org/apache/shardingsphere/single/util/SingleTableLoadUtils.java

@@ -25,16 +25,11 @@
 import org.apache.shardingsphere.infra.database.core.type.DatabaseTypeRegistry;
 import org.apache.shardingsphere.infra.datanode.DataNode;
 import org.apache.shardingsphere.infra.rule.ShardingSphereRule;
-import org.apache.shardingsphere.infra.rule.attribute.datasource.DataSourceMapperRuleAttribute;
 import org.apache.shardingsphere.infra.rule.attribute.table.TableMapperRuleAttribute;
 import org.apache.shardingsphere.single.constant.SingleTableConstants;
 
-import javax.sql.DataSource;
 import java.util.Collection;
-import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
-import java.util.Map;
-import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.TreeSet;
 
@@ -46,37 +41,6 @@ public final class SingleTableLoadUtils {
 
     private static final String DELIMITER = ",";
 
-    /**
-     * Get aggregated data source map.
-     *
-     * @param dataSourceMap data source map
-     * @param builtRules built rules
-     * @return aggregated data source map
-     */
-    public static Map<String, DataSource> getAggregatedDataSourceMap(final Map<String, DataSource> dataSourceMap, final Collection<ShardingSphereRule> builtRules) {
-        Map<String, DataSource> result = new LinkedHashMap<>(dataSourceMap);
-        for (ShardingSphereRule each : builtRules) {
-            Optional<DataSourceMapperRuleAttribute> ruleAttribute = each.getAttributes().findAttribute(DataSourceMapperRuleAttribute.class);
-            if (ruleAttribute.isPresent()) {
-                result = getAggregatedDataSourceMap(result, ruleAttribute.get());
-            }
-        }
-        return result;
-    }
-
-    private static Map<String, DataSource> getAggregatedDataSourceMap(final Map<String, DataSource> dataSourceMap, final DataSourceMapperRuleAttribute ruleAttribute) {
-        Map<String, DataSource> result = new LinkedHashMap<>();
-        for (Entry<String, Collection<String>> entry : ruleAttribute.getDataSourceMapper().entrySet()) {
-            for (String each : entry.getValue()) {
-                if (dataSourceMap.containsKey(each)) {
-                    result.putIfAbsent(entry.getKey(), dataSourceMap.remove(each));
-                }
-            }
-        }
-        result.putAll(dataSourceMap);
-        return result;
-    }
-
     /**
      * Get excluded tables.
      *

kernel/single/distsql/handler/src/main/java/org/apache/shardingsphere/single/distsql/handler/query/ShowUnloadedSingleTablesExecutor.java

@@ -24,6 +24,7 @@
 import org.apache.shardingsphere.infra.datanode.DataNode;
 import org.apache.shardingsphere.infra.merge.result.impl.local.LocalDataQueryResultRow;
 import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
+import org.apache.shardingsphere.infra.metadata.database.resource.PhysicalResourceAggregator;
 import org.apache.shardingsphere.infra.metadata.database.resource.ResourceMetaData;
 import org.apache.shardingsphere.infra.rule.attribute.table.TableMapperRuleAttribute;
 import org.apache.shardingsphere.mode.manager.ContextManager;
@@ -66,7 +67,7 @@ public Collection<LocalDataQueryResultRow> getRows(final ShowUnloadedSingleTable
 
     private Map<String, Collection<DataNode>> getActualDataNodes(final ShardingSphereDatabase database) {
         ResourceMetaData resourceMetaData = database.getResourceMetaData();
-        Map<String, DataSource> aggregateDataSourceMap = SingleTableLoadUtils.getAggregatedDataSourceMap(
+        Map<String, DataSource> aggregateDataSourceMap = PhysicalResourceAggregator.getAggregatedResources(
                 resourceMetaData.getStorageUnits().entrySet().stream()
                         .collect(Collectors.toMap(Entry::getKey, entry -> entry.getValue().getDataSource(), (oldValue, currentValue) -> oldValue, LinkedHashMap::new)),
                 database.getRuleMetaData().getRules());

kernel/single/distsql/handler/src/main/java/org/apache/shardingsphere/single/distsql/handler/update/LoadSingleTableExecutor.java

@@ -30,6 +30,7 @@
 import org.apache.shardingsphere.infra.exception.kernel.metadata.resource.storageunit.InvalidStorageUnitStatusException;
 import org.apache.shardingsphere.infra.exception.kernel.metadata.resource.storageunit.MissingRequiredStorageUnitsException;
 import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
+import org.apache.shardingsphere.infra.metadata.database.resource.PhysicalResourceAggregator;
 import org.apache.shardingsphere.infra.metadata.database.schema.model.ShardingSphereSchema;
 import org.apache.shardingsphere.infra.rule.attribute.datasource.DataSourceMapperRuleAttribute;
 import org.apache.shardingsphere.single.config.SingleRuleConfiguration;
@@ -38,7 +39,6 @@
 import org.apache.shardingsphere.single.distsql.segment.SingleTableSegment;
 import org.apache.shardingsphere.single.distsql.statement.rdl.LoadSingleTableStatement;
 import org.apache.shardingsphere.single.rule.SingleRule;
-import org.apache.shardingsphere.single.util.SingleTableLoadUtils;
 
 import javax.sql.DataSource;
 import java.util.Collection;
@@ -113,7 +113,7 @@ private void checkTableNodeFormat(final boolean isSchemaSupportedDatabaseType, f
     private void checkShouldExistActualTables(final LoadSingleTableStatement sqlStatement, final Collection<String> storageUnitNames, final String defaultSchemaName) {
         Map<String, DataSource> dataSourceMap = database.getResourceMetaData().getStorageUnits().entrySet()
                 .stream().collect(Collectors.toMap(Entry::getKey, entry -> entry.getValue().getDataSource()));
-        Map<String, DataSource> aggregatedDataSourceMap = SingleTableLoadUtils.getAggregatedDataSourceMap(dataSourceMap, database.getRuleMetaData().getRules());
+        Map<String, DataSource> aggregatedDataSourceMap = PhysicalResourceAggregator.getAggregatedResources(dataSourceMap, database.getRuleMetaData().getRules());
         Collection<String> invalidDataSources = storageUnitNames.stream().filter(each -> !aggregatedDataSourceMap.containsKey(each)).collect(Collectors.toList());
         ShardingSpherePreconditions.checkState(invalidDataSources.isEmpty(), () -> new InvalidStorageUnitStatusException(String.format("`%s` is invalid, please use `%s`",
                 String.join(",", invalidDataSources), String.join(",", aggregatedDataSourceMap.keySet()))));

kernel/single/distsql/handler/src/test/java/org/apache/shardingsphere/single/distsql/handler/update/LoadSingleTableExecutorTest.java

@@ -25,6 +25,7 @@
 import org.apache.shardingsphere.infra.exception.kernel.metadata.datanode.InvalidDataNodeFormatException;
 import org.apache.shardingsphere.infra.exception.kernel.metadata.resource.storageunit.MissingRequiredStorageUnitsException;
 import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
+import org.apache.shardingsphere.infra.metadata.database.resource.PhysicalResourceAggregator;
 import org.apache.shardingsphere.infra.metadata.database.resource.unit.StorageUnit;
 import org.apache.shardingsphere.infra.metadata.database.rule.RuleMetaData;
 import org.apache.shardingsphere.infra.metadata.database.schema.model.ShardingSphereSchema;
@@ -65,7 +66,7 @@
 import static org.mockito.Mockito.when;
 
 @ExtendWith(AutoMockExtension.class)
-@StaticMockSettings({SingleTableDataNodeLoader.class, SingleTableLoadUtils.class})
+@StaticMockSettings({SingleTableDataNodeLoader.class, SingleTableLoadUtils.class, PhysicalResourceAggregator.class})
 @MockitoSettings(strictness = Strictness.LENIENT)
 class LoadSingleTableExecutorTest {
 
@@ -124,7 +125,7 @@ void assertExecuteUpdateWithNotExistedActualTables() {
         StorageUnit storageUnit = mock(StorageUnit.class);
         when(storageUnit.getDataSource()).thenReturn(new MockedDataSource());
         when(database.getResourceMetaData().getStorageUnits()).thenReturn(Collections.singletonMap("foo_ds", storageUnit));
-        when(SingleTableLoadUtils.getAggregatedDataSourceMap(any(), any())).thenReturn(Collections.singletonMap("foo_ds", new MockedDataSource()));
+        when(PhysicalResourceAggregator.getAggregatedResources(any(), any())).thenReturn(Collections.singletonMap("foo_ds", new MockedDataSource()));
         LoadSingleTableStatement sqlStatement = new LoadSingleTableStatement(Collections.singleton(new SingleTableSegment("foo_ds", "foo_tbl")));
         assertThrows(TableNotFoundException.class, () -> new DistSQLUpdateExecuteEngine(sqlStatement, "foo_db", mockContextManager(mock(SingleRule.class))).executeUpdate());
     }