Decouple StorageUnit and Map<StorageNodeName, DataSource> (#28619)

* Refactor StorageUnit * Refactor StorageUnit * Decouple StorageUnit and Map<StorageNodeName, DataSource>
apache · Sep 28, 2023 · 972abd3 · 972abd3
1 parent 0e9af37
commit 972abd3
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 43 deletions.
diff --git a/...ain/java/org/apache/shardingsphere/infra/metadata/database/resource/unit/StorageUnit.java b/...ain/java/org/apache/shardingsphere/infra/metadata/database/resource/unit/StorageUnit.java
@@ -27,71 +27,47 @@
 import org.apache.shardingsphere.infra.datasource.pool.props.creator.DataSourcePoolPropertiesCreator;
 import org.apache.shardingsphere.infra.datasource.pool.props.domain.DataSourcePoolProperties;
 import org.apache.shardingsphere.infra.metadata.database.resource.node.StorageNode;
-import org.apache.shardingsphere.infra.metadata.database.resource.node.StorageNodeName;
 import org.apache.shardingsphere.infra.state.datasource.DataSourceStateManager;
 
 import javax.sql.DataSource;
 import java.util.Collections;
-import java.util.LinkedHashMap;
 import java.util.Map;
-import java.util.Map.Entry;
-import java.util.Optional;
-import java.util.stream.Collectors;
 
 /**
  * Storage unit.
  */
 @Getter
 public final class StorageUnit {
 
+    private final DataSource dataSource;
+
     private final DataSourcePoolProperties dataSourcePoolProperties;
 
     private final StorageNode storageNode;
 
-    private final DataSource dataSource;
-
     private final DatabaseType storageType;
 
     private final ConnectionProperties connectionProperties;
 
-    public StorageUnit(final String databaseName, final Map<StorageNodeName, DataSource> storageNodeDataSources,
-                       final DataSourcePoolProperties props, final StorageNode storageNode) {
-        this.dataSourcePoolProperties = props;
+    public StorageUnit(final String databaseName, final DataSource dataSource, final DataSourcePoolProperties dataSourcePoolProperties, final StorageNode storageNode) {
+        this.dataSource = new CatalogSwitchableDataSource(dataSource, storageNode.getCatalog(), storageNode.getUrl());
+        this.dataSourcePoolProperties = dataSourcePoolProperties;
         this.storageNode = storageNode;
-        dataSource = getStorageUnitDataSource(storageNodeDataSources);
-        Map<StorageNodeName, DataSource> enabledStorageNodeDataSources = getEnabledStorageNodeDataSources(databaseName, storageNodeDataSources);
-        storageType = createStorageType(enabledStorageNodeDataSources);
-        connectionProperties = createConnectionProperties(enabledStorageNodeDataSources).orElse(null);
-    }
-
-    private DataSource getStorageUnitDataSource(final Map<StorageNodeName, DataSource> storageNodeDataSources) {
-        DataSource dataSource = storageNodeDataSources.get(storageNode.getName());
-        return new CatalogSwitchableDataSource(dataSource, storageNode.getCatalog(), storageNode.getUrl());
-    }
-
-    private Map<StorageNodeName, DataSource> getEnabledStorageNodeDataSources(final String databaseName, final Map<StorageNodeName, DataSource> storageNodeDataSources) {
-        Map<String, DataSource> toBeCheckedDataSources = new LinkedHashMap<>(storageNodeDataSources.size(), 1F);
-        for (Entry<StorageNodeName, DataSource> entry : storageNodeDataSources.entrySet()) {
-            toBeCheckedDataSources.put(entry.getKey().getName(), entry.getValue());
-        }
-        Map<String, DataSource> enabledDataSources = DataSourceStateManager.getInstance().getEnabledDataSources(databaseName, toBeCheckedDataSources);
-        return storageNodeDataSources.entrySet().stream()
-                .filter(entry -> enabledDataSources.containsKey(entry.getKey().getName())).collect(Collectors.toMap(Entry::getKey, Entry::getValue));
+        boolean isDataSourceEnabled = !DataSourceStateManager.getInstance().getEnabledDataSources(databaseName, Collections.singletonMap(storageNode.getName().getName(), dataSource)).isEmpty();
+        storageType = createStorageType(isDataSourceEnabled);
+        connectionProperties = createConnectionProperties(isDataSourceEnabled);
     }
 
-    private DatabaseType createStorageType(final Map<StorageNodeName, DataSource> enabledStorageNodeDataSources) {
-        return DatabaseTypeEngine.getStorageType(enabledStorageNodeDataSources.containsKey(storageNode.getName())
-                ? Collections.singleton(enabledStorageNodeDataSources.get(storageNode.getName()))
-                : Collections.emptyList());
+    private DatabaseType createStorageType(final boolean isDataSourceEnabled) {
+        return DatabaseTypeEngine.getStorageType(isDataSourceEnabled ? Collections.singleton(dataSource) : Collections.emptyList());
     }
 
-    private Optional<ConnectionProperties> createConnectionProperties(final Map<StorageNodeName, DataSource> enabledStorageNodeDataSources) {
-        if (!enabledStorageNodeDataSources.containsKey(storageNode.getName())) {
-            return Optional.empty();
+    private ConnectionProperties createConnectionProperties(final boolean isDataSourceEnabled) {
+        if (!isDataSourceEnabled) {
+            return null;
         }
-        Map<String, Object> standardProps = DataSourcePoolPropertiesCreator.create(
-                enabledStorageNodeDataSources.get(storageNode.getName())).getConnectionPropertySynonyms().getStandardProperties();
+        Map<String, Object> standardProps = DataSourcePoolPropertiesCreator.create(dataSource).getConnectionPropertySynonyms().getStandardProperties();
         ConnectionPropertiesParser parser = DatabaseTypedSPILoader.getService(ConnectionPropertiesParser.class, storageType);
-        return Optional.of(parser.parse(standardProps.get("url").toString(), standardProps.get("username").toString(), storageNode.getCatalog()));
+        return parser.parse(standardProps.get("url").toString(), standardProps.get("username").toString(), storageNode.getCatalog());
     }
 }
diff --git a/.../org/apache/shardingsphere/infra/metadata/database/resource/unit/StorageUnitMetaData.java b/.../org/apache/shardingsphere/infra/metadata/database/resource/unit/StorageUnitMetaData.java
@@ -45,13 +45,13 @@ public final class StorageUnitMetaData {
     private final Map<String, DataSource> dataSources;
 
     public StorageUnitMetaData(final String databaseName, final Map<String, StorageNode> storageNodes, final Map<String, DataSourcePoolProperties> dataSourcePoolPropertiesMap,
-                               final Map<StorageNodeName, DataSource> storageNodeDataSources) {
+                               final Map<StorageNodeName, DataSource> dataSources) {
         this.storageNodes = storageNodes;
         this.dataSourcePoolPropertiesMap = dataSourcePoolPropertiesMap;
         storageUnits = storageNodes.entrySet().stream().collect(
-                Collectors.toMap(Entry::getKey, entry -> new StorageUnit(databaseName, storageNodeDataSources, dataSourcePoolPropertiesMap.get(entry.getKey()), entry.getValue()),
+                Collectors.toMap(Entry::getKey, entry -> new StorageUnit(databaseName, dataSources.get(entry.getValue().getName()), dataSourcePoolPropertiesMap.get(entry.getKey()), entry.getValue()),
                         (oldValue, currentValue) -> currentValue, () -> new LinkedHashMap<>(this.storageNodes.size(), 1F)));
-        dataSources = storageUnits.entrySet().stream().collect(
+        this.dataSources = storageUnits.entrySet().stream().collect(
                 Collectors.toMap(Entry::getKey, entry -> entry.getValue().getDataSource(), (oldValue, currentValue) -> currentValue, () -> new LinkedHashMap<>(storageUnits.size(), 1F)));
     }
 }
diff --git a/mode/core/src/test/java/org/apache/shardingsphere/mode/manager/ContextManagerTest.java b/mode/core/src/test/java/org/apache/shardingsphere/mode/manager/ContextManagerTest.java
@@ -258,7 +258,7 @@ private ResourceMetaData createOriginalResource() {
         Map<String, StorageUnit> storageUnits = new LinkedHashMap<>(2, 1F);
         Map<String, StorageNode> storageUnitNodeMap = StorageUnitNodeMapUtils.fromDataSources(originalDataSources);
         for (Entry<String, StorageNode> entry : storageUnitNodeMap.entrySet()) {
-            storageUnits.put(entry.getKey(), new StorageUnit("foo_db", storageNodeDataSourceMap, mock(DataSourcePoolProperties.class), entry.getValue()));
+            storageUnits.put(entry.getKey(), new StorageUnit("foo_db", storageNodeDataSourceMap.get(entry.getValue().getName()), mock(DataSourcePoolProperties.class), entry.getValue()));
         }
         when(result.getStorageUnitMetaData().getStorageUnits()).thenReturn(storageUnits);
         when(result.getStorageUnitMetaData().getStorageNodes()).thenReturn(storageUnitNodeMap);