Skip to content

Commit

Permalink
Split EncryptInsertSelectProjectionTokenGenerator and EncryptSelectPr…
Browse files Browse the repository at this point in the history
…ojectionTokenGenerator (#32285)

* Split EncryptInsertSelectProjectionTokenGenerator and EncryptSelectProjectionTokenGenerator

* Split EncryptInsertSelectProjectionTokenGenerator and EncryptSelectProjectionTokenGenerator
  • Loading branch information
terrymanu authored Jul 26, 2024
1 parent 0efdc32 commit bfeb5ff
Show file tree
Hide file tree
Showing 7 changed files with 139 additions and 42 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -26,17 +26,18 @@
import org.apache.shardingsphere.encrypt.rewrite.token.generator.EncryptAlterTableTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.EncryptAssignmentTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.EncryptCreateTableTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.insert.EncryptInsertDefaultColumnsTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.EncryptGroupByItemTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.EncryptIndexColumnTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.EncryptOrderByItemTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.EncryptPredicateColumnTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.EncryptPredicateRightValueTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.EncryptProjectionTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.insert.EncryptInsertCipherNameTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.insert.EncryptInsertDefaultColumnsTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.insert.EncryptInsertDerivedColumnsTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.insert.EncryptInsertOnUpdateTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.insert.EncryptInsertValuesTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.projection.EncryptInsertSelectProjectionTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.projection.EncryptSelectProjectionTokenGenerator;
import org.apache.shardingsphere.encrypt.rule.EncryptRule;
import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
import org.apache.shardingsphere.infra.rewrite.sql.token.generator.SQLTokenGenerator;
Expand All @@ -62,7 +63,8 @@ public final class EncryptTokenGenerateBuilder implements SQLTokenGeneratorBuild
@Override
public Collection<SQLTokenGenerator> getSQLTokenGenerators() {
Collection<SQLTokenGenerator> result = new LinkedList<>();
addSQLTokenGenerator(result, new EncryptProjectionTokenGenerator());
addSQLTokenGenerator(result, new EncryptSelectProjectionTokenGenerator());
addSQLTokenGenerator(result, new EncryptInsertSelectProjectionTokenGenerator());
addSQLTokenGenerator(result, new EncryptAssignmentTokenGenerator());
addSQLTokenGenerator(result, new EncryptPredicateColumnTokenGenerator());
addSQLTokenGenerator(result, new EncryptPredicateRightValueTokenGenerator());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ public final class InsertSelectColumnsEncryptorComparator {
* @param insertColumns insert columns
* @param projections projections
* @param encryptRule encrypt rule
* @return same encryptors or not
* @return same encryptors or not
*/
public static boolean isSame(final Collection<ColumnSegment> insertColumns, final Collection<Projection> projections, final EncryptRule encryptRule) {
Iterator<ColumnSegment> insertColumnsIterator = insertColumns.iterator();
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package org.apache.shardingsphere.encrypt.rewrite.token.generator.projection;

import lombok.Setter;
import org.apache.shardingsphere.encrypt.rewrite.aware.DatabaseTypeAware;
import org.apache.shardingsphere.encrypt.rewrite.aware.EncryptRuleAware;
import org.apache.shardingsphere.encrypt.rule.EncryptRule;
import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
import org.apache.shardingsphere.infra.binder.context.statement.dml.InsertStatementContext;
import org.apache.shardingsphere.infra.database.core.type.DatabaseType;
import org.apache.shardingsphere.infra.rewrite.sql.token.generator.CollectionSQLTokenGenerator;
import org.apache.shardingsphere.infra.rewrite.sql.token.generator.aware.PreviousSQLTokensAware;
import org.apache.shardingsphere.infra.rewrite.sql.token.pojo.SQLToken;

import java.util.Collection;
import java.util.List;

/**
* Projection token generator for encrypt.
*/
@Setter
public final class EncryptInsertSelectProjectionTokenGenerator implements CollectionSQLTokenGenerator<InsertStatementContext>, PreviousSQLTokensAware, EncryptRuleAware, DatabaseTypeAware {

private List<SQLToken> previousSQLTokens;

private EncryptRule encryptRule;

private DatabaseType databaseType;

@Override
public boolean isGenerateSQLToken(final SQLStatementContext sqlStatementContext) {
return sqlStatementContext instanceof InsertStatementContext && null != ((InsertStatementContext) sqlStatementContext).getInsertSelectContext();
}

@Override
public Collection<SQLToken> generateSQLTokens(final InsertStatementContext sqlStatementContext) {
return new EncryptProjectionTokenGenerator(previousSQLTokens, encryptRule, databaseType).generateSQLTokens(sqlStatementContext.getInsertSelectContext().getSelectStatementContext());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,9 @@
* limitations under the License.
*/

package org.apache.shardingsphere.encrypt.rewrite.token.generator;
package org.apache.shardingsphere.encrypt.rewrite.token.generator.projection;

import lombok.Setter;
import org.apache.shardingsphere.encrypt.rewrite.aware.DatabaseTypeAware;
import org.apache.shardingsphere.encrypt.rewrite.aware.EncryptRuleAware;
import lombok.RequiredArgsConstructor;
import org.apache.shardingsphere.encrypt.rule.EncryptRule;
import org.apache.shardingsphere.encrypt.rule.column.EncryptColumn;
import org.apache.shardingsphere.encrypt.rule.table.EncryptTable;
Expand All @@ -29,15 +27,11 @@
import org.apache.shardingsphere.infra.binder.context.segment.select.projection.ProjectionsContext;
import org.apache.shardingsphere.infra.binder.context.segment.select.projection.impl.ColumnProjection;
import org.apache.shardingsphere.infra.binder.context.segment.select.projection.impl.ShorthandProjection;
import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
import org.apache.shardingsphere.infra.binder.context.statement.dml.InsertStatementContext;
import org.apache.shardingsphere.infra.binder.context.statement.dml.SelectStatementContext;
import org.apache.shardingsphere.infra.database.core.metadata.database.enums.QuoteCharacter;
import org.apache.shardingsphere.infra.database.core.type.DatabaseType;
import org.apache.shardingsphere.infra.exception.core.ShardingSpherePreconditions;
import org.apache.shardingsphere.infra.exception.generic.UnsupportedSQLOperationException;
import org.apache.shardingsphere.infra.rewrite.sql.token.generator.CollectionSQLTokenGenerator;
import org.apache.shardingsphere.infra.rewrite.sql.token.generator.aware.PreviousSQLTokensAware;
import org.apache.shardingsphere.infra.rewrite.sql.token.pojo.SQLToken;
import org.apache.shardingsphere.infra.rewrite.sql.token.pojo.generic.SubstitutableColumnNameToken;
import org.apache.shardingsphere.sql.parser.statement.core.enums.SubqueryType;
Expand All @@ -59,37 +53,28 @@
/**
* Projection token generator for encrypt.
*/
@Setter
public final class EncryptProjectionTokenGenerator implements CollectionSQLTokenGenerator<SQLStatementContext>, PreviousSQLTokensAware, EncryptRuleAware, DatabaseTypeAware {
@RequiredArgsConstructor
public final class EncryptProjectionTokenGenerator {

private List<SQLToken> previousSQLTokens;
private final List<SQLToken> previousSQLTokens;

private EncryptRule encryptRule;
private final EncryptRule encryptRule;

private DatabaseType databaseType;
private final DatabaseType databaseType;

@Override
public boolean isGenerateSQLToken(final SQLStatementContext sqlStatementContext) {
return sqlStatementContext instanceof SelectStatementContext && !((SelectStatementContext) sqlStatementContext).getTablesContext().getSimpleTables().isEmpty()
|| sqlStatementContext instanceof InsertStatementContext && null != ((InsertStatementContext) sqlStatementContext).getInsertSelectContext();
}

@Override
public Collection<SQLToken> generateSQLTokens(final SQLStatementContext sqlStatementContext) {
/**
* Generate SQL tokens.
*
* @param selectStatementContext select statement context
* @return SQL tokens
*/
public Collection<SQLToken> generateSQLTokens(final SelectStatementContext selectStatementContext) {
Collection<SQLToken> result = new LinkedHashSet<>();
if (sqlStatementContext instanceof SelectStatementContext) {
generateSQLTokens((SelectStatementContext) sqlStatementContext, result);
} else if (sqlStatementContext instanceof InsertStatementContext && null != ((InsertStatementContext) sqlStatementContext).getInsertSelectContext()) {
generateSQLTokens(((InsertStatementContext) sqlStatementContext).getInsertSelectContext().getSelectStatementContext(), result);
}
return result;
}

private void generateSQLTokens(final SelectStatementContext selectStatementContext, final Collection<SQLToken> sqlTokens) {
addGenerateSQLTokens(sqlTokens, selectStatementContext);
addGenerateSQLTokens(result, selectStatementContext);
for (SelectStatementContext each : selectStatementContext.getSubqueryContexts().values()) {
addGenerateSQLTokens(sqlTokens, each);
addGenerateSQLTokens(result, each);
}
return result;
}

private void addGenerateSQLTokens(final Collection<SQLToken> sqlTokens, final SelectStatementContext selectStatementContext) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package org.apache.shardingsphere.encrypt.rewrite.token.generator.projection;

import lombok.Setter;
import org.apache.shardingsphere.encrypt.rewrite.aware.DatabaseTypeAware;
import org.apache.shardingsphere.encrypt.rewrite.aware.EncryptRuleAware;
import org.apache.shardingsphere.encrypt.rule.EncryptRule;
import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
import org.apache.shardingsphere.infra.binder.context.statement.dml.SelectStatementContext;
import org.apache.shardingsphere.infra.database.core.type.DatabaseType;
import org.apache.shardingsphere.infra.rewrite.sql.token.generator.CollectionSQLTokenGenerator;
import org.apache.shardingsphere.infra.rewrite.sql.token.generator.aware.PreviousSQLTokensAware;
import org.apache.shardingsphere.infra.rewrite.sql.token.pojo.SQLToken;

import java.util.Collection;
import java.util.List;

/**
* Projection token generator for encrypt.
*/
@Setter
public final class EncryptSelectProjectionTokenGenerator implements CollectionSQLTokenGenerator<SelectStatementContext>, PreviousSQLTokensAware, EncryptRuleAware, DatabaseTypeAware {

private List<SQLToken> previousSQLTokens;

private EncryptRule encryptRule;

private DatabaseType databaseType;

@Override
public boolean isGenerateSQLToken(final SQLStatementContext sqlStatementContext) {
return sqlStatementContext instanceof SelectStatementContext && !((SelectStatementContext) sqlStatementContext).getTablesContext().getSimpleTables().isEmpty();
}

@Override
public Collection<SQLToken> generateSQLTokens(final SelectStatementContext sqlStatementContext) {
return new EncryptProjectionTokenGenerator(previousSQLTokens, encryptRule, databaseType).generateSQLTokens(sqlStatementContext);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
import lombok.SneakyThrows;
import org.apache.shardingsphere.encrypt.rewrite.aware.EncryptRuleAware;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.EncryptOrderByItemTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.EncryptProjectionTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.projection.EncryptSelectProjectionTokenGenerator;
import org.apache.shardingsphere.encrypt.rule.EncryptRule;
import org.apache.shardingsphere.infra.binder.context.segment.select.orderby.OrderByItem;
import org.apache.shardingsphere.infra.binder.context.statement.dml.SelectStatementContext;
Expand Down Expand Up @@ -64,7 +64,7 @@ void assertGetSQLTokenGenerators() {
assertThat(sqlTokenGenerators.size(), is(2));
Iterator<SQLTokenGenerator> iterator = sqlTokenGenerators.iterator();
SQLTokenGenerator item1 = iterator.next();
assertThat(item1, instanceOf(EncryptProjectionTokenGenerator.class));
assertThat(item1, instanceOf(EncryptSelectProjectionTokenGenerator.class));
assertSQLTokenGenerator(item1);
SQLTokenGenerator item2 = iterator.next();
assertThat(item2, instanceOf(EncryptOrderByItemTokenGenerator.class));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

package org.apache.shardingsphere.encrypt.rewrite.token.generator;

import org.apache.shardingsphere.encrypt.rewrite.token.generator.projection.EncryptProjectionTokenGenerator;
import org.apache.shardingsphere.encrypt.rule.EncryptRule;
import org.apache.shardingsphere.encrypt.rule.column.EncryptColumn;
import org.apache.shardingsphere.encrypt.rule.table.EncryptTable;
Expand Down Expand Up @@ -59,14 +60,13 @@

class EncryptProjectionTokenGeneratorTest {

private final EncryptProjectionTokenGenerator generator = new EncryptProjectionTokenGenerator();

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

private EncryptProjectionTokenGenerator generator;

@BeforeEach
void setup() {
generator.setEncryptRule(mockEncryptRule());
generator.setDatabaseType(databaseType);
generator = new EncryptProjectionTokenGenerator(Collections.emptyList(), mockEncryptRule(), databaseType);
}

private EncryptRule mockEncryptRule() {
Expand Down

0 comments on commit bfeb5ff

Please sign in to comment.