From 65e12bc010fa596728e6d840e6d14c5daf34db5b Mon Sep 17 00:00:00 2001
From: Anton Standrik <astandrik@gmail.com>
Date: Mon, 27 Jan 2025 15:14:42 +0300
Subject: [PATCH] fix: add traceId logics

---
 .../components/TraceButton/TraceButton.tsx    |  7 ++++--
 src/services/api/viewer.ts                    | 24 ++++++++++++-------
 src/store/reducers/query/query.ts             |  1 +
 src/types/store/streaming.ts                  |  3 +++
 4 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/src/containers/Tenant/Query/QueryResult/components/TraceButton/TraceButton.tsx b/src/containers/Tenant/Query/QueryResult/components/TraceButton/TraceButton.tsx
index 948b75c7c1..4e3a1822c0 100644
--- a/src/containers/Tenant/Query/QueryResult/components/TraceButton/TraceButton.tsx
+++ b/src/containers/Tenant/Query/QueryResult/components/TraceButton/TraceButton.tsx
@@ -13,7 +13,10 @@ interface TraceUrlButtonProps {
     isTraceReady?: true;
 }
 
-export function TraceButton({traceId, isTraceReady}: TraceUrlButtonProps) {
+export const TraceButton = React.memo(function TraceButton({
+    traceId,
+    isTraceReady,
+}: TraceUrlButtonProps) {
     const {traceCheck, traceView} = useClusterBaseInfo();
 
     const checkTraceUrl = traceCheck?.url ? replaceParams(traceCheck.url, {traceId}) : '';
@@ -47,4 +50,4 @@ export function TraceButton({traceId, isTraceReady}: TraceUrlButtonProps) {
             </Button.Icon>
         </Button>
     );
-}
+});
diff --git a/src/services/api/viewer.ts b/src/services/api/viewer.ts
index 85d2056db4..2c303c7d59 100644
--- a/src/services/api/viewer.ts
+++ b/src/services/api/viewer.ts
@@ -38,6 +38,7 @@ import type {DescribeTopicResult} from '../../types/api/topic';
 import type {TEvVDiskStateResponse} from '../../types/api/vdisk';
 import type {TUserToken} from '../../types/api/whoami';
 import type {QuerySyntax, TransactionMode} from '../../types/store/query';
+import type {StreamingChunk} from '../../types/store/streaming';
 import {BINARY_DATA_IN_PLAIN_TEXT_DISPLAY} from '../../utils/constants';
 import type {Nullable} from '../../utils/typecheckers';
 import {parseMultipart} from '../parsers/parseMultipart';
@@ -387,7 +388,13 @@ export class ViewerAPI extends BaseYdbAPI {
             limit_rows?: number;
             output_chunk_max_size?: number;
         },
-        {concurrentId, signal, onChunk}: AxiosOptions & {onChunk?: (chunk: any) => void} = {},
+        {
+            concurrentId,
+            signal,
+            onChunk,
+        }: AxiosOptions & {
+            onChunk: (chunk: StreamingChunk) => void;
+        },
     ) {
         const base64 = !settingsManager.readUserSettingsValue(
             BINARY_DATA_IN_PLAIN_TEXT_DISPLAY,
@@ -426,26 +433,27 @@ export class ViewerAPI extends BaseYdbAPI {
                     let mergedChunk = null;
 
                     for (const chunk of chunks) {
-                        if (isSessionChunk(chunk) || isQueryResponseChunk(chunk)) {
-                            // First dispatch any accumulated data chunk
+                        if (isSessionChunk(chunk)) {
+                            const traceId = response
+                                .getResponseHeader('traceresponse')
+                                ?.split('-')[1];
+                            chunk.meta.trace_id = traceId;
+                            onChunk(chunk);
+                        } else if (isQueryResponseChunk(chunk)) {
                             if (mergedChunk) {
                                 onChunk?.(mergedChunk);
                                 mergedChunk = null;
                             }
-                            // Then dispatch control chunk
-                            onChunk?.(chunk);
+                            onChunk(chunk);
                         } else if (isStreamDataChunk(chunk)) {
                             if (mergedChunk) {
-                                // Merge rows from subsequent chunks
                                 mergedChunk.result.rows.push(...chunk.result.rows);
                             } else {
-                                // First data chunk - use as base with columns
                                 mergedChunk = chunk;
                             }
                         }
                     }
 
-                    // Dispatch any remaining merged chunk
                     if (mergedChunk) {
                         onChunk?.(mergedChunk);
                     }
diff --git a/src/store/reducers/query/query.ts b/src/store/reducers/query/query.ts
index d6c0bd0937..07f445b064 100644
--- a/src/store/reducers/query/query.ts
+++ b/src/store/reducers/query/query.ts
@@ -144,6 +144,7 @@ const slice = createSlice({
 
             if (isSessionChunk(chunk)) {
                 state.result.queryId = chunk.meta.query_id;
+                state.result.data.traceId = chunk.meta.trace_id;
             } else if (isStreamDataChunk(chunk)) {
                 const {
                     result: {columns, rows},
diff --git a/src/types/store/streaming.ts b/src/types/store/streaming.ts
index 25c70952f6..3deb3e027e 100644
--- a/src/types/store/streaming.ts
+++ b/src/types/store/streaming.ts
@@ -6,6 +6,9 @@ export interface SessionChunk {
         node_id: number;
         query_id: string;
         session_id: string;
+
+        // Custom client-set property.
+        trace_id?: string;
     };
 }