diff --git a/sharding-core/sharding-core-merge/src/test/java/org/apache/shardingsphere/core/merge/MergeEngineFactoryTest.java b/sharding-core/sharding-core-merge/src/test/java/org/apache/shardingsphere/core/merge/MergeEngineFactoryTest.java index 66518d4f926de9126e2fcfd669d742ebf7d4450e..d70d193265eed1396b730519e08684a71db52ae3 100644 --- a/sharding-core/sharding-core-merge/src/test/java/org/apache/shardingsphere/core/merge/MergeEngineFactoryTest.java +++ b/sharding-core/sharding-core-merge/src/test/java/org/apache/shardingsphere/core/merge/MergeEngineFactoryTest.java @@ -26,7 +26,6 @@ import org.apache.shardingsphere.core.merge.fixture.TestQueryResult; import org.apache.shardingsphere.core.optimize.encrypt.condition.EncryptCondition; import org.apache.shardingsphere.core.optimize.encrypt.statement.EncryptTransparentOptimizedStatement; import org.apache.shardingsphere.core.optimize.sharding.segment.condition.ShardingCondition; -import org.apache.shardingsphere.core.optimize.sharding.segment.insert.ShardingInsertColumns; import org.apache.shardingsphere.core.optimize.sharding.segment.select.groupby.GroupBy; import org.apache.shardingsphere.core.optimize.sharding.segment.select.item.SelectItem; import org.apache.shardingsphere.core.optimize.sharding.segment.select.item.SelectItems; @@ -90,9 +89,7 @@ public final class MergeEngineFactoryTest { @Test public void assertNewInstanceWithOtherStatement() throws SQLException { - ShardingInsertColumns insertColumns = mock(ShardingInsertColumns.class); - when(insertColumns.getRegularColumnNames()).thenReturn(Collections.emptySet()); - SQLRouteResult routeResult = new SQLRouteResult(new ShardingInsertOptimizedStatement(new InsertStatement(), Collections.emptyList(), insertColumns, null), + SQLRouteResult routeResult = new SQLRouteResult(new ShardingInsertOptimizedStatement(new InsertStatement(), Collections.emptyList(), Collections.emptySet(), null), new EncryptTransparentOptimizedStatement(new InsertStatement())); assertThat(MergeEngineFactory.newInstance(DatabaseTypes.getActualDatabaseType("MySQL"), null, routeResult, null, queryResults), instanceOf(TransparentMergeEngine.class)); } diff --git a/sharding-core/sharding-core-optimize/src/main/java/org/apache/shardingsphere/core/optimize/sharding/engnie/dml/ShardingInsertOptimizeEngine.java b/sharding-core/sharding-core-optimize/src/main/java/org/apache/shardingsphere/core/optimize/sharding/engnie/dml/ShardingInsertOptimizeEngine.java index 221490cfebb64f71c9f624c6fb1b0bcbea641e13..91d0f44ea20b1825b7a015a21426b44407674169 100644 --- a/sharding-core/sharding-core-optimize/src/main/java/org/apache/shardingsphere/core/optimize/sharding/engnie/dml/ShardingInsertOptimizeEngine.java +++ b/sharding-core/sharding-core-optimize/src/main/java/org/apache/shardingsphere/core/optimize/sharding/engnie/dml/ShardingInsertOptimizeEngine.java @@ -25,7 +25,6 @@ import org.apache.shardingsphere.core.optimize.sharding.engnie.ShardingOptimizeE import org.apache.shardingsphere.core.optimize.sharding.segment.condition.ShardingCondition; import org.apache.shardingsphere.core.optimize.sharding.segment.condition.engine.InsertClauseShardingConditionEngine; import org.apache.shardingsphere.core.optimize.sharding.segment.insert.GeneratedKey; -import org.apache.shardingsphere.core.optimize.sharding.segment.insert.ShardingInsertColumns; import org.apache.shardingsphere.core.optimize.sharding.statement.dml.ShardingInsertOptimizedStatement; import org.apache.shardingsphere.core.parse.sql.segment.dml.column.ColumnSegment; import org.apache.shardingsphere.core.parse.sql.segment.dml.column.OnDuplicateKeyColumnsSegment; @@ -53,42 +52,56 @@ public final class ShardingInsertOptimizeEngine implements ShardingOptimizeEngin final TableMetas tableMetas, final String sql, final List parameters, final InsertStatement sqlStatement) { InsertClauseShardingConditionEngine shardingConditionEngine = new InsertClauseShardingConditionEngine(shardingRule); String tableName = sqlStatement.getTable().getTableName(); - ShardingInsertColumns insertColumns = new ShardingInsertColumns(shardingRule, tableMetas, sqlStatement); - Collection columnNames = new LinkedHashSet<>(sqlStatement.useDefaultColumns() ? tableMetas.getAllColumnNames(tableName) : sqlStatement.getColumnNames()); + Collection columnNames = sqlStatement.useDefaultColumns() ? tableMetas.getAllColumnNames(tableName) : sqlStatement.getColumnNames(); Optional generatedKey = GeneratedKey.getGenerateKey(shardingRule, parameters, sqlStatement, columnNames); + Optional generateKeyColumnName = shardingRule.findGenerateKeyColumnName(tableName); boolean isGeneratedValue = generatedKey.isPresent() && generatedKey.get().isGenerated(); - Iterator> generatedValues = isGeneratedValue ? generatedKey.get().getGeneratedValues().iterator() : null; - Collection allColumnNames = new LinkedHashSet<>(insertColumns.getRegularColumnNames()); - if (generatedKey.isPresent()) { - allColumnNames.add(generatedKey.get().getColumnName()); + Iterator> generatedValues = null; + if (isGeneratedValue) { + columnNames.remove(generatedKey.get().getColumnName()); + generatedValues = generatedKey.get().getGeneratedValues().iterator(); } - allColumnNames.addAll(shardingRule.getEncryptRule().getAssistedQueryAndPlainColumns(sqlStatement.getTable().getTableName())); + Collection derivedColumnNames = getDerivedColumnNames(shardingRule, tableName, generatedKey.orNull()); + Collection allColumnNames = new LinkedHashSet<>(columnNames); + allColumnNames.addAll(derivedColumnNames); int derivedColumnsCount = getDerivedColumnsCount(shardingRule, tableName, isGeneratedValue); List shardingConditions = shardingConditionEngine.createShardingConditions(sqlStatement, parameters, allColumnNames, generatedKey.orNull()); - ShardingInsertOptimizedStatement result = new ShardingInsertOptimizedStatement(sqlStatement, shardingConditions, insertColumns, generatedKey.orNull()); + ShardingInsertOptimizedStatement result = new ShardingInsertOptimizedStatement(sqlStatement, shardingConditions, columnNames, generatedKey.orNull()); checkDuplicateKeyForShardingKey(shardingRule, sqlStatement, tableName); int parametersCount = 0; for (Collection each : sqlStatement.getAllValueExpressions()) { - String generateKeyColumnName = shardingRule.findGenerateKeyColumnName(sqlStatement.getTable().getTableName()).orNull(); - Collection encryptDerivedColumnNames = shardingRule.getEncryptRule().getAssistedQueryAndPlainColumns(sqlStatement.getTable().getTableName()); - InsertValue insertValue = result.createInsertValue( - generateKeyColumnName, encryptDerivedColumnNames, each, derivedColumnsCount, parameters, parametersCount); + Collection encryptDerivedColumnNames = shardingRule.getEncryptRule().getAssistedQueryAndPlainColumns(tableName); + InsertValue insertValue = result.createInsertValue(generateKeyColumnName.orNull(), encryptDerivedColumnNames, each, derivedColumnsCount, parameters, parametersCount); result.addInsertValue(insertValue); Object[] currentParameters = insertValue.getParameters(); if (isGeneratedValue) { insertValue.appendValue(generatedValues.next(), Arrays.asList(currentParameters)); } if (shardingRule.getEncryptRule().containsQueryAssistedColumn(tableName)) { - fillAssistedQueryInsertValue(shardingRule, tableName, insertColumns.getRegularColumnNames(), insertValue, Arrays.asList(currentParameters)); + fillAssistedQueryInsertValue(shardingRule, tableName, columnNames, insertValue, Arrays.asList(currentParameters)); } if (shardingRule.getEncryptRule().containsPlainColumn(tableName)) { - fillPlainInsertValue(shardingRule, tableName, insertColumns.getRegularColumnNames(), insertValue, Arrays.asList(currentParameters)); + fillPlainInsertValue(shardingRule, tableName, columnNames, insertValue, Arrays.asList(currentParameters)); } parametersCount += insertValue.getParametersCount(); } return result; } + private Collection getDerivedColumnNames(final ShardingRule shardingRule, final String tableName, final GeneratedKey generatedKey) { + Collection result = new LinkedHashSet<>(); + if (null != generatedKey) { + result.add(generatedKey.getColumnName()); + } + result.addAll(shardingRule.getEncryptRule().getAssistedQueryAndPlainColumns(tableName)); + return result; + } + + private int getDerivedColumnsCount(final ShardingRule shardingRule, final String tableName, final boolean isGeneratedValue) { + int assistedQueryAndPlainColumnsCount = shardingRule.getEncryptRule().getAssistedQueryAndPlainColumnCount(tableName); + return isGeneratedValue ? assistedQueryAndPlainColumnsCount + 1 : assistedQueryAndPlainColumnsCount; + } + private void checkDuplicateKeyForShardingKey(final ShardingRule shardingRule, final InsertStatement sqlStatement, final String tableName) { Optional onDuplicateKeyColumnsSegment = sqlStatement.findSQLSegment(OnDuplicateKeyColumnsSegment.class); if (onDuplicateKeyColumnsSegment.isPresent() && isUpdateShardingKey(shardingRule, onDuplicateKeyColumnsSegment.get(), tableName)) { @@ -105,11 +118,6 @@ public final class ShardingInsertOptimizeEngine implements ShardingOptimizeEngin return false; } - private int getDerivedColumnsCount(final ShardingRule shardingRule, final String tableName, final boolean isGeneratedValue) { - int assistedQueryAndPlainColumnsCount = shardingRule.getEncryptRule().getAssistedQueryAndPlainColumnCount(tableName); - return isGeneratedValue ? assistedQueryAndPlainColumnsCount + 1 : assistedQueryAndPlainColumnsCount; - } - private void fillAssistedQueryInsertValue(final ShardingRule shardingRule, final String tableName, final Collection columnNames, final InsertValue insertValue, final List parameters) { for (String each : columnNames) { diff --git a/sharding-core/sharding-core-optimize/src/main/java/org/apache/shardingsphere/core/optimize/sharding/statement/dml/ShardingInsertOptimizedStatement.java b/sharding-core/sharding-core-optimize/src/main/java/org/apache/shardingsphere/core/optimize/sharding/statement/dml/ShardingInsertOptimizedStatement.java index 54c539dc0fc18675661b7f422cad0a0ea66afdc0..9ee66bedababc952301813b718c3a5b17d526c6a 100644 --- a/sharding-core/sharding-core-optimize/src/main/java/org/apache/shardingsphere/core/optimize/sharding/statement/dml/ShardingInsertOptimizedStatement.java +++ b/sharding-core/sharding-core-optimize/src/main/java/org/apache/shardingsphere/core/optimize/sharding/statement/dml/ShardingInsertOptimizedStatement.java @@ -28,7 +28,6 @@ import org.apache.shardingsphere.core.optimize.encrypt.condition.EncryptConditio import org.apache.shardingsphere.core.optimize.sharding.segment.condition.ShardingCondition; import org.apache.shardingsphere.core.optimize.sharding.segment.condition.ShardingConditions; import org.apache.shardingsphere.core.optimize.sharding.segment.insert.GeneratedKey; -import org.apache.shardingsphere.core.optimize.sharding.segment.insert.ShardingInsertColumns; import org.apache.shardingsphere.core.parse.sql.segment.dml.expr.ExpressionSegment; import org.apache.shardingsphere.core.parse.sql.statement.SQLStatement; @@ -48,17 +47,17 @@ public final class ShardingInsertOptimizedStatement extends ShardingConditionOpt private final Tables tables; - private final ShardingInsertColumns insertColumns; + private final Collection columnNames; private final GeneratedKey generatedKey; private final List insertValues = new LinkedList<>(); public ShardingInsertOptimizedStatement(final SQLStatement sqlStatement, - final List shardingConditions, final ShardingInsertColumns insertColumns, final GeneratedKey generatedKey) { + final List shardingConditions, final Collection columnNames, final GeneratedKey generatedKey) { super(sqlStatement, new ShardingConditions(shardingConditions), new EncryptConditions(Collections.emptyList())); tables = new Tables(sqlStatement); - this.insertColumns = insertColumns; + this.columnNames = columnNames; this.generatedKey = generatedKey; } @@ -75,7 +74,7 @@ public final class ShardingInsertOptimizedStatement extends ShardingConditionOpt */ public InsertValue createInsertValue(final String generateKeyColumnName, final Collection derivedColumnNames, final Collection assignments, final int derivedColumnsCount, final List parameters, final int startIndexOfAppendedParameters) { - List allColumnNames = new LinkedList<>(insertColumns.getRegularColumnNames()); + List allColumnNames = new LinkedList<>(columnNames); allColumnNames.add(generateKeyColumnName); allColumnNames.addAll(derivedColumnNames); return new InsertValue(allColumnNames, assignments, derivedColumnsCount, parameters, startIndexOfAppendedParameters); diff --git a/sharding-core/sharding-core-optimize/src/test/java/org/apache/shardingsphere/core/optimize/sharding/segment/insert/InsertOptimizedStatementTest.java b/sharding-core/sharding-core-optimize/src/test/java/org/apache/shardingsphere/core/optimize/sharding/segment/insert/InsertOptimizedStatementTest.java index 6b2424d31288925dd3fe42fbb1ecc7fbe66e5b06..64118388cc5475d018fe6b92d58c929a269a94da 100644 --- a/sharding-core/sharding-core-optimize/src/test/java/org/apache/shardingsphere/core/optimize/sharding/segment/insert/InsertOptimizedStatementTest.java +++ b/sharding-core/sharding-core-optimize/src/test/java/org/apache/shardingsphere/core/optimize/sharding/segment/insert/InsertOptimizedStatementTest.java @@ -34,16 +34,13 @@ import java.util.List; import static org.hamcrest.CoreMatchers.is; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThat; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; public final class InsertOptimizedStatementTest { @Test public void assertAddInsertValueWithSet() { - ShardingInsertColumns insertColumns = mock(ShardingInsertColumns.class); - when(insertColumns.getRegularColumnNames()).thenReturn(Arrays.asList("id", "value", "status")); - ShardingInsertOptimizedStatement insertClauseOptimizedStatement = new ShardingInsertOptimizedStatement(new InsertStatement(), Collections.emptyList(), insertColumns, null); + ShardingInsertOptimizedStatement insertClauseOptimizedStatement = new ShardingInsertOptimizedStatement( + new InsertStatement(), Collections.emptyList(), Arrays.asList("id", "value", "status"), null); ExpressionSegment assignment1 = new LiteralExpressionSegment(0, 0, 1); ExpressionSegment assignment2 = new ParameterMarkerExpressionSegment(0, 0, 1); ExpressionSegment assignment3 = new LiteralExpressionSegment(0, 0, "test"); diff --git a/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/SQLRewriteEngine.java b/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/SQLRewriteEngine.java index db503c411e296c33a1bd9def650d8ec235382e05..1c7792f45b726e032c6ced75ffbcc7ff89e58c81 100644 --- a/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/SQLRewriteEngine.java +++ b/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/SQLRewriteEngine.java @@ -108,7 +108,7 @@ public final class SQLRewriteEngine { private void encryptInsertOptimizedStatement(final EncryptRule encryptRule, final ShardingInsertOptimizedStatement insertOptimizedStatement) { for (InsertValue insertValue : insertOptimizedStatement.getInsertValues()) { - for (String each : insertOptimizedStatement.getInsertColumns().getRegularColumnNames()) { + for (String each : insertOptimizedStatement.getColumnNames()) { encryptInsertValue(encryptRule, insertValue, insertOptimizedStatement.getTables().getSingleTableName(), each); } } diff --git a/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertGeneratedKeyNameTokenGenerator.java b/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertGeneratedKeyNameTokenGenerator.java index ac8cd45dd077b439d3badfcd880e83cf2fe8d595..4003b91b95be842899f8320005fe962531c75ae6 100644 --- a/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertGeneratedKeyNameTokenGenerator.java +++ b/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertGeneratedKeyNameTokenGenerator.java @@ -46,7 +46,7 @@ public final class InsertGeneratedKeyNameTokenGenerator implements OptionalSQLTo final ShardingInsertOptimizedStatement optimizedStatement, final InsertColumnsSegment segment, final ShardingRule shardingRule) { String tableName = optimizedStatement.getTables().getSingleTableName(); Optional generatedKeyColumnName = shardingRule.findGenerateKeyColumnName(tableName); - return generatedKeyColumnName.isPresent() && !optimizedStatement.getInsertColumns().getRegularColumnNames().contains(generatedKeyColumnName.get()) + return generatedKeyColumnName.isPresent() && !optimizedStatement.getColumnNames().contains(generatedKeyColumnName.get()) ? Optional.of(new InsertGeneratedKeyNameToken(segment.getStopIndex(), generatedKeyColumnName.get(), isToAddCloseParenthesis(tableName, segment, shardingRule))) : Optional.absent(); } @@ -54,5 +54,4 @@ public final class InsertGeneratedKeyNameTokenGenerator implements OptionalSQLTo private boolean isToAddCloseParenthesis(final String tableName, final InsertColumnsSegment segment, final ShardingRule shardingRule) { return segment.getColumns().isEmpty() && 0 == shardingRule.getEncryptRule().getAssistedQueryAndPlainColumnCount(tableName); } - } diff --git a/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertQueryAndPlainNamesTokenGenerator.java b/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertQueryAndPlainNamesTokenGenerator.java index d70fcffe65e68797704d7330cb3df9522172a061..41a267e13a4782ce9648cb57cc79a614fa2ee584 100644 --- a/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertQueryAndPlainNamesTokenGenerator.java +++ b/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertQueryAndPlainNamesTokenGenerator.java @@ -72,8 +72,7 @@ public final class InsertQueryAndPlainNamesTokenGenerator implements OptionalSQL private Collection getAssistedQueryColumns(final InsertOptimizedStatement optimizedStatement, final EncryptRule encryptRule) { Collection result = new LinkedList<>(); Collection columnNames = optimizedStatement instanceof ShardingInsertOptimizedStatement - ? ((ShardingInsertOptimizedStatement) optimizedStatement).getInsertColumns().getRegularColumnNames() - : ((EncryptInsertOptimizedStatement) optimizedStatement).getColumnNames(); + ? ((ShardingInsertOptimizedStatement) optimizedStatement).getColumnNames() : ((EncryptInsertOptimizedStatement) optimizedStatement).getColumnNames(); for (String each : columnNames) { Optional assistedQueryColumn = encryptRule.getAssistedQueryColumn(optimizedStatement.getTables().getSingleTableName(), each); if (assistedQueryColumn.isPresent()) { @@ -86,8 +85,7 @@ public final class InsertQueryAndPlainNamesTokenGenerator implements OptionalSQL private Collection getPlainColumns(final InsertOptimizedStatement optimizedStatement, final EncryptRule encryptRule) { Collection result = new LinkedList<>(); Collection columnNames = optimizedStatement instanceof ShardingInsertOptimizedStatement - ? ((ShardingInsertOptimizedStatement) optimizedStatement).getInsertColumns().getRegularColumnNames() - : ((EncryptInsertOptimizedStatement) optimizedStatement).getColumnNames(); + ? ((ShardingInsertOptimizedStatement) optimizedStatement).getColumnNames() : ((EncryptInsertOptimizedStatement) optimizedStatement).getColumnNames(); for (String each : columnNames) { Optional plainColumn = encryptRule.getPlainColumn(optimizedStatement.getTables().getSingleTableName(), each); if (plainColumn.isPresent()) { diff --git a/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertRegularNamesTokenGenerator.java b/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertRegularNamesTokenGenerator.java index ebb3ad33104f080bce5fbedea6adc616a4859769..87b916098cf0374309d86c2a2d00195adb45d267 100644 --- a/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertRegularNamesTokenGenerator.java +++ b/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertRegularNamesTokenGenerator.java @@ -68,8 +68,7 @@ public final class InsertRegularNamesTokenGenerator implements OptionalSQLTokenG Collection result = new LinkedList<>(); Map logicAndCipherColumns = getEncryptRule(baseRule).getLogicAndCipherColumns(optimizedStatement.getTables().getSingleTableName()); Collection columnNames = optimizedStatement instanceof ShardingInsertOptimizedStatement - ? ((ShardingInsertOptimizedStatement) optimizedStatement).getInsertColumns().getRegularColumnNames() - : ((EncryptInsertOptimizedStatement) optimizedStatement).getColumnNames(); + ? ((ShardingInsertOptimizedStatement) optimizedStatement).getColumnNames() : ((EncryptInsertOptimizedStatement) optimizedStatement).getColumnNames(); for (String each : columnNames) { result.add(getCipherColumn(each, logicAndCipherColumns)); } @@ -96,8 +95,7 @@ public final class InsertRegularNamesTokenGenerator implements OptionalSQLTokenG private boolean isNeedToAppendGeneratedKey(final InsertOptimizedStatement optimizedStatement, final ShardingRule shardingRule) { Optional generateKeyColumnName = shardingRule.findGenerateKeyColumnName(optimizedStatement.getTables().getSingleTableName()); Collection columnNames = optimizedStatement instanceof ShardingInsertOptimizedStatement - ? ((ShardingInsertOptimizedStatement) optimizedStatement).getInsertColumns().getRegularColumnNames() - : ((EncryptInsertOptimizedStatement) optimizedStatement).getColumnNames(); + ? ((ShardingInsertOptimizedStatement) optimizedStatement).getColumnNames() : ((EncryptInsertOptimizedStatement) optimizedStatement).getColumnNames(); return generateKeyColumnName.isPresent() && !columnNames.contains(generateKeyColumnName.get()); } diff --git a/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertSetGeneratedKeyColumnTokenGenerator.java b/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertSetGeneratedKeyColumnTokenGenerator.java index 32d543bc0a8bd7d049faf0e4375141977261bc08..7d3e1d4e17fbcd1930804db255dd24f813ea9fa5 100644 --- a/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertSetGeneratedKeyColumnTokenGenerator.java +++ b/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertSetGeneratedKeyColumnTokenGenerator.java @@ -64,6 +64,6 @@ public final class InsertSetGeneratedKeyColumnTokenGenerator implements Optional private Optional getGeneratedKeyColumn(final ShardingInsertOptimizedStatement optimizedStatement, final ShardingRule shardingRule) { Optional generateKeyColumn = shardingRule.findGenerateKeyColumnName(optimizedStatement.getTables().getSingleTableName()); - return generateKeyColumn.isPresent() && !optimizedStatement.getInsertColumns().getRegularColumnNames().contains(generateKeyColumn.get()) ? generateKeyColumn : Optional.absent(); + return generateKeyColumn.isPresent() && !optimizedStatement.getColumnNames().contains(generateKeyColumn.get()) ? generateKeyColumn : Optional.absent(); } } diff --git a/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertSetQueryAndPlainColumnsTokenGenerator.java b/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertSetQueryAndPlainColumnsTokenGenerator.java index f15e23ba4705e547285b8f0a61a33bbbd9d034aa..8dc701454473c6356fb0bac8e8f639c495291d76 100644 --- a/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertSetQueryAndPlainColumnsTokenGenerator.java +++ b/sharding-core/sharding-core-rewrite/src/main/java/org/apache/shardingsphere/core/rewrite/token/generator/InsertSetQueryAndPlainColumnsTokenGenerator.java @@ -75,8 +75,7 @@ public final class InsertSetQueryAndPlainColumnsTokenGenerator implements Option private List getAssistedQueryColumnNames(final InsertOptimizedStatement optimizedStatement, final EncryptRule encryptRule) { List result = new LinkedList<>(); Collection columnNames = optimizedStatement instanceof ShardingInsertOptimizedStatement - ? ((ShardingInsertOptimizedStatement) optimizedStatement).getInsertColumns().getRegularColumnNames() - : ((EncryptInsertOptimizedStatement) optimizedStatement).getColumnNames(); + ? ((ShardingInsertOptimizedStatement) optimizedStatement).getColumnNames() : ((EncryptInsertOptimizedStatement) optimizedStatement).getColumnNames(); for (String each : columnNames) { Optional assistedQueryColumn = encryptRule.getAssistedQueryColumn(optimizedStatement.getTables().getSingleTableName(), each); if (assistedQueryColumn.isPresent()) { @@ -89,8 +88,7 @@ public final class InsertSetQueryAndPlainColumnsTokenGenerator implements Option private List getPlainColumnNames(final InsertOptimizedStatement optimizedStatement, final EncryptRule encryptRule) { List result = new LinkedList<>(); Collection columnNames = optimizedStatement instanceof ShardingInsertOptimizedStatement - ? ((ShardingInsertOptimizedStatement) optimizedStatement).getInsertColumns().getRegularColumnNames() - : ((EncryptInsertOptimizedStatement) optimizedStatement).getColumnNames(); + ? ((ShardingInsertOptimizedStatement) optimizedStatement).getColumnNames() : ((EncryptInsertOptimizedStatement) optimizedStatement).getColumnNames(); for (String each : columnNames) { Optional plainColumn = encryptRule.getPlainColumn(optimizedStatement.getTables().getSingleTableName(), each); if (plainColumn.isPresent()) { diff --git a/sharding-core/sharding-core-rewrite/src/test/java/org/apache/shardingsphere/core/rewrite/rewriter/ShardingSQLRewriteEngineTest.java b/sharding-core/sharding-core-rewrite/src/test/java/org/apache/shardingsphere/core/rewrite/rewriter/ShardingSQLRewriteEngineTest.java index 35f6ae2360d141342afecedd2b21bef045e1e37e..9a5dd3e48458395896650c81c9fccd97e300f06c 100644 --- a/sharding-core/sharding-core-rewrite/src/test/java/org/apache/shardingsphere/core/rewrite/rewriter/ShardingSQLRewriteEngineTest.java +++ b/sharding-core/sharding-core-rewrite/src/test/java/org/apache/shardingsphere/core/rewrite/rewriter/ShardingSQLRewriteEngineTest.java @@ -25,7 +25,6 @@ import org.apache.shardingsphere.core.optimize.encrypt.condition.EncryptConditio import org.apache.shardingsphere.core.optimize.encrypt.statement.EncryptTransparentOptimizedStatement; import org.apache.shardingsphere.core.optimize.sharding.segment.condition.ShardingCondition; import org.apache.shardingsphere.core.optimize.sharding.segment.condition.ShardingConditions; -import org.apache.shardingsphere.core.optimize.sharding.segment.insert.ShardingInsertColumns; import org.apache.shardingsphere.core.optimize.sharding.segment.select.groupby.GroupBy; import org.apache.shardingsphere.core.optimize.sharding.segment.select.item.AggregationDistinctSelectItem; import org.apache.shardingsphere.core.optimize.sharding.segment.select.item.AggregationSelectItem; @@ -93,7 +92,6 @@ import java.util.Map; import static org.hamcrest.CoreMatchers.is; import static org.junit.Assert.assertThat; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; public final class ShardingSQLRewriteEngineTest { @@ -220,9 +218,7 @@ public final class ShardingSQLRewriteEngineTest { insertStatement.getAllSQLSegments().add(new InsertColumnsSegment(30, 30, Collections.singleton(mock(ColumnSegment.class)))); insertStatement.getAllSQLSegments().add(new InsertValuesSegment(39, 44, Collections.emptyList())); insertStatement.getAllSQLSegments().add(new TableSegment(12, 18, "table_x")); - ShardingInsertColumns insertColumns = mock(ShardingInsertColumns.class); - when(insertColumns.getRegularColumnNames()).thenReturn(Arrays.asList("name", "age")); - ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), insertColumns, null); + ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), Arrays.asList("name", "age"), null); Collection assignments = Arrays.asList( new ParameterMarkerExpressionSegment(0, 0, 0), new ParameterMarkerExpressionSegment(0, 0, 1), new ParameterMarkerExpressionSegment(0, 0, 2)); List parameters = Arrays.asList("x", 1, 1); @@ -248,9 +244,7 @@ public final class ShardingSQLRewriteEngineTest { insertStatement.getAllSQLSegments().add(new InsertValuesSegment(29, 31, Collections.emptyList())); insertStatement.setTable(new TableSegment(0, 0, "table_x")); insertStatement.getAllSQLSegments().add(new TableSegment(12, 20, "`table_x`")); - ShardingInsertColumns insertColumns = mock(ShardingInsertColumns.class); - when(insertColumns.getRegularColumnNames()).thenReturn(Collections.singletonList("name")); - ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), insertColumns, null); + ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), Collections.singletonList("name"), null); Collection assignments = Arrays.asList(new ParameterMarkerExpressionSegment(0, 0, 0), new ParameterMarkerExpressionSegment(0, 0, 1)); List parameters = Arrays.asList("Bill", 1); InsertValue insertValue = shardingStatement.createInsertValue("id", Collections.emptyList(), assignments, 0, parameters, 0); @@ -275,9 +269,7 @@ public final class ShardingSQLRewriteEngineTest { insertStatement.getAllSQLSegments().add(new InsertValuesSegment(29, 32, Collections.emptyList())); insertStatement.setTable(new TableSegment(0, 0, "table_x")); insertStatement.getAllSQLSegments().add(new TableSegment(12, 20, "`table_x`")); - ShardingInsertColumns insertColumns = mock(ShardingInsertColumns.class); - when(insertColumns.getRegularColumnNames()).thenReturn(Collections.singletonList("name")); - ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), insertColumns, null); + ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), Collections.singletonList("name"), null); Collection assignments = Arrays.asList(new LiteralExpressionSegment(0, 0, 10), new LiteralExpressionSegment(0, 0, 1)); List parameters = Collections.emptyList(); InsertValue insertValue = shardingStatement.createInsertValue("id", Collections.emptyList(), assignments, 0, parameters, 0); @@ -302,9 +294,7 @@ public final class ShardingSQLRewriteEngineTest { insertStatement.getAllSQLSegments().add(new InsertValuesSegment(29, 32, Collections.emptyList())); insertStatement.setTable(new TableSegment(0, 0, "table_x")); insertStatement.getAllSQLSegments().add(new TableSegment(12, 20, "`table_x`")); - ShardingInsertColumns insertColumns = mock(ShardingInsertColumns.class); - when(insertColumns.getRegularColumnNames()).thenReturn(Collections.singletonList("name")); - ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), insertColumns, null); + ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), Collections.singletonList("name"), null); Collection assignments = Arrays.asList(new LiteralExpressionSegment(0, 0, 10), new LiteralExpressionSegment(0, 0, 1)); List parameters = Collections.emptyList(); InsertValue insertValue = shardingStatement.createInsertValue("id", Collections.emptyList(), assignments, 0, parameters, 0); @@ -328,9 +318,7 @@ public final class ShardingSQLRewriteEngineTest { insertStatement.getColumns().add(new ColumnSegment(0, 0, "id")); insertStatement.setTable(new TableSegment(0, 0, "table_x")); insertStatement.getAllSQLSegments().add(new TableSegment(12, 20, "`table_x`")); - ShardingInsertColumns insertColumns = mock(ShardingInsertColumns.class); - when(insertColumns.getRegularColumnNames()).thenReturn(Arrays.asList("name", "id")); - ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), insertColumns, null); + ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), Arrays.asList("name", "id"), null); Collection assignments = Arrays.asList(new LiteralExpressionSegment(0, 0, 10), new LiteralExpressionSegment(0, 0, 1)); List parameters = Collections.emptyList(); InsertValue insertValue = shardingStatement.createInsertValue("id", Collections.emptyList(), assignments, 1, parameters, 0); @@ -355,9 +343,7 @@ public final class ShardingSQLRewriteEngineTest { insertStatement.getAllSQLSegments().add(new TableSegment(12, 20, "`table_x`")); insertStatement.getAllSQLSegments().add(new InsertColumnsSegment(21, 21, Collections.emptyList())); insertStatement.getAllSQLSegments().add(new InsertValuesSegment(29, 35, Collections.emptyList())); - ShardingInsertColumns insertColumns = mock(ShardingInsertColumns.class); - when(insertColumns.getRegularColumnNames()).thenReturn(Arrays.asList("name", "id")); - ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), insertColumns, null); + ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), Arrays.asList("name", "id"), null); Collection assignments = Arrays.asList(new LiteralExpressionSegment(0, 0, 10), new LiteralExpressionSegment(0, 0, 1)); List parameters = Collections.emptyList(); InsertValue insertValue = shardingStatement.createInsertValue("id", Collections.emptyList(), assignments, 0, parameters, 0); @@ -381,9 +367,7 @@ public final class ShardingSQLRewriteEngineTest { insertStatement.getAllSQLSegments().add(new InsertValuesSegment(29, 34, Collections.emptyList())); insertStatement.setTable(new TableSegment(0, 0, "table_x")); insertStatement.getAllSQLSegments().add(new TableSegment(12, 20, "`table_x`")); - ShardingInsertColumns insertColumns = mock(ShardingInsertColumns.class); - when(insertColumns.getRegularColumnNames()).thenReturn(Collections.singletonList("name")); - ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), insertColumns, null); + ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), Collections.singletonList("name"), null); Collection assignments = Arrays.asList(new ParameterMarkerExpressionSegment(0, 0, 0), new ParameterMarkerExpressionSegment(0, 0, 1)); List parameters = Arrays.asList("x", 1); InsertValue insertValue = shardingStatement.createInsertValue("id", Collections.emptyList(), assignments, 0, parameters, 0); @@ -850,10 +834,8 @@ public final class ShardingSQLRewriteEngineTest { private SQLRouteResult createSQLRouteResultForTableTokenWithSchemaForInsert() { InsertStatement insertStatement = new InsertStatement(); insertStatement.getAllSQLSegments().add(new TableSegment(12, 30, "table_x")); - ShardingInsertColumns insertColumns = mock(ShardingInsertColumns.class); - when(insertColumns.getRegularColumnNames()).thenReturn(Arrays.asList("order_id", "user_id", "status", "id")); - SQLRouteResult result = new SQLRouteResult( - new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), insertColumns, null), new EncryptTransparentOptimizedStatement(insertStatement)); + SQLRouteResult result = new SQLRouteResult(new ShardingInsertOptimizedStatement( + insertStatement, Collections.emptyList(), Arrays.asList("order_id", "user_id", "status", "id"), null), new EncryptTransparentOptimizedStatement(insertStatement)); result.setRoutingResult(new RoutingResult()); return result; } @@ -989,9 +971,7 @@ public final class ShardingSQLRewriteEngineTest { insertStatement.getAllSQLSegments().add(new SetAssignmentsSegment(26, 34, Collections.singleton(new AssignmentSegment(26, 34, columnSegment, expressionSegment)))); insertStatement.getAllSQLSegments().add(new TableSegment(12, 20, "`table_w`")); insertStatement.setTable(new TableSegment(0, 0, "table_w")); - ShardingInsertColumns insertColumns = mock(ShardingInsertColumns.class); - when(insertColumns.getRegularColumnNames()).thenReturn(Collections.singletonList("name")); - ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), insertColumns, null); + ShardingInsertOptimizedStatement shardingStatement = new ShardingInsertOptimizedStatement(insertStatement, Collections.emptyList(), Collections.singletonList("name"), null); Collection assignments = Arrays.asList( new LiteralExpressionSegment(0, 0, 10), new LiteralExpressionSegment(0, 0, 1), new LiteralExpressionSegment(0, 0, 1), new LiteralExpressionSegment(0, 0, 10)); List parameters = Collections.emptyList();