[Discover] Restore Unified Search UI state on tab change (query, time range)

src/platform/plugins/shared/discover/public/application/main/components/top_nav/discover_topnav.tsx

@@ -9,7 +9,7 @@
 
 import React, { useCallback, useEffect, useMemo, useRef } from 'react';
 import { DataViewType } from '@kbn/data-views-plugin/public';
-import type { DataViewPickerProps } from '@kbn/unified-search-plugin/public';
+import type { DataViewPickerProps, UnifiedSearchDraft } from '@kbn/unified-search-plugin/public';
 import { DiscoverFlyouts, dismissAllFlyoutsExceptFor } from '@kbn/discover-utils';
 import type { EuiHeaderLinksProps } from '@elastic/eui';
 import { useSavedSearchInitial } from '../../state_management/discover_state_provider';
@@ -25,10 +25,13 @@ import { ESQLToDataViewTransitionModal } from './esql_dataview_transition';
 import {
   internalStateActions,
   useCurrentDataView,
+  useCurrentTabAction,
+  useCurrentTabSelector,
   useDataViewsForPicker,
   useInternalStateDispatch,
   useInternalStateSelector,
 } from '../../state_management/redux';
+import { TABS_ENABLED } from '../../../../constants';
 
 export interface DiscoverTopNavProps {
   savedQuery?: string;
@@ -211,6 +214,19 @@ export const DiscoverTopNav = ({
     [searchBarCustomization?.CustomSearchBar, navigation.ui.AggregateQueryTopNavMenu]
   );
 
+  const searchDraftUiState = useCurrentTabSelector((state) => state.uiState.searchDraft);
+  const setSearchDraftUiState = useCurrentTabAction(internalStateActions.setSearchDraftUiState);
+  const onDraftChange = useCallback(
+    (newSearchDraftUiState: UnifiedSearchDraft | undefined) => {
+      dispatch(
+        setSearchDraftUiState({
+          searchDraftUiState: newSearchDraftUiState,
+        })
+      );
+    },
+    [dispatch, setSearchDraftUiState]
+  );
+
   const shouldHideDefaultDataviewPicker =
     !!searchBarCustomization?.CustomDataViewPicker || !!searchBarCustomization?.hideDataViewPicker;
 
@@ -248,6 +264,8 @@ export const DiscoverTopNav = ({
           ) : undefined
         }
         onESQLDocsFlyoutVisibilityChanged={onESQLDocsFlyoutVisibilityChanged}
+        draft={searchDraftUiState}
+        onDraftChange={TABS_ENABLED ? onDraftChange : undefined}
       />
       {isESQLToDataViewTransitionModalVisible && (
         <ESQLToDataViewTransitionModal onClose={onESQLToDataViewTransitionModalClose} />

src/platform/plugins/shared/discover/public/application/main/state_management/redux/internal_state.ts

@@ -245,6 +245,14 @@ export const internalStateSlice = createSlice({
       withTab(state, action, (tab) => {
         tab.uiState.layout = action.payload.layoutUiState;
       }),
+
+    setSearchDraftUiState: (
+      state,
+      action: TabAction<{ searchDraftUiState: Partial<TabState['uiState']['searchDraft']> }>
+    ) =>
+      withTab(state, action, (tab) => {
+        tab.uiState.searchDraft = action.payload.searchDraftUiState;
+      }),
   },
   extraReducers: (builder) => {
     builder.addCase(loadDataViewList.fulfilled, (state, action) => {

src/platform/plugins/shared/discover/public/application/main/state_management/redux/types.ts

@@ -13,6 +13,7 @@ import type { DataTableRecord } from '@kbn/discover-utils';
 import type { Filter, TimeRange } from '@kbn/es-query';
 import type { UnifiedDataTableRestorableState } from '@kbn/unified-data-table';
 import type { UnifiedFieldListRestorableState } from '@kbn/unified-field-list';
+import type { UnifiedSearchDraft } from '@kbn/unified-search-plugin/public';
 import type { UnifiedHistogramVisContext } from '@kbn/unified-histogram';
 import type { TabItem } from '@kbn/unified-tabs';
 import type { DiscoverAppState } from '../discover_app_state_container';
@@ -80,6 +81,7 @@ export interface TabState extends TabItem {
     dataGrid?: Partial<UnifiedDataTableRestorableState>;
     fieldList?: Partial<UnifiedFieldListRestorableState>;
     layout?: Partial<DiscoverLayoutRestorableState>;
+    searchDraft?: Partial<UnifiedSearchDraft>;
   };
 }
 

src/platform/plugins/shared/unified_search/public/index.ts

@@ -25,6 +25,7 @@ export type {
   UnifiedSearchPublicPluginStart,
   UnifiedSearchPluginSetup,
   IUnifiedSearchPluginServices,
+  UnifiedSearchDraft,
 } from './types';
 export type { FilterItemsProps } from './filter_bar';
 export type { DataViewPickerProps } from './dataview_picker';

src/platform/plugins/shared/unified_search/public/query_string_input/query_bar_top_row.test.tsx

@@ -404,6 +404,70 @@ describe('QueryBarTopRowTopRow', () => {
 
     expect(component.find(CANCEL_BUTTON_SELECTOR).length).toBe(0);
   });
+
+  describe('draft', () => {
+    it('should call onDraftChange when in dirty state', () => {
+      const onDraftChange = jest.fn();
+      const state = {
+        query: kqlQuery,
+        dateRangeFrom: 'now-7d',
+        dateRangeTo: 'now',
+      };
+      const { getByText } = render(
+        wrapQueryBarTopRowInContext({
+          isDirty: true,
+          onDraftChange,
+          ...state,
+        })
+      );
+
+      expect(getByText(kqlQuery.query)).toBeInTheDocument();
+      expect(onDraftChange).toHaveBeenCalledWith(state);
+    });
+
+    it('should call onDraftChange when in dirty state and no date picker', () => {
+      const onDraftChange = jest.fn();
+      const state = {
+        query: kqlQuery,
+        dateRangeFrom: 'now-7d',
+        dateRangeTo: 'now',
+      };
+      const { getByText } = render(
+        wrapQueryBarTopRowInContext({
+          isDirty: true,
+          showDatePicker: false,
+          onDraftChange,
+          ...state,
+        })
+      );
+
+      expect(getByText(kqlQuery.query)).toBeInTheDocument();
+      expect(onDraftChange).toHaveBeenCalledWith({
+        query: state.query,
+        dateRangeFrom: undefined,
+        dateRangeTo: undefined,
+      });
+    });
+
+    it('should call onDraftChange with empty draft when in normal state', () => {
+      const onDraftChange = jest.fn();
+      const state = {
+        query: kqlQuery,
+        dateRangeFrom: 'now-7d',
+        dateRangeTo: 'now',
+      };
+      const { getByText } = render(
+        wrapQueryBarTopRowInContext({
+          isDirty: false,
+          onDraftChange,
+          ...state,
+        })
+      );
+
+      expect(getByText(kqlQuery.query)).toBeInTheDocument();
+      expect(onDraftChange).toHaveBeenCalledWith(undefined);
+    });
+  });
 });
 
 describe('SharingMetaFields', () => {

src/platform/plugins/shared/unified_search/public/query_string_input/query_bar_top_row.tsx

@@ -46,7 +46,7 @@ import type { DataView } from '@kbn/data-views-plugin/public';
 import type { PersistedLog } from '@kbn/data-plugin/public';
 import { useKibana } from '@kbn/kibana-react-plugin/public';
 import { UI_SETTINGS } from '@kbn/data-plugin/common';
-import type { IUnifiedSearchPluginServices } from '../types';
+import type { IUnifiedSearchPluginServices, UnifiedSearchDraft } from '../types';
 import { QueryStringInput } from './query_string_input';
 import { NoDataPopover } from './no_data_popover';
 import { shallowEqual } from '../utils/shallow_equal';
@@ -146,6 +146,7 @@ export interface QueryBarTopRowProps<QT extends Query | AggregateQuery = Query>
   onRefreshChange?: (options: { isPaused: boolean; refreshInterval: number }) => void;
   onSubmit: (payload: { dateRange: TimeRange; query?: Query | QT }) => void;
   onCancel?: () => void;
+  onDraftChange?: (draft: UnifiedSearchDraft | undefined) => void;
   placeholder?: string;
   prepend?: React.ComponentProps<typeof EuiFieldText>['prepend'];
   query?: Query | QT;
@@ -445,6 +446,36 @@ export const QueryBarTopRow = React.memo(
       [onSubmit]
     );
 
+    const {
+      onDraftChange,
+      isDirty: draftIsDirty,
+      query: draftQuery,
+      dateRangeFrom: draftDateRangeFrom,
+      dateRangeTo: draftDateRangeTo,
+    } = props;
+    const draft = useMemo(
+      () =>
+        onDraftChange && draftIsDirty
+          ? {
+              query: draftQuery,
+              dateRangeFrom: showDatePicker ? draftDateRangeFrom : undefined,
+              dateRangeTo: showDatePicker ? draftDateRangeTo : undefined,
+            }
+          : undefined,
+      [
+        onDraftChange,
+        draftIsDirty,
+        draftQuery,
+        draftDateRangeFrom,
+        draftDateRangeTo,
+        showDatePicker,
+      ]
+    );
+
+    useEffect(() => {
+      onDraftChange?.(draft);
+    }, [onDraftChange, draft]);
+
     function shouldRenderQueryInput(): boolean {
       return Boolean(showQueryInput && props.query && storage);
     }

src/platform/plugins/shared/unified_search/public/search_bar/create_search_bar.tsx

@@ -245,6 +245,8 @@ export function createSearchBar({
             onCancel={props.onCancel}
             filters={filters}
             query={query}
+            draft={props.draft}
+            onDraftChange={props.onDraftChange}
             onFiltersUpdated={defaultFiltersUpdated(data.query, props.onFiltersUpdated)}
             onRefreshChange={
               !props.isAutoRefreshDisabled

src/platform/plugins/shared/unified_search/public/search_bar/search_bar.test.tsx

@@ -409,4 +409,50 @@ describe('SearchBar', () => {
       });
     });
   });
+
+  describe('draft', () => {
+    it('should prefill with the draft query if provided', () => {
+      const draft = {
+        query: { language: 'kuery', query: 'test_draft' },
+        dateRangeFrom: 'now-30m',
+        dateRangeTo: 'now-10m',
+      };
+      const onDraftChange = jest.fn();
+      const component = mount(
+        wrapSearchBarInContext({
+          indexPatterns: [stubIndexPattern],
+          query: kqlQuery,
+          dateRangeTo: 'now',
+          dateRangeFrom: 'now-15m',
+          draft,
+          onDraftChange,
+        })
+      );
+
+      expect(onDraftChange).toHaveBeenCalledWith(draft);
+      expect(component.find('textarea').prop('value')).toEqual(draft.query.query);
+    });
+
+    it('should check for query type mismatch', () => {
+      const draft = {
+        query: esqlQuery,
+        dateRangeFrom: 'now-30m',
+        dateRangeTo: 'now-10m',
+      };
+      const onDraftChange = jest.fn();
+      const component = mount(
+        wrapSearchBarInContext({
+          indexPatterns: [stubIndexPattern],
+          query: kqlQuery,
+          dateRangeTo: 'now',
+          dateRangeFrom: 'now-15m',
+          draft,
+          onDraftChange,
+        })
+      );
+
+      expect(onDraftChange).toHaveBeenCalledWith(undefined);
+      expect(component.find('textarea').prop('value')).toEqual(kqlQuery.query);
+    });
+  });
 });

src/platform/plugins/shared/unified_search/public/search_bar/search_bar.tsx

@@ -36,7 +36,7 @@ import { DataView } from '@kbn/data-views-plugin/public';
 
 import { i18n } from '@kbn/i18n';
 import { AdditionalQueryBarMenuItems } from '../query_string_input/query_bar_menu_panels';
-import type { IUnifiedSearchPluginServices } from '../types';
+import type { IUnifiedSearchPluginServices, UnifiedSearchDraft } from '../types';
 import { SavedQueryMeta, SaveQueryForm } from '../saved_query_form';
 import { SavedQueryManagementList } from '../saved_query_management';
 import { QueryBarMenu, QueryBarMenuProps } from '../query_string_input/query_bar_menu';
@@ -96,6 +96,9 @@ export interface SearchBarOwnProps<QT extends AggregateQuery | Query = Query> {
     payload: { dateRange: TimeRange; query?: QT | Query },
     isUpdate?: boolean
   ) => void;
+  // To initialize with a predefined query which has not been submitted yet (in dirty state)
+  draft?: UnifiedSearchDraft;
+  onDraftChange?: QueryBarTopRowProps<QT>['onDraftChange'];
   // User has saved the current state as a saved query
   onSaved?: (savedQuery: SavedQuery) => void;
   // User has modified the saved query, your app should persist the update
@@ -251,22 +254,46 @@ export class SearchBarUI<QT extends (Query | AggregateQuery) | Query = Query> ex
     return nextState;
   }
 
+  private prefillWithInitialDraftState = (state: SearchBarState<QT>): SearchBarState<QT> => {
+    const { draft, query } = this.props;
+
+    if (!draft) {
+      return state;
+    }
+
+    if (
+      query &&
+      draft.query &&
+      isOfAggregateQueryType(query) !== isOfAggregateQueryType(draft.query)
+    ) {
+      // safeguard against query type mismatch
+      return state;
+    }
+
+    return {
+      ...state,
+      query: draft.query ? ({ ...draft.query } as SearchBarState<QT>['query']) : state.query,
+      dateRangeFrom: draft.dateRangeFrom || state.dateRangeFrom,
+      dateRangeTo: draft.dateRangeTo || state.dateRangeTo,
+    };
+  };
+
   /*
    Keep the "draft" value in local state until the user actually submits the query. There are a couple advantages:
 
    Each app doesn't have to maintain its own "draft" value if it wants to put off updating the query in app state
    until the user manually submits their changes. Some apps have watches on the query value in app state so we don't
    want to trigger those on every keypress.
   */
-  public state = {
+  public state = this.prefillWithInitialDraftState({
     isFiltersVisible: true,
     openQueryBarMenu: false,
     showSavedQueryPopover: false,
     currentProps: this.props,
     query: this.props.query ? { ...this.props.query } : undefined,
     dateRangeFrom: get(this.props, 'dateRangeFrom', 'now-15m'),
     dateRangeTo: get(this.props, 'dateRangeTo', 'now'),
-  } as SearchBarState<QT>;
+  } as SearchBarState<QT>);
 
   public isDirty = () => {
     if (!this.props.showDatePicker && this.state.query && this.props.query) {
@@ -636,6 +663,7 @@ export class SearchBarUI<QT extends (Query | AggregateQuery) | Query = Query> ex
           onRefreshChange={this.props.onRefreshChange}
           onCancel={this.props.onCancel}
           onChange={this.onQueryBarChange}
+          onDraftChange={this.props.onDraftChange}
           isDirty={this.isDirty()}
           customSubmitButton={
             this.props.customSubmitButton ? this.props.customSubmitButton : undefined

src/platform/plugins/shared/unified_search/public/types.ts

@@ -23,6 +23,12 @@ import type { IndexPatternSelectProps, QueryStringInputProps, StatefulSearchBarP
 import type { FiltersBuilderProps } from './filters_builder/filters_builder';
 import { StatefulSearchBarDeps } from './search_bar/create_search_bar';
 
+export interface UnifiedSearchDraft {
+  query?: AggregateQuery | Query;
+  dateRangeFrom?: string;
+  dateRangeTo?: string;
+}
+
 export interface UnifiedSearchSetupDependencies {
   uiActions: UiActionsSetup;
   data: DataPublicPluginStart;