Skip to content

Commit

Permalink
Add more test cases on ShardingDDLResultMerger (#33726)
Browse files Browse the repository at this point in the history
  • Loading branch information
@terrymanu
terrymanu authored Nov 19, 2024
1 parent 6943fdd commit 68a9804
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 58 deletions.
18 changes: 7 additions & 11 deletions ...e/src/main/java/org/apache/shardingsphere/sharding/merge/ddl/ShardingDDLResultMerger.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
import com.cedarsoftware.util.CaseInsensitiveMap;
import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
import org.apache.shardingsphere.infra.binder.context.statement.ddl.FetchStatementContext;
import org.apache.shardingsphere.infra.binder.context.type.TableAvailable;
import org.apache.shardingsphere.infra.database.core.type.DatabaseTypeRegistry;
import org.apache.shardingsphere.infra.executor.sql.execute.result.query.QueryResult;
import org.apache.shardingsphere.infra.merge.engine.merger.ResultMerger;
Expand Down Expand Up @@ -52,16 +51,8 @@ public MergedResult merge(final List<QueryResult> queryResults, final SQLStateme
return new IteratorStreamMergedResult(queryResults);
}
FetchStatementContext fetchStatementContext = (FetchStatementContext) sqlStatementContext;
Map<String, Integer> columnLabelIndexMap = getColumnLabelIndexMap(queryResults.get(0));
fetchStatementContext.getCursorStatementContext().getSelectStatementContext().setIndexes(columnLabelIndexMap);
return new FetchStreamMergedResult(queryResults, fetchStatementContext, getSchema(sqlStatementContext, database), connectionContext);
}

private ShardingSphereSchema getSchema(final SQLStatementContext sqlStatementContext, final ShardingSphereDatabase database) {
String defaultSchemaName = new DatabaseTypeRegistry(sqlStatementContext.getDatabaseType()).getDefaultSchemaName(database.getName());
return sqlStatementContext instanceof TableAvailable
? ((TableAvailable) sqlStatementContext).getTablesContext().getSchemaName().map(database::getSchema).orElseGet(() -> database.getSchema(defaultSchemaName))
: database.getSchema(defaultSchemaName);
fetchStatementContext.getCursorStatementContext().getSelectStatementContext().setIndexes(getColumnLabelIndexMap(queryResults.get(0)));
return new FetchStreamMergedResult(queryResults, fetchStatementContext, getSchema(fetchStatementContext, database), connectionContext);
}

private Map<String, Integer> getColumnLabelIndexMap(final QueryResult queryResult) throws SQLException {
Expand All @@ -71,4 +62,9 @@ private Map<String, Integer> getColumnLabelIndexMap(final QueryResult queryResul
}
return result;
}

private ShardingSphereSchema getSchema(final FetchStatementContext fetchStatementContext, final ShardingSphereDatabase database) {
String defaultSchemaName = new DatabaseTypeRegistry(fetchStatementContext.getDatabaseType()).getDefaultSchemaName(database.getName());
return fetchStatementContext.getTablesContext().getSchemaName().map(database::getSchema).orElseGet(() -> database.getSchema(defaultSchemaName));
}
}
89 changes: 42 additions & 47 deletions ...c/test/java/org/apache/shardingsphere/sharding/merge/ddl/ShardingDDLResultMergerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,32 +21,33 @@
import org.apache.shardingsphere.infra.binder.context.statement.ddl.FetchStatementContext;
import org.apache.shardingsphere.infra.binder.context.statement.dml.SelectStatementContext;
import org.apache.shardingsphere.infra.config.props.ConfigurationProperties;
import org.apache.shardingsphere.infra.metadata.database.resource.ResourceMetaData;
import org.apache.shardingsphere.infra.session.connection.ConnectionContext;
import org.apache.shardingsphere.infra.session.connection.cursor.CursorConnectionContext;
import org.apache.shardingsphere.infra.database.core.DefaultDatabase;
import org.apache.shardingsphere.infra.database.core.type.DatabaseType;
import org.apache.shardingsphere.infra.executor.sql.execute.result.query.QueryResult;
import org.apache.shardingsphere.infra.merge.result.impl.transparent.TransparentMergedResult;
import org.apache.shardingsphere.infra.metadata.ShardingSphereMetaData;
import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
import org.apache.shardingsphere.infra.metadata.database.resource.ResourceMetaData;
import org.apache.shardingsphere.infra.metadata.database.rule.RuleMetaData;
import org.apache.shardingsphere.infra.metadata.database.schema.model.ShardingSphereSchema;
import org.apache.shardingsphere.infra.session.connection.ConnectionContext;
import org.apache.shardingsphere.infra.session.connection.cursor.CursorConnectionContext;
import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader;
import org.apache.shardingsphere.sharding.merge.common.IteratorStreamMergedResult;
import org.apache.shardingsphere.sharding.merge.ddl.fetch.FetchStreamMergedResult;
import org.apache.shardingsphere.sql.parser.statement.core.segment.ddl.cursor.CursorNameSegment;
import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.item.ProjectionsSegment;
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.ddl.FetchStatement;
import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.SelectStatement;
import org.apache.shardingsphere.sql.parser.statement.core.value.identifier.IdentifierValue;
import org.apache.shardingsphere.sql.parser.statement.mysql.dml.MySQLSelectStatement;
import org.apache.shardingsphere.sql.parser.statement.opengauss.ddl.OpenGaussFetchStatement;
import org.junit.jupiter.api.Test;

