Fixes #31773, support logic data source when set default single table…

… storage unit.
apache · Jun 20, 2024 · 2a1e111 · 2a1e111
1 parent 75c5b58
commit 2a1e111
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 4 deletions.
diff --git a/...hardingsphere/single/distsql/handler/update/SetDefaultSingleTableStorageUnitExecutor.java b/...hardingsphere/single/distsql/handler/update/SetDefaultSingleTableStorageUnitExecutor.java
@@ -23,12 +23,15 @@
 import org.apache.shardingsphere.distsql.handler.engine.update.rdl.rule.spi.database.DatabaseRuleCreateExecutor;
 import org.apache.shardingsphere.infra.exception.core.ShardingSpherePreconditions;
 import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
+import org.apache.shardingsphere.infra.rule.attribute.datasource.DataSourceMapperRuleAttribute;
 import org.apache.shardingsphere.single.api.config.SingleRuleConfiguration;
 import org.apache.shardingsphere.single.distsql.statement.rdl.SetDefaultSingleTableStorageUnitStatement;
 import org.apache.shardingsphere.single.rule.SingleRule;
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashSet;
+import java.util.stream.Collectors;
 
 /**
  * Set default single table storage unit executor.
@@ -47,12 +50,18 @@ public void checkBeforeUpdate(final SetDefaultSingleTableStorageUnitStatement sq
 
     private void checkStorageUnitExist(final SetDefaultSingleTableStorageUnitStatement sqlStatement) {
         if (!Strings.isNullOrEmpty(sqlStatement.getDefaultStorageUnit())) {
-            Collection<String> storageUnitNames = database.getResourceMetaData().getStorageUnits().keySet();
-            ShardingSpherePreconditions.checkContains(storageUnitNames, sqlStatement.getDefaultStorageUnit(),
+            Collection<String> dataSourceNames = new HashSet<>(database.getResourceMetaData().getStorageUnits().keySet());
+            dataSourceNames.addAll(getLogicDataSourceNames());
+            ShardingSpherePreconditions.checkContains(dataSourceNames, sqlStatement.getDefaultStorageUnit(),
                     () -> new MissingRequiredStorageUnitsException(database.getName(), Collections.singleton(sqlStatement.getDefaultStorageUnit())));
         }
     }
 
+    private Collection<String> getLogicDataSourceNames() {
+        return database.getRuleMetaData().getAttributes(DataSourceMapperRuleAttribute.class).stream()
+                .flatMap(each -> each.getDataSourceMapper().keySet().stream()).collect(Collectors.toSet());
+    }
+
     @Override
     public SingleRuleConfiguration buildToBeCreatedRuleConfiguration(final SetDefaultSingleTableStorageUnitStatement sqlStatement) {
         SingleRuleConfiguration result = new SingleRuleConfiguration();

diff --git a/...ingsphere/single/distsql/handler/update/SetDefaultSingleTableStorageUnitExecutorTest.java b/...ingsphere/single/distsql/handler/update/SetDefaultSingleTableStorageUnitExecutorTest.java
@@ -19,18 +19,23 @@
 
 import org.apache.shardingsphere.infra.exception.kernel.metadata.resource.storageunit.MissingRequiredStorageUnitsException;
 import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
+import org.apache.shardingsphere.infra.rule.attribute.datasource.DataSourceMapperRuleAttribute;
 import org.apache.shardingsphere.single.api.config.SingleRuleConfiguration;
 import org.apache.shardingsphere.single.distsql.statement.rdl.SetDefaultSingleTableStorageUnitStatement;
 import org.apache.shardingsphere.single.rule.SingleRule;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.junit.jupiter.MockitoExtension;
 
+import java.util.Collections;
+
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -41,11 +46,23 @@ class SetDefaultSingleTableStorageUnitExecutorTest {
     private final SetDefaultSingleTableStorageUnitExecutor executor = new SetDefaultSingleTableStorageUnitExecutor();
 
     @Test
-    void assertCheckWithInvalidResource() {
-        executor.setDatabase(mock(ShardingSphereDatabase.class, RETURNS_DEEP_STUBS));
+    void assertCheckWithInvalidDataSource() {
+        ShardingSphereDatabase database = mock(ShardingSphereDatabase.class, RETURNS_DEEP_STUBS);
+        when(database.getRuleMetaData().getAttributes(any())).thenReturn(Collections.emptyList());
+        executor.setDatabase(database);
         assertThrows(MissingRequiredStorageUnitsException.class, () -> executor.checkBeforeUpdate(new SetDefaultSingleTableStorageUnitStatement("bar_ds")));
     }
 
+    @Test
+    void assertCheckWithLogicDataSource() {
+        ShardingSphereDatabase database = mock(ShardingSphereDatabase.class, RETURNS_DEEP_STUBS);
+        DataSourceMapperRuleAttribute ruleAttribute = mock(DataSourceMapperRuleAttribute.class, RETURNS_DEEP_STUBS);
+        when(ruleAttribute.getDataSourceMapper().keySet()).thenReturn(Collections.singleton("logic_ds"));
+        when(database.getRuleMetaData().getAttributes(any())).thenReturn(Collections.singleton(ruleAttribute));
+        executor.setDatabase(database);
+        assertDoesNotThrow(() -> executor.checkBeforeUpdate(new SetDefaultSingleTableStorageUnitStatement("logic_ds")));
+    }
+
     @Test
     void assertBuild() {
         SingleRule rule = mock(SingleRule.class);