[Discover] Allow save query to load correctly in Discover (#5951)

* [Discover] Allow save query to load correctly in Discover

* add save query logic in Discover
* add save query logic in VisBuilder
* remove double render

Issue Resolve
#5942

---------

Signed-off-by: Anan Zhuang <ananzh@amazon.com>

src/plugins/discover/public/application/utils/state_management/discover_slice.test.tsx

@@ -4,6 +4,7 @@
  */
 
 import { discoverSlice, DiscoverState } from './discover_slice';
+import { SortOrder } from '../../../saved_searches/types';
 
 describe('discoverSlice', () => {
   let initialState: DiscoverState;
@@ -134,4 +135,40 @@ describe('discoverSlice', () => {
     const result = discoverSlice.reducer(initialState, action);
     expect(result.columns).toEqual(['column1', 'column2', 'column3']);
   });
+
+  it('should set the savedQuery when a valid id is provided', () => {
+    const savedQueryId = 'some-query-id';
+    const action = { type: 'discover/setSavedQuery', payload: savedQueryId };
+    const result = discoverSlice.reducer(initialState, action);
+    expect(result.savedQuery).toEqual(savedQueryId);
+  });
+
+  it('should remove the savedQuery from state when payload is undefined', () => {
+    // pre-set the savedQuery in the initialState
+    const initialStateWithSavedQuery = {
+      ...initialState,
+      savedQuery: 'existing-query-id',
+    };
+
+    const action = { type: 'discover/setSavedQuery', payload: undefined };
+    const result = discoverSlice.reducer(initialStateWithSavedQuery, action);
+
+    // Check that savedQuery is not in the resulting state
+    expect(result.savedQuery).toBeUndefined();
+  });
+
+  it('should not affect other state properties when setting savedQuery', () => {
+    const initialStateWithOtherProperties = {
+      ...initialState,
+      columns: ['column1', 'column2'],
+      sort: [['field1', 'asc']] as SortOrder[],
+    };
+    const savedQueryId = 'new-query-id';
+    const action = { type: 'discover/setSavedQuery', payload: savedQueryId };
+    const result = discoverSlice.reducer(initialStateWithOtherProperties, action);
+    // check that other properties remain unchanged
+    expect(result.columns).toEqual(['column1', 'column2']);
+    expect(result.sort).toEqual([['field1', 'asc']] as SortOrder[]);
+    expect(result.savedQuery).toEqual(savedQueryId);
+  });
 });

src/plugins/discover/public/application/utils/state_management/discover_slice.tsx

@@ -46,6 +46,7 @@ export interface DiscoverState {
   /**
    * Metadata for the view
    */
+  savedQuery?: string;
   metadata?: {
     /**
      * Number of lines to display per row
@@ -188,6 +189,18 @@ export const discoverSlice = createSlice({
         },
       };
     },
+    setSavedQuery(state, action: PayloadAction<string | undefined>) {
+      if (action.payload === undefined) {
+        // if the payload is undefined, remove the savedQuery property
+        const { savedQuery, ...restState } = state;
+        return restState;
+      } else {
+        return {
+          ...state,
+          savedQuery: action.payload,
+        };
+      }
+    },
   },
 });
 
@@ -204,5 +217,6 @@ export const {
   updateState,
   setSavedSearchId,
   setMetadata,
+  setSavedQuery,
 } = discoverSlice.actions;
 export const { reducer } = discoverSlice;

src/plugins/discover/public/application/view_components/canvas/discover_table.tsx

@@ -22,6 +22,7 @@ import { DocViewFilterFn } from '../../doc_views/doc_views_types';
 import { SortOrder } from '../../../saved_searches/types';
 import { OpenSearchSearchHit } from '../../doc_views/doc_views_types';
 import { popularizeField } from '../../helpers/popularize_field';
+import { buildColumns } from '../../utils/columns';
 
 interface Props {
   rows?: OpenSearchSearchHit[];
@@ -39,7 +40,20 @@ export const DiscoverTable = ({ rows, scrollToTop }: Props) => {
   } = services;
 
   const { refetch$, indexPattern, savedSearch } = useDiscoverContext();
-  const { columns, sort } = useSelector((state) => state.discover);
+  const { columns } = useSelector((state) => {
+    const stateColumns = state.discover.columns;
+    // check if state columns is not undefined, otherwise use buildColumns
+    return {
+      columns: stateColumns !== undefined ? stateColumns : buildColumns([]),
+    };
+  });
+  const { sort } = useSelector((state) => {
+    const stateSort = state.discover.sort;
+    // check if state sort is not undefined, otherwise assign an empty array
+    return {
+      sort: stateSort !== undefined ? stateSort : [],
+    };
+  });
   const dispatch = useDispatch();
   const onAddColumn = (col: string) => {
     if (indexPattern && capabilities.discover?.save) {

src/plugins/discover/public/application/view_components/canvas/index.tsx

@@ -21,6 +21,7 @@ import { useOpenSearchDashboards } from '../../../../../opensearch_dashboards_re
 import { filterColumns } from '../utils/filter_columns';
 import { DEFAULT_COLUMNS_SETTING, MODIFY_COLUMNS_ON_SWITCH } from '../../../../common';
 import { OpenSearchSearchHit } from '../../../application/doc_views/doc_views_types';
+import { buildColumns } from '../../utils/columns';
 import './discover_canvas.scss';
 import { getNewDiscoverSetting, setNewDiscoverSetting } from '../../components/utils/local_storage';
 
@@ -29,9 +30,16 @@ export default function DiscoverCanvas({ setHeaderActionMenu, history }: ViewPro
   const panelRef = useRef<HTMLDivElement>(null);
   const { data$, refetch$, indexPattern } = useDiscoverContext();
   const {
-    services: { uiSettings, storage },
+    services: { uiSettings, storage, capabilities },
   } = useOpenSearchDashboards<DiscoverViewServices>();
-  const { columns } = useSelector((state) => state.discover);
+  const { columns } = useSelector((state) => {
+    const stateColumns = state.discover.columns;
+
+    // check if stateColumns is not undefined, otherwise use buildColumns
+    return {
+      columns: stateColumns !== undefined ? stateColumns : buildColumns([]),
+    };
+  });
   const filteredColumns = filterColumns(
     columns,
     indexPattern,
@@ -97,6 +105,7 @@ export default function DiscoverCanvas({ setHeaderActionMenu, history }: ViewPro
       panelRef.current.scrollTop = 0;
     }
   };
+  const showSaveQuery = !!capabilities.discover?.saveQuery;
 
   const [isOptionsOpen, setOptionsOpen] = useState(false);
   const [useLegacy, setUseLegacy] = useState(!getNewDiscoverSetting(storage));
@@ -159,6 +168,7 @@ export default function DiscoverCanvas({ setHeaderActionMenu, history }: ViewPro
           setHeaderActionMenu,
           onQuerySubmit,
         }}
+        showSaveQuery={showSaveQuery}
       />
       {fetchState.status === ResultStatus.NO_RESULTS && (
         <DiscoverNoResults timeFieldName={timeField} queryLanguage={''} />

src/plugins/discover/public/application/view_components/canvas/top_nav.tsx

@@ -14,18 +14,22 @@ import { IndexPattern } from '../../../opensearch_dashboards_services';
 import { getTopNavLinks } from '../../components/top_nav/get_top_nav_links';
 import { getRootBreadcrumbs } from '../../helpers/breadcrumbs';
 import { useDiscoverContext } from '../context';
+import { useDispatch, setSavedQuery, useSelector } from '../../utils/state_management';
 
 export interface TopNavProps {
   opts: {
     setHeaderActionMenu: AppMountParameters['setHeaderActionMenu'];
     onQuerySubmit: (payload: { dateRange: TimeRange; query?: Query }, isUpdate?: boolean) => void;
   };
+  showSaveQuery: boolean;
 }
 
-export const TopNav = ({ opts }: TopNavProps) => {
+export const TopNav = ({ opts, showSaveQuery }: TopNavProps) => {
   const { services } = useOpenSearchDashboards<DiscoverViewServices>();
   const { inspectorAdapters, savedSearch, indexPattern } = useDiscoverContext();
   const [indexPatterns, setIndexPatterns] = useState<IndexPattern[] | undefined>(undefined);
+  const state = useSelector((s) => s.discover);
+  const dispatch = useDispatch();
 
   const {
     navigation: {
@@ -79,13 +83,17 @@ export const TopNav = ({ opts }: TopNavProps) => {
     indexPattern,
   ]);
 
+  const updateSavedQueryId = (newSavedQueryId: string | undefined) => {
+    dispatch(setSavedQuery(newSavedQueryId));
+  };
+
   return (
     <TopNavMenu
       appName={PLUGIN_ID}
       config={topNavLinks}
       showSearchBar
       showDatePicker={showDatePicker}
-      showSaveQuery
+      showSaveQuery={showSaveQuery}
       useDefaultBehaviors
       setMenuMountPoint={opts.setHeaderActionMenu}
       indexPatterns={indexPattern ? [indexPattern] : indexPatterns}
@@ -94,6 +102,8 @@ export const TopNav = ({ opts }: TopNavProps) => {
       // is ported to main, pass dataSource to TopNavMenu by picking
       // commit 328e08e688c again.
       onQuerySubmit={opts.onQuerySubmit}
+      savedQueryId={state.savedQuery}
+      onSavedQueryIdChange={updateSavedQueryId}
     />
   );
 };

src/plugins/discover/public/application/view_components/context/index.tsx

@@ -10,7 +10,7 @@ import {
   useOpenSearchDashboards,
 } from '../../../../../opensearch_dashboards_react/public';
 import { getServices } from '../../../opensearch_dashboards_services';
-import { useSearch, SearchContextValue } from '../utils/use_search';
+import { useSearch, SearchContextValue, ResultStatus } from '../utils/use_search';
 
 const SearchContext = React.createContext<SearchContextValue>({} as SearchContextValue);
 
@@ -22,6 +22,9 @@ export default function DiscoverContext({ children }: React.PropsWithChildren<Vi
     ...deServices,
     ...services,
   });
+  searchParams.data$.next({
+    status: ResultStatus.LOADING,
+  });
 
   return (
     <OpenSearchDashboardsContextProvider services={services}>

src/plugins/discover/public/application/view_components/panel/index.tsx

@@ -37,9 +37,13 @@ export default function DiscoverPanel(props: ViewProps) {
   const { data$, indexPattern } = useDiscoverContext();
   const [fetchState, setFetchState] = useState<SearchData>(data$.getValue());
 
-  const { columns } = useSelector((state) => ({
-    columns: state.discover.columns,
-  }));
+  const { columns } = useSelector((state) => {
+    const stateColumns = state.discover.columns;
+    // check if state columns is not undefined, otherwise use buildColumns
+    return {
+      columns: stateColumns !== undefined ? stateColumns : buildColumns([]),
+    };
+  });
 
   const prevColumns = useRef(columns);
   const dispatch = useDispatch();
@@ -49,6 +53,7 @@ export default function DiscoverPanel(props: ViewProps) {
     if (columns !== prevColumns.current) {
       let updatedColumns = buildColumns(columns);
       if (
+        columns &&
         timeFieldname &&
         !prevColumns.current.includes(timeFieldname) &&
         columns.includes(timeFieldname)

src/plugins/discover/public/url_generator.ts

@@ -78,6 +78,10 @@ export interface DiscoverUrlGeneratorState {
    * whether to hash the data in the url to avoid url length issues.
    */
   useHash?: boolean;
+  /**
+   * Saved query Id
+   */
+  savedQuery?: string;
 }
 
 interface Params {
@@ -99,12 +103,14 @@ export class DiscoverUrlGenerator
     savedSearchId,
     timeRange,
     useHash = this.params.useHash,
+    savedQuery,
   }: DiscoverUrlGeneratorState): Promise<string> => {
     const savedSearchPath = savedSearchId ? encodeURIComponent(savedSearchId) : '';
     const appState: {
       query?: Query;
       filters?: Filter[];
       index?: string;
+      savedQuery?: string;
     } = {};
     const queryState: QueryState = {};
 
@@ -117,6 +123,7 @@ export class DiscoverUrlGenerator
     if (filters && filters.length)
       queryState.filters = filters?.filter((f) => opensearchFilters.isFilterPinned(f));
     if (refreshInterval) queryState.refreshInterval = refreshInterval;
+    if (savedQuery) appState.savedQuery = savedQuery;
 
     let url = `${this.params.appBasePath}#/${savedSearchPath}`;
     url = setStateToOsdUrl<QueryState>('_g', queryState, { useHash }, url);

src/plugins/vis_builder/public/application/components/top_nav.tsx

@@ -14,11 +14,13 @@ import { VisBuilderServices } from '../../types';
 import './top_nav.scss';
 import { useIndexPatterns, useSavedVisBuilderVis } from '../utils/use';
 import { useTypedSelector, useTypedDispatch } from '../utils/state_management';
+import { setSavedQuery } from '../utils/state_management/visualization_slice';
 import { setEditorState } from '../utils/state_management/metadata_slice';
 import { useCanSave } from '../utils/use/use_can_save';
 import { saveStateToSavedObject } from '../../saved_visualizations/transforms';
 import { TopNavMenuData } from '../../../../navigation/public';
 import { opensearchFilters, connectStorageToQueryState } from '../../../../data/public';
+import { RootState } from '../../../../data_explorer/public';
 
 function useDeepEffect(callback, dependencies) {
   const currentDepsRef = useRef(dependencies);
@@ -40,8 +42,9 @@ export const TopNav = () => {
       ui: { TopNavMenu },
     },
     appName,
+    capabilities,
   } = services;
-  const rootState = useTypedSelector((state) => state);
+  const rootState = useTypedSelector((state: RootState) => state);
   const dispatch = useTypedDispatch();
 
   useDeepEffect(() => {
@@ -93,6 +96,11 @@ export const TopNav = () => {
     dispatch(setEditorState({ state: 'loading' }));
   });
 
+  const updateSavedQueryId = (newSavedQueryId: string | undefined) => {
+    dispatch(setSavedQuery(newSavedQueryId));
+  };
+  const showSaveQuery = !!capabilities['visualization-visbuilder']?.saveQuery;
+
   return (
     <div className="vbTopNav">
       <TopNavMenu
@@ -102,8 +110,10 @@ export const TopNav = () => {
         indexPatterns={indexPattern ? [indexPattern] : []}
         showDatePicker={!!indexPattern?.timeFieldName ?? true}
         showSearchBar
-        showSaveQuery
+        showSaveQuery={showSaveQuery}
         useDefaultBehaviors
+        savedQueryId={rootState.visualization.savedQuery}
+        onSavedQueryIdChange={updateSavedQueryId}
       />
     </div>
   );