Support encrypt rewrite for SQL Server OPENQUERY UPDATE with narrow SELECT ... FROM shape

RELEASE-NOTES.md

@@ -76,6 +76,7 @@
 1. Proxy: Support basic Firebird batch operations: create, send, execute, cancel, and release - [#38605](https://github.com/apache/shardingsphere/pull/38605)
 1. JDBC & Proxy: Add a check to verify database name naming conventions. - [#38883](https://github.com/apache/shardingsphere/pull/38883)
 1. Encrypt: Support SqlServer update statement for Updating data in a remote table by using a linked server when use encrypt feature - [#39122](https://github.com/apache/shardingsphere/pull/39122)
+1. Encrypt: Support encrypt rewrite for SQL Server OPENQUERY UPDATE with narrow SELECT ... FROM shape - [#39156](https://github.com/apache/shardingsphere/pull/39156)
 1. Encrypt: Support SqlServer update statement for Specifying a table alias as the target object when use encrypt feature - [#38733](https://github.com/apache/shardingsphere/pull/38733)
 1. Encrypt: Support SqlServer update statement for Specifying a view as the target object when use encrypt feature - [#38896](https://github.com/apache/shardingsphere/pull/38896)
 1. Encrypt: Support SqlServer for Using the UPDATE statement with information from another table  when use encrypt feature - [#38926](https://github.com/apache/shardingsphere/pull/38926)

docs/document/content/features/encrypt/limitations.cn.md

@@ -10,3 +10,26 @@ weight = 2
 - 加密字段无法支持计算操作，如：AVG、SUM 以及计算表达式；
 - 不支持使用 `;` 分隔的多条 SQL 同时执行；
 - 当投影子查询中包含加密字段时，必须使用别名。
+
+## SQL Server OPENQUERY 加密功能
+
+`OPENQUERY` 的加密改写仅支持如下窄形态透传查询：
+
+```sql
+UPDATE OPENQUERY (linked_server, 'SELECT <columns> FROM [<schema>.]<table> [WHERE ...]')
+SET <encrypt_column> = <literal_or_parameter>
+```
+
+不支持以下场景：
+
+- `SELECT` 列表中的字符串字面量、数字字面量、关键字表达式（例如 `NULL`）或表达式；
+- 括号标识符中包含空格，例如 `[Human Resources]`；
+- 三部分表名，例如 `db.schema.table`；
+- 逗号分隔的多表源；
+- `JOIN`、`CROSS APPLY`、`OUTER APPLY`；
+- `UNION`、`UNION ALL`、`EXCEPT`、`INTERSECT`；
+- 表引用后的 `ORDER BY`、`GROUP BY`、`HAVING` 等额外子句；
+- 使用 `;` 分隔的多条语句；
+- `WHERE` 后引用加密列；
+- 非字面量、非参数的赋值表达式，例如 `SET col = UPPER('x')`；
+- 物理列名包含 `]`。

docs/document/content/features/encrypt/limitations.en.md

@@ -10,3 +10,26 @@ weight = 2
 - Calculation operations are not supported for encrypted fields, such as `AVG`, `SUM`, and computation expressions.
 - Not support simultaneous execution of multiple SQL statements separated by `;`.
 - When projection subquery contains encrypt column, you must use alias.
+
+## SQL Server OPENQUERY encryption
+
+Encrypt rewrite for `OPENQUERY` only supports a narrow pass-through shape:
+
+```sql
+UPDATE OPENQUERY (linked_server, 'SELECT <columns> FROM [<schema>.]<table> [WHERE ...]')
+SET <encrypt_column> = <literal_or_parameter>
+```
+
+The following are not supported:
+
+- `SELECT` list items that are string literals, numeric literals, keyword expressions such as `NULL`, or other expressions.
+- Identifiers that contain spaces inside brackets, such as `[Human Resources]`.
+- Three-part table names, such as `db.schema.table`.
+- Comma-separated table sources.
+- `JOIN`, `CROSS APPLY`, `OUTER APPLY`.
+- `UNION`, `UNION ALL`, `EXCEPT`, `INTERSECT`.
+- Additional trailing clauses after the table reference, such as `ORDER BY`, `GROUP BY`, and `HAVING`.
+- Multiple statements separated by `;`.
+- Predicates after `WHERE` that reference encrypted columns.
+- Assignment expressions other than literals or parameter markers, such as `SET col = UPPER('x')`.
+- Physical column names that contain `]`.

features/encrypt/core/src/main/java/org/apache/shardingsphere/encrypt/rewrite/context/EncryptSQLRewriteContextDecorator.java

@@ -22,12 +22,14 @@
 import org.apache.shardingsphere.encrypt.rewrite.condition.EncryptConditionEngine;
 import org.apache.shardingsphere.encrypt.rewrite.parameter.EncryptParameterRewritersRegistry;
 import org.apache.shardingsphere.encrypt.rewrite.token.EncryptTokenGenerateBuilder;
+import org.apache.shardingsphere.encrypt.rewrite.token.generator.assignment.EncryptOpenQueryUtils;
 import org.apache.shardingsphere.encrypt.rule.EncryptRule;
 import org.apache.shardingsphere.infra.annotation.HighFrequencyInvocation;
 import org.apache.shardingsphere.infra.binder.context.available.WhereContextAvailable;
 import org.apache.shardingsphere.infra.binder.context.extractor.SQLStatementContextExtractor;
 import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
 import org.apache.shardingsphere.infra.binder.context.statement.type.dml.SelectStatementContext;
+import org.apache.shardingsphere.infra.binder.context.statement.type.dml.UpdateStatementContext;
 import org.apache.shardingsphere.infra.config.props.ConfigurationProperties;
 import org.apache.shardingsphere.infra.rewrite.context.SQLRewriteContext;
 import org.apache.shardingsphere.infra.rewrite.context.SQLRewriteContextDecorator;
@@ -37,6 +39,7 @@
 import org.apache.shardingsphere.infra.route.context.RouteContext;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.predicate.WhereSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.generic.table.SimpleTableSegment;
+import org.apache.shardingsphere.sql.parser.statement.core.statement.type.dml.UpdateStatement;
 
 import java.util.Collection;
 import java.util.Collections;
@@ -69,7 +72,16 @@ private boolean containsEncryptTable(final EncryptRule rule, final SQLStatementC
                 return true;
             }
         }
-        return false;
+        return containsOpenQueryEncryptTable(rule, sqlStatementContext);
+    }
+
+    private boolean containsOpenQueryEncryptTable(final EncryptRule rule, final SQLStatementContext sqlStatementContext) {
+        if (!(sqlStatementContext instanceof UpdateStatementContext)) {
+            return false;
+        }
+        UpdateStatement updateStatement = ((UpdateStatementContext) sqlStatementContext).getSqlStatement();
+        return EncryptOpenQueryUtils.isOpenQueryFunctionTable(updateStatement.getTable())
+                && EncryptOpenQueryUtils.findEncryptTable(rule, updateStatement.getTable()).isPresent();
     }
 
     private Collection<EncryptCondition> createEncryptConditions(final EncryptRule rule, final SQLStatementContext sqlStatementContext) {

features/encrypt/core/src/main/java/org/apache/shardingsphere/encrypt/rewrite/parameter/rewriter/EncryptAssignmentParameterRewriter.java

@@ -20,8 +20,10 @@
 import com.google.common.base.Preconditions;
 import lombok.RequiredArgsConstructor;
 import org.apache.shardingsphere.database.connector.core.type.DatabaseTypeRegistry;
+import org.apache.shardingsphere.encrypt.rewrite.token.generator.assignment.EncryptOpenQueryUtils;
 import org.apache.shardingsphere.encrypt.rule.EncryptRule;
 import org.apache.shardingsphere.encrypt.rule.column.EncryptColumn;
+import org.apache.shardingsphere.encrypt.rule.table.EncryptTable;
 import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
 import org.apache.shardingsphere.infra.binder.context.statement.type.dml.InsertStatementContext;
 import org.apache.shardingsphere.infra.binder.context.statement.type.dml.UpdateStatementContext;
@@ -33,6 +35,7 @@
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.assignment.SetAssignmentSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.expr.ExpressionSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.expr.simple.ParameterMarkerExpressionSegment;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.generic.table.TableSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.SQLStatement;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.type.dml.InsertStatement;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.type.dml.UpdateStatement;
@@ -66,23 +69,54 @@ public boolean isNeedRewrite(final SQLStatementContext sqlStatementContext) {
 
     @Override
     public void rewrite(final ParameterBuilder paramBuilder, final SQLStatementContext sqlStatementContext, final List<Object> params) {
-        String schemaName = sqlStatementContext.getTablesContext().getSchemaName()
-                .orElseGet(() -> new DatabaseTypeRegistry(sqlStatementContext.getSqlStatement().getDatabaseType()).getDefaultSchemaName(databaseName));
+        TableSegment openQueryTarget = findOpenQueryTarget(sqlStatementContext);
         for (ColumnAssignmentSegment each : getSetAssignmentSegment(sqlStatementContext.getSqlStatement()).getAssignments()) {
             String columnName = each.getColumns().get(0).getIdentifier().getValue();
-            String tableName = each.getColumns().get(0).getColumnBoundInfo().getOriginalTable().getValue();
-            if (!rule.findEncryptTable(tableName).map(optional -> optional.isEncryptColumn(columnName)).orElse(false)) {
+            String originalTableName = each.getColumns().get(0).getColumnBoundInfo().getOriginalTable().getValue();
+            EncryptTable encryptTable = resolveEncryptTable(originalTableName, columnName, openQueryTarget);
+            if (null == encryptTable) {
                 continue;
             }
-            EncryptColumn encryptColumn = rule.getEncryptTable(tableName).getEncryptColumn(columnName);
-            StandardParameterBuilder standardParamBuilder = paramBuilder instanceof StandardParameterBuilder
-                    ? (StandardParameterBuilder) paramBuilder
-                    : ((GroupedParameterBuilder) paramBuilder).getParameterBuilders().get(0);
-            ExpressionSegment valueExpression = each.getValue();
-            if (valueExpression instanceof ParameterMarkerExpressionSegment) {
-                encryptParameters(standardParamBuilder, schemaName, tableName, encryptColumn, ((ParameterMarkerExpressionSegment) valueExpression).getParameterMarkerIndex(), params);
-            }
+            rewriteEncryptAssignment(paramBuilder, sqlStatementContext, openQueryTarget, encryptTable, columnName, each, params);
+        }
+    }
+
+    private void rewriteEncryptAssignment(final ParameterBuilder paramBuilder, final SQLStatementContext sqlStatementContext, final TableSegment openQueryTarget,
+                                          final EncryptTable encryptTable, final String columnName, final ColumnAssignmentSegment assignmentSegment, final List<Object> params) {
+        String schemaName = resolveSchemaName(sqlStatementContext, openQueryTarget);
+        EncryptColumn encryptColumn = encryptTable.getEncryptColumn(columnName);
+        StandardParameterBuilder standardParamBuilder = paramBuilder instanceof StandardParameterBuilder
+                ? (StandardParameterBuilder) paramBuilder
+                : ((GroupedParameterBuilder) paramBuilder).getParameterBuilders().get(0);
+        ExpressionSegment valueExpression = assignmentSegment.getValue();
+        if (valueExpression instanceof ParameterMarkerExpressionSegment) {
+            encryptParameters(standardParamBuilder, schemaName, encryptTable.getTable(), encryptColumn, ((ParameterMarkerExpressionSegment) valueExpression).getParameterMarkerIndex(), params);
+        }
+    }
+
+    private String resolveSchemaName(final SQLStatementContext sqlStatementContext, final TableSegment openQueryTarget) {
+        String defaultSchemaName = sqlStatementContext.getTablesContext().getSchemaName()
+                .orElseGet(() -> new DatabaseTypeRegistry(sqlStatementContext.getSqlStatement().getDatabaseType()).getDefaultSchemaName(databaseName));
+        return null == openQueryTarget ? defaultSchemaName : EncryptOpenQueryUtils.findSchemaName(openQueryTarget).orElse(defaultSchemaName);
+    }
+
+    private TableSegment findOpenQueryTarget(final SQLStatementContext sqlStatementContext) {
+        if (!(sqlStatementContext instanceof UpdateStatementContext)) {
+            return null;
+        }
+        TableSegment table = ((UpdateStatementContext) sqlStatementContext).getSqlStatement().getTable();
+        return EncryptOpenQueryUtils.isOpenQueryFunctionTable(table) ? table : null;
+    }
+
+    private EncryptTable resolveEncryptTable(final String originalTableName, final String columnName, final TableSegment openQueryTarget) {
+        Optional<EncryptTable> fromOriginal = rule.findEncryptTable(originalTableName).filter(t -> t.isEncryptColumn(columnName));
+        if (fromOriginal.isPresent()) {
+            return fromOriginal.get();
+        }
+        if (null == openQueryTarget) {
+            return null;
         }
+        return EncryptOpenQueryUtils.findEncryptTable(rule, openQueryTarget).filter(t -> t.isEncryptColumn(columnName)).orElse(null);
     }
 
     private SetAssignmentSegment getSetAssignmentSegment(final SQLStatement sqlStatement) {