From 8e8dc5d0baf955697f2ee512bd73eb3eeb618118 Mon Sep 17 00:00:00 2001
From: Ivan Lapa <lin@arenadata.io>
Date: Tue, 2 Jun 2026 15:24:42 +0300
Subject: [PATCH] ADH-5827: Use raw Hive query id in Tez caller context

Tez UI resolves Hive query details by matching the Tez callerId with the HIVE_QUERY_ID entity emitted by ATSHook. HIVE-26789 appended the user name to caller context for audit logging, but applying the same format to the Tez DAG callerId breaks that lookup.

Keep the user-enriched caller context for Hadoop audit logging paths and pass the raw query id to Tez. Add a regression test for the Tez caller context.
---
 .../hadoop/hive/ql/exec/tez/TezTask.java      | 12 ++++--
 .../ql/exec/tez/TestTezTaskCallerContext.java | 38 +++++++++++++++++++
 2 files changed, 46 insertions(+), 4 deletions(-)
 create mode 100644 ql/src/test/org/apache/hadoop/hive/ql/exec/tez/TestTezTaskCallerContext.java

diff --git a/ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezTask.java b/ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezTask.java
index 6bd801dd64c3..baba623567d2 100644
--- a/ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezTask.java
+++ b/ql/src/java/org/apache/hadoop/hive/ql/exec/tez/TezTask.java
@@ -94,8 +94,6 @@
 
 import com.google.common.annotations.VisibleForTesting;
 
-import static org.apache.hadoop.hive.shims.HadoopShims.USER_ID;
-
 /**
  *
  * TezTask handles the execution of TezWork. Currently it executes a graph of map and reduce work
@@ -214,8 +212,7 @@ public int execute() {
       // TODO: we could perhaps reuse the same directory for HiveResources?
       Path scratchDir = utils.createTezDir(ctx.getMRScratchDir(), conf);
       CallerContext callerContext =
-          CallerContext.create("HIVE", String.format(USER_ID, queryPlan.getQueryId(), userName), "HIVE_QUERY_ID",
-              queryPlan.getQueryStr());
+          createCallerContext(queryPlan.getQueryId(), queryPlan.getQueryStr());
 
       perfLogger.perfLogBegin(CLASS_NAME, PerfLogger.TEZ_GET_SESSION);
       session = sessionRef.value = WorkloadManagerFederation.getSession(
@@ -464,6 +461,13 @@ private void logResources(List<LocalResource> additionalLr) {
     }
   }
 
+  @VisibleForTesting
+  static CallerContext createCallerContext(String queryId, String queryStr) {
+    // The callerId must match the HIVE_QUERY_ID entity id emitted by ATSHook.
+    // User information belongs to the Hadoop caller context used by audit logs.
+    return CallerContext.create("HIVE", queryId, "HIVE_QUERY_ID", queryStr);
+  }
+
   /**
    * Ensures that the Tez Session is open and the AM has all necessary jars configured.
    */
diff --git a/ql/src/test/org/apache/hadoop/hive/ql/exec/tez/TestTezTaskCallerContext.java b/ql/src/test/org/apache/hadoop/hive/ql/exec/tez/TestTezTaskCallerContext.java
new file mode 100644
index 000000000000..447ab5e80e4a
--- /dev/null
+++ b/ql/src/test/org/apache/hadoop/hive/ql/exec/tez/TestTezTaskCallerContext.java
@@ -0,0 +1,38 @@
+/*
+ * 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.
+ */
+package org.apache.hadoop.hive.ql.exec.tez;
+
+import org.apache.tez.client.CallerContext;
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+public class TestTezTaskCallerContext {
+
+  @Test
+  public void testCreateCallerContextUsesHiveQueryId() {
+    String queryId = "hive_20260601151408_438d1789-d603-412b-bb1d-5401effba17c";
+
+    CallerContext callerContext = TezTask.createCallerContext(queryId, "select 1");
+
+    assertEquals("HIVE", callerContext.getContext());
+    assertEquals("HIVE_QUERY_ID", callerContext.getCallerType());
+    assertEquals(queryId, callerContext.getCallerId());
+    assertEquals("select 1", callerContext.getBlob());
+  }
+}