From c80bfea88ef267e1e0ab28ad93e89dd6cca868bc Mon Sep 17 00:00:00 2001
From: Vlad Kuznecov <va-kuznecov@ydb.tech>
Date: Wed, 10 Jan 2024 10:10:18 +0000
Subject: [PATCH] Get right TraceId in SessionActor

---
 ydb/core/kqp/common/events/query.h               | 7 +++++++
 ydb/core/kqp/session_actor/kqp_query_state.h     | 4 ++--
 ydb/core/kqp/session_actor/kqp_session_actor.cpp | 3 +--
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/ydb/core/kqp/common/events/query.h b/ydb/core/kqp/common/events/query.h
index 2f269d85f3fd..38fa0f63d934 100644
--- a/ydb/core/kqp/common/events/query.h
+++ b/ydb/core/kqp/common/events/query.h
@@ -161,6 +161,13 @@ struct TEvQueryRequest: public NActors::TEventLocal<TEvQueryRequest, TKqpEvents:
         return Record.GetTraceId();
     }
 
+    NWilson::TTraceId GetWilsonTraceId() const {
+        if (RequestCtx) {
+            return RequestCtx->GetWilsonTraceId();
+        }
+        return {};
+    }
+
     const TString& GetRequestType() const {
         if (RequestCtx) {
             if (!RequestType) {
diff --git a/ydb/core/kqp/session_actor/kqp_query_state.h b/ydb/core/kqp/session_actor/kqp_query_state.h
index 0eb2c0ee046b..0d360e922556 100644
--- a/ydb/core/kqp/session_actor/kqp_query_state.h
+++ b/ydb/core/kqp/session_actor/kqp_query_state.h
@@ -35,7 +35,7 @@ class TKqpQueryState : public TNonCopyable {
     TKqpQueryState(TEvKqp::TEvQueryRequest::TPtr& ev, ui64 queryId, const TString& database,
         const TString& cluster, TKqpDbCountersPtr dbCounters, bool longSession,
         const NKikimrConfig::TTableServiceConfig& tableServiceConfig, const NKikimrConfig::TQueryServiceConfig& queryServiceConfig,
-        NWilson::TTraceId&& traceId, const TString& sessionId, TMonotonic startedAt)
+        const TString& sessionId, TMonotonic startedAt)
         : QueryId(queryId)
         , Database(database)
         , Cluster(cluster)
@@ -62,7 +62,7 @@ class TKqpQueryState : public TNonCopyable {
         SetQueryDeadlines(tableServiceConfig, queryServiceConfig);
         auto action = GetAction();
         KqpSessionSpan = NWilson::TSpan(
-            TWilsonKqp::KqpSession, std::move(traceId),
+            TWilsonKqp::KqpSession, std::move(RequestEv->GetWilsonTraceId()),
             "Session.query." + NKikimrKqp::EQueryAction_Name(action), NWilson::EFlags::AUTO_END);
         if (RequestEv->GetUserRequestContext()) {
             UserRequestContext = RequestEv->GetUserRequestContext();
diff --git a/ydb/core/kqp/session_actor/kqp_session_actor.cpp b/ydb/core/kqp/session_actor/kqp_session_actor.cpp
index 368ae0f043b9..399740fbbc3e 100644
--- a/ydb/core/kqp/session_actor/kqp_session_actor.cpp
+++ b/ydb/core/kqp/session_actor/kqp_session_actor.cpp
@@ -213,8 +213,7 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
         ev->Get()->SetClientLostAction(selfId, as);
         QueryState = std::make_shared<TKqpQueryState>(
             ev, QueryId, Settings.Database, Settings.Cluster, Settings.DbCounters, Settings.LongSession,
-            Settings.TableService, Settings.QueryService, std::move(ev->TraceId), SessionId,
-            AppData()->MonotonicTimeProvider->Now());
+            Settings.TableService, Settings.QueryService, SessionId, AppData()->MonotonicTimeProvider->Now());
         if (QueryState->UserRequestContext->TraceId.empty()) {
             QueryState->UserRequestContext->TraceId = UlidGen.Next().ToString();
         }