feat: keep plan representation tab selection between query executions (#2625)

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: adameat <[email protected]>
Co-authored-by: Alexey Efimov <[email protected]>
Co-authored-by: artemmufazalov <[email protected]>

Author: 3 people

src/containers/Tenant/Query/QueryResult/QueryResultViewer.tsx

@@ -1,4 +1,5 @@
 import React from 'react';
+// Query result viewer with tab persistence functionality
 
 import type {Settings} from '@gravity-ui/react-data-table';
 import type {ControlGroupOption} from '@gravity-ui/uikit';
@@ -10,13 +11,14 @@ import {Illustration} from '../../../../components/Illustration';
 import {LoaderWrapper} from '../../../../components/LoaderWrapper/LoaderWrapper';
 import {QueryExecutionStatus} from '../../../../components/QueryExecutionStatus';
 import {disableFullscreen} from '../../../../store/reducers/fullscreen';
+import {selectResultTab, setResultTab} from '../../../../store/reducers/query/query';
 import type {QueryResult} from '../../../../store/reducers/query/types';
 import type {ValueOf} from '../../../../types/common';
 import type {QueryAction} from '../../../../types/store/query';
 import {cn} from '../../../../utils/cn';
 import {USE_SHOW_PLAN_SVG_KEY} from '../../../../utils/constants';
 import {getStringifiedData} from '../../../../utils/dataFormatters/dataFormatters';
-import {useSetting, useTypedDispatch} from '../../../../utils/hooks';
+import {useSetting, useTypedDispatch, useTypedSelector} from '../../../../utils/hooks';
 import {PaneVisibilityToggleButtons} from '../../utils/paneVisibilityToggleHelpers';
 import {QuerySettingsBanner} from '../QuerySettingsBanner/QuerySettingsBanner';
 import {QueryStoppedBanner} from '../QueryStoppedBanner/QueryStoppedBanner';
@@ -98,27 +100,43 @@ export function QueryResultViewer({
     onExpandResults,
 }: ExecuteResultProps) {
     const dispatch = useTypedDispatch();
+    const selectedTabs = useTypedSelector(selectResultTab);
 
     const isExecute = resultType === 'execute';
     const isExplain = resultType === 'explain';
 
     const [selectedResultSet, setSelectedResultSet] = React.useState(0);
-    const [activeSection, setActiveSection] = React.useState<SectionID>(() => {
-        return isExecute ? RESULT_OPTIONS_IDS.result : RESULT_OPTIONS_IDS.schema;
-    });
     const [useShowPlanToSvg] = useSetting<boolean>(USE_SHOW_PLAN_SVG_KEY);
 
+    // Get the saved tab for the current query type, or use default
+    const getDefaultSection = (): SectionID => {
+        return isExecute ? RESULT_OPTIONS_IDS.result : RESULT_OPTIONS_IDS.schema;
+    };
+
+    const activeSection: SectionID = React.useMemo(() => {
+        const savedTab = selectedTabs?.[resultType];
+        if (savedTab) {
+            // Validate that the saved tab is valid for the current result type
+            const validSections = isExecute ? EXECUTE_SECTIONS : EXPLAIN_SECTIONS;
+            if (validSections.includes(savedTab as SectionID)) {
+                return savedTab as SectionID;
+            }
+        }
+        return getDefaultSection();
+    }, [selectedTabs, resultType, isExecute]);
+
     const {error, isLoading, data = {}} = result;
     const {preparedPlan, simplifiedPlan, stats, resultSets, ast} = data;
 
     React.useEffect(() => {
-        if (resultType === 'execute' && !EXECUTE_SECTIONS.includes(activeSection)) {
-            setActiveSection('result');
-        }
-        if (resultType === 'explain' && !EXPLAIN_SECTIONS.includes(activeSection)) {
-            setActiveSection('schema');
-        }
-    }, [activeSection, resultType]);
+        return () => {
+            dispatch(disableFullscreen());
+        };
+    }, [dispatch]);
+
+    const onSelectSection = (value: SectionID) => {
+        dispatch(setResultTab({queryType: resultType, tabId: value}));
+    };
 
     const radioButtonOptions: ControlGroupOption<SectionID>[] = React.useMemo(() => {
         let sections: SectionID[] = [];
@@ -137,16 +155,6 @@ export function QueryResultViewer({
         });
     }, [isExecute, isExplain]);
 
