Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Code optimization for shardingsphere-shadow project. (#11282) #11285

Merged
merged 6 commits into from
Jul 13, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -17,30 +17,45 @@

package org.apache.shardingsphere.shadow.rewrite.context;

import org.apache.shardingsphere.infra.binder.statement.SQLStatementContext;
import org.apache.shardingsphere.infra.config.properties.ConfigurationProperties;
import org.apache.shardingsphere.infra.rewrite.context.SQLRewriteContext;
import org.apache.shardingsphere.infra.rewrite.context.SQLRewriteContextDecorator;
import org.apache.shardingsphere.infra.rewrite.parameter.builder.ParameterBuilder;
import org.apache.shardingsphere.infra.rewrite.parameter.rewriter.ParameterRewriter;
import org.apache.shardingsphere.infra.route.context.RouteContext;
import org.apache.shardingsphere.shadow.constant.ShadowOrder;
import org.apache.shardingsphere.shadow.rewrite.parameter.ShadowParameterRewriterBuilder;
import org.apache.shardingsphere.shadow.rewrite.token.ShadowTokenGenerateBuilder;
import org.apache.shardingsphere.shadow.rule.ShadowRule;

import java.util.List;

/**
* SQL rewrite context decorator for shadow.
*/
public final class ShadowSQLRewriteContextDecorator implements SQLRewriteContextDecorator<ShadowRule> {

@SuppressWarnings("unchecked")
@Override
public void decorate(final ShadowRule shadowRule, final ConfigurationProperties props, final SQLRewriteContext sqlRewriteContext, final RouteContext routeContext) {
doShadowDecorate(shadowRule, sqlRewriteContext);
}

private void doShadowDecorate(final ShadowRule shadowRule, final SQLRewriteContext sqlRewriteContext) {
doParameterRewriter(shadowRule, sqlRewriteContext);
sqlRewriteContext.addSQLTokenGenerators(new ShadowTokenGenerateBuilder(shadowRule).getSQLTokenGenerators());
}

@SuppressWarnings("unchecked")
private void doParameterRewriter(final ShadowRule shadowRule, final SQLRewriteContext sqlRewriteContext) {
List<Object> parameters = sqlRewriteContext.getParameters();
SQLStatementContext<?> sqlStatementContext = sqlRewriteContext.getSqlStatementContext();
for (ParameterRewriter each : new ShadowParameterRewriterBuilder(shadowRule).getParameterRewriters(sqlRewriteContext.getSchema())) {
if (!sqlRewriteContext.getParameters().isEmpty() && each.isNeedRewrite(sqlRewriteContext.getSqlStatementContext())) {
each.rewrite(sqlRewriteContext.getParameterBuilder(), sqlRewriteContext.getSqlStatementContext(), sqlRewriteContext.getParameters());
if (!parameters.isEmpty() && each.isNeedRewrite(sqlStatementContext)) {
ParameterBuilder parameterBuilder = sqlRewriteContext.getParameterBuilder();
each.rewrite(parameterBuilder, sqlStatementContext, parameters);
}
}
sqlRewriteContext.addSQLTokenGenerators(new ShadowTokenGenerateBuilder(shadowRule).getSQLTokenGenerators());
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,4 +41,8 @@ public final boolean isNeedRewrite(final SQLStatementContext sqlStatementContext
}

protected abstract boolean isNeedRewriteForShadow(SQLStatementContext sqlStatementContext);

protected String getShadowColumn() {
return shadowRule.getColumn();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need another function?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Each subclass is called multiple times, and protected methods are extracted from the parent class.

}
}
Original file line number Diff line number Diff line change
Expand Up @@ -19,12 +19,11 @@

import org.apache.shardingsphere.infra.rewrite.parameter.builder.ParameterBuilder;
import org.apache.shardingsphere.infra.rewrite.parameter.builder.impl.GroupedParameterBuilder;
import org.apache.shardingsphere.infra.rewrite.parameter.builder.impl.StandardParameterBuilder;
import org.apache.shardingsphere.shadow.rewrite.parameter.ShadowParameterRewriter;
import org.apache.shardingsphere.infra.binder.statement.SQLStatementContext;
import org.apache.shardingsphere.infra.binder.statement.dml.InsertStatementContext;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

/**
Expand All @@ -34,31 +33,40 @@ public final class ShadowInsertValueParameterRewriter extends ShadowParameterRew

@Override
protected boolean isNeedRewriteForShadow(final SQLStatementContext sqlStatementContext) {
return sqlStatementContext instanceof InsertStatementContext && ((InsertStatementContext) sqlStatementContext).getInsertColumnNames().contains(getShadowRule().getColumn());
if (sqlStatementContext instanceof InsertStatementContext) {
InsertStatementContext insertStatementContext = (InsertStatementContext) sqlStatementContext;
return insertStatementContext.getInsertColumnNames().contains(getShadowColumn());
}
return false;
}

@Override
public void rewrite(final ParameterBuilder parameterBuilder, final InsertStatementContext insertStatementContext, final List<Object> parameters) {
String columnName = getShadowRule().getColumn();
int columnIndex = getColumnIndex((GroupedParameterBuilder) parameterBuilder, insertStatementContext, columnName);
doShadowRewrite(parameterBuilder, insertStatementContext);
}

private void doShadowRewrite(final ParameterBuilder parameterBuilder, final InsertStatementContext insertStatementContext) {
GroupedParameterBuilder groupedParameterBuilder;
if (parameterBuilder instanceof GroupedParameterBuilder) {
groupedParameterBuilder = (GroupedParameterBuilder) parameterBuilder;
int columnIndex = getShadowColumnIndex(groupedParameterBuilder, insertStatementContext);
addRemovedParametersForShadow(groupedParameterBuilder, insertStatementContext, columnIndex);
}
}

private void addRemovedParametersForShadow(final GroupedParameterBuilder groupedParameterBuilder, final InsertStatementContext insertStatementContext, final int columnIndex) {
int count = 0;
for (List<Object> each : insertStatementContext.getGroupedParameters()) {
if (!each.isEmpty()) {
StandardParameterBuilder standardParameterBuilder = ((GroupedParameterBuilder) parameterBuilder).getParameterBuilders().get(count);
standardParameterBuilder.addRemovedParameters(columnIndex);
groupedParameterBuilder.getParameterBuilders().get(count).addRemovedParameters(columnIndex);
}
count++;
}
}

private int getColumnIndex(final GroupedParameterBuilder parameterBuilder, final InsertStatementContext insertStatementContext, final String shadowColumnName) {
List<String> columnNames;
if (parameterBuilder.getDerivedColumnName().isPresent()) {
columnNames = new ArrayList<>(insertStatementContext.getColumnNames());
columnNames.remove(parameterBuilder.getDerivedColumnName().get());
} else {
columnNames = insertStatementContext.getColumnNames();
}
return columnNames.indexOf(shadowColumnName);

private int getShadowColumnIndex(final GroupedParameterBuilder groupedParameterBuilder, final InsertStatementContext insertStatementContext) {
List<String> columnNames = new LinkedList<>(insertStatementContext.getColumnNames());
groupedParameterBuilder.getDerivedColumnName().ifPresent(columnNames::remove);
return columnNames.indexOf(getShadowColumn());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -39,4 +39,8 @@ public final boolean isGenerateSQLToken(final SQLStatementContext sqlStatementCo
}

protected abstract boolean isGenerateSQLTokenForShadow(SQLStatementContext sqlStatementContext);

protected String getShadowColumn() {
return shadowRule.getColumn();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@

package org.apache.shardingsphere.shadow.rewrite.token.generator.impl;

import com.google.common.base.Preconditions;
import org.apache.shardingsphere.infra.rewrite.sql.token.generator.CollectionSQLTokenGenerator;
import org.apache.shardingsphere.infra.rewrite.sql.token.pojo.generic.RemoveToken;
import org.apache.shardingsphere.shadow.rewrite.token.generator.BaseShadowSQLTokenGenerator;
Expand All @@ -38,28 +37,33 @@ public final class RemoveShadowColumnTokenGenerator extends BaseShadowSQLTokenGe

@Override
protected boolean isGenerateSQLTokenForShadow(final SQLStatementContext sqlStatementContext) {
if (!(sqlStatementContext instanceof InsertStatementContext)) {
return false;
InsertStatementContext insertStatementContext;
if (sqlStatementContext instanceof InsertStatementContext) {
insertStatementContext = (InsertStatementContext) sqlStatementContext;
Optional<InsertColumnsSegment> insertColumnsSegment = insertStatementContext.getSqlStatement().getInsertColumns();
return insertColumnsSegment.isPresent() && !insertColumnsSegment.get().getColumns().isEmpty();
}
Optional<InsertColumnsSegment> insertColumnsSegment = (((InsertStatementContext) sqlStatementContext).getSqlStatement()).getInsertColumns();
return insertColumnsSegment.isPresent() && !insertColumnsSegment.get().getColumns().isEmpty();
return false;
}

@Override
public Collection<RemoveToken> generateSQLTokens(final InsertStatementContext insertStatementContext) {
Optional<InsertColumnsSegment> sqlSegment = insertStatementContext.getSqlStatement().getInsertColumns();
Preconditions.checkState(sqlSegment.isPresent());
Collection<RemoveToken> result = new LinkedList<>();
List<ColumnSegment> columns = (LinkedList<ColumnSegment>) sqlSegment.get().getColumns();
for (int i = 0; i < columns.size(); i++) {
if (getShadowRule().getColumn().equals(columns.get(i).getIdentifier().getValue())) {
if (i == 0) {
result.add(new RemoveToken(columns.get(0).getStartIndex(), columns.get(i + 1).getStartIndex() - 1));
} else {
result.add(new RemoveToken(columns.get(i - 1).getStopIndex() + 1, columns.get(i).getStopIndex()));
}
sqlSegment.ifPresent(insertColumnsSegment -> generateRemoveTokenForShadow(insertColumnsSegment, result));
return result;
}

private void generateRemoveTokenForShadow(final InsertColumnsSegment insertColumnsSegment, final Collection<RemoveToken> removeTokens) {
List<ColumnSegment> columnSegments = (LinkedList<ColumnSegment>) insertColumnsSegment.getColumns();
String shadowColumn = getShadowColumn();
for (int i = 0; i < columnSegments.size(); i++) {
ColumnSegment columnSegment = columnSegments.get(i);
if (shadowColumn.equals(columnSegment.getIdentifier().getValue())) {
RemoveToken removeToken = i == 0 ? new RemoveToken(columnSegments.get(i).getStartIndex(), columnSegments.get(i + 1).getStartIndex() - 1)
: new RemoveToken(columnSegments.get(i - 1).getStopIndex() + 1, columnSegments.get(i).getStopIndex());
removeTokens.add(removeToken);
}
}
return result;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,12 @@ public final class ShadowInsertValuesTokenGenerator extends BaseShadowSQLTokenGe

@Override
protected boolean isGenerateSQLTokenForShadow(final SQLStatementContext sqlStatementContext) {
return sqlStatementContext instanceof InsertStatementContext && ((InsertStatementContext) sqlStatementContext).getInsertColumnNames().contains(getShadowRule().getColumn());
InsertStatementContext insertStatementContext;
if (sqlStatementContext instanceof InsertStatementContext) {
insertStatementContext = (InsertStatementContext) sqlStatementContext;
return insertStatementContext.getInsertColumnNames().contains(getShadowColumn());
}
return false;
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -30,12 +30,14 @@
import org.apache.shardingsphere.shadow.route.judge.impl.PreparedShadowDataSourceJudgeEngine;
import org.apache.shardingsphere.shadow.route.judge.impl.SimpleShadowDataSourceJudgeEngine;
import org.apache.shardingsphere.shadow.rule.ShadowRule;
import org.apache.shardingsphere.sql.parser.sql.common.statement.SQLStatement;
import org.apache.shardingsphere.sql.parser.sql.common.statement.dml.DMLStatement;

import java.util.Collection;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

/**
* Shadow SQL router.
Expand All @@ -44,22 +46,42 @@ public final class ShadowSQLRouter implements SQLRouter<ShadowRule> {

@Override
public RouteContext createRouteContext(final LogicSQL logicSQL, final ShardingSphereMetaData metaData, final ShadowRule rule, final ConfigurationProperties props) {
RouteContext result = new RouteContext();
if (!(logicSQL.getSqlStatementContext().getSqlStatement() instanceof DMLStatement)) {
rule.getShadowMappings().forEach((key, value) -> {
result.getRouteUnits().add(new RouteUnit(new RouteMapper(key, key), Collections.emptyList()));
result.getRouteUnits().add(new RouteUnit(new RouteMapper(value, value), Collections.emptyList()));
});
return result;
}
if (isShadow(logicSQL.getSqlStatementContext(), logicSQL.getParameters(), rule)) {
rule.getShadowMappings().values().forEach(each -> result.getRouteUnits().add(new RouteUnit(new RouteMapper(each, each), Collections.emptyList())));
SQLStatement sqlStatement = logicSQL.getSqlStatementContext().getSqlStatement();
return sqlStatement instanceof DMLStatement ? createRouteContextInDML(logicSQL, rule) : createRouteContextWithoutDML(rule);
}

private RouteContext createRouteContextWithoutDML(final ShadowRule rule) {
final RouteContext result = new RouteContext();
rule.getShadowMappings().forEach((key, value) -> {
result.getRouteUnits().add(createRouteUnit(key, key));
result.getRouteUnits().add(createRouteUnit(value, value));
});
return result;
}

private RouteContext createRouteContextInDML(final LogicSQL logicSQL, final ShadowRule rule) {
final RouteContext result = new RouteContext();
Map<String, String> shadowMappings = rule.getShadowMappings();
if (isShadow(logicSQL, rule)) {
shadowMappings.values().forEach(each -> result.getRouteUnits().add(createRouteUnit(each, each)));
} else {
rule.getShadowMappings().keySet().forEach(each -> result.getRouteUnits().add(new RouteUnit(new RouteMapper(each, each), Collections.emptyList())));
shadowMappings.keySet().forEach(each -> result.getRouteUnits().add(createRouteUnit(each, each)));
}
return result;
}

private boolean isShadow(final LogicSQL logicSQL, final ShadowRule rule) {
final SQLStatementContext<?> sqlStatementContext = logicSQL.getSqlStatementContext();
final List<Object> parameters = logicSQL.getParameters();
ShadowDataSourceJudgeEngine shadowDataSourceRouter = parameters.isEmpty()
? new SimpleShadowDataSourceJudgeEngine(rule, sqlStatementContext) : new PreparedShadowDataSourceJudgeEngine(rule, sqlStatementContext, parameters);
return shadowDataSourceRouter.isShadow();
}

private RouteUnit createRouteUnit(final String logicName, final String actualName) {
return new RouteUnit(new RouteMapper(logicName, actualName), Collections.emptyList());
}

@Override
public void decorateRouteContext(final RouteContext routeContext,
final LogicSQL logicSQL, final ShardingSphereMetaData metaData, final ShadowRule rule, final ConfigurationProperties props) {
Expand All @@ -73,7 +95,7 @@ public void decorateRouteContext(final RouteContext routeContext,
return;
}
Collection<RouteUnit> toBeRemoved = new LinkedList<>();
if (isShadow(logicSQL.getSqlStatementContext(), logicSQL.getParameters(), rule)) {
if (isShadow(logicSQL, rule)) {
for (RouteUnit each : routeContext.getRouteUnits()) {
toBeRemoved.add(each);
String shadowDataSourceName = rule.getShadowMappings().get(each.getDataSourceMapper().getActualName());
Expand All @@ -84,12 +106,6 @@ public void decorateRouteContext(final RouteContext routeContext,
routeContext.getRouteUnits().addAll(toBeAdded);
}

private boolean isShadow(final SQLStatementContext<?> sqlStatementContext, final List<Object> parameters, final ShadowRule rule) {
ShadowDataSourceJudgeEngine shadowDataSourceRouter = parameters.isEmpty()
? new SimpleShadowDataSourceJudgeEngine(rule, sqlStatementContext) : new PreparedShadowDataSourceJudgeEngine(rule, sqlStatementContext, parameters);
return shadowDataSourceRouter.isShadow();
}

@Override
public int getOrder() {
return ShadowOrder.ORDER;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,18 @@

package org.apache.shardingsphere.shadow.condition;

import org.apache.shardingsphere.infra.binder.statement.SQLStatementContext;
import org.apache.shardingsphere.shadow.rule.ShadowRule;
import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.expr.simple.LiteralExpressionSegment;
import org.junit.Test;

import java.util.Collections;
import java.util.List;

import static org.hamcrest.CoreMatchers.is;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertThat;
import static org.mockito.Mockito.mock;

public final class ShadowConditionEngineTest {

Expand All @@ -34,4 +38,9 @@ public void assertGetConditionValues() {
assertThat(actual.size(), is(1));
assertThat(actual.get(0), is(1));
}

@Test
public void assertCreateShadowCondition() {
assertFalse(new ShadowConditionEngine(mock(ShadowRule.class)).createShadowCondition(mock(SQLStatementContext.class)).isPresent());
}
}