This is an automated email from the ASF dual-hosted git repository.
yiguolei pushed a commit to branch branch-4.1
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/branch-4.1 by this push:
new c656fb9e30c branch-4.1: [Opt](load) Support using `AES_ENCRYPT` with
internal custom ENCRYPTKEY in stream load and routine load #60503 (#65243)
c656fb9e30c is described below
commit c656fb9e30c62824563448742d9abdf9b01d13a1
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Tue Jul 7 13:31:33 2026 +0800
branch-4.1: [Opt](load) Support using `AES_ENCRYPT` with internal custom
ENCRYPTKEY in stream load and routine load #60503 (#65243)
Cherry-picked from #60503
Co-authored-by: bobhan1 <[email protected]>
---
.../doris/load/routineload/RoutineLoadJob.java | 2 +
.../doris/nereids/load/NereidsLoadUtils.java | 32 ++++
.../test_routine_load_with_aes_encrypt.out | 22 +++
.../stream_load/test_stream_load_aes_encrypt.csv | 3 +
.../test_stream_load_with_aes_encrypt.out | 17 ++
.../test_routine_load_with_aes_encrypt.groovy | 175 +++++++++++++++++++++
.../test_stream_load_with_aes_encrypt.groovy | 116 ++++++++++++++
7 files changed, 367 insertions(+)
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/load/routineload/RoutineLoadJob.java
b/fe/fe-core/src/main/java/org/apache/doris/load/routineload/RoutineLoadJob.java
index 37c9a6a1c52..a06a0e1044a 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/load/routineload/RoutineLoadJob.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/load/routineload/RoutineLoadJob.java
@@ -1059,6 +1059,8 @@ public abstract class RoutineLoadJob
ConnectContext.get().setCurrentUserIdentity(this.getUserIdentity());
} else {
setComputeGroup();
+ // Set user identity for privilege check in expression rewrite
(e.g., EncryptKeyRef)
+
ConnectContext.get().setCurrentUserIdentity(this.getUserIdentity());
}
if (ConnectContext.get().getEnv() == null) {
ConnectContext.get().setEnv(Env.getCurrentEnv());
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/nereids/load/NereidsLoadUtils.java
b/fe/fe-core/src/main/java/org/apache/doris/nereids/load/NereidsLoadUtils.java
index 96a728ac1f8..d088bdce6bc 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/nereids/load/NereidsLoadUtils.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/nereids/load/NereidsLoadUtils.java
@@ -37,6 +37,8 @@ import org.apache.doris.nereids.rules.RuleType;
import org.apache.doris.nereids.rules.analysis.BindExpression;
import org.apache.doris.nereids.rules.analysis.BindSink;
import org.apache.doris.nereids.rules.expression.ExpressionNormalization;
+import org.apache.doris.nereids.rules.expression.ExpressionRewrite;
+import org.apache.doris.nereids.rules.expression.rules.FoldConstantRuleOnFE;
import org.apache.doris.nereids.rules.rewrite.MergeProjects;
import org.apache.doris.nereids.rules.rewrite.OneRewriteRuleFactory;
import org.apache.doris.nereids.trees.expressions.Alias;
@@ -45,6 +47,7 @@ import
org.apache.doris.nereids.trees.expressions.NamedExpression;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator;
+import
org.apache.doris.nereids.trees.expressions.functions.scalar.EncryptKeyRef;
import
org.apache.doris.nereids.trees.expressions.functions.scalar.JsonbParseErrorToNull;
import
org.apache.doris.nereids.trees.expressions.functions.scalar.JsonbParseErrorToValue;
import org.apache.doris.nereids.trees.expressions.literal.Literal;
@@ -236,6 +239,11 @@ public class NereidsLoadUtils {
// the NereidsLoadPlanInfoCollector will not generate
slot by id#0,
// so we must use MergeProjects here
new MergeProjects(),
+ // RewriteEncryptKeyRef must be placed before
ExpressionNormalization,
+ // because setDebugSkipFoldConstant(true) will skip
FoldConstantRule which
+ // is responsible for folding EncryptKeyRef to
StringLiteral.
+ // We need to handle EncryptKeyRef separately to support
KEY syntax in stream load.
+ new RewriteEncryptKeyRef(),
new ExpressionNormalization())
)).execute();
Rewriter.getWholeTreeRewriterWithCustomJobs(cascadesContext,
ImmutableList.of()).execute();
@@ -371,4 +379,28 @@ public class NereidsLoadUtils {
}).toRule(RuleType.ADD_POST_PROJECT_FOR_LOAD);
}
}
+
+ /**
+ * RewriteEncryptKeyRef
+ * This rule rewrites EncryptKeyRef to StringLiteral in stream load.
+ * Since setDebugSkipFoldConstant(true) is set during stream load planning,
+ * FoldConstantRule will be skipped and EncryptKeyRef won't be folded.
+ * This rule handles EncryptKeyRef separately to support KEY syntax in
stream load columns parameter.
+ */
+ private static class RewriteEncryptKeyRef extends ExpressionRewrite {
+ private static final FoldConstantRuleOnFE FOLD_ENCRYPT_KEY_REF =
FoldConstantRuleOnFE.VISITOR_INSTANCE;
+
+ public RewriteEncryptKeyRef() {
+ super(((expression, context) -> {
+ // Use rewriteUp to traverse the expression tree bottom-up and
only fold
+ // EncryptKeyRef nodes to StringLiteral, leaving all other
expressions unchanged.
+ return expression.rewriteUp(e -> {
+ if (e instanceof EncryptKeyRef) {
+ return e.accept(FOLD_ENCRYPT_KEY_REF, context);
+ }
+ return e;
+ });
+ }));
+ }
+ }
}
diff --git
a/regression-test/data/load_p0/routine_load/test_routine_load_with_aes_encrypt.out
b/regression-test/data/load_p0/routine_load/test_routine_load_with_aes_encrypt.out
new file mode 100644
index 00000000000..88041b9e162
--- /dev/null
+++
b/regression-test/data/load_p0/routine_load/test_routine_load_with_aes_encrypt.out
@@ -0,0 +1,22 @@
+-- This file is automatically generated. You should know what you did if you
want to edit this
+-- !sql_count_direct --
+3
+
+-- !sql_decrypt_direct --
+1 Alice hello world
+2 Bob doris is great
+3 Charlie test encryption
+
+-- !sql_count_encryptkey --
+3
+
+-- !sql_decrypt_with_encryptkey --
+1 Alice hello world
+2 Bob doris is great
+3 Charlie test encryption
+
+-- !sql_decrypt_with_direct_key --
+1 Alice hello world
+2 Bob doris is great
+3 Charlie test encryption
+
diff --git
a/regression-test/data/load_p0/stream_load/test_stream_load_aes_encrypt.csv
b/regression-test/data/load_p0/stream_load/test_stream_load_aes_encrypt.csv
new file mode 100644
index 00000000000..75288201e1a
--- /dev/null
+++ b/regression-test/data/load_p0/stream_load/test_stream_load_aes_encrypt.csv
@@ -0,0 +1,3 @@
+1,Alice,hello world
+2,Bob,doris is great
+3,Charlie,test encryption
diff --git
a/regression-test/data/load_p0/stream_load/test_stream_load_with_aes_encrypt.out
b/regression-test/data/load_p0/stream_load/test_stream_load_with_aes_encrypt.out
new file mode 100644
index 00000000000..3be01d99d26
--- /dev/null
+++
b/regression-test/data/load_p0/stream_load/test_stream_load_with_aes_encrypt.out
@@ -0,0 +1,17 @@
+-- This file is automatically generated. You should know what you did if you
want to edit this
+-- !sql_count --
+3
+
+-- !sql_decrypt_direct --
+1 Alice hello world
+2 Bob doris is great
+3 Charlie test encryption
+
+-- !sql_count_key --
+3
+
+-- !sql_decrypt_with_key --
+1 Alice hello world
+2 Bob doris is great
+3 Charlie test encryption
+
diff --git
a/regression-test/suites/load_p0/routine_load/test_routine_load_with_aes_encrypt.groovy
b/regression-test/suites/load_p0/routine_load/test_routine_load_with_aes_encrypt.groovy
new file mode 100644
index 00000000000..2d87a9b3626
--- /dev/null
+++
b/regression-test/suites/load_p0/routine_load/test_routine_load_with_aes_encrypt.groovy
@@ -0,0 +1,175 @@
+// 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.
+
+import org.apache.doris.regression.util.RoutineLoadTestUtils
+import org.apache.kafka.clients.producer.ProducerRecord
+
+suite("test_routine_load_with_aes_encrypt", "p0") {
+ def tableName1 = "test_routine_load_aes_encrypt_direct_key"
+ def tableName2 = "test_routine_load_aes_encrypt_with_encryptkey"
+ def kafkaTopic1 = "test_routine_load_aes_encrypt_direct"
+ def kafkaTopic2 = "test_routine_load_aes_encrypt_key"
+ def jobName1 = "routine_load_aes_direct_key_job"
+ def jobName2 = "routine_load_aes_encryptkey_job"
+ def dbName = context.dbName
+ def encryptKeyName = "routine_load_aes_key"
+ def aesKey = "F3229A0B371ED2D9441B830D21A390C3"
+
+ if (RoutineLoadTestUtils.isKafkaTestEnabled(context)) {
+ def runSql = { String q -> sql q }
+ def kafka_broker = RoutineLoadTestUtils.getKafkaBroker(context)
+ def producer = RoutineLoadTestUtils.createKafkaProducer(kafka_broker)
+
+ // Test data
+ def testData = [
+ "1,Alice,hello world",
+ "2,Bob,doris is great",
+ "3,Charlie,test encryption"
+ ]
+
+ // Send test data to both Kafka topics
+ [kafkaTopic1, kafkaTopic2].each { topic ->
+ testData.each { line ->
+ logger.info("Sending to Kafka topic ${topic}: ${line}")
+ def record = new ProducerRecord<>(topic, null, line)
+ producer.send(record).get()
+ }
+ }
+ producer.flush()
+
+ // ============================================================
+ // Test 1: AES_ENCRYPT with direct key string
+ // ============================================================
+ log.info("Test 1: Routine load with AES_ENCRYPT using direct key
string")
+
+ sql """ DROP TABLE IF EXISTS ${tableName1} """
+ sql """
+ CREATE TABLE IF NOT EXISTS ${tableName1} (
+ `id` int(11) NULL,
+ `name` string NULL,
+ `encrypted_data` string NULL
+ ) ENGINE=OLAP
+ DUPLICATE KEY(`id`)
+ DISTRIBUTED BY HASH(`id`) BUCKETS 1
+ PROPERTIES ("replication_allocation" = "tag.location.default: 1");
+ """
+
+ try {
+ sql """
+ CREATE ROUTINE LOAD ${jobName1} ON ${tableName1}
+ COLUMNS TERMINATED BY ",",
+ COLUMNS(id, name, tmp_data,
encrypted_data=TO_BASE64(AES_ENCRYPT(tmp_data, '${aesKey}')))
+ PROPERTIES
+ (
+ "max_batch_interval" = "5",
+ "max_batch_rows" = "300000",
+ "max_batch_size" = "209715200"
+ )
+ FROM KAFKA
+ (
+ "kafka_broker_list" = "${kafka_broker}",
+ "kafka_topic" = "${kafkaTopic1}",
+ "property.kafka_default_offsets" = "OFFSET_BEGINNING"
+ );
+ """
+
+ RoutineLoadTestUtils.waitForTaskFinish(runSql, jobName1,
tableName1, 2)
+
+ // Verify data count
+ qt_sql_count_direct """ SELECT count(*) FROM ${tableName1} """
+
+ // Verify encrypted data can be decrypted
+ qt_sql_decrypt_direct """
+ SELECT id, name, AES_DECRYPT(FROM_BASE64(encrypted_data),
'${aesKey}') as decrypted_data
+ FROM ${tableName1}
+ ORDER BY id
+ """
+ } finally {
+ sql "STOP ROUTINE LOAD FOR ${jobName1}"
+ sql """ DROP TABLE IF EXISTS ${tableName1} """
+ }
+
+ // ============================================================
+ // Test 2: AES_ENCRYPT with ENCRYPTKEY (KEY syntax)
+ // ============================================================
+ log.info("Test 2: Routine load with AES_ENCRYPT using ENCRYPTKEY (KEY
syntax)")
+
+ sql """ DROP TABLE IF EXISTS ${tableName2} """
+ sql """
+ CREATE TABLE IF NOT EXISTS ${tableName2} (
+ `id` int(11) NULL,
+ `name` string NULL,
+ `encrypted_data` string NULL
+ ) ENGINE=OLAP
+ DUPLICATE KEY(`id`)
+ DISTRIBUTED BY HASH(`id`) BUCKETS 1
+ PROPERTIES ("replication_allocation" = "tag.location.default: 1");
+ """
+
+ // Create encryptkey
+ try_sql """ DROP ENCRYPTKEY IF EXISTS ${encryptKeyName} """
+ sql """ CREATE ENCRYPTKEY ${encryptKeyName} AS "${aesKey}" """
+
+ // Verify encryptkey was created
+ def keyRes = sql """ SHOW ENCRYPTKEYS FROM ${dbName} """
+ log.info("Encryptkeys: ${keyRes}")
+ assertTrue(keyRes.size() >= 1, "Encryptkey should be created")
+
+ try {
+ sql """
+ CREATE ROUTINE LOAD ${jobName2} ON ${tableName2}
+ COLUMNS TERMINATED BY ",",
+ COLUMNS(id, name, tmp_data,
encrypted_data=TO_BASE64(AES_ENCRYPT(tmp_data, KEY
${dbName}.${encryptKeyName})))
+ PROPERTIES
+ (
+ "max_batch_interval" = "5",
+ "max_batch_rows" = "300000",
+ "max_batch_size" = "209715200"
+ )
+ FROM KAFKA
+ (
+ "kafka_broker_list" = "${kafka_broker}",
+ "kafka_topic" = "${kafkaTopic2}",
+ "property.kafka_default_offsets" = "OFFSET_BEGINNING"
+ );
+ """
+
+ RoutineLoadTestUtils.waitForTaskFinish(runSql, jobName2,
tableName2, 2)
+
+ // Verify data count
+ qt_sql_count_encryptkey """ SELECT count(*) FROM ${tableName2} """
+
+ // Verify encrypted data can be decrypted using ENCRYPTKEY
+ qt_sql_decrypt_with_encryptkey """
+ SELECT id, name, AES_DECRYPT(FROM_BASE64(encrypted_data), KEY
${dbName}.${encryptKeyName}) as decrypted_data
+ FROM ${tableName2}
+ ORDER BY id
+ """
+
+ // Verify encrypted data can also be decrypted using direct key
string (same key value)
+ qt_sql_decrypt_with_direct_key """
+ SELECT id, name, AES_DECRYPT(FROM_BASE64(encrypted_data),
'${aesKey}') as decrypted_data
+ FROM ${tableName2}
+ ORDER BY id
+ """
+ } finally {
+ sql "STOP ROUTINE LOAD FOR ${jobName2}"
+ sql """ DROP TABLE IF EXISTS ${tableName2} """
+ try_sql """ DROP ENCRYPTKEY IF EXISTS ${encryptKeyName} """
+ }
+ }
+}
diff --git
a/regression-test/suites/load_p0/stream_load/test_stream_load_with_aes_encrypt.groovy
b/regression-test/suites/load_p0/stream_load/test_stream_load_with_aes_encrypt.groovy
new file mode 100644
index 00000000000..1ef961529a0
--- /dev/null
+++
b/regression-test/suites/load_p0/stream_load/test_stream_load_with_aes_encrypt.groovy
@@ -0,0 +1,116 @@
+// 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.
+
+suite("test_stream_load_with_aes_encrypt", "p0") {
+ def tableName = "test_stream_load_with_aes_encrypt"
+ def dbName = context.dbName
+ def encryptKeyName = "my_stream_load_key"
+ def aesKey = "F3229A0B371ED2D9441B830D21A390C3"
+
+ sql """ DROP TABLE IF EXISTS ${tableName} """
+ sql """ CREATE TABLE ${tableName} (
+ id int,
+ name string,
+ encrypted_data string
+ ) ENGINE=OLAP
+ DUPLICATE KEY (`id`)
+ DISTRIBUTED BY HASH(`id`) BUCKETS 1
+ PROPERTIES (
+ "replication_num" = "1"
+ );
+ """
+
+ // Test 1: Stream load with AES_ENCRYPT using direct key string
+ log.info("Test 1: Stream load with AES_ENCRYPT using direct key string")
+ streamLoad {
+ table "${tableName}"
+
+ set 'column_separator', ','
+ set 'columns', """ id, name, tmp_data,
encrypted_data=TO_BASE64(AES_ENCRYPT(tmp_data, '${aesKey}')) """
+ file 'test_stream_load_aes_encrypt.csv'
+ check { result, exception, startTime, endTime ->
+ if (exception != null) {
+ throw exception
+ }
+ log.info("Stream load result: ${result}".toString())
+ def json = parseJson(result)
+ assertEquals("success", json.Status.toLowerCase())
+ assertEquals(3, json.NumberLoadedRows)
+ }
+ }
+
+ sql """sync"""
+
+ // Verify data was loaded
+ qt_sql_count """ SELECT count(*) FROM ${tableName} """
+
+ // Verify encrypted data can be decrypted back to original using direct key
+ qt_sql_decrypt_direct """
+ SELECT id, name, AES_DECRYPT(FROM_BASE64(encrypted_data), '${aesKey}')
as decrypted_data
+ FROM ${tableName}
+ ORDER BY id
+ """
+
+ // Clean table for next test
+ sql """ TRUNCATE TABLE ${tableName} """
+
+ // Test 2: Stream load with AES_ENCRYPT using ENCRYPTKEY (KEY syntax)
+ log.info("Test 2: Stream load with AES_ENCRYPT using ENCRYPTKEY (KEY
syntax)")
+
+ // Create encryptkey with the same key value
+ try_sql """ DROP ENCRYPTKEY IF EXISTS ${encryptKeyName} """
+ sql """ CREATE ENCRYPTKEY ${encryptKeyName} AS "${aesKey}" """
+
+ // Verify encryptkey was created
+ def keyRes = sql """ SHOW ENCRYPTKEYS FROM ${dbName} """
+ log.info("Encryptkeys: ${keyRes}")
+ assertTrue(keyRes.size() >= 1, "Encryptkey should be created")
+
+ // Stream load using KEY syntax to reference encryptkey
+ streamLoad {
+ table "${tableName}"
+
+ set 'column_separator', ','
+ set 'columns', """ id, name, tmp_data,
encrypted_data=TO_BASE64(AES_ENCRYPT(tmp_data, KEY
${dbName}.${encryptKeyName})) """
+ file 'test_stream_load_aes_encrypt.csv'
+ check { result, exception, startTime, endTime ->
+ if (exception != null) {
+ throw exception
+ }
+ log.info("Stream load result with ENCRYPTKEY:
${result}".toString())
+ def json = parseJson(result)
+ assertEquals("success", json.Status.toLowerCase())
+ assertEquals(3, json.NumberLoadedRows)
+ }
+ }
+
+ sql """sync"""
+
+ // Verify data count
+ qt_sql_count_key """ SELECT count(*) FROM ${tableName} """
+
+ // Verify encrypted data can be decrypted using ENCRYPTKEY
+ qt_sql_decrypt_with_key """
+ SELECT id, name, AES_DECRYPT(FROM_BASE64(encrypted_data), KEY
${dbName}.${encryptKeyName}) as decrypted_data
+ FROM ${tableName}
+ ORDER BY id
+ """
+
+ // Cleanup
+ sql """ DROP TABLE IF EXISTS ${tableName} """
+ try_sql """ DROP ENCRYPTKEY IF EXISTS ${encryptKeyName} """
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
