Support encrypt insert select sql rewrite (#28425)

apache · Sep 13, 2023 · 05d5a9b · 05d5a9b
1 parent 6c04da1
commit 05d5a9b
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 11 deletions.
diff --git a/...pache/shardingsphere/encrypt/rewrite/token/generator/EncryptProjectionTokenGenerator.java b/...pache/shardingsphere/encrypt/rewrite/token/generator/EncryptProjectionTokenGenerator.java
@@ -17,7 +17,6 @@
 
 package org.apache.shardingsphere.encrypt.rewrite.token.generator;
 
-import com.google.common.base.Preconditions;
 import lombok.Setter;
 import org.apache.shardingsphere.encrypt.rewrite.aware.DatabaseTypeAware;
 import org.apache.shardingsphere.encrypt.rewrite.aware.EncryptRuleAware;
@@ -30,6 +29,7 @@
 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.DialectDatabaseMetaData;
 import org.apache.shardingsphere.infra.database.core.metadata.database.enums.QuoteCharacter;
@@ -69,21 +69,28 @@ public final class EncryptProjectionTokenGenerator implements CollectionSQLToken
 
     @Override
     public boolean isGenerateSQLToken(final SQLStatementContext sqlStatementContext) {
-        return sqlStatementContext instanceof SelectStatementContext && !((SelectStatementContext) sqlStatementContext).getAllTables().isEmpty();
+        return sqlStatementContext instanceof SelectStatementContext && !((SelectStatementContext) sqlStatementContext).getAllTables().isEmpty()
+                || sqlStatementContext instanceof InsertStatementContext && null != ((InsertStatementContext) sqlStatementContext).getInsertSelectContext();
     }
 
     @Override
     public Collection<SQLToken> generateSQLTokens(final SQLStatementContext sqlStatementContext) {
-        Preconditions.checkState(sqlStatementContext instanceof SelectStatementContext);
         Collection<SQLToken> result = new LinkedHashSet<>();
-        SelectStatementContext selectStatementContext = (SelectStatementContext) sqlStatementContext;
-        addGenerateSQLTokens(result, selectStatementContext);
-        for (SelectStatementContext each : selectStatementContext.getSubqueryContexts().values()) {
-            addGenerateSQLTokens(result, each);
+        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);
+        for (SelectStatementContext each : selectStatementContext.getSubqueryContexts().values()) {
+            addGenerateSQLTokens(sqlTokens, each);
+        }
+    }
+
     private void addGenerateSQLTokens(final Collection<SQLToken> sqlTokens, final SelectStatementContext selectStatementContext) {
         for (ProjectionSegment each : selectStatementContext.getSqlStatement().getProjections().getProjections()) {
             SubqueryType subqueryType = selectStatementContext.getSubqueryType();
@@ -153,8 +160,11 @@ private Collection<Projection> generateProjections(final EncryptColumn encryptCo
             return generateProjectionsInTableSegmentSubquery(encryptColumn, columnProjection, shorthandProjection, subqueryType);
         } else if (SubqueryType.PREDICATE_SUBQUERY == subqueryType) {
             return Collections.singleton(generateProjectionInPredicateSubquery(encryptColumn, columnProjection, shorthandProjection));
+        } else if (SubqueryType.INSERT_SELECT_SUBQUERY == subqueryType) {
+            return generateProjectionsInInsertSelectSubquery(encryptColumn, columnProjection, shorthandProjection);
         }
-        throw new UnsupportedSQLOperationException("Projections not in simple select, table subquery, join subquery and predicate subquery are not supported in encrypt feature.");
+        throw new UnsupportedSQLOperationException(
+                "Projections not in simple select, table subquery, join subquery, predicate subquery and insert select subquery are not supported in encrypt feature.");
     }
 
     private ColumnProjection generateProjection(final EncryptColumn encryptColumn, final ColumnProjection columnProjection, final boolean shorthandProjection) {
@@ -186,6 +196,18 @@ private ColumnProjection generateProjectionInPredicateSubquery(final EncryptColu
                         databaseType));
     }
 
+    private Collection<Projection> generateProjectionsInInsertSelectSubquery(final EncryptColumn encryptColumn, final ColumnProjection columnProjection, final boolean shorthandProjection) {
+        QuoteCharacter quoteCharacter = columnProjection.getName().getQuoteCharacter();
+        IdentifierValue columnName = new IdentifierValue(encryptColumn.getCipher().getName(), quoteCharacter);
+        Collection<Projection> result = new LinkedList<>();
+        IdentifierValue encryptColumnOwner = shorthandProjection ? columnProjection.getOwner().orElse(null) : null;
+        result.add(new ColumnProjection(encryptColumnOwner, columnName, null, databaseType));
+        IdentifierValue assistedColumOwner = columnProjection.getOwner().orElse(null);
+        encryptColumn.getAssistedQuery().ifPresent(optional -> result.add(new ColumnProjection(assistedColumOwner, new IdentifierValue(optional.getName(), quoteCharacter), null, databaseType)));
+        encryptColumn.getLikeQuery().ifPresent(optional -> result.add(new ColumnProjection(assistedColumOwner, new IdentifierValue(optional.getName(), quoteCharacter), null, databaseType)));
+        return result;
+    }
+
     private ShorthandProjection getShorthandProjection(final ShorthandProjectionSegment segment, final ProjectionsContext projectionsContext) {
         Optional<String> owner = segment.getOwner().isPresent() ? Optional.of(segment.getOwner().get().getIdentifier().getValue()) : Optional.empty();
         for (Projection each : projectionsContext.getProjections()) {

diff --git a/.../org/apache/shardingsphere/infra/binder/context/statement/dml/InsertStatementContext.java b/.../org/apache/shardingsphere/infra/binder/context/statement/dml/InsertStatementContext.java
@@ -18,8 +18,6 @@
 package org.apache.shardingsphere.infra.binder.context.statement.dml;
 
 import lombok.Getter;
-import org.apache.shardingsphere.infra.exception.dialect.exception.syntax.database.NoDatabaseSelectedException;
-import org.apache.shardingsphere.infra.exception.dialect.exception.syntax.database.UnknownDatabaseException;
 import org.apache.shardingsphere.infra.binder.context.aware.ParameterAware;
 import org.apache.shardingsphere.infra.binder.context.segment.insert.keygen.GeneratedKeyContext;
 import org.apache.shardingsphere.infra.binder.context.segment.insert.keygen.engine.GeneratedKeyContextEngine;
@@ -30,10 +28,13 @@
 import org.apache.shardingsphere.infra.binder.context.statement.CommonSQLStatementContext;
 import org.apache.shardingsphere.infra.binder.context.type.TableAvailable;
 import org.apache.shardingsphere.infra.database.core.type.DatabaseTypeRegistry;
+import org.apache.shardingsphere.infra.exception.core.ShardingSpherePreconditions;
+import org.apache.shardingsphere.infra.exception.dialect.exception.syntax.database.NoDatabaseSelectedException;
+import org.apache.shardingsphere.infra.exception.dialect.exception.syntax.database.UnknownDatabaseException;
 import org.apache.shardingsphere.infra.metadata.ShardingSphereMetaData;
 import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
 import org.apache.shardingsphere.infra.metadata.database.schema.model.ShardingSphereSchema;
-import org.apache.shardingsphere.infra.exception.core.ShardingSpherePreconditions;
+import org.apache.shardingsphere.sql.parser.sql.common.enums.SubqueryType;
 import org.apache.shardingsphere.sql.parser.sql.common.extractor.TableExtractor;
 import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.assignment.AssignmentSegment;
 import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.assignment.InsertValuesSegment;
@@ -129,6 +130,7 @@ private Optional<InsertSelectContext> getInsertSelectContext(final ShardingSpher
         }
         SubquerySegment insertSelectSegment = getSqlStatement().getInsertSelect().get();
         SelectStatementContext selectStatementContext = new SelectStatementContext(metaData, params, insertSelectSegment.getSelect(), defaultDatabaseName);
+        selectStatementContext.setSubqueryType(SubqueryType.INSERT_SELECT_SUBQUERY);
         InsertSelectContext insertSelectContext = new InsertSelectContext(selectStatementContext, params, paramsOffset.get());
         paramsOffset.addAndGet(insertSelectContext.getParameterCount());
         return Optional.of(insertSelectContext);