You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@shardingsphere.apache.org by ji...@apache.org on 2024/11/01 16:20:58 UTC
(shardingsphere) branch master updated: Add more test cases on BroadcastRouteEngineFactory (#33502)
This is an automated email from the ASF dual-hosted git repository.
jianglongtao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shardingsphere.git
The following commit(s) were added to refs/heads/master by this push:
new d7b77dbc91d Add more test cases on BroadcastRouteEngineFactory (#33502)
d7b77dbc91d is described below
commit d7b77dbc91d9836d102a897187f856ad96cc9dc6
Author: Liang Zhang <zh...@apache.org>
AuthorDate: Sat Nov 2 00:20:50 2024 +0800
Add more test cases on BroadcastRouteEngineFactory (#33502)
* Add more test cases on BroadcastRouteEngineFactory
* Add more test cases on BroadcastRouteEngineFactory
---
.../route/engine/BroadcastRouteEngineFactory.java | 82 ++--------
.../engine/BroadcastRouteEngineFactoryTest.java | 182 ++++++++++++++++++++-
.../extractor/SQLStatementContextExtractor.java | 67 ++++++++
3 files changed, 261 insertions(+), 70 deletions(-)
diff --git a/features/broadcast/core/src/main/java/org/apache/shardingsphere/broadcast/route/engine/BroadcastRouteEngineFactory.java b/features/broadcast/core/src/main/java/org/apache/shardingsphere/broadcast/route/engine/BroadcastRouteEngineFactory.java
index 3f4488ef8fd..1b9a5ce0baa 100644
--- a/features/broadcast/core/src/main/java/org/apache/shardingsphere/broadcast/route/engine/BroadcastRouteEngineFactory.java
+++ b/features/broadcast/core/src/main/java/org/apache/shardingsphere/broadcast/route/engine/BroadcastRouteEngineFactory.java
@@ -26,18 +26,14 @@ import org.apache.shardingsphere.broadcast.route.engine.type.ignore.BroadcastIgn
import org.apache.shardingsphere.broadcast.route.engine.type.unicast.BroadcastUnicastRoutingEngine;
import org.apache.shardingsphere.broadcast.rule.BroadcastRule;
import org.apache.shardingsphere.infra.annotation.HighFrequencyInvocation;
+import org.apache.shardingsphere.infra.binder.context.extractor.SQLStatementContextExtractor;
import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
import org.apache.shardingsphere.infra.binder.context.statement.ddl.CloseStatementContext;
import org.apache.shardingsphere.infra.binder.context.type.CursorAvailable;
-import org.apache.shardingsphere.infra.binder.context.type.IndexAvailable;
import org.apache.shardingsphere.infra.binder.context.type.TableAvailable;
-import org.apache.shardingsphere.infra.database.core.type.DatabaseType;
import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
-import org.apache.shardingsphere.infra.metadata.database.schema.QualifiedTable;
-import org.apache.shardingsphere.infra.metadata.database.schema.util.IndexMetaDataUtils;
import org.apache.shardingsphere.infra.session.connection.ConnectionContext;
import org.apache.shardingsphere.infra.session.query.QueryContext;
-import org.apache.shardingsphere.sql.parser.statement.core.segment.ddl.index.IndexSegment;
import org.apache.shardingsphere.sql.parser.statement.core.statement.SQLStatement;
import org.apache.shardingsphere.sql.parser.statement.core.statement.dal.DALStatement;
import org.apache.shardingsphere.sql.parser.statement.core.statement.dcl.DCLStatement;
@@ -48,7 +44,6 @@ import org.apache.shardingsphere.sql.parser.statement.mysql.dal.MySQLUseStatemen
import java.util.Collection;
import java.util.Collections;
-import java.util.LinkedList;
import java.util.stream.Collectors;
/**
@@ -75,95 +70,54 @@ public final class BroadcastRouteEngineFactory {
if (sqlStatement instanceof DDLStatement) {
return sqlStatementContext instanceof CursorAvailable
? getCursorRouteEngine(rule, sqlStatementContext, queryContext.getConnectionContext())
- : getDDLRoutingEngine(rule, database, queryContext);
+ : getDDLRoutingEngine(rule, database, sqlStatementContext);
}
if (sqlStatement instanceof DALStatement) {
- return getDALRoutingEngine(rule, queryContext);
+ return getDALRoutingEngine(rule, sqlStatementContext);
}
if (sqlStatement instanceof DCLStatement) {
- return getDCLRoutingEngine(rule, queryContext);
+ return getDCLRoutingEngine(rule, sqlStatementContext);
}
- return getDQLRoutingEngine(rule, queryContext);
+ return getDMLRoutingEngine(rule, sqlStatementContext, queryContext.getConnectionContext());
}
private static BroadcastRouteEngine getCursorRouteEngine(final BroadcastRule rule, final SQLStatementContext sqlStatementContext, final ConnectionContext connectionContext) {
if (sqlStatementContext instanceof CloseStatementContext && ((CloseStatementContext) sqlStatementContext).getSqlStatement().isCloseAll()) {
return new BroadcastDatabaseBroadcastRoutingEngine();
}
- Collection<String> tableNames = sqlStatementContext instanceof TableAvailable
- ? ((TableAvailable) sqlStatementContext).getTablesContext().getSimpleTables().stream().map(each -> each.getTableName().getIdentifier().getValue()).collect(Collectors.toSet())
- : Collections.emptyList();
- return rule.isAllBroadcastTables(tableNames) ? new BroadcastUnicastRoutingEngine(sqlStatementContext, tableNames, connectionContext) : new BroadcastIgnoreRoutingEngine();
- }
-
- private static BroadcastRouteEngine getDDLRoutingEngine(final BroadcastRule rule, final ShardingSphereDatabase database, final QueryContext queryContext) {
- SQLStatementContext sqlStatementContext = queryContext.getSqlStatementContext();
- Collection<String> tableNames = getTableNames(database, sqlStatementContext);
- if (rule.isAllBroadcastTables(tableNames)) {
- return new BroadcastTableBroadcastRoutingEngine(tableNames);
+ if (sqlStatementContext instanceof TableAvailable) {
+ Collection<String> tableNames = ((TableAvailable) sqlStatementContext)
+ .getTablesContext().getSimpleTables().stream().map(each -> each.getTableName().getIdentifier().getValue()).collect(Collectors.toSet());
+ return rule.isAllBroadcastTables(tableNames) ? new BroadcastUnicastRoutingEngine(sqlStatementContext, tableNames, connectionContext) : new BroadcastIgnoreRoutingEngine();
}
return new BroadcastIgnoreRoutingEngine();
}
- private static Collection<String> getTableNames(final ShardingSphereDatabase database, final SQLStatementContext sqlStatementContext) {
- Collection<String> tableNames = sqlStatementContext instanceof TableAvailable
- ? ((TableAvailable) sqlStatementContext).getTablesContext().getSimpleTables().stream().map(each -> each.getTableName().getIdentifier().getValue()).collect(Collectors.toSet())
- : Collections.emptyList();
- if (!tableNames.isEmpty()) {
- return tableNames;
- }
- return sqlStatementContext instanceof IndexAvailable
- ? getTableNames(database, sqlStatementContext.getDatabaseType(), ((IndexAvailable) sqlStatementContext).getIndexes())
- : Collections.emptyList();
+ private static BroadcastRouteEngine getDDLRoutingEngine(final BroadcastRule rule, final ShardingSphereDatabase database, final SQLStatementContext sqlStatementContext) {
+ Collection<String> tableNames = SQLStatementContextExtractor.getTableNames(database, sqlStatementContext);
+ return rule.isAllBroadcastTables(tableNames) ? new BroadcastTableBroadcastRoutingEngine(tableNames) : new BroadcastIgnoreRoutingEngine();
}
- private static Collection<String> getTableNames(final ShardingSphereDatabase database, final DatabaseType databaseType, final Collection<IndexSegment> indexes) {
- Collection<String> result = new LinkedList<>();
- for (QualifiedTable each : IndexMetaDataUtils.getTableNames(database, databaseType, indexes)) {
- result.add(each.getTableName());
- }
- return result;
- }
-
- private static BroadcastRouteEngine getDALRoutingEngine(final BroadcastRule rule, final QueryContext queryContext) {
- SQLStatementContext sqlStatementContext = queryContext.getSqlStatementContext();
+ private static BroadcastRouteEngine getDALRoutingEngine(final BroadcastRule rule, final SQLStatementContext sqlStatementContext) {
SQLStatement sqlStatement = sqlStatementContext.getSqlStatement();
if (sqlStatement instanceof MySQLUseStatement) {
return new BroadcastIgnoreRoutingEngine();
}
Collection<String> tableNames = sqlStatementContext instanceof TableAvailable ? ((TableAvailable) sqlStatementContext).getTablesContext().getTableNames() : Collections.emptyList();
- Collection<String> broadcastTableNames = rule.filterBroadcastTableNames(tableNames);
- if (rule.isAllBroadcastTables(broadcastTableNames)) {
- return new BroadcastTableBroadcastRoutingEngine(broadcastTableNames);
- }
- return new BroadcastIgnoreRoutingEngine();
+ return tableNames.isEmpty() ? new BroadcastIgnoreRoutingEngine() : new BroadcastTableBroadcastRoutingEngine(rule.filterBroadcastTableNames(tableNames));
}
- private static BroadcastRouteEngine getDCLRoutingEngine(final BroadcastRule rule, final QueryContext queryContext) {
- SQLStatementContext sqlStatementContext = queryContext.getSqlStatementContext();
+ private static BroadcastRouteEngine getDCLRoutingEngine(final BroadcastRule rule, final SQLStatementContext sqlStatementContext) {
Collection<String> tableNames = sqlStatementContext instanceof TableAvailable ? ((TableAvailable) sqlStatementContext).getTablesContext().getTableNames() : Collections.emptyList();
Collection<String> broadcastTableNames = rule.filterBroadcastTableNames(tableNames);
- if (isDCLForSingleTable(sqlStatementContext) && !broadcastTableNames.isEmpty() || rule.isAllBroadcastTables(broadcastTableNames)) {
- return new BroadcastTableBroadcastRoutingEngine(broadcastTableNames);
- }
- return new BroadcastIgnoreRoutingEngine();
+ return broadcastTableNames.isEmpty() ? new BroadcastIgnoreRoutingEngine() : new BroadcastTableBroadcastRoutingEngine(broadcastTableNames);
}
- private static boolean isDCLForSingleTable(final SQLStatementContext sqlStatementContext) {
- if (sqlStatementContext instanceof TableAvailable) {
- TableAvailable tableSegmentsAvailable = (TableAvailable) sqlStatementContext;
- return 1 == tableSegmentsAvailable.getTablesContext().getSimpleTables().size()
- && !"*".equals(tableSegmentsAvailable.getTablesContext().getSimpleTables().iterator().next().getTableName().getIdentifier().getValue());
- }
- return false;
- }
-
- private static BroadcastRouteEngine getDQLRoutingEngine(final BroadcastRule rule, final QueryContext queryContext) {
- SQLStatementContext sqlStatementContext = queryContext.getSqlStatementContext();
+ private static BroadcastRouteEngine getDMLRoutingEngine(final BroadcastRule rule, final SQLStatementContext sqlStatementContext, final ConnectionContext connectionContext) {
Collection<String> tableNames = sqlStatementContext instanceof TableAvailable ? ((TableAvailable) sqlStatementContext).getTablesContext().getTableNames() : Collections.emptyList();
if (rule.isAllBroadcastTables(tableNames)) {
return sqlStatementContext.getSqlStatement() instanceof SelectStatement
- ? new BroadcastUnicastRoutingEngine(sqlStatementContext, tableNames, queryContext.getConnectionContext())
+ ? new BroadcastUnicastRoutingEngine(sqlStatementContext, tableNames, connectionContext)
: new BroadcastDatabaseBroadcastRoutingEngine();
}
return new BroadcastIgnoreRoutingEngine();
diff --git a/features/broadcast/core/src/test/java/org/apache/shardingsphere/broadcast/route/engine/BroadcastRouteEngineFactoryTest.java b/features/broadcast/core/src/test/java/org/apache/shardingsphere/broadcast/route/engine/BroadcastRouteEngineFactoryTest.java
index f1c6c0a2ecb..d4e886bb1c5 100644
--- a/features/broadcast/core/src/test/java/org/apache/shardingsphere/broadcast/route/engine/BroadcastRouteEngineFactoryTest.java
+++ b/features/broadcast/core/src/test/java/org/apache/shardingsphere/broadcast/route/engine/BroadcastRouteEngineFactoryTest.java
@@ -18,38 +18,208 @@
package org.apache.shardingsphere.broadcast.route.engine;
import org.apache.shardingsphere.broadcast.route.engine.type.broadcast.BroadcastDatabaseBroadcastRoutingEngine;
+import org.apache.shardingsphere.broadcast.route.engine.type.broadcast.BroadcastTableBroadcastRoutingEngine;
+import org.apache.shardingsphere.broadcast.route.engine.type.ignore.BroadcastIgnoreRoutingEngine;
+import org.apache.shardingsphere.broadcast.route.engine.type.unicast.BroadcastUnicastRoutingEngine;
import org.apache.shardingsphere.broadcast.rule.BroadcastRule;
+import org.apache.shardingsphere.infra.binder.context.extractor.SQLStatementContextExtractor;
+import org.apache.shardingsphere.infra.binder.context.segment.table.TablesContext;
import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
+import org.apache.shardingsphere.infra.binder.context.statement.ddl.CloseStatementContext;
+import org.apache.shardingsphere.infra.binder.context.type.CursorAvailable;
+import org.apache.shardingsphere.infra.binder.context.type.TableAvailable;
+import org.apache.shardingsphere.infra.database.core.type.DatabaseType;
import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
import org.apache.shardingsphere.infra.session.query.QueryContext;
+import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.generic.table.SimpleTableSegment;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.generic.table.TableNameSegment;
+import org.apache.shardingsphere.sql.parser.statement.core.statement.dal.DALStatement;
+import org.apache.shardingsphere.sql.parser.statement.core.statement.dcl.DCLStatement;
+import org.apache.shardingsphere.sql.parser.statement.core.statement.ddl.DDLStatement;
+import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.DMLStatement;
+import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.SelectStatement;
+import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.UpdateStatement;
import org.apache.shardingsphere.sql.parser.statement.core.statement.tcl.TCLStatement;
+import org.apache.shardingsphere.sql.parser.statement.core.value.identifier.IdentifierValue;
+import org.apache.shardingsphere.test.mock.AutoMockExtension;
+import org.apache.shardingsphere.test.mock.StaticMockSettings;
+import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Answers;
import org.mockito.Mock;
-import org.mockito.junit.jupiter.MockitoExtension;
+import org.mockito.junit.jupiter.MockitoSettings;
+import org.mockito.quality.Strictness;
+
+import java.util.Collections;
import static org.hamcrest.CoreMatchers.instanceOf;
import static org.hamcrest.MatcherAssert.assertThat;
+import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.withSettings;
-@ExtendWith(MockitoExtension.class)
+@ExtendWith(AutoMockExtension.class)
+@MockitoSettings(strictness = Strictness.LENIENT)
+@StaticMockSettings(SQLStatementContextExtractor.class)
class BroadcastRouteEngineFactoryTest {
+ private final DatabaseType databaseType = TypedSPILoader.getService(DatabaseType.class, "FIXTURE");
+
@Mock
- private BroadcastRule broadcastRule;
+ private BroadcastRule rule;
@Mock
private ShardingSphereDatabase database;
- @Mock
+ @Mock(answer = Answers.RETURNS_DEEP_STUBS)
private QueryContext queryContext;
+ @BeforeEach
+ void setUp() {
+ when(rule.filterBroadcastTableNames(Collections.singleton("foo_tbl"))).thenReturn(Collections.singleton("foo_tbl"));
+ when(rule.isAllBroadcastTables(Collections.singleton("foo_tbl"))).thenReturn(true);
+ }
+
@Test
void assertNewInstanceWithTCLStatement() {
+ when(queryContext.getSqlStatementContext().getSqlStatement()).thenReturn(mock(TCLStatement.class));
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastDatabaseBroadcastRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithCloseStatementContextAndCloseAll() {
+ CloseStatementContext sqlStatementContext = mock(CloseStatementContext.class, RETURNS_DEEP_STUBS);
+ when(sqlStatementContext.getSqlStatement().isCloseAll()).thenReturn(true);
+ when(queryContext.getSqlStatementContext()).thenReturn(sqlStatementContext);
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastDatabaseBroadcastRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithCursorAvailableAndIsNotAllBroadcastTables() {
+ CloseStatementContext sqlStatementContext = mock(CloseStatementContext.class, RETURNS_DEEP_STUBS);
+ when(queryContext.getSqlStatementContext()).thenReturn(sqlStatementContext);
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastIgnoreRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithCursorAvailableAndIsAllBroadcastTables() {
+ CloseStatementContext sqlStatementContext = mock(CloseStatementContext.class, RETURNS_DEEP_STUBS);
+ when(sqlStatementContext.getTablesContext().getSimpleTables()).thenReturn(Collections.singletonList(new SimpleTableSegment(new TableNameSegment(0, 0, new IdentifierValue("foo_tbl")))));
+ when(queryContext.getSqlStatementContext()).thenReturn(sqlStatementContext);
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastUnicastRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithCursorAvailableAndNotTableAvailable() {
+ SQLStatementContext sqlStatementContext = mock(SQLStatementContext.class, withSettings().extraInterfaces(CursorAvailable.class));
+ when(sqlStatementContext.getSqlStatement()).thenReturn(mock(DDLStatement.class));
+ when(queryContext.getSqlStatementContext()).thenReturn(sqlStatementContext);
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastIgnoreRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithDDLStatementAndIsAllBroadcastTables() {
+ SQLStatementContext sqlStatementContext = mock(SQLStatementContext.class);
+ when(sqlStatementContext.getSqlStatement()).thenReturn(mock(DDLStatement.class));
+ when(queryContext.getSqlStatementContext()).thenReturn(sqlStatementContext);
+ when(SQLStatementContextExtractor.getTableNames(database, sqlStatementContext)).thenReturn(Collections.singleton("foo_tbl"));
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastTableBroadcastRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithDDLStatementAndIsNotAllBroadcastTables() {
SQLStatementContext sqlStatementContext = mock(SQLStatementContext.class);
- when(sqlStatementContext.getSqlStatement()).thenReturn(mock(TCLStatement.class));
+ when(sqlStatementContext.getSqlStatement()).thenReturn(mock(DDLStatement.class));
+ when(queryContext.getSqlStatementContext()).thenReturn(sqlStatementContext);
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastIgnoreRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithDALStatementAndNotTableAvailable() {
+ when(queryContext.getSqlStatementContext().getSqlStatement()).thenReturn(mock(DALStatement.class));
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastIgnoreRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithDALStatementAndEmptyTables() {
+ SQLStatementContext sqlStatementContext = mock(SQLStatementContext.class, withSettings().extraInterfaces(TableAvailable.class));
+ when(((TableAvailable) sqlStatementContext).getTablesContext()).thenReturn(new TablesContext(Collections.emptyList(), databaseType, null));
+ when(sqlStatementContext.getSqlStatement()).thenReturn(mock(DALStatement.class));
+ when(queryContext.getSqlStatementContext()).thenReturn(sqlStatementContext);
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastIgnoreRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithDALStatementAndTables() {
+ SQLStatementContext sqlStatementContext = mock(SQLStatementContext.class, withSettings().extraInterfaces(TableAvailable.class));
+ TablesContext tablesContext = new TablesContext(Collections.singleton(new SimpleTableSegment(new TableNameSegment(0, 0, new IdentifierValue("foo_tbl")))), databaseType, null);
+ when(((TableAvailable) sqlStatementContext).getTablesContext()).thenReturn(tablesContext);
+ when(sqlStatementContext.getSqlStatement()).thenReturn(mock(DALStatement.class));
+ when(queryContext.getSqlStatementContext()).thenReturn(sqlStatementContext);
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastTableBroadcastRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithDCLStatementAndNotTableAvailable() {
+ when(queryContext.getSqlStatementContext().getSqlStatement()).thenReturn(mock(DCLStatement.class));
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastIgnoreRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithDCLStatementAndEmptyTables() {
+ SQLStatementContext sqlStatementContext = mock(SQLStatementContext.class, withSettings().extraInterfaces(TableAvailable.class));
+ when(((TableAvailable) sqlStatementContext).getTablesContext()).thenReturn(new TablesContext(Collections.emptyList(), databaseType, null));
+ when(sqlStatementContext.getSqlStatement()).thenReturn(mock(DCLStatement.class));
+ when(queryContext.getSqlStatementContext()).thenReturn(sqlStatementContext);
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastIgnoreRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithDCLStatementAndTables() {
+ SQLStatementContext sqlStatementContext = mock(SQLStatementContext.class, withSettings().extraInterfaces(TableAvailable.class));
+ TablesContext tablesContext = new TablesContext(Collections.singleton(new SimpleTableSegment(new TableNameSegment(0, 0, new IdentifierValue("foo_tbl")))), databaseType, null);
+ when(((TableAvailable) sqlStatementContext).getTablesContext()).thenReturn(tablesContext);
+ when(sqlStatementContext.getSqlStatement()).thenReturn(mock(DCLStatement.class));
+ when(queryContext.getSqlStatementContext()).thenReturn(sqlStatementContext);
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastTableBroadcastRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithDMLStatementAndNotTableAvailable() {
+ when(queryContext.getSqlStatementContext().getSqlStatement()).thenReturn(mock(DMLStatement.class));
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastIgnoreRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithDMLStatementAndIsNotAllBroadcastTables() {
+ SQLStatementContext sqlStatementContext = mock(SQLStatementContext.class, withSettings().extraInterfaces(TableAvailable.class));
+ TablesContext tablesContext = new TablesContext(Collections.emptyList(), databaseType, null);
+ when(((TableAvailable) sqlStatementContext).getTablesContext()).thenReturn(tablesContext);
+ when(sqlStatementContext.getSqlStatement()).thenReturn(mock(DMLStatement.class));
+ when(queryContext.getSqlStatementContext()).thenReturn(sqlStatementContext);
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastIgnoreRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithDMLSelectStatementAndIsAllBroadcastTables() {
+ SQLStatementContext sqlStatementContext = mock(SQLStatementContext.class, withSettings().extraInterfaces(TableAvailable.class));
+ TablesContext tablesContext = new TablesContext(Collections.singleton(new SimpleTableSegment(new TableNameSegment(0, 0, new IdentifierValue("foo_tbl")))), databaseType, null);
+ when(((TableAvailable) sqlStatementContext).getTablesContext()).thenReturn(tablesContext);
+ when(sqlStatementContext.getSqlStatement()).thenReturn(mock(SelectStatement.class));
+ when(queryContext.getSqlStatementContext()).thenReturn(sqlStatementContext);
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastUnicastRoutingEngine.class));
+ }
+
+ @Test
+ void assertNewInstanceWithDMLUpdateStatementAndIsAllBroadcastTables() {
+ SQLStatementContext sqlStatementContext = mock(SQLStatementContext.class, withSettings().extraInterfaces(TableAvailable.class));
+ TablesContext tablesContext = new TablesContext(Collections.singleton(new SimpleTableSegment(new TableNameSegment(0, 0, new IdentifierValue("foo_tbl")))), databaseType, null);
+ when(((TableAvailable) sqlStatementContext).getTablesContext()).thenReturn(tablesContext);
+ when(sqlStatementContext.getSqlStatement()).thenReturn(mock(UpdateStatement.class));
when(queryContext.getSqlStatementContext()).thenReturn(sqlStatementContext);
- assertThat(BroadcastRouteEngineFactory.newInstance(broadcastRule, database, queryContext), instanceOf(BroadcastDatabaseBroadcastRoutingEngine.class));
+ assertThat(BroadcastRouteEngineFactory.newInstance(rule, database, queryContext), instanceOf(BroadcastDatabaseBroadcastRoutingEngine.class));
}
}
diff --git a/infra/binder/src/main/java/org/apache/shardingsphere/infra/binder/context/extractor/SQLStatementContextExtractor.java b/infra/binder/src/main/java/org/apache/shardingsphere/infra/binder/context/extractor/SQLStatementContextExtractor.java
new file mode 100644
index 00000000000..8d1c48a317b
--- /dev/null
+++ b/infra/binder/src/main/java/org/apache/shardingsphere/infra/binder/context/extractor/SQLStatementContextExtractor.java
@@ -0,0 +1,67 @@
+/*
+ * 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.binder.context.extractor;
+
+import lombok.AccessLevel;
+import lombok.NoArgsConstructor;
+import org.apache.shardingsphere.infra.annotation.HighFrequencyInvocation;
+import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
+import org.apache.shardingsphere.infra.binder.context.type.IndexAvailable;
+import org.apache.shardingsphere.infra.binder.context.type.TableAvailable;
+import org.apache.shardingsphere.infra.database.core.type.DatabaseType;
+import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
+import org.apache.shardingsphere.infra.metadata.database.schema.QualifiedTable;
+import org.apache.shardingsphere.infra.metadata.database.schema.util.IndexMetaDataUtils;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.ddl.index.IndexSegment;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.stream.Collectors;
+
+@HighFrequencyInvocation
+@NoArgsConstructor(access = AccessLevel.PRIVATE)
+public final class SQLStatementContextExtractor {
+
+ /**
+ * Get table names.
+ *
+ * @param database database
+ * @param sqlStatementContext SQL statement context
+ * @return table names
+ */
+ public static Collection<String> getTableNames(final ShardingSphereDatabase database, final SQLStatementContext sqlStatementContext) {
+ Collection<String> tableNames = sqlStatementContext instanceof TableAvailable
+ ? ((TableAvailable) sqlStatementContext).getTablesContext().getSimpleTables().stream().map(each -> each.getTableName().getIdentifier().getValue()).collect(Collectors.toSet())
+ : Collections.emptyList();
+ if (!tableNames.isEmpty()) {
+ return tableNames;
+ }
+ return sqlStatementContext instanceof IndexAvailable
+ ? getTableNames(database, sqlStatementContext.getDatabaseType(), ((IndexAvailable) sqlStatementContext).getIndexes())
+ : Collections.emptyList();
+ }
+
+ private static Collection<String> getTableNames(final ShardingSphereDatabase database, final DatabaseType databaseType, final Collection<IndexSegment> indexes) {
+ Collection<String> result = new LinkedList<>();
+ for (QualifiedTable each : IndexMetaDataUtils.getTableNames(database, databaseType, indexes)) {
+ result.add(each.getTableName());
+ }
+ return result;
+ }
+}