From d017fa05cec9a5495ffa52764177c80112d88725 Mon Sep 17 00:00:00 2001 From: deardeng Date: Mon, 3 Aug 2026 14:22:48 +0800 Subject: [PATCH] [fix](fe) Preserve effective Compute Group for audit logs (#66281) Problem Summary: In Cloud mode, a per-query SET_VAR(cloud_cluster=...) temporarily changes the session variable used by query scheduling. StmtExecutor reverted this value before AuditLogHelper built the audit event, so ComputeGroupName and related per-Compute-Group query, error, and latency metrics were attributed to the session's original Compute Group even though the query ran on the hinted group. Preserve the query-scoped effective Compute Group before reverting SET_VAR and use it for both audit events and metrics, with fallback to the existing session resolution path. --- .../org/apache/doris/qe/AuditLogHelper.java | 12 +++- .../org/apache/doris/qe/ConnectContext.java | 11 ++++ .../org/apache/doris/qe/StmtExecutor.java | 6 ++ .../apache/doris/qe/AuditLogHelperTest.java | 9 +++ ...audit_log_hint_compute_group_docker.groovy | 64 +++++++++++++++++++ 5 files changed, 99 insertions(+), 3 deletions(-) create mode 100644 regression-test/suites/audit/test_audit_log_hint_compute_group_docker.groovy diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/AuditLogHelper.java b/fe/fe-core/src/main/java/org/apache/doris/qe/AuditLogHelper.java index e07012fd75ed55..d6c992412c2ec5 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/AuditLogHelper.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/AuditLogHelper.java @@ -189,7 +189,7 @@ private static void logAuditLogImpl(ConnectContext ctx, String origStmt, Stateme String cloudCluster = ""; try { if (Config.isCloudMode()) { - cloudCluster = ctx.getCloudCluster(false); + cloudCluster = getCloudClusterForAudit(ctx); } } catch (ComputeGroupException e) { LOG.warn("Failed to get cloud cluster", e); @@ -372,6 +372,13 @@ private static long getQueueTimeMs(ConnectContext ctx) { return queueToken == null ? -1 : queueToken.getQueueEndTime() - queueToken.getQueueStartTime(); } + static String getCloudClusterForAudit(ConnectContext ctx) throws ComputeGroupException { + if (!Strings.isNullOrEmpty(ctx.getEffectiveCloudCluster())) { + return ctx.getEffectiveCloudCluster(); + } + return ctx.getCloudCluster(false); + } + /** * Update query metrics without writing audit log. This is used when * enable_prepared_stmt_audit_log is disabled, to ensure QPS metrics @@ -407,7 +414,7 @@ private static void updateMetricsImpl(ConnectContext ctx) { String physicalClusterName = ""; try { if (Config.isCloudMode()) { - cloudCluster = ctx.getCloudCluster(false); + cloudCluster = getCloudClusterForAudit(ctx); physicalClusterName = ((CloudSystemInfoService) Env.getCurrentSystemInfo()) .getPhysicalCluster(cloudCluster); if (!cloudCluster.equals(physicalClusterName)) { @@ -471,4 +478,3 @@ private static String getStmtType(StatementBase stmt) { } } } - diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java index ce119805151714..4efec5e7242197 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java @@ -220,6 +220,9 @@ public enum ConnectType { // cloud cluster name protected volatile String cloudCluster = null; + // The compute group selected for the statement currently being executed. Unlike cloudCluster, + // this value is query-scoped and remains available after a per-query SET_VAR is reverted. + protected volatile String effectiveCloudCluster = null; // If set to true, the nondeterministic function will not be rewrote to constant. private boolean notEvalNondeterministicFunction = false; @@ -1449,6 +1452,14 @@ public void setCloudCluster(String cluster) { this.getSessionVariable().setCloudCluster(cluster); } + public String getEffectiveCloudCluster() { + return effectiveCloudCluster; + } + + public void setEffectiveCloudCluster(String cluster) { + this.effectiveCloudCluster = cluster; + } + public String getCloudCluster() throws ComputeGroupException { return getCloudCluster(true); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java index 4c75df883a8279..4a36237e212ce1 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java @@ -588,6 +588,7 @@ boolean shouldDisableCloudVersionCacheOnRetry(String errorMessage) { public void execute(TUniqueId queryId) throws Exception { SessionVariable sessionVariable = context.getSessionVariable(); + context.setEffectiveCloudCluster(null); if (context.getConnectType() == ConnectType.ARROW_FLIGHT_SQL) { context.setReturnResultFromLocal(true); } @@ -614,6 +615,11 @@ public void execute(TUniqueId queryId) throws Exception { throw e; } } finally { + // Preserve the effective per-query compute group before SET_VAR values are reverted. + // Audit logging runs after this method returns and otherwise sees the session value. + if (Config.isCloudMode()) { + context.setEffectiveCloudCluster(sessionVariable.getCloudCluster()); + } // Snapshot changed session variables (including SET_VAR hint values) BEFORE revert, // so the audit log (logged after execute() returns, i.e. after the revert below) can // reflect what was actually in effect for this statement instead of the reverted values. diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/AuditLogHelperTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/AuditLogHelperTest.java index e1b6b8c961cb9e..82afc4a02a1727 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/AuditLogHelperTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/AuditLogHelperTest.java @@ -101,4 +101,13 @@ public void testUpdateMetricsDebugModeShortCircuit() { Config.enable_bdbje_debug_mode = original; } } + + @Test + public void testGetCloudClusterForAuditPrefersEffectiveCluster() throws Exception { + ConnectContext ctx = createMockContext(true, false); + ctx.getSessionVariable().setCloudCluster("session_cluster"); + ctx.setEffectiveCloudCluster("hint_cluster"); + + Assert.assertEquals("hint_cluster", AuditLogHelper.getCloudClusterForAudit(ctx)); + } } diff --git a/regression-test/suites/audit/test_audit_log_hint_compute_group_docker.groovy b/regression-test/suites/audit/test_audit_log_hint_compute_group_docker.groovy new file mode 100644 index 00000000000000..ce223bd110fa3f --- /dev/null +++ b/regression-test/suites/audit/test_audit_log_hint_compute_group_docker.groovy @@ -0,0 +1,64 @@ +// 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.suite.ClusterOptions + +suite("test_audit_log_hint_compute_group_docker", "docker") { + def options = new ClusterOptions(cloudMode: true, feNum: 1, beNum: 1, msNum: 1) + options.feConfigs += ['cloud_cluster_check_interval_second=1'] + + docker(options) { + def hintComputeGroup = "audit_hint_compute_group" + cluster.addBackend(1, hintComputeGroup) + + def computeGroups = sql_return_maparray "SHOW CLUSTERS" + assertEquals(2, computeGroups.size()) + def sessionComputeGroup = computeGroups.find { it.is_current == "TRUE" }.cluster + assertNotNull(sessionComputeGroup) + assertTrue(computeGroups.any { it.cluster == hintComputeGroup }) + + try { + sql "set global enable_audit_plugin = true" + sql "use @${sessionComputeGroup}" + sql "truncate table __internal_schema.audit_log" + + def marker = "audit_hint_cg_docker_marker_7F3A2B" + sql """select /*+ SET_VAR(cloud_cluster = '${hintComputeGroup}') */ + 1, '${marker}'""" + + def retry = 60 + def query = """select count(*) + from __internal_schema.audit_log + where stmt like '%${marker}%' + and compute_group = '${hintComputeGroup}'""" + def found = (sql "${query}")[0][0] as long + while (found == 0) { + if (retry-- < 0) { + throw new RuntimeException("audit_log row for the hint query was not found in the " + + "hint Compute Group") + } + sleep(3000) + sql "call flush_audit_log()" + found = (sql "${query}")[0][0] as long + } + + assertTrue(found >= 1) + } finally { + sql "set global enable_audit_plugin = false" + } + } +}