feat: multiple responses

Author: astandrik

src/services/api/viewer.ts

@@ -38,7 +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 type {StreamDataChunk, 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';
@@ -386,7 +386,8 @@ export class ViewerAPI extends BaseYdbAPI {
             transaction_mode?: TransactionMode;
             timeout?: Timeout;
             limit_rows?: number;
-            output_chunk_max_size?: number;
+            output_chunk_max_size: number;
+            concurrent_results: number;
         },
         {
             concurrentId,
@@ -409,7 +410,6 @@ export class ViewerAPI extends BaseYdbAPI {
                 ...params,
                 base64,
                 schema: 'multipart',
-                output_chunk_max_size: params.output_chunk_max_size || 1000,
             },
             {
                 concurrentId,
@@ -430,32 +430,38 @@ export class ViewerAPI extends BaseYdbAPI {
                         lastProcessedLength,
                     });
                     lastProcessedLength = currentLength;
-                    let mergedChunk = null;
+                    const mergedChunks = new Map<number, StreamDataChunk>();
 
                     for (const chunk of chunks) {
                         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;
+                            for (const [index, mergedChunk] of mergedChunks) {
+                                onChunk(mergedChunk);
+                                mergedChunks.delete(index);
                             }
+
                             onChunk(chunk);
                         } else if (isStreamDataChunk(chunk)) {
+                            const resultIndex = chunk.meta.result_index;
+                            const mergedChunk = mergedChunks.get(resultIndex);
+
                             if (mergedChunk) {
                                 mergedChunk.result.rows.push(...chunk.result.rows);
                             } else {
-                                mergedChunk = chunk;
+                                mergedChunks.set(resultIndex, chunk);
                             }
                         }
                     }
 
-                    if (mergedChunk) {
-                        onChunk?.(mergedChunk);
+                    for (const [index, mergedChunk] of mergedChunks) {
+                        onChunk(mergedChunk);
+                        mergedChunks.delete(index);
                     }
                 },
             },

src/store/reducers/query/query.ts

@@ -148,26 +148,31 @@ const slice = createSlice({
             } else if (isStreamDataChunk(chunk)) {
                 const {
                     result: {columns, rows},
+                    meta: {result_index: resultIndex},
                 } = chunk;
 
-                if (columns && !state.result.data.resultSets?.[0]?.columns?.length) {
-                    if (state.result.data.resultSets) {
-                        state.result.data.resultSets[0].columns = columns;
+                if (columns && !state.result.data.resultSets?.[resultIndex]?.columns?.length) {
+                    if (!state.result.data.resultSets) {
+                        state.result.data.resultSets = [];
+                    }
+
+                    if (state.result.data.resultSets[resultIndex]) {
+                        state.result.data.resultSets[resultIndex].columns = columns;
                     } else {
-                        state.result.data.resultSets = [{columns, result: []}];
+                        state.result.data.resultSets[resultIndex] = {columns, result: []};
                     }
                 }
                 // Convert and append new rows
-                if (rows.length > 0 && state.result.data.resultSets?.[0]?.columns) {
+                if (rows.length > 0 && state.result.data.resultSets?.[resultIndex]?.columns) {
                     const formattedRows = parseResult(
                         rows,
-                        state.result.data.resultSets[0].columns,
+                        state.result.data.resultSets[resultIndex].columns,
                     );
 
-                    if (state.result.data.resultSets[0].result) {
-                        state.result.data.resultSets[0].result.push(...formattedRows);
+                    if (state.result.data.resultSets[resultIndex].result) {
+                        state.result.data.resultSets[resultIndex].result.push(...formattedRows);
                     } else {
-                        state.result.data.resultSets[0].result = formattedRows;
+                        state.result.data.resultSets[resultIndex].result = formattedRows;
                     }
                 }
             } else if (isQueryResponseChunk(chunk)) {
@@ -239,6 +244,9 @@ interface QueryStats {
     endTime?: string | number;
 }
 
+const DEFAULT_STREAM_CHUNK_SIZE = 2;
+const DEFAULT_CONCURRENT_RESULTS = 1;
+
 export const queryApi = api.injectEndpoints({
     endpoints: (build) => ({
         useStreamQuery: build.mutation<null, SendQueryParams>({
@@ -273,7 +281,8 @@ export const queryApi = api.injectEndpoints({
                             timeout: isNumeric(querySettings.timeout)
                                 ? Number(querySettings.timeout) * 1000
                                 : undefined,
-                            output_chunk_max_size: 1000,
+                            output_chunk_max_size: DEFAULT_STREAM_CHUNK_SIZE,
+                            concurrent_results: DEFAULT_CONCURRENT_RESULTS,
                         },
                         {
                             signal,