feat: save progress

astandrik · astandrik · commit 471b72da5408 · 2025-01-22T14:22:19.000+03:00
diff --git a/src/containers/Tenant/Query/QueryEditor/QueryEditor.tsx b/src/containers/Tenant/Query/QueryEditor/QueryEditor.tsx
@@ -8,7 +8,10 @@ import {v4 as uuidv4} from 'uuid';
 
 import {MonacoEditor} from '../../../../components/MonacoEditor/MonacoEditor';
 import SplitPane from '../../../../components/SplitPane';
-import {useTracingLevelOptionAvailable} from '../../../../store/reducers/capabilities/hooks';
+import {
+    useStreamingAvailable,
+    useTracingLevelOptionAvailable,
+} from '../../../../store/reducers/capabilities/hooks';
 import {
     goToNextQuery,
     goToPreviousQuery,
@@ -137,6 +140,8 @@ export default function QueryEditor(props: QueryEditorProps) {
         return historyQueries[historyQueries.length - 1].queryText;
     });
 
+    const isStreamingSupported = useStreamingAvailable();
+
     const handleSendExecuteClick = useEventHandler((text?: string) => {
         const query = text ?? input;
 
@@ -155,6 +160,7 @@ export default function QueryEditor(props: QueryEditorProps) {
             querySettings,
             enableTracingLevel,
             queryId,
+            isStreaming: isStreamingSupported,
         });
 
         dispatch(setShowPreview(false));
diff --git a/src/containers/Tenant/Query/QueryResult/QueryResultViewer.tsx b/src/containers/Tenant/Query/QueryResult/QueryResultViewer.tsx
@@ -107,8 +107,9 @@ export function QueryResultViewer({
     });
     const [useShowPlanToSvg] = useSetting<boolean>(USE_SHOW_PLAN_SVG_KEY);
 
-    const {error, isLoading, queryId, data = {}} = result;
+    const {error, isLoading, queryId, data = {}, streaming} = result;
     const {preparedPlan, simplifiedPlan, stats, resultSets, ast} = data;
+    const isStreaming = streaming?.active;
 
     React.useEffect(() => {
         if (resultType === 'execute' && !EXECUTE_SECTIONS.includes(activeSection)) {
@@ -217,9 +218,17 @@ export function QueryResultViewer({
             return <QueryResultError error={error} />;
         }
         if (activeSection === RESULT_OPTIONS_IDS.result) {
-            return (
+            return isStreaming ? (
                 <ResultSetsViewer
-                    resultSets={resultSets}
+                    type="streaming"
+                    chunks={streaming?.chunks || []}
+                    selectedResultSet={selectedResultSet}
+                    setSelectedResultSet={setSelectedResultSet}
+                />
+            ) : (
+                <ResultSetsViewer
+                    type="regular"
+                    resultSets={resultSets || []}
                     selectedResultSet={selectedResultSet}
                     setSelectedResultSet={setSelectedResultSet}
                 />
@@ -263,8 +272,8 @@ export function QueryResultViewer({
     const renderLeftControls = () => {
         return (
             <div className={b('controls-left')}>
-                <QueryExecutionStatus error={error} loading={isLoading} />
-                {!error && !isLoading && (
+                <QueryExecutionStatus error={error} loading={isLoading || isStreaming} />
+                {!error && !isLoading && !isStreaming && (
                     <React.Fragment>
                         {valueIsDefined(stats?.DurationUs) ? (
                             <QueryDuration duration={Number(stats.DurationUs)} />
diff --git a/src/containers/Tenant/Query/QueryResult/components/ResultSetsViewer/ResultSetsViewer.tsx b/src/containers/Tenant/Query/QueryResult/components/ResultSetsViewer/ResultSetsViewer.tsx
@@ -1,6 +1,7 @@
 import {Tabs, Text} from '@gravity-ui/uikit';
 
 import {QueryResultTable} from '../../../../../../components/QueryResultTable';
+import type {ExecuteQueryResponse} from '../../../../../../types/api/query';
 import type {ParsedResultSet} from '../../../../../../types/store/query';
 import {getArray} from '../../../../../../utils';
 import {cn} from '../../../../../../utils/cn';
@@ -10,25 +11,57 @@ import './ResultSetsViewer.scss';
 
 const b = cn('ydb-query-result-sets-viewer');
 
-interface ResultSetsViewerProps {
-    resultSets?: ParsedResultSet[];
+interface BaseResultSetsViewerProps {
     selectedResultSet: number;
     setSelectedResultSet: (resultSet: number) => void;
 }
 
-export function ResultSetsViewer({
-    resultSets,
-    selectedResultSet,
-    setSelectedResultSet,
-}: ResultSetsViewerProps) {
-    const resultsSetsCount = resultSets?.length || 0;
-    const currentResult = resultSets?.[selectedResultSet];
+interface RegularResultSetsProps extends BaseResultSetsViewerProps {
+    type: 'regular';
+    resultSets: ParsedResultSet[];
+}
+
+interface StreamingResultSetsProps extends BaseResultSetsViewerProps {
+    type: 'streaming';
+    chunks: ExecuteQueryResponse[];
+}
+
+type ResultSetsViewerProps = RegularResultSetsProps | StreamingResultSetsProps;
+
+export function ResultSetsViewer(props: ResultSetsViewerProps) {
+    const {selectedResultSet, setSelectedResultSet} = props;
+
+    const resultsSetsCount =
+        props.type === 'streaming' ? props.chunks.length : props.resultSets.length;
+
+    const currentResult =
+        props.type === 'streaming'
+            ? {
+                  result:
+                      props.chunks[selectedResultSet]?.result?.[0]?.rows?.map((row) => {
+                          if (!Array.isArray(row)) {
+                              return row || {};
+                          }
+                          return row.reduce<Record<string, any>>((obj, val, idx) => {
+                              const columnName =
+                                  props.chunks[0]?.result?.[0]?.columns?.[idx]?.name ||
+                                  `column${idx}`;
+                              return {...obj, [columnName]: val};
+                          }, {});
+                      }) || [],
+                  columns: props.chunks[0]?.result?.[0]?.columns,
+                  truncated: false,
+              }
+            : props.resultSets[selectedResultSet];
 
     const renderTabs = () => {
         if (resultsSetsCount > 1) {
             const tabsItems = getArray(resultsSetsCount).map((item) => ({
                 id: String(item),
-                title: `Result #${item + 1}${resultSets?.[item]?.truncated ? ' (T)' : ''}`,
+                title:
+                    props.type === 'streaming'
+                        ? `Chunk #${item + 1}`
+                        : `Result #${item + 1}${props.resultSets[item]?.truncated ? ' (T)' : ''}`,
             }));
 
             return (
diff --git a/src/services/api/viewer.ts b/src/services/api/viewer.ts
@@ -35,6 +35,7 @@ import type {TUserToken} from '../../types/api/whoami';
 import type {QuerySyntax, TransactionMode} from '../../types/store/query';
 import {BINARY_DATA_IN_PLAIN_TEXT_DISPLAY} from '../../utils/constants';
 import type {Nullable} from '../../utils/typecheckers';
+import {multipartParser} from '../parsers/multipart';
 import {settingsManager} from '../settings';
 
 import {BaseYdbAPI} from './base';
@@ -368,6 +369,61 @@ export class ViewerAPI extends BaseYdbAPI {
         );
     }
 
+    streamQuery<Action extends Actions>(
+        params: {
+            query?: string;
+            database?: string;
+            action?: Action;
+            syntax?: QuerySyntax;
+            stats?: Stats;
+            tracingLevel?: TracingLevel;
+            transaction_mode?: TransactionMode;
+            timeout?: Timeout;
+            limit_rows?: number;
+            output_chunk_max_size?: number;
+        },
+        {concurrentId, signal, onChunk}: AxiosOptions & {onChunk?: (chunk: any) => void} = {},
+    ) {
+        const base64 = !settingsManager.readUserSettingsValue(
+            BINARY_DATA_IN_PLAIN_TEXT_DISPLAY,
+            true,
+        );
+
+        // Reset parser state for new request
+        multipartParser.reset();
+
+        return this.get<string>(
+            this.getPath('/viewer/query'),
+            {
+                ...params,
+                base64,
+                schema: 'multipart',
+                output_chunk_max_size: params.output_chunk_max_size || 1000,
+            },
+            {
+                concurrentId,
+                timeout: params.timeout,
+                requestConfig: {
+                    signal,
+                    'axios-retry': {retries: 0}, // No retries for streaming
+                    responseType: 'text',
+                },
+                headers: {
+                    ...(params.tracingLevel ? {'X-Trace-Verbosity': params.tracingLevel} : {}),
+                    Accept: 'multipart/x-mixed-replace',
+                },
+                onDownloadProgress: (progressEvent) => {
+                    console.log('progressEvent', progressEvent.event.target);
+                    const response = progressEvent.event.target as XMLHttpRequest;
+                    multipartParser.processNewData<QueryAPIResponse<Action>>(
+                        response.responseText,
+                        onChunk,
+                    );
+                },
+            },
+        );
+    }
+
     getHotKeys(
         {path, database, enableSampling}: {path: string; database: string; enableSampling: boolean},
         {concurrentId, signal}: AxiosOptions = {},
diff --git a/src/services/parsers/multipart.ts b/src/services/parsers/multipart.ts
@@ -0,0 +1,74 @@
+interface MultipartChunk<T = any> {
+    part_number: number;
+    total_parts: number;
+    content: T;
+}
+
+export class MultipartParser {
+    private boundary = 'boundary';
+    private lastProcessedLength = 0;
+
+    parseChunk<T>(part: string): MultipartChunk<T> | null {
+        try {
+            // Normalize line endings and split headers and body
+            const normalizedPart = part.replace(/\r\n/g, '\n');
+            const sections = normalizedPart.split('\n\n');
+
+            if (sections.length < 2) {
+                console.log('Invalid chunk format: no body found');
+                return null;
+            }
+
+            // Get the body (last section after headers)
+            const body = sections[sections.length - 1].trim();
+            if (!body) {
+                console.log('Invalid chunk format: empty body');
+                return null;
+            }
+
+            // Parse the JSON content
+            const content = JSON.parse(body);
+
+            return {
+                part_number: content.Counter || 0,
+                total_parts: 0, // We don't know the total in advance
+                content: content as T,
+            };
+        } catch (error) {
+            console.error('Failed to parse chunk:', error);
+            return null;
+        }
+    }
+
+    processNewData<T>(responseText: string, onChunk?: (chunk: MultipartChunk<T>) => void): void {
+        // Get only the new data
+        const newData = responseText.slice(this.lastProcessedLength);
+
+        if (newData) {
+            // Split on boundary with double dashes
+            const boundaryStr = `--${this.boundary}`;
+            const parts = newData.split(boundaryStr).filter(Boolean);
+
+            for (const part of parts) {
+                // Skip the final boundary marker
+                if (part.trim() === '--') {
+                    continue;
+                }
+
+                const chunk = this.parseChunk<T>(part);
+                if (chunk && onChunk) {
+                    onChunk(chunk);
+                }
+            }
+        }
+
+        // Update the processed length
+        this.lastProcessedLength = responseText.length;
+    }
+
+    reset(): void {
+        this.lastProcessedLength = 0;
+    }
+}
+
+export const multipartParser = new MultipartParser();
diff --git a/src/store/reducers/capabilities/hooks.ts b/src/store/reducers/capabilities/hooks.ts
@@ -72,3 +72,7 @@ export const useFeatureFlagsAvailable = () => {
 export const useClusterDashboardAvailable = () => {
     return useGetFeatureVersion('/viewer/cluster') > 4;
 };
+
+export const useStreamingAvailable = () => {
+    return useGetFeatureVersion('/viewer/query') >= 7;
+};
diff --git a/src/store/reducers/query/query.ts b/src/store/reducers/query/query.ts
diff --git a/src/store/reducers/query/types.ts b/src/store/reducers/query/types.ts
