[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
opensearch-project#5942

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

CHANGELOG.md

@@ -42,6 +42,8 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - [Workspace] Add workspace id in basePath ([#6060](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6060))
 
 ### 🐛 Bug Fixes
+- [BUG][Discover] Allow saved sort from search embeddable to load in Dashboard ([#5934](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/5934))
+- [BUG][Discover] Allow save query to load correctly in Discover ([#5951](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/5951))
 
 - [BUG][Discover] Allow saved sort from search embeddable to load in Dashboard ([#5934](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/5934))
 - [BUG][Discover] Add key to index pattern options for support deplicate index pattern names([#5946](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/5946))

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 string 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
@@ -110,6 +111,9 @@ export const discoverSlice = createSlice({
     setState(state, action: PayloadAction<DiscoverState>) {
       return action.payload;
     },
+    getState(state, action: PayloadAction<DiscoverState>) {
+      return state;
+    },
     addColumn(state, action: PayloadAction<{ column: string; index?: number }>) {
       const columns = utils.addColumn(state.columns || [], action.payload);
       return { ...state, columns: buildColumns(columns) };
@@ -188,6 +192,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,
+        };
+      }
+    },
   },
 });
 
@@ -201,8 +217,10 @@ export const {
   setSort,
   setInterval,
   setState,
+  getState,
   updateState,
   setSavedSearchId,
   setMetadata,
+  setSavedQuery,
 } = discoverSlice.actions;
 export const { reducer } = discoverSlice;

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

@@ -23,6 +23,7 @@ import { SortOrder } from '../../../saved_searches/types';
 import { DOC_HIDE_TIME_COLUMN_SETTING } from '../../../../common';
 import { OpenSearchSearchHit } from '../../doc_views/doc_views_types';
 import { popularizeField } from '../../helpers/popularize_field';
+import { buildColumns } from '../../utils/columns';
 
 interface Props {
   rows?: OpenSearchSearchHit[];
@@ -41,7 +42,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

@@ -20,16 +20,24 @@ 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';
 
 // eslint-disable-next-line import/no-default-export
 export default function DiscoverCanvas({ setHeaderActionMenu, history }: ViewProps) {
   const panelRef = useRef<HTMLDivElement>(null);
   const { data$, refetch$, indexPattern } = useDiscoverContext();
   const {
-    services: { uiSettings },
+    services: { uiSettings, 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,
@@ -95,6 +103,7 @@ export default function DiscoverCanvas({ setHeaderActionMenu, history }: ViewPro
       panelRef.current.scrollTop = 0;
     }
   };
+  const showSaveQuery = !!capabilities.discover?.saveQuery;
 
   return (
     <EuiPanel
@@ -110,6 +119,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 { getTopNavLinks } from '../../components/top_nav/get_top_nav_links';
 import { useDiscoverContext } from '../context';
 import { getRootBreadcrumbs } from '../../helpers/breadcrumbs';
 import { opensearchFilters, connectStorageToQueryState } from '../../../../../data/public';
+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: {
@@ -36,16 +40,10 @@ export const TopNav = ({ opts }: TopNavProps) => {
     },
     data,
     chrome,
-    osdUrlStateStorage,
   } = services;
 
   const topNavLinks = savedSearch ? getTopNavLinks(services, inspectorAdapters, savedSearch) : [];
 
-  connectStorageToQueryState(services.data.query, osdUrlStateStorage, {
-    filters: opensearchFilters.FilterStateStore.APP_STATE,
-    query: true,
-  });
-
   useEffect(() => {
     let isMounted = true;
     const getDefaultIndexPattern = async () => {
@@ -79,17 +77,23 @@ 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}
       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

@@ -35,9 +35,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();
@@ -47,6 +51,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

@@ -13,11 +13,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';
 
 export const TopNav = () => {
   // id will only be set for the edit route
@@ -29,8 +31,9 @@ export const TopNav = () => {
       ui: { TopNavMenu },
     },
     appName,
+    capabilities,
   } = services;
-  const rootState = useTypedSelector((state) => state);
+  const rootState = useTypedSelector((state: RootState) => state);
   const dispatch = useTypedDispatch();
 
   const saveDisabledReason = useCanSave();
@@ -78,6 +81,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
@@ -87,8 +95,10 @@ export const TopNav = () => {
         indexPatterns={indexPattern ? [indexPattern] : []}
         showDatePicker={!!indexPattern?.timeFieldName ?? true}
         showSearchBar
-        showSaveQuery
+        showSaveQuery={showSaveQuery}
         useDefaultBehaviors
+        savedQueryId={rootState.visualization.savedQuery}
+        onSavedQueryIdChange={updateSavedQueryId}
       />
     </div>
   );