import java.sql.SQLException;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.Optional;

import static org.hamcrest.CoreMatchers.instanceOf;
import static org.hamcrest.MatcherAssert.assertThat;
Expand All @@ -56,66 +57,65 @@

class ShardingDDLResultMergerTest {

private final DatabaseType databaseType = TypedSPILoader.getService(DatabaseType.class, "FIXTURE");

private final ShardingDDLResultMerger merger = new ShardingDDLResultMerger();

@Test
void assertBuildIteratorStreamMergedResult() throws SQLException {
ShardingDDLResultMerger merger = new ShardingDDLResultMerger();
assertThat(merger.merge(createSingleQueryResults(), mock(FetchStatementContext.class), mock(ShardingSphereDatabase.class), mock(ConnectionContext.class)),
void assertMergeWithIteratorStreamMergedResult() throws SQLException {
assertThat(merger.merge(Collections.singletonList(createQueryResult()), mock(FetchStatementContext.class), mock(ShardingSphereDatabase.class), mock(ConnectionContext.class)),
instanceOf(IteratorStreamMergedResult.class));
}

@Test
void assertBuildFetchStreamMergedResult() throws SQLException {
ShardingDDLResultMerger merger = new ShardingDDLResultMerger();
void assertMergeWithTransparentMergedResult() throws SQLException {
assertThat(merger.merge(createQueryResults(), mock(SelectStatementContext.class), mock(ShardingSphereDatabase.class), mock(ConnectionContext.class)),
instanceOf(TransparentMergedResult.class));
}

@Test
void assertMergeWithFetchStreamMergedResult() throws SQLException {
ShardingSphereDatabase database = mock(ShardingSphereDatabase.class, RETURNS_DEEP_STUBS);
when(database.getSchema(DefaultDatabase.LOGIC_NAME)).thenReturn(mock(ShardingSphereSchema.class));
when(database.getSchema("foo_db")).thenReturn(mock(ShardingSphereSchema.class));
ConnectionContext connectionContext = mock(ConnectionContext.class);
when(connectionContext.getCursorContext()).thenReturn(new CursorConnectionContext());
assertThat(merger.merge(createMultiQueryResults(), createFetchStatementContext(database), mock(ShardingSphereDatabase.class), connectionContext),
assertThat(merger.merge(createQueryResults(), createFetchStatementContext(database), mock(ShardingSphereDatabase.class), connectionContext),
instanceOf(FetchStreamMergedResult.class));
}

@Test
void assertBuildTransparentMergedResult() throws SQLException {
ShardingDDLResultMerger merger = new ShardingDDLResultMerger();
assertThat(merger.merge(createMultiQueryResults(), mock(SelectStatementContext.class), mock(ShardingSphereDatabase.class), mock(ConnectionContext.class)),
instanceOf(TransparentMergedResult.class));
}

private FetchStatementContext createFetchStatementContext(final ShardingSphereDatabase database) {
OpenGaussFetchStatement fetchStatement = createFetchStatement();
FetchStatementContext result = new FetchStatementContext(fetchStatement, DefaultDatabase.LOGIC_NAME);
CursorStatementContext cursorStatementContext = createCursorStatementContext(database);
result.setCursorStatementContext(cursorStatementContext);
FetchStatementContext result = new FetchStatementContext(mockFetchStatement(), "foo_db");
result.setCursorStatementContext(createCursorStatementContext(database));
return result;
}

private CursorStatementContext createCursorStatementContext(final ShardingSphereDatabase database) {
CursorStatementContext result = mock(CursorStatementContext.class, RETURNS_DEEP_STUBS);
SelectStatement selectStatement = createSelectStatement();
selectStatement.setProjections(new ProjectionsSegment(0, 0));
SelectStatementContext selectStatementContext = new SelectStatementContext(createShardingSphereMetaData(database), Collections.emptyList(),
selectStatement, DefaultDatabase.LOGIC_NAME, Collections.emptyList());
SelectStatementContext selectStatementContext = new SelectStatementContext(
createShardingSphereMetaData(database), Collections.emptyList(), selectStatement, "foo_db", Collections.emptyList());
when(result.getSelectStatementContext()).thenReturn(selectStatementContext);
when(result.getSqlStatement().getSelect()).thenReturn(selectStatement);
return result;
}

private ShardingSphereMetaData createShardingSphereMetaData(final ShardingSphereDatabase database) {
return new ShardingSphereMetaData(Collections.singletonMap(DefaultDatabase.LOGIC_NAME, database), mock(ResourceMetaData.class),
mock(RuleMetaData.class), mock(ConfigurationProperties.class));
private SelectStatement createSelectStatement() {
SelectStatement result = mock(SelectStatement.class, RETURNS_DEEP_STUBS);
when(result.getFrom()).thenReturn(Optional.of(new SimpleTableSegment(new TableNameSegment(10, 13, new IdentifierValue("tbl")))));
when(result.getProjections()).thenReturn(new ProjectionsSegment(0, 0));
when(result.getDatabaseType()).thenReturn(databaseType);
return result;
}

private List<QueryResult> createSingleQueryResults() throws SQLException {
List<QueryResult> result = new LinkedList<>();
QueryResult queryResult = createQueryResult();
result.add(queryResult);
return result;
private ShardingSphereMetaData createShardingSphereMetaData(final ShardingSphereDatabase database) {
return new ShardingSphereMetaData(Collections.singletonMap("foo_db", database), mock(ResourceMetaData.class),
mock(RuleMetaData.class), mock(ConfigurationProperties.class));
}

private List<QueryResult> createMultiQueryResults() throws SQLException {
private List<QueryResult> createQueryResults() throws SQLException {
List<QueryResult> result = new LinkedList<>();
QueryResult queryResult = createQueryResult();
result.add(queryResult);
result.add(createQueryResult());
result.add(mock(QueryResult.class, RETURNS_DEEP_STUBS));
result.add(mock(QueryResult.class, RETURNS_DEEP_STUBS));
result.add(mock(QueryResult.class, RETURNS_DEEP_STUBS));
Expand All @@ -130,16 +130,11 @@ private QueryResult createQueryResult() throws SQLException {
return result;
}

private OpenGaussFetchStatement createFetchStatement() {
OpenGaussFetchStatement result = new OpenGaussFetchStatement();
result.setCursorName(new CursorNameSegment(0, 0, new IdentifierValue("t_order_cursor")));
return result;
}

private SelectStatement createSelectStatement() {
SelectStatement result = new MySQLSelectStatement();
result.setFrom(new SimpleTableSegment(new TableNameSegment(10, 13, new IdentifierValue("tbl"))));
result.setProjections(new ProjectionsSegment(0, 0));
private FetchStatement mockFetchStatement() {
FetchStatement result = mock(FetchStatement.class);
when(result.getCursorName()).thenReturn(new CursorNameSegment(0, 0, new IdentifierValue("foo_cursor")));
when(result.getDirection()).thenReturn(Optional.empty());
when(result.getDatabaseType()).thenReturn(databaseType);
return result;
}
}

0 comments on commit 68a9804

Please sign in to comment.