Refactor StorageUnit

apache · Sep 28, 2023 · 49bcf13 · 49bcf13
1 parent 3b8acde
commit 49bcf13
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 10 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,7 +27,6 @@
 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;
@@ -50,20 +49,15 @@ public final class StorageUnit {
 
     private final ConnectionProperties connectionProperties;
 
-    public StorageUnit(final String databaseName, final Map<StorageNodeName, DataSource> storageNodeDataSources, final DataSourcePoolProperties props, final StorageNode storageNode) {
+    public StorageUnit(final String databaseName, final DataSource dataSource, final DataSourcePoolProperties props, final StorageNode storageNode) {
         this.dataSourcePoolProperties = props;
         this.storageNode = storageNode;
-        dataSource = getStorageUnitDataSource(storageNodeDataSources);
+        this.dataSource = new CatalogSwitchableDataSource(dataSource, storageNode.getCatalog(), storageNode.getUrl());
         boolean isDataSourceEnabled = !DataSourceStateManager.getInstance().getEnabledDataSources(databaseName, Collections.singletonMap(storageNode.getName().getName(), dataSource)).isEmpty();
         storageType = createStorageType(isDataSourceEnabled);
         connectionProperties = createConnectionProperties(isDataSourceEnabled);
     }
 
-    private DataSource getStorageUnitDataSource(final Map<StorageNodeName, DataSource> storageNodeDataSources) {
-        DataSource dataSource = storageNodeDataSources.get(storageNode.getName());
-        return new CatalogSwitchableDataSource(dataSource, storageNode.getCatalog(), storageNode.getUrl());
-    }
-
     private DatabaseType createStorageType(final boolean isDataSourceEnabled) {
         return DatabaseTypeEngine.getStorageType(isDataSourceEnabled ? Collections.singleton(dataSource) : Collections.emptyList());
     }

diff --git a/.../org/apache/shardingsphere/infra/metadata/database/resource/unit/StorageUnitMetaData.java b/.../org/apache/shardingsphere/infra/metadata/database/resource/unit/StorageUnitMetaData.java
@@ -49,7 +49,7 @@ public StorageUnitMetaData(final String databaseName, final Map<String, StorageN
         this.storageNodes = storageNodes;
         this.dataSourcePoolPropertiesMap = dataSourcePoolPropertiesMap;
         storageUnits = storageNodes.entrySet().stream().collect(
-                Collectors.toMap(Entry::getKey, entry -> new StorageUnit(databaseName, dataSources, 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)));
         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);