Skip to content

Commit b3cc174

Browse files
maxiaoguang64maxiaoguang64
authored
Merge pull request #1 from shardingjdbc/dev
merge from sharding jdbc
2 parents 4d04eef + 84f0d43 commit b3cc174

File tree

11 files changed

+110
-103
lines changed

11 files changed

+110
-103
lines changed

‎sharding-core/src/main/java/io/shardingjdbc/core/parsing/SQLParsingEngine.java

+1-13
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,6 @@
2424
import io.shardingjdbc.core.parsing.lexer.LexerEngineFactory;
2525
import io.shardingjdbc.core.parsing.parser.sql.SQLParserFactory;
2626
import io.shardingjdbc.core.parsing.parser.sql.SQLStatement;
27-
import io.shardingjdbc.core.parsing.parser.token.GeneratedKeyToken;
28-
import io.shardingjdbc.core.parsing.parser.token.SQLToken;
2927
import io.shardingjdbc.core.rule.ShardingRule;
3028
import lombok.RequiredArgsConstructor;
3129

@@ -57,8 +55,7 @@ public SQLStatement parse(final boolean useCache) {
5755
LexerEngine lexerEngine = LexerEngineFactory.newInstance(dbType, sql);
5856
lexerEngine.nextToken();
5957
SQLStatement result = SQLParserFactory.newInstance(dbType, lexerEngine.getCurrentToken().getType(), shardingRule, lexerEngine).parse();
60-
// TODO cannot cache InsertStatement here by generate key, should not modify original InsertStatement on router.
61-
if (useCache && !findGeneratedKeyToken(result)) {
58+
if (useCache) {
6259
ParsingResultCache.getInstance().put(sql, result);
6360
}
6461
return result;
@@ -67,13 +64,4 @@ public SQLStatement parse(final boolean useCache) {
6764
private Optional<SQLStatement> getSQLStatementFromCache(final boolean useCache) {
6865
return useCache ? Optional.fromNullable(ParsingResultCache.getInstance().getSQLStatement(sql)) : Optional.<SQLStatement>absent();
6966
}
70-
71-
private boolean findGeneratedKeyToken(final SQLStatement sqlStatement) {
72-
for (SQLToken each : sqlStatement.getSqlTokens()) {
73-
if (each instanceof GeneratedKeyToken) {
74-
return true;
75-
}
76-
}
77-
return false;
78-
}
7967
}

‎sharding-core/src/main/java/io/shardingjdbc/core/parsing/parser/sql/dml/insert/InsertStatement.java

+4-42
Original file line numberDiff line numberDiff line change
@@ -18,19 +18,12 @@
1818
package io.shardingjdbc.core.parsing.parser.sql.dml.insert;
1919

2020
import com.google.common.base.Optional;
21-
import io.shardingjdbc.core.parsing.lexer.token.Symbol;
2221
import io.shardingjdbc.core.parsing.parser.context.GeneratedKey;
2322
import io.shardingjdbc.core.parsing.parser.context.condition.Column;
24-
import io.shardingjdbc.core.parsing.parser.context.condition.Condition;
2523
import io.shardingjdbc.core.parsing.parser.context.condition.Conditions;
26-
import io.shardingjdbc.core.parsing.parser.expression.SQLNumberExpression;
27-
import io.shardingjdbc.core.parsing.parser.expression.SQLPlaceholderExpression;
2824
import io.shardingjdbc.core.parsing.parser.sql.dml.DMLStatement;
2925
import io.shardingjdbc.core.parsing.parser.token.GeneratedKeyToken;
30-
import io.shardingjdbc.core.parsing.parser.token.ItemsToken;
3126
import io.shardingjdbc.core.parsing.parser.token.SQLToken;
32-
import io.shardingjdbc.core.rule.ShardingRule;
33-
import io.shardingjdbc.core.rule.TableRule;
3427
import lombok.Getter;
3528
import lombok.Setter;
3629
import lombok.ToString;
@@ -64,42 +57,11 @@ public final class InsertStatement extends DMLStatement {
6457
private GeneratedKey generatedKey;
6558

6659
/**
67-
* Append generate key token.
68-
*
69-
* @param shardingRule databases and tables sharding rule
60+
* Find generated key token.
61+
*
62+
* @return generated key token
7063
*/
71-
public void appendGenerateKeyToken(final ShardingRule shardingRule) {
72-
if (null != generatedKey) {
73-
return;
74-
}
75-
Optional<TableRule> tableRule = shardingRule.tryFindTableRuleByLogicTable(getTables().getSingleTableName());
76-
if (!tableRule.isPresent()) {
77-
return;
78-
}
79-
Optional<GeneratedKeyToken> generatedKeysToken = findGeneratedKeyToken();
80-
if (!generatedKeysToken.isPresent()) {
81-
return;
82-
}
83-
ItemsToken valuesToken = new ItemsToken(generatedKeysToken.get().getBeginPosition());
84-
appendGenerateKeyToken(shardingRule, tableRule.get(), valuesToken);
85-
getSqlTokens().remove(generatedKeysToken.get());
86-
getSqlTokens().add(valuesToken);
87-
}
88-
89-
private void appendGenerateKeyToken(final ShardingRule shardingRule, final TableRule tableRule, final ItemsToken valuesToken) {
90-
if (0 == getParametersIndex()) {
91-
Number generatedKey = shardingRule.generateKey(tableRule.getLogicTable());
92-
valuesToken.getItems().add(generatedKey.toString());
93-
getConditions().add(new Condition(new Column(tableRule.getGenerateKeyColumn(), tableRule.getLogicTable()), new SQLNumberExpression(generatedKey)), shardingRule);
94-
this.generatedKey = new GeneratedKey(tableRule.getLogicTable(), -1, generatedKey);
95-
} else {
96-
valuesToken.getItems().add(Symbol.QUESTION.getLiterals());
97-
getConditions().add(new Condition(new Column(tableRule.getGenerateKeyColumn(), tableRule.getLogicTable()), new SQLPlaceholderExpression(getParametersIndex())), shardingRule);
98-
generatedKey = new GeneratedKey(tableRule.getGenerateKeyColumn(), getParametersIndex(), null);
99-
}
100-
}
101-
102-
private Optional<GeneratedKeyToken> findGeneratedKeyToken() {
64+
public Optional<GeneratedKeyToken> findGeneratedKeyToken() {
10365
for (SQLToken each : getSqlTokens()) {
10466
if (each instanceof GeneratedKeyToken) {
10567
return Optional.of((GeneratedKeyToken) each);

‎sharding-core/src/main/java/io/shardingjdbc/core/rewrite/SQLRewriteEngine.java

+20-1
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,13 @@
2121
import com.google.common.base.Strings;
2222
import io.shardingjdbc.core.constant.DatabaseType;
2323
import io.shardingjdbc.core.parsing.lexer.token.DefaultKeyword;
24+
import io.shardingjdbc.core.parsing.lexer.token.Symbol;
25+
import io.shardingjdbc.core.parsing.parser.context.GeneratedKey;
2426
import io.shardingjdbc.core.parsing.parser.context.OrderItem;
2527
import io.shardingjdbc.core.parsing.parser.context.limit.Limit;
2628
import io.shardingjdbc.core.parsing.parser.sql.SQLStatement;
2729
import io.shardingjdbc.core.parsing.parser.sql.dql.select.SelectStatement;
30+
import io.shardingjdbc.core.parsing.parser.token.GeneratedKeyToken;
2831
import io.shardingjdbc.core.parsing.parser.token.IndexToken;
2932
import io.shardingjdbc.core.parsing.parser.token.ItemsToken;
3033
import io.shardingjdbc.core.parsing.parser.token.OffsetToken;
@@ -68,19 +71,23 @@ public final class SQLRewriteEngine {
6871

6972
private final SQLStatement sqlStatement;
7073

74+
private final GeneratedKey generatedKey;
75+
7176
/**
7277
* Constructs SQL rewrite engine.
7378
*
7479
* @param shardingRule databases and tables sharding rule
7580
* @param originalSQL original SQL
7681
* @param databaseType database type
7782
* @param sqlStatement SQL statement
83+
* @param generatedKey generated key
7884
*/
79-
public SQLRewriteEngine(final ShardingRule shardingRule, final String originalSQL, final DatabaseType databaseType, final SQLStatement sqlStatement) {
85+
public SQLRewriteEngine(final ShardingRule shardingRule, final String originalSQL, final DatabaseType databaseType, final SQLStatement sqlStatement, final GeneratedKey generatedKey) {
8086
this.shardingRule = shardingRule;
8187
this.originalSQL = originalSQL;
8288
this.databaseType = databaseType;
8389
this.sqlStatement = sqlStatement;
90+
this.generatedKey = generatedKey;
8491
sqlTokens.addAll(sqlStatement.getSqlTokens());
8592
}
8693

@@ -110,6 +117,8 @@ public SQLBuilder rewrite(final boolean isRewriteLimit) {
110117
appendIndexPlaceholder(result, (IndexToken) each, count, sqlTokens);
111118
} else if (each instanceof ItemsToken) {
112119
appendItemsToken(result, (ItemsToken) each, count, sqlTokens);
120+
} else if (each instanceof GeneratedKeyToken) {
121+
appendGenerateKeyToken(result, (GeneratedKeyToken) each, count, sqlTokens);
113122
} else if (each instanceof RowCountToken) {
114123
appendLimitRowCount(result, (RowCountToken) each, count, sqlTokens, isRewriteLimit);
115124
} else if (each instanceof OffsetToken) {
@@ -164,6 +173,16 @@ private void appendItemsToken(final SQLBuilder sqlBuilder, final ItemsToken item
164173
appendRest(sqlBuilder, count, sqlTokens, beginPosition);
165174
}
166175

176+
private void appendGenerateKeyToken(final SQLBuilder sqlBuilder, final GeneratedKeyToken generatedKeyToken, final int count, final List<SQLToken> sqlTokens) {
177+
ItemsToken valuesToken = new ItemsToken(generatedKeyToken.getBeginPosition());
178+
if (0 == sqlStatement.getParametersIndex()) {
179+
valuesToken.getItems().add(generatedKey.getValue().toString());
180+
} else {
181+
valuesToken.getItems().add(Symbol.QUESTION.getLiterals());
182+
}
183+
appendItemsToken(sqlBuilder, valuesToken, count, sqlTokens);
184+
}
185+
167186
private void appendLimitRowCount(final SQLBuilder sqlBuilder, final RowCountToken rowCountToken, final int count, final List<SQLToken> sqlTokens, final boolean isRewrite) {
168187
SelectStatement selectStatement = (SelectStatement) sqlStatement;
169188
Limit limit = selectStatement.getLimit();

‎sharding-core/src/main/java/io/shardingjdbc/core/routing/SQLRouteResult.java

+3
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
package io.shardingjdbc.core.routing;
1919

20+
import io.shardingjdbc.core.parsing.parser.context.GeneratedKey;
2021
import io.shardingjdbc.core.parsing.parser.sql.SQLStatement;
2122
import lombok.Getter;
2223
import lombok.RequiredArgsConstructor;
@@ -38,6 +39,8 @@ public final class SQLRouteResult {
3839

3940
private final SQLStatement sqlStatement;
4041

42+
private final GeneratedKey generatedKey;
43+
4144
private final Set<SQLExecutionUnit> executionUnits = new LinkedHashSet<>();
4245

4346
private final List<Number> generatedKeys = new LinkedList<>();

‎sharding-core/src/main/java/io/shardingjdbc/core/routing/router/DatabaseHintSQLRouter.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ public SQLStatement parse(final String logicSQL, final boolean useCache) {
5151
@Override
5252
// TODO insert SQL need parse gen key
5353
public SQLRouteResult route(final String logicSQL, final List<Object> parameters, final SQLStatement sqlStatement) {
54-
SQLRouteResult result = new SQLRouteResult(sqlStatement);
54+
SQLRouteResult result = new SQLRouteResult(sqlStatement, null);
5555
RoutingResult routingResult = new DatabaseHintRoutingEngine(shardingRule.getDataSourceNames(), (HintShardingStrategy) shardingRule.getDefaultDatabaseShardingStrategy()).route();
5656
for (TableUnit each : routingResult.getTableUnits().getTableUnits()) {
5757
result.getExecutionUnits().add(new SQLExecutionUnit(each.getDataSourceName(), logicSQL));

‎sharding-core/src/main/java/io/shardingjdbc/core/routing/router/ParsingSQLRouter.java

+37-16
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
package io.shardingjdbc.core.routing.router;
1919

20+
import com.google.common.base.Optional;
21+
import com.google.common.base.Preconditions;
2022
import io.shardingjdbc.core.constant.DatabaseType;
2123
import io.shardingjdbc.core.parsing.SQLParsingEngine;
2224
import io.shardingjdbc.core.parsing.parser.context.GeneratedKey;
@@ -28,6 +30,7 @@
2830
import io.shardingjdbc.core.parsing.parser.sql.ddl.DDLStatement;
2931
import io.shardingjdbc.core.parsing.parser.sql.dml.insert.InsertStatement;
3032
import io.shardingjdbc.core.parsing.parser.sql.dql.select.SelectStatement;
33+
import io.shardingjdbc.core.parsing.parser.token.GeneratedKeyToken;
3134
import io.shardingjdbc.core.rewrite.SQLBuilder;
3235
import io.shardingjdbc.core.rewrite.SQLRewriteEngine;
3336
import io.shardingjdbc.core.routing.SQLExecutionUnit;
@@ -45,6 +48,7 @@
4548
import io.shardingjdbc.core.routing.type.standard.StandardRoutingEngine;
4649
import io.shardingjdbc.core.routing.type.unicast.UnicastRoutingEngine;
4750
import io.shardingjdbc.core.rule.ShardingRule;
51+
import io.shardingjdbc.core.rule.TableRule;
4852
import io.shardingjdbc.core.util.SQLLogger;
4953
import lombok.RequiredArgsConstructor;
5054

@@ -70,22 +74,21 @@ public final class ParsingSQLRouter implements SQLRouter {
7074

7175
@Override
7276
public SQLStatement parse(final String logicSQL, final boolean useCache) {
73-
SQLParsingEngine parsingEngine = new SQLParsingEngine(databaseType, logicSQL, shardingRule);
74-
SQLStatement result = parsingEngine.parse(useCache);
75-
if (result instanceof InsertStatement) {
76-
((InsertStatement) result).appendGenerateKeyToken(shardingRule);
77-
}
78-
return result;
77+
return new SQLParsingEngine(databaseType, logicSQL, shardingRule).parse(useCache);
7978
}
8079

8180
@Override
8281
public SQLRouteResult route(final String logicSQL, final List<Object> parameters, final SQLStatement sqlStatement) {
83-
SQLRouteResult result = new SQLRouteResult(sqlStatement);
84-
if (sqlStatement instanceof InsertStatement && null != ((InsertStatement) sqlStatement).getGeneratedKey()) {
85-
processGeneratedKey(parameters, (InsertStatement) sqlStatement, result);
82+
GeneratedKey generatedKey = null;
83+
if (sqlStatement instanceof InsertStatement) {
84+
generatedKey = getGenerateKey(shardingRule, (InsertStatement) sqlStatement);
8685
}
87-
RoutingResult routingResult = route(parameters, sqlStatement);
88-
SQLRewriteEngine rewriteEngine = new SQLRewriteEngine(shardingRule, logicSQL, databaseType, sqlStatement);
86+
SQLRouteResult result = new SQLRouteResult(sqlStatement, generatedKey);
87+
if (null != generatedKey) {
88+
processGeneratedKey(parameters, generatedKey, sqlStatement.getTables().getSingleTableName(), result);
89+
}
90+
RoutingResult routingResult = route(parameters, sqlStatement, generatedKey);
91+
SQLRewriteEngine rewriteEngine = new SQLRewriteEngine(shardingRule, logicSQL, databaseType, sqlStatement, generatedKey);
8992
boolean isSingleRouting = routingResult.isSingleRouting();
9093
if (sqlStatement instanceof SelectStatement && null != ((SelectStatement) sqlStatement).getLimit()) {
9194
processLimit(parameters, (SelectStatement) sqlStatement, isSingleRouting);
@@ -108,7 +111,7 @@ public SQLRouteResult route(final String logicSQL, final List<Object> parameters
108111
return result;
109112
}
110113

111-
private RoutingResult route(final List<Object> parameters, final SQLStatement sqlStatement) {
114+
private RoutingResult route(final List<Object> parameters, final SQLStatement sqlStatement, final GeneratedKey generatedKey) {
112115
Collection<String> tableNames = sqlStatement.getTables().getTableNames();
113116
RoutingEngine routingEngine;
114117
if (sqlStatement instanceof UseStatement) {
@@ -124,20 +127,38 @@ private RoutingResult route(final List<Object> parameters, final SQLStatement sq
124127
} else if (tableNames.isEmpty()) {
125128
routingEngine = new DatabaseBroadcastRoutingEngine(shardingRule);
126129
} else if (1 == tableNames.size() || shardingRule.isAllBindingTables(tableNames) || shardingRule.isAllInDefaultDataSource(tableNames)) {
127-
routingEngine = new StandardRoutingEngine(shardingRule, parameters, tableNames.iterator().next(), sqlStatement);
130+
routingEngine = new StandardRoutingEngine(shardingRule, parameters, tableNames.iterator().next(), sqlStatement, generatedKey);
128131
} else {
129132
// TODO config for cartesian set
130133
routingEngine = new ComplexRoutingEngine(shardingRule, parameters, tableNames, sqlStatement);
131134
}
132135
return routingEngine.route();
133136
}
134137

135-
private void processGeneratedKey(final List<Object> parameters, final InsertStatement insertStatement, final SQLRouteResult sqlRouteResult) {
136-
GeneratedKey generatedKey = insertStatement.getGeneratedKey();
138+
private GeneratedKey getGenerateKey(final ShardingRule shardingRule, final InsertStatement insertStatement) {
139+
if (null != insertStatement.getGeneratedKey()) {
140+
return insertStatement.getGeneratedKey();
141+
}
142+
Optional<TableRule> tableRule = shardingRule.tryFindTableRuleByLogicTable(insertStatement.getTables().getSingleTableName());
143+
if (!tableRule.isPresent()) {
144+
return null;
145+
}
146+
Optional<GeneratedKeyToken> generatedKeysToken = insertStatement.findGeneratedKeyToken();
147+
if (!generatedKeysToken.isPresent()) {
148+
return null;
149+
}
150+
String logicTableName = insertStatement.getTables().getSingleTableName();
151+
Optional<String> generateKeyColumn = shardingRule.getGenerateKeyColumn(logicTableName);
152+
Preconditions.checkState(generateKeyColumn.isPresent());
153+
return 0 == insertStatement.getParametersIndex()
154+
? new GeneratedKey(generateKeyColumn.get(), -1, shardingRule.generateKey(logicTableName)) : new GeneratedKey(generateKeyColumn.get(), insertStatement.getParametersIndex(), null);
155+
}
156+
157+
private void processGeneratedKey(final List<Object> parameters, final GeneratedKey generatedKey, final String logicTableName, final SQLRouteResult sqlRouteResult) {
137158
if (parameters.isEmpty()) {
138159
sqlRouteResult.getGeneratedKeys().add(generatedKey.getValue());
139160
} else if (parameters.size() == generatedKey.getIndex()) {
140-
Number key = shardingRule.generateKey(insertStatement.getTables().getSingleTableName());
161+
Number key = shardingRule.generateKey(logicTableName);
141162
parameters.add(key);
142163
setGeneratedKeys(sqlRouteResult, key);
143164
} else if (-1 != generatedKey.getIndex()) {

‎sharding-core/src/main/java/io/shardingjdbc/core/routing/type/complex/ComplexRoutingEngine.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ public RoutingResult route() {
6262
Optional<TableRule> tableRule = shardingRule.tryFindTableRuleByLogicTable(each);
6363
if (tableRule.isPresent()) {
6464
if (!bindingTableNames.contains(each)) {
65-
result.add(new StandardRoutingEngine(shardingRule, parameters, tableRule.get().getLogicTable(), sqlStatement).route());
65+
result.add(new StandardRoutingEngine(shardingRule, parameters, tableRule.get().getLogicTable(), sqlStatement, null).route());
6666
}
6767
Optional<BindingTableRule> bindingTableRule = shardingRule.findBindingTableRule(each);
6868
if (bindingTableRule.isPresent()) {

‎sharding-core/src/main/java/io/shardingjdbc/core/routing/type/standard/StandardRoutingEngine.java

+8
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,11 @@
1919

2020
import com.google.common.base.Optional;
2121
import com.google.common.base.Preconditions;
22+
import io.shardingjdbc.core.api.algorithm.sharding.ListShardingValue;
2223
import io.shardingjdbc.core.api.algorithm.sharding.ShardingValue;
2324
import io.shardingjdbc.core.hint.HintManagerHolder;
2425
import io.shardingjdbc.core.hint.ShardingKey;
26+
import io.shardingjdbc.core.parsing.parser.context.GeneratedKey;
2527
import io.shardingjdbc.core.parsing.parser.context.condition.Column;
2628
import io.shardingjdbc.core.parsing.parser.context.condition.Condition;
2729
import io.shardingjdbc.core.parsing.parser.sql.SQLStatement;
@@ -36,6 +38,7 @@
3638

3739
import java.util.ArrayList;
3840
import java.util.Collection;
41+
import java.util.Collections;
3942
import java.util.LinkedList;
4043
import java.util.List;
4144

@@ -55,6 +58,8 @@ public final class StandardRoutingEngine implements RoutingEngine {
5558

5659
private final SQLStatement sqlStatement;
5760

61+
private final GeneratedKey generatedKey;
62+
5863
@Override
5964
public RoutingResult route() {
6065
TableRule tableRule = shardingRule.getTableRule(logicTableName);
@@ -106,6 +111,9 @@ private List<ShardingValue> getShardingValues(final Collection<String> shardingC
106111
Optional<Condition> condition = sqlStatement.getConditions().find(new Column(each, logicTableName));
107112
if (condition.isPresent()) {
108113
result.add(condition.get().getShardingValue(parameters));
114+
} else if (null != generatedKey && each.equals(generatedKey.getColumn())) {
115+
Comparable key = null == generatedKey.getValue() ? (Comparable) parameters.get(generatedKey.getIndex()) : (Comparable) generatedKey.getValue();
116+
result.add(new ListShardingValue<>(sqlStatement.getTables().getSingleTableName(), each, Collections.singletonList(key)));
109117
}
110118
}
111119
return result;

0 commit comments

Comments
 (0)