-    React.useEffect(() => {
-        return () => {
-            dispatch(disableFullscreen());
-        };
-    }, [dispatch]);
-
-    const onSelectSection = (value: SectionID) => {
-        setActiveSection(value);
-    };
-
     const getStatsToCopy = () => {
         switch (activeSection) {
             case RESULT_OPTIONS_IDS.result: {

src/store/reducers/query/__test__/tabPersistence.test.tsx

@@ -0,0 +1,79 @@
+import queryReducer, {setResultTab} from '../query';
+import type {QueryState} from '../types';
+
+describe('QueryResultViewer tab persistence integration', () => {
+    const initialState: QueryState = {
+        input: '',
+        history: {
+            queries: [],
+            currentIndex: -1,
+        },
+    };
+
+    it('should save and retrieve tab selection for explain queries', () => {
+        // Test that we can set and get the tab preference
+        let state = queryReducer(initialState, setResultTab({queryType: 'explain', tabId: 'json'}));
+
+        expect(state.selectedResultTab).toEqual({
+            explain: 'json',
+        });
+
+        // Test updating the same query type
+        state = queryReducer(state, setResultTab({queryType: 'explain', tabId: 'ast'}));
+
+        expect(state.selectedResultTab).toEqual({
+            explain: 'ast',
+        });
+    });
+
+    it('should save and retrieve tab selection for execute queries', () => {
+        const state = queryReducer(
+            initialState,
+            setResultTab({queryType: 'execute', tabId: 'stats'}),
+        );
+
+        expect(state.selectedResultTab).toEqual({
+            execute: 'stats',
+        });
+    });
+
+    it('should maintain separate preferences for different query types', () => {
+        let state = initialState;
+
+        // Set explain tab
+        state = queryReducer(state, setResultTab({queryType: 'explain', tabId: 'json'}));
+        expect(state.selectedResultTab).toEqual({
+            explain: 'json',
+        });
+
+        // Set execute tab - should not override explain
+        state = queryReducer(state, setResultTab({queryType: 'execute', tabId: 'stats'}));
+        expect(state.selectedResultTab).toEqual({
+            explain: 'json',
+            execute: 'stats',
+        });
+
+        // Update explain tab - should not affect execute
+        state = queryReducer(state, setResultTab({queryType: 'explain', tabId: 'ast'}));
+        expect(state.selectedResultTab).toEqual({
+            explain: 'ast',
+            execute: 'stats',
+        });
+    });
+
+    it('should handle multiple updates to the same query type', () => {
+        let state = initialState;
+
+        // Set initial value
+        state = queryReducer(state, setResultTab({queryType: 'explain', tabId: 'schema'}));
+        expect(state.selectedResultTab?.explain).toBe('schema');
+
+        // Update to different tab
+        state = queryReducer(state, setResultTab({queryType: 'explain', tabId: 'json'}));
+        expect(state.selectedResultTab?.explain).toBe('json');
+
+        // Update again
+        state = queryReducer(state, setResultTab({queryType: 'explain', tabId: 'ast'}));
+        expect(state.selectedResultTab?.explain).toBe('ast');
+    });
+});

src/store/reducers/query/query.ts

@@ -128,6 +128,16 @@ const slice = createSlice({
         setQueryHistoryFilter: (state, action: PayloadAction<string>) => {
             state.history.filter = action.payload;
         },
+        setResultTab: (
+            state,
+            action: PayloadAction<{queryType: 'execute' | 'explain'; tabId: string}>,
+        ) => {
+            const {queryType, tabId} = action.payload;
+            if (!state.selectedResultTab) {
+                state.selectedResultTab = {};
+            }
+            state.selectedResultTab[queryType] = tabId;
+        },
         setStreamSession: setStreamSessionReducer,
         addStreamingChunks: addStreamingChunksReducer,
         setStreamQueryResponse: setStreamQueryResponseReducer,
@@ -149,6 +159,7 @@ const slice = createSlice({
         selectUserInput: (state) => state.input,
         selectIsDirty: (state) => state.isDirty,
         selectQueriesHistoryCurrentIndex: (state) => state.history?.currentIndex,
+        selectResultTab: (state) => state.selectedResultTab,
     },
 });
 
@@ -177,6 +188,7 @@ export const {
     setStreamQueryResponse,
     setStreamSession,
     setIsDirty,
+    setResultTab,
 } = slice.actions;
 
 export const {
@@ -187,6 +199,7 @@ export const {
     selectResult,
     selectUserInput,
     selectIsDirty,
+    selectResultTab,
 } = slice.selectors;
 
 interface SendQueryParams extends QueryRequestParams {

src/store/reducers/query/types.ts

@@ -68,4 +68,8 @@ export interface QueryState {
         filter?: string;
     };
     tenantPath?: string;
+    selectedResultTab?: {
+        execute?: string;
+        explain?: string;
+    };
 }