From 111d22f49c2088a86792fa06b0d19431715932e0 Mon Sep 17 00:00:00 2001 From: duanzhengqiang Date: Fri, 20 Dec 2024 12:43:29 +0800 Subject: [PATCH] Use sql bind info in EncryptInsertPredicateColumnTokenGenerator to avoid wrong column table mapping --- ...ptInsertPredicateColumnTokenGenerator.java | 11 +------- .../EncryptPredicateColumnTokenGenerator.java | 27 ++++++------------- .../EncryptGeneratorFixtureBuilder.java | 12 ++++++--- ...ryptPredicateColumnTokenGeneratorTest.java | 3 --- .../cases/dql/e2e-dql-select-sub-query.xml | 8 ++++++ 5 files changed, 25 insertions(+), 36 deletions(-) diff --git a/features/encrypt/core/src/main/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/predicate/EncryptInsertPredicateColumnTokenGenerator.java b/features/encrypt/core/src/main/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/predicate/EncryptInsertPredicateColumnTokenGenerator.java index 609c050a10d90..d93fc96219038 100644 --- a/features/encrypt/core/src/main/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/predicate/EncryptInsertPredicateColumnTokenGenerator.java +++ b/features/encrypt/core/src/main/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/predicate/EncryptInsertPredicateColumnTokenGenerator.java @@ -23,13 +23,10 @@ import org.apache.shardingsphere.infra.annotation.HighFrequencyInvocation; import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext; import org.apache.shardingsphere.infra.binder.context.statement.dml.InsertStatementContext; -import org.apache.shardingsphere.infra.metadata.database.schema.model.ShardingSphereSchema; import org.apache.shardingsphere.infra.rewrite.sql.token.common.generator.CollectionSQLTokenGenerator; -import org.apache.shardingsphere.infra.rewrite.sql.token.common.generator.aware.SchemaMetaDataAware; import org.apache.shardingsphere.infra.rewrite.sql.token.common.pojo.SQLToken; import java.util.Collection; -import java.util.Map; /** * Insert predicate column token generator for encrypt. @@ -37,14 +34,10 @@ @HighFrequencyInvocation @RequiredArgsConstructor @Setter -public final class EncryptInsertPredicateColumnTokenGenerator implements CollectionSQLTokenGenerator, SchemaMetaDataAware { +public final class EncryptInsertPredicateColumnTokenGenerator implements CollectionSQLTokenGenerator { private final EncryptRule rule; - private Map schemas; - - private ShardingSphereSchema defaultSchema; - @Override public boolean isGenerateSQLToken(final SQLStatementContext sqlStatementContext) { return sqlStatementContext instanceof InsertStatementContext && null != ((InsertStatementContext) sqlStatementContext).getInsertSelectContext() @@ -54,8 +47,6 @@ public boolean isGenerateSQLToken(final SQLStatementContext sqlStatementContext) @Override public Collection generateSQLTokens(final SQLStatementContext sqlStatementContext) { EncryptPredicateColumnTokenGenerator generator = new EncryptPredicateColumnTokenGenerator(rule); - generator.setSchemas(schemas); - generator.setDefaultSchema(defaultSchema); return generator.generateSQLTokens(((InsertStatementContext) sqlStatementContext).getInsertSelectContext().getSelectStatementContext()); } } diff --git a/features/encrypt/core/src/main/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/predicate/EncryptPredicateColumnTokenGenerator.java b/features/encrypt/core/src/main/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/predicate/EncryptPredicateColumnTokenGenerator.java index 191380747b5c8..af41bb328b9e6 100644 --- a/features/encrypt/core/src/main/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/predicate/EncryptPredicateColumnTokenGenerator.java +++ b/features/encrypt/core/src/main/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/predicate/EncryptPredicateColumnTokenGenerator.java @@ -30,27 +30,23 @@ import org.apache.shardingsphere.infra.binder.context.segment.select.projection.impl.ColumnProjection; import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext; import org.apache.shardingsphere.infra.binder.context.statement.dml.SelectStatementContext; -import org.apache.shardingsphere.infra.binder.context.type.TableAvailable; import org.apache.shardingsphere.infra.binder.context.type.WhereAvailable; 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.metadata.database.schema.model.ShardingSphereSchema; import org.apache.shardingsphere.infra.rewrite.sql.token.common.generator.CollectionSQLTokenGenerator; -import org.apache.shardingsphere.infra.rewrite.sql.token.common.generator.aware.SchemaMetaDataAware; import org.apache.shardingsphere.infra.rewrite.sql.token.common.pojo.SQLToken; import org.apache.shardingsphere.infra.rewrite.sql.token.common.pojo.generic.SubstitutableColumnNameToken; +import org.apache.shardingsphere.sql.parser.statement.core.extractor.ExpressionExtractor; import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.column.ColumnSegment; import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.expr.BinaryOperationExpression; import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.expr.ExpressionSegment; import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.predicate.AndPredicate; import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.predicate.WhereSegment; -import org.apache.shardingsphere.sql.parser.statement.core.extractor.ExpressionExtractor; import org.apache.shardingsphere.sql.parser.statement.core.value.identifier.IdentifierValue; import java.util.Collection; import java.util.Collections; import java.util.LinkedList; -import java.util.Map; import java.util.Optional; /** @@ -59,14 +55,10 @@ @HighFrequencyInvocation @RequiredArgsConstructor @Setter -public final class EncryptPredicateColumnTokenGenerator implements CollectionSQLTokenGenerator, SchemaMetaDataAware { +public final class EncryptPredicateColumnTokenGenerator implements CollectionSQLTokenGenerator { private final EncryptRule rule; - private Map schemas; - - private ShardingSphereSchema defaultSchema; - @Override public boolean isGenerateSQLToken(final SQLStatementContext sqlStatementContext) { return sqlStatementContext instanceof WhereAvailable && !((WhereAvailable) sqlStatementContext).getWhereSegments().isEmpty(); @@ -77,19 +69,16 @@ public Collection generateSQLTokens(final SQLStatementContext sqlState Collection allSubqueryContexts = SQLStatementContextExtractor.getAllSubqueryContexts(sqlStatementContext); Collection whereSegments = SQLStatementContextExtractor.getWhereSegments((WhereAvailable) sqlStatementContext, allSubqueryContexts); Collection columnSegments = SQLStatementContextExtractor.getColumnSegments((WhereAvailable) sqlStatementContext, allSubqueryContexts); - ShardingSphereSchema schema = ((TableAvailable) sqlStatementContext).getTablesContext().getSchemaName().map(schemas::get).orElseGet(() -> defaultSchema); - Map columnExpressionTableNames = ((TableAvailable) sqlStatementContext).getTablesContext().findTableNames(columnSegments, schema); - return generateSQLTokens(columnSegments, columnExpressionTableNames, whereSegments, sqlStatementContext.getDatabaseType()); + return generateSQLTokens(columnSegments, whereSegments, sqlStatementContext.getDatabaseType()); } - private Collection generateSQLTokens(final Collection columnSegments, final Map columnExpressionTableNames, - final Collection whereSegments, final DatabaseType databaseType) { + private Collection generateSQLTokens(final Collection columnSegments, final Collection whereSegments, final DatabaseType databaseType) { Collection result = new LinkedList<>(); for (ColumnSegment each : columnSegments) { - String tableName = columnExpressionTableNames.getOrDefault(each.getExpression(), ""); - Optional encryptTable = rule.findEncryptTable(tableName); - if (encryptTable.isPresent() && encryptTable.get().isEncryptColumn(each.getIdentifier().getValue())) { - result.add(buildSubstitutableColumnNameToken(encryptTable.get().getEncryptColumn(each.getIdentifier().getValue()), each, whereSegments, databaseType)); + Optional encryptTable = rule.findEncryptTable(each.getColumnBoundInfo().getOriginalTable().getValue()); + if (encryptTable.isPresent() && encryptTable.get().isEncryptColumn(each.getColumnBoundInfo().getOriginalColumn().getValue())) { + EncryptColumn encryptColumn = encryptTable.get().getEncryptColumn(each.getColumnBoundInfo().getOriginalColumn().getValue()); + result.add(buildSubstitutableColumnNameToken(encryptColumn, each, whereSegments, databaseType)); } } return result; diff --git a/features/encrypt/core/src/test/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/fixture/EncryptGeneratorFixtureBuilder.java b/features/encrypt/core/src/test/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/fixture/EncryptGeneratorFixtureBuilder.java index 5befae4ed0535..5ee1eb212fbee 100644 --- a/features/encrypt/core/src/test/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/fixture/EncryptGeneratorFixtureBuilder.java +++ b/features/encrypt/core/src/test/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/fixture/EncryptGeneratorFixtureBuilder.java @@ -168,10 +168,14 @@ public static UpdateStatementContext createUpdateStatementContext() { } private static WhereSegment createWhereSegment() { - BinaryOperationExpression nameExpression = new BinaryOperationExpression(10, 24, - new ColumnSegment(10, 13, new IdentifierValue("name")), new LiteralExpressionSegment(18, 22, "LiLei"), "=", "name = 'LiLei'"); - BinaryOperationExpression pwdExpression = new BinaryOperationExpression(30, 44, - new ColumnSegment(30, 32, new IdentifierValue("pwd")), new LiteralExpressionSegment(40, 45, "123456"), "=", "pwd = '123456'"); + ColumnSegment nameColumnSegment = new ColumnSegment(10, 13, new IdentifierValue("name")); + nameColumnSegment.setColumnBoundInfo( + new ColumnSegmentBoundInfo(new TableSegmentBoundInfo(new IdentifierValue("foo_db"), new IdentifierValue("foo_db")), new IdentifierValue("t_user"), new IdentifierValue("name"))); + BinaryOperationExpression nameExpression = new BinaryOperationExpression(10, 24, nameColumnSegment, new LiteralExpressionSegment(18, 22, "LiLei"), "=", "name = 'LiLei'"); + ColumnSegment pwdColumnSegment = new ColumnSegment(30, 32, new IdentifierValue("pwd")); + pwdColumnSegment.setColumnBoundInfo( + new ColumnSegmentBoundInfo(new TableSegmentBoundInfo(new IdentifierValue("foo_db"), new IdentifierValue("foo_db")), new IdentifierValue("t_user"), new IdentifierValue("pwd"))); + BinaryOperationExpression pwdExpression = new BinaryOperationExpression(30, 44, pwdColumnSegment, new LiteralExpressionSegment(40, 45, "123456"), "=", "pwd = '123456'"); return new WhereSegment(0, 0, new BinaryOperationExpression(0, 0, nameExpression, pwdExpression, "AND", "name = 'LiLei' AND pwd = '123456'")); } diff --git a/features/encrypt/core/src/test/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/predicate/EncryptPredicateColumnTokenGeneratorTest.java b/features/encrypt/core/src/test/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/predicate/EncryptPredicateColumnTokenGeneratorTest.java index 5d9dd2c6a185a..1599403278ce1 100644 --- a/features/encrypt/core/src/test/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/predicate/EncryptPredicateColumnTokenGeneratorTest.java +++ b/features/encrypt/core/src/test/java/org/apache/shardingsphere/encrypt/rewrite/token/generator/predicate/EncryptPredicateColumnTokenGeneratorTest.java @@ -24,7 +24,6 @@ import org.junit.jupiter.api.Test; import java.util.Collection; -import java.util.Collections; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; @@ -41,13 +40,11 @@ void setup() { @Test void assertIsGenerateSQLToken() { - generator.setSchemas(Collections.emptyMap()); assertTrue(generator.isGenerateSQLToken(EncryptGeneratorFixtureBuilder.createUpdateStatementContext())); } @Test void assertGenerateSQLTokenFromGenerateNewSQLToken() { - generator.setSchemas(Collections.emptyMap()); Collection substitutableColumnNameTokens = generator.generateSQLTokens(EncryptGeneratorFixtureBuilder.createUpdateStatementContext()); assertThat(substitutableColumnNameTokens.size(), is(1)); assertThat(((SubstitutableColumnNameToken) substitutableColumnNameTokens.iterator().next()).toString(null), is("pwd_assist")); diff --git a/test/e2e/sql/src/test/resources/cases/dql/e2e-dql-select-sub-query.xml b/test/e2e/sql/src/test/resources/cases/dql/e2e-dql-select-sub-query.xml index 09b2c507e1a98..0c89bae490a99 100644 --- a/test/e2e/sql/src/test/resources/cases/dql/e2e-dql-select-sub-query.xml +++ b/test/e2e/sql/src/test/resources/cases/dql/e2e-dql-select-sub-query.xml @@ -86,4 +86,12 @@ scenario-comments="Test join contains some encrypt columns in multi tables when use encrypt feature."> + + + + + + + +