feat: support multipart responses in query (#1865)

Author: astandrik

config-overrides.js

@@ -58,7 +58,7 @@ module.exports = {
         // By default jest does not transform anything in node_modules
         // So this override excludes node_modules except @gravity-ui
         // see https://github.com/timarney/react-app-rewired/issues/241
-        config.transformIgnorePatterns = ['node_modules/(?!(@gravity-ui)/)'];
+        config.transformIgnorePatterns = ['node_modules/(?!(@gravity-ui|@mjackson)/)'];
 
         // Add .github directory to roots
         config.roots = ['<rootDir>/src', '<rootDir>/.github'];

package-lock.json

@@ -21,12 +21,13 @@
     "@gravity-ui/icons": "^2.12.0",
     "@gravity-ui/navigation": "^2.30.0",
     "@gravity-ui/paranoid": "^2.0.1",
-    "@gravity-ui/react-data-table": "^2.1.1",
+    "@gravity-ui/react-data-table": "^2.2.1",
     "@gravity-ui/table": "^1.7.0",
     "@gravity-ui/uikit": "^6.40.0",
     "@gravity-ui/unipika": "^5.2.1",
     "@gravity-ui/websql-autocomplete": "^13.7.0",
     "@hookform/resolvers": "^3.10.0",
+    "@mjackson/multipart-parser": "^0.8.2",
     "@reduxjs/toolkit": "^2.5.0",
     "@tanstack/react-table": "^8.20.6",
     "@ydb-platform/monaco-ghost": "^0.6.1",
@@ -163,10 +164,10 @@
     "typescript": "^5.7.3"
   },
   "peerDependencies": {
+    "monaco-yql-languages": ">=1.3.0",
     "prop-types": "^15.8.1",
     "react": "^18.3.1",
-    "react-dom": "^18.3.1",
-    "monaco-yql-languages": ">=1.3.0"
+    "react-dom": "^18.3.1"
   },
   "overrides": {
     "fork-ts-checker-webpack-plugin": "^9.0.2",

package.json

@@ -9,4 +9,10 @@
     &__message {
         padding: 15px 10px;
     }
+
+    // Must have fixed height for frequent data updates (to avoid interface shaking).
+    // Will be fixed after migration to new @gravity-ui/table.
+    &__table-wrapper {
+        height: 0px;
+    }
 }

src/components/QueryResultTable/QueryResultTable.scss

@@ -7,7 +7,7 @@ import type {ColumnType, KeyValueRow} from '../../types/api/query';
 import {cn} from '../../utils/cn';
 import {DEFAULT_TABLE_SETTINGS} from '../../utils/constants';
 import {getColumnWidth} from '../../utils/getColumnWidth';
-import {getColumnType, prepareQueryResponse} from '../../utils/query';
+import {getColumnType} from '../../utils/query';
 import {isNumeric} from '../../utils/utils';
 import type {ResizeableDataTableProps} from '../ResizeableDataTable/ResizeableDataTable';
 import {ResizeableDataTable} from '../ResizeableDataTable/ResizeableDataTable';
@@ -49,8 +49,8 @@ const prepareTypedColumns = (columns: ColumnType[], data?: KeyValueRow[]) => {
     });
 };
 
