[Discover] Allow save query to load correctly in Discover

* add save query logic in Discover
* add save query logic in VisBuilder

Issue Resolve
opensearch-project#5942

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

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,12 @@ export const discoverSlice = createSlice({
         },
       };
     },
+    setSavedQuery(state, action: PayloadAction<string>) {
+      return {
+        ...state,
+        savedQuery: action.payload,
+      }
+    }
   },
 });
 
@@ -201,8 +211,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

@@ -24,6 +24,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,24 @@ export const DiscoverTable = ({ rows }: Props) => {
   } = services;
 
   const { refetch$, indexPattern, savedSearch } = useDiscoverContext();
-  const { columns, sort } = useSelector((state) => state.discover);
+  const { columns } = useSelector(state => {
+    // Extract columns from the state
+    const stateColumns = state.discover.columns;
+
+    // Check if stateColumns is not undefined, otherwise use buildColumns
+    return {
+      columns: stateColumns !== undefined ? stateColumns : buildColumns([])
+    };
+  });
+  const { sort } = useSelector(state => {
+    // Extract columns from the state
+    const stateSort = state.discover.sort;
+
+    // Check if stateColumns is not undefined, otherwise use buildColumns
+    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,15 +20,21 @@ 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 { data$, refetch$, indexPattern } = useDiscoverContext();
-  const {
-    services: { uiSettings },
-  } = useOpenSearchDashboards<DiscoverViewServices>();
-  const { columns } = useSelector((state) => state.discover);
+  const { services: {uiSettings, capabilities} } = useOpenSearchDashboards<DiscoverViewServices>();
+  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,
@@ -89,6 +95,7 @@ export default function DiscoverCanvas({ setHeaderActionMenu, history }: ViewPro
   }, [dispatch, filteredColumns, indexPattern]);
 
   const timeField = indexPattern?.timeFieldName ? indexPattern.timeFieldName : undefined;
+  const showSaveQuery = !!capabilities.discover?.saveQuery;
 
   return (
     <EuiPanel
@@ -103,6 +110,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,23 @@ 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, setState } 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((state) => state.discover);
+  const dispatch = useDispatch();
 
   const {
     navigation: {
@@ -79,17 +84,30 @@ export const TopNav = ({ opts }: TopNavProps) => {
     indexPattern,
   ]);
 
+  const updateSavedQueryId = (newSavedQueryId: string | undefined) => {
+    if (newSavedQueryId) {
+      dispatch(setSavedQuery(newSavedQueryId));
+    } else {
+      // remove savedQueryId from state
+      const newState = state;
+      delete newState.savedQuery;
+      dispatch(setState(newState));
+    }
+  };
+
   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/panel/index.tsx

@@ -35,9 +35,15 @@ 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 => {
+    // Extract columns from the state
+    const stateColumns = state.discover.columns;
+
+    // Check if stateColumns is not undefined, otherwise use buildColumns
+    return {
+      columns: stateColumns !== undefined ? stateColumns : buildColumns([])
+    };
+  });
 
   const prevColumns = useRef(columns);
   const dispatch = useDispatch();
@@ -47,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

@@ -13,6 +13,7 @@ import { VisBuilderServices } from '../../types';
 import './top_nav.scss';
 import { useIndexPatterns, useSavedVisBuilderVis } from '../utils/use';
 import { useTypedSelector, useTypedDispatch } from '../utils/state_management';
+import { setSavedQuery, setState } 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';
@@ -29,6 +30,7 @@ export const TopNav = () => {
       ui: { TopNavMenu },
     },
     appName,
+    capabilities,
   } = services;
   const rootState = useTypedSelector((state) => state);
   const dispatch = useTypedDispatch();
@@ -78,6 +80,18 @@ export const TopNav = () => {
     dispatch(setEditorState({ state: 'loading' }));
   });
 
+  const updateSavedQueryId = (newSavedQueryId: string | undefined) => {
+    if (newSavedQueryId) {
+      dispatch(setSavedQuery(newSavedQueryId));
+    } else {
+      // remove savedQueryId from state
+      const newState = rootState;
+      delete newState.visualization.savedQuery;
+      dispatch(setState(newState.visualization));
+    }
+  };
+  const showSaveQuery=!!capabilities['visualization-visbuilder']?.saveQuery;
+
   return (
     <div className="vbTopNav">
       <TopNavMenu
@@ -87,8 +101,10 @@ export const TopNav = () => {
         indexPatterns={indexPattern ? [indexPattern] : []}
         showDatePicker={!!indexPattern?.timeFieldName ?? true}
         showSearchBar
-        showSaveQuery
+        showSaveQuery={showSaveQuery}
         useDefaultBehaviors
+        savedQueryId={rootState.visualization.savedQuery}
+        onSavedQueryIdChange={updateSavedQueryId}
       />
     </div>
   );

src/plugins/vis_builder/public/application/utils/state_management/visualization_slice.ts

@@ -16,6 +16,7 @@ export interface VisualizationState {
     aggConfigParams: CreateAggConfigParams[];
     draftAgg?: CreateAggConfigParams;
   };
+  savedQuery?: string
 }
 
 const initialState: VisualizationState = {
@@ -115,6 +116,12 @@ export const slice = createSlice({
         [action.payload.paramName]: action.payload.value,
       };
     },
+    setSavedQuery(state, action: PayloadAction<string>) {
+      return {
+        ...state,
+        savedQuery: action.payload,
+      }
+    },
     setState: (_state, action: PayloadAction<VisualizationState>) => {
       return action.payload;
     },
@@ -138,5 +145,6 @@ export const {
   updateAggConfigParams,
   setAggParamValue,
   reorderAgg,
+  setSavedQuery,
   setState,
 } = slice.actions;

src/plugins/vis_builder/public/plugin.ts

@@ -159,6 +159,7 @@ export class VisBuilderPlugin
           embeddable: pluginsStart.embeddable,
           dashboard: pluginsStart.dashboard,
           uiActions: pluginsStart.uiActions,
+          capabilities: coreStart.application.capabilities,
         };
 
         // Instantiate the store

src/plugins/vis_builder/public/types.ts

@@ -17,6 +17,7 @@ import { AppMountParameters, CoreStart, ToastsStart, ScopedHistory } from '../..
 import { IOsdUrlStateStorage } from '../../opensearch_dashboards_utils/public';
 import { DataPublicPluginSetup } from '../../data/public';
 import { UiActionsStart } from '../../ui_actions/public';
+import { Capabilities } from '../../../core/public';
 
 export type VisBuilderSetup = TypeServiceSetup;
 export interface VisBuilderStart extends TypeServiceStart {
@@ -54,6 +55,7 @@ export interface VisBuilderServices extends CoreStart {
   osdUrlStateStorage: IOsdUrlStateStorage;
   dashboard: DashboardStart;
   uiActions: UiActionsStart;
+  capabilities: Capabilities;
 }
 
 export interface ISavedVis {

src/plugins/vis_builder/server/capabilities_provider.ts

@@ -12,6 +12,6 @@ export const capabilitiesProvider = () => ({
     show: true,
     // showWriteControls: true,
     // save: true,
-    // saveQuery: true,
+    saveQuery: true,
   },
 });