branch-4.1: [fix](fe) Preserve effective Compute Group for audit logs #66281

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) {
         }
     }
 }
-

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);
     }

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.

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));
+    }
 }

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"
+        }
+    }
+}