-const prepareGenericColumns = (data: KeyValueRow[]) => {
-    if (!data.length) {
+const prepareGenericColumns = (data?: KeyValueRow[]) => {
+    if (!data?.length) {
         return [];
     }
 
@@ -77,35 +77,42 @@ interface QueryResultTableProps
     extends Omit<ResizeableDataTableProps<KeyValueRow>, 'data' | 'columns'> {
     data?: KeyValueRow[];
     columns?: ColumnType[];
+    settings?: Partial<Settings>;
 }
 
 export const QueryResultTable = (props: QueryResultTableProps) => {
-    const {columns: rawColumns, data: rawData, ...restProps} = props;
+    const {columns, data, settings: propsSettings} = props;
 
-    const data = React.useMemo(() => prepareQueryResponse(rawData), [rawData]);
-    const columns = React.useMemo(() => {
-        return rawColumns ? prepareTypedColumns(rawColumns, data) : prepareGenericColumns(data);
-    }, [data, rawColumns]);
+    const preparedColumns = React.useMemo(() => {
+        return columns ? prepareTypedColumns(columns, data) : prepareGenericColumns(data);
+    }, [data, columns]);
+
+    const settings = React.useMemo(() => {
+        return {
+            ...TABLE_SETTINGS,
+            ...propsSettings,
+        };
+    }, [propsSettings]);
 
     // empty data is expected to be be an empty array
     // undefined data is not rendered at all
-    if (!Array.isArray(rawData)) {
+    if (!Array.isArray(data)) {
         return null;
     }
 
-    if (!columns.length) {
+    if (!preparedColumns.length) {
         return <div className={b('message')}>{i18n('empty')}</div>;
     }
 
     return (
         <ResizeableDataTable
             data={data}
-            columns={columns}
-            settings={TABLE_SETTINGS}
+            columns={preparedColumns}
+            settings={settings}
             // prevent accessing row.id in case it is present but is not the PK (i.e. may repeat)
             rowKey={getRowIndex}
             visibleRowIndex={getVisibleRowIndex}
-            {...restProps}
+            wrapperClassName={b('table-wrapper')}
         />
     );
 };

src/components/QueryResultTable/QueryResultTable.tsx

@@ -1,9 +1,6 @@
-import React from 'react';
-
 import {StopFill} from '@gravity-ui/icons';
 import {Button, Icon} from '@gravity-ui/uikit';
 
-import {cancelQueryApi} from '../../../../store/reducers/cancelQuery';
 import {cn} from '../../../../utils/cn';
 import i18n from '../i18n';
 
@@ -12,22 +9,17 @@ import './CancelQueryButton.scss';
 const b = cn('cancel-query-button');
 
 interface CancelQueryButtonProps {
-    queryId: string;
-    tenantName: string;
+    isLoading: boolean;
+    isError: boolean;
+    onClick?: VoidFunction;
 }
 
-export function CancelQueryButton({queryId, tenantName}: CancelQueryButtonProps) {
-    const [sendCancelQuery, cancelQueryResponse] = cancelQueryApi.useCancelQueryMutation();
-
-    const onStopButtonClick = React.useCallback(() => {
-        sendCancelQuery({queryId, database: tenantName});
-    }, [queryId, sendCancelQuery, tenantName]);
-
+export function CancelQueryButton({isLoading, isError, onClick}: CancelQueryButtonProps) {
     return (
         <Button
-            loading={cancelQueryResponse.isLoading}
-            onClick={onStopButtonClick}
-            className={b('stop-button', {error: Boolean(cancelQueryResponse.error)})}
+            loading={isLoading}
+            onClick={onClick}
+            className={b('stop-button', {error: isError})}
         >
             <Icon data={StopFill} size={16} />
             {i18n('action.stop')}

src/containers/Tenant/Query/CancelQueryButton/CancelQueryButton.tsx

@@ -1,10 +1,15 @@
 import React from 'react';
 
+import type {Settings} from '@gravity-ui/react-data-table';
 import {isEqual} from 'lodash';
 import {v4 as uuidv4} from 'uuid';
 
 import SplitPane from '../../../../components/SplitPane';
-import {useTracingLevelOptionAvailable} from '../../../../store/reducers/capabilities/hooks';
+import {cancelQueryApi} from '../../../../store/reducers/cancelQuery';
+import {
+    useStreamingAvailable,
+    useTracingLevelOptionAvailable,
+} from '../../../../store/reducers/capabilities/hooks';
 import {
     queryApi,
     saveQueryToHistory,
@@ -23,6 +28,7 @@ import {cn} from '../../../../utils/cn';
 import {
     DEFAULT_IS_QUERY_RESULT_COLLAPSED,
     DEFAULT_SIZE_RESULT_PANE_KEY,
+    ENABLE_QUERY_STREAMING,
     LAST_USED_QUERY_ACTION_KEY,
 } from '../../../../utils/constants';
 import {
@@ -34,7 +40,7 @@ import {
 } from '../../../../utils/hooks';
 import {useChangedQuerySettings} from '../../../../utils/hooks/useChangedQuerySettings';
 import {useLastQueryExecutionSettings} from '../../../../utils/hooks/useLastQueryExecutionSettings';
-import {QUERY_ACTIONS} from '../../../../utils/query';
+import {DEFAULT_QUERY_SETTINGS, QUERY_ACTIONS} from '../../../../utils/query';
 import type {InitialPaneState} from '../../utils/paneVisibilityToggleHelpers';
 import {
     PaneVisibilityActionTypes,
@@ -86,8 +92,24 @@ export default function QueryEditor(props: QueryEditorProps) {
         LAST_USED_QUERY_ACTION_KEY,
     );
     const [lastExecutedQueryText, setLastExecutedQueryText] = React.useState<string>('');
+    const [isQueryStreamingEnabled] = useSetting(ENABLE_QUERY_STREAMING);
+    const isStreamingEnabled = useStreamingAvailable() && isQueryStreamingEnabled;
 
     const [sendQuery] = queryApi.useUseSendQueryMutation();
+    const [streamQuery] = queryApi.useUseStreamQueryMutation();
+    const [sendCancelQuery, cancelQueryResponse] = cancelQueryApi.useCancelQueryMutation();
+
+    const runningQueryRef = React.useRef<{abort: VoidFunction} | null>(null);
+
+    const tableSettings = React.useMemo(() => {
+        return isStreamingEnabled
+            ? {
+                  displayIndices: {
+                      maxIndex: (querySettings.limitRows || DEFAULT_QUERY_SETTINGS.limitRows) + 1,
+                  },
+              }
+            : undefined;
+    }, [isStreamingEnabled, querySettings.limitRows]);
 
     React.useEffect(() => {
         if (savedPath !== tenantName) {
@@ -121,14 +143,25 @@ export default function QueryEditor(props: QueryEditorProps) {
         }
         const queryId = uuidv4();
 
-        sendQuery({
-            actionType: 'execute',
-            query: text,
-            database: tenantName,
-            querySettings,
-            enableTracingLevel,
-            queryId,
-        });
+        if (isStreamingEnabled) {
+            runningQueryRef.current = streamQuery({
+                actionType: 'execute',
+                query: text,
+                database: tenantName,
+                querySettings,
+                enableTracingLevel,
+                queryId,
+            });
+        } else {
+            runningQueryRef.current = sendQuery({
+                actionType: 'execute',
+                query: text,
+                database: tenantName,
+                querySettings,
+                enableTracingLevel,
+                queryId,
+            });
+        }
 
         dispatch(setShowPreview(false));
 
@@ -155,7 +188,7 @@ export default function QueryEditor(props: QueryEditorProps) {
 
         const queryId = uuidv4();
 
-        sendQuery({
+        runningQueryRef.current = sendQuery({
             actionType: 'explain',
             query: text,
             database: tenantName,
@@ -169,6 +202,14 @@ export default function QueryEditor(props: QueryEditorProps) {
         dispatchResultVisibilityState(PaneVisibilityActionTypes.triggerExpand);
     });
 
+    const handleCancelRunningQuery = React.useCallback(() => {
+        if (isStreamingEnabled && runningQueryRef.current) {
+            runningQueryRef.current.abort();
+        } else if (result?.queryId) {
+            sendCancelQuery({queryId: result?.queryId, database: tenantName});
+        }
+    }, [isStreamingEnabled, result?.queryId, sendCancelQuery, tenantName]);
+
     const onCollapseResultHandler = () => {
         dispatchResultVisibilityState(PaneVisibilityActionTypes.triggerCollapse);
     };
@@ -229,10 +270,13 @@ export default function QueryEditor(props: QueryEditorProps) {
                         theme={theme}
                         key={result?.queryId}
                         result={result}
+                        cancelQueryResponse={cancelQueryResponse}
                         tenantName={tenantName}
                         path={path}
                         showPreview={showPreview}
                         queryText={lastExecutedQueryText}
+                        onCancelRunningQuery={handleCancelRunningQuery}
+                        tableSettings={tableSettings}
                     />
                 </div>
             </SplitPane>
@@ -248,13 +292,17 @@ interface ResultProps {
     type?: EPathType;
     theme: string;
     result?: QueryResult;
+    cancelQueryResponse?: Pick<QueryResult, 'isLoading' | 'error'>;
     tenantName: string;
     path: string;
     showPreview?: boolean;
     queryText: string;
+    tableSettings?: Partial<Settings>;
+    onCancelRunningQuery: VoidFunction;
 }
 function Result({
     resultVisibilityState,
+    cancelQueryResponse,
     onExpandResultHandler,
     onCollapseResultHandler,
     type,
@@ -264,6 +312,8 @@ function Result({
     path,
     showPreview,
     queryText,
+    tableSettings,
+    onCancelRunningQuery,
 }: ResultProps) {
     if (showPreview) {
         return <Preview database={tenantName} path={path} type={type} />;
@@ -277,9 +327,13 @@ function Result({
                 theme={theme}
                 tenantName={tenantName}
                 isResultsCollapsed={resultVisibilityState.collapsed}
+                isCancelError={Boolean(cancelQueryResponse?.error)}
+                isCancelling={Boolean(cancelQueryResponse?.isLoading)}
+                tableSettings={tableSettings}
                 onExpandResults={onExpandResultHandler}
                 onCollapseResults={onCollapseResultHandler}
                 queryText={queryText}
+                onCancelRunningQuery={onCancelRunningQuery}
             />
         );
     }