Skip to content

Commit

Permalink
refactor SQLRewriteEngine
Browse files Browse the repository at this point in the history
  • Loading branch information
terrymanu committed Jan 2, 2019
1 parent 2a5cecb commit 3dbe777
Show file tree
Hide file tree
Showing 4 changed files with 73 additions and 75 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import io.shardingsphere.core.merger.QueryResult;
import lombok.AccessLevel;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;

import java.io.InputStream;
Expand All @@ -44,6 +45,7 @@
*
* @author panjuan
*/
@RequiredArgsConstructor
@Getter(AccessLevel.PROTECTED)
public class DistinctQueryResult implements QueryResult {

Expand All @@ -53,11 +55,6 @@ public class DistinctQueryResult implements QueryResult {

private QueryRow currentRow;

protected DistinctQueryResult(final Multimap<String, Integer> columnLabelAndIndexMap, final Iterator<QueryRow> resultData) {
this.columnLabelAndIndexMap = columnLabelAndIndexMap;
this.resultData = resultData;
}

@SneakyThrows
public DistinctQueryResult(final Collection<QueryResult> queryResults, final List<String> distinctColumnLabels) {
this.columnLabelAndIndexMap = getColumnLabelAndIndexMap(queryResults.iterator().next());
Expand All @@ -75,18 +72,18 @@ private Multimap<String, Integer> getColumnLabelAndIndexMap(final QueryResult qu

@SneakyThrows
private Iterator<QueryRow> getResultData(final Collection<QueryResult> queryResults, final List<String> distinctColumnLabels) {
Set<QueryRow> resultData = new LinkedHashSet<>();
Set<QueryRow> result = new LinkedHashSet<>();
List<Integer> distinctColumnIndexes = Lists.transform(distinctColumnLabels, new Function<String, Integer>() {

@Override
public Integer apply(final String input) {
return getColumnIndex(input);
}
});
for (QueryResult each : queryResults) {
fill(resultData, each, distinctColumnIndexes);
fill(result, each, distinctColumnIndexes);
}
return resultData.iterator();
return result.iterator();
}

private void fill(final Set<QueryRow> resultData, final QueryResult queryResult, final List<Integer> distinctColumnIndexes) throws SQLException {
Expand All @@ -106,7 +103,7 @@ private void fill(final Set<QueryRow> resultData, final QueryResult queryResult,
*/
public List<DistinctQueryResult> divide() {
return Lists.newArrayList(Iterators.transform(resultData, new Function<QueryRow, DistinctQueryResult>() {

@Override
public DistinctQueryResult apply(final QueryRow row) {
Set<QueryRow> resultData = new LinkedHashSet<>();
Expand All @@ -117,7 +114,7 @@ public DistinctQueryResult apply(final QueryRow row) {
}

@Override
public boolean next() {
public final boolean next() {
if (resultData.hasNext()) {
currentRow = resultData.next();
return true;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -107,73 +107,74 @@ public SQLRewriteEngine(final ShardingRule shardingRule, final String originalSQ
/**
* rewrite SQL.
*
* @param isRewrite is rewrite
* @param isSingleRouting is rewrite
* @return SQL builder
*/
public SQLBuilder rewrite(final boolean isRewrite) {
public SQLBuilder rewrite(final boolean isSingleRouting) {
SQLBuilder result = new SQLBuilder(parameters);
if (sqlTokens.isEmpty()) {
return appendOriginalLiterals(result);
}
appendInitialLiterals(isRewrite, result);
appendTokensAndPlaceholders(isRewrite, result);
appendInitialLiterals(!isSingleRouting, result);
appendTokensAndPlaceholders(!isSingleRouting, result);
return result;
}

private SQLBuilder appendOriginalLiterals(final SQLBuilder result) {
result.appendLiterals(originalSQL);
return result;
private SQLBuilder appendOriginalLiterals(final SQLBuilder sqlBuilder) {
sqlBuilder.appendLiterals(originalSQL);
return sqlBuilder;
}

private void appendInitialLiterals(final boolean isRewrite, final SQLBuilder result) {
if (isRewrite && isContainAggregationDistinctToken()) {
appendDistinctLiteral(result);
private void appendInitialLiterals(final boolean isRewrite, final SQLBuilder sqlBuilder) {
if (isRewrite && isContainsAggregationDistinctToken()) {
appendAggregationDistinctLiteral(sqlBuilder);
} else {
result.appendLiterals(originalSQL.substring(0, sqlTokens.get(0).getBeginPosition()));
sqlBuilder.appendLiterals(originalSQL.substring(0, sqlTokens.get(0).getBeginPosition()));
}
}

private boolean isContainAggregationDistinctToken() {
private boolean isContainsAggregationDistinctToken() {
return Iterators.tryFind(sqlTokens.iterator(), new Predicate<SQLToken>() {

@Override
public boolean apply(final SQLToken input) {
return input instanceof AggregationDistinctToken;
}
}).isPresent();
}

private void appendDistinctLiteral(final SQLBuilder result) {
private void appendAggregationDistinctLiteral(final SQLBuilder sqlBuilder) {
int firstSelectItemStartPosition = ((SelectStatement) sqlStatement).getFirstSelectItemStartPosition();
result.appendLiterals(originalSQL.substring(0, firstSelectItemStartPosition));
result.appendLiterals("DISTINCT ");
result.appendLiterals(originalSQL.substring(firstSelectItemStartPosition, sqlTokens.get(0).getBeginPosition()));
sqlBuilder.appendLiterals(originalSQL.substring(0, firstSelectItemStartPosition));
sqlBuilder.appendLiterals("DISTINCT ");
sqlBuilder.appendLiterals(originalSQL.substring(firstSelectItemStartPosition, sqlTokens.get(0).getBeginPosition()));
}

private void appendTokensAndPlaceholders(final boolean isRewrite, final SQLBuilder result) {
private void appendTokensAndPlaceholders(final boolean isRewrite, final SQLBuilder sqlBuilder) {
int count = 0;
for (SQLToken each : sqlTokens) {
if (each instanceof TableToken) {
appendTablePlaceholder(result, (TableToken) each, count);
appendTablePlaceholder(sqlBuilder, (TableToken) each, count);
} else if (each instanceof SchemaToken) {
appendSchemaPlaceholder(result, (SchemaToken) each, count);
appendSchemaPlaceholder(sqlBuilder, (SchemaToken) each, count);
} else if (each instanceof IndexToken) {
appendIndexPlaceholder(result, (IndexToken) each, count);
appendIndexPlaceholder(sqlBuilder, (IndexToken) each, count);
} else if (each instanceof ItemsToken) {
appendItemsToken(result, (ItemsToken) each, count, isRewrite);
appendItemsToken(sqlBuilder, (ItemsToken) each, count, isRewrite);
} else if (each instanceof InsertValuesToken) {
appendInsertValuesToken(result, (InsertValuesToken) each, count);
appendInsertValuesToken(sqlBuilder, (InsertValuesToken) each, count);
} else if (each instanceof RowCountToken) {
appendLimitRowCount(result, (RowCountToken) each, count, isRewrite);
appendLimitRowCount(sqlBuilder, (RowCountToken) each, count, isRewrite);
} else if (each instanceof OffsetToken) {
appendLimitOffsetToken(result, (OffsetToken) each, count, isRewrite);
appendLimitOffsetToken(sqlBuilder, (OffsetToken) each, count, isRewrite);
} else if (each instanceof OrderByToken) {
appendOrderByToken(result, count, isRewrite);
appendOrderByToken(sqlBuilder, count, isRewrite);
} else if (each instanceof InsertColumnToken) {
appendSymbolToken(result, (InsertColumnToken) each, count);
appendSymbolToken(sqlBuilder, (InsertColumnToken) each, count);
} else if (each instanceof AggregationDistinctToken) {
appendAggregationDistinctPlaceholder(result, (AggregationDistinctToken) each, count, isRewrite);
appendAggregationDistinctPlaceholder(sqlBuilder, (AggregationDistinctToken) each, count, isRewrite);
} else if (each instanceof RemoveToken) {
appendRest(result, count, ((RemoveToken) each).getEndPosition());
appendRest(sqlBuilder, count, ((RemoveToken) each).getEndPosition());
}
count++;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ public SQLRouteResult route(final String logicSQL, final List<Object> parameters
if (sqlStatement instanceof SelectStatement && null != ((SelectStatement) sqlStatement).getLimit()) {
processLimit(parameters, (SelectStatement) sqlStatement);
}
SQLBuilder sqlBuilder = rewriteEngine.rewrite(!routingResult.isSingleRouting());
SQLBuilder sqlBuilder = rewriteEngine.rewrite(routingResult.isSingleRouting());
for (TableUnit each : routingResult.getTableUnits().getTableUnits()) {
result.getRouteUnits().add(new RouteUnit(each.getDataSourceName(), rewriteEngine.generateSQL(each, sqlBuilder, shardingMetaData.getDataSource())));
}
Expand Down
Loading

0 comments on commit 3dbe777

Please sign in to comment.