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 9174afc2fcc92f..34b090cc435dbb 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 f9114313cd814e..7d7db3df7df766 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 @@ -214,6 +214,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; @@ -1430,6 +1433,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 0037fc28aa5fd7..8f591c58a19513 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 @@ -589,6 +589,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); } @@ -615,6 +616,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.