Moves ui state to separate slice

Signed-off-by: Ashwin P Chandran <ashwinpc@amazon.com>

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

@@ -122,6 +122,6 @@ const OptionItem = ({ icon, title }: { icon: IconType; title: string }) => (
   </>
 );
 
-// The app uses EuiResizableContainer that triggers a rerender for ever mouseover action.
+// The app uses EuiResizableContainer that triggers a rerender for every mouseover action.
 // To prevent this child component from unnecessarily rerendering in that instance, it needs to be memoized
 export const RightNav = React.memo(RightNavUI);

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

@@ -10,7 +10,7 @@ import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react
 import { IExpressionLoaderParams } from '../../../../expressions/public';
 import { VisBuilderServices } from '../../types';
 import { validateSchemaState, validateAggregations } from '../utils/validations';
-import { useTypedDispatch, useTypedSelector } from '../utils/state_management';
+import { useTypedDispatch, useTypedSelector, setUIStateState } from '../utils/state_management';
 import { useAggs, useVisualizationType } from '../utils/use';
 import { PersistedState } from '../../../../visualizations/public';
 
@@ -20,7 +20,6 @@ import fields_bg from '../../assets/fields_bg.svg';
 import './workspace.scss';
 import { ExperimentalInfo } from './experimental_info';
 import { handleVisEvent } from '../utils/handle_vis_event';
-import { setUiState } from '../utils/state_management/metadata_slice';
 
 export const WorkspaceUI = () => {
   const {
@@ -43,11 +42,11 @@ export const WorkspaceUI = () => {
   const dispatch = useTypedDispatch();
   // Visualizations require the uiState object to persist even when the expression changes
   // eslint-disable-next-line react-hooks/exhaustive-deps
-  const uiState = useMemo(() => new PersistedState(rootState.metadata.uiState), []);
+  const uiState = useMemo(() => new PersistedState(rootState.ui), []);
 
   useEffect(() => {
     if (rootState.metadata.editor.state === 'loaded') {
-      uiState.setSilent(rootState.metadata.uiState);
+      uiState.setSilent(rootState.ui);
     }
     // To update uiState once saved object data is loaded
     // eslint-disable-next-line react-hooks/exhaustive-deps
@@ -56,7 +55,7 @@ export const WorkspaceUI = () => {
   useEffect(() => {
     uiState.on('change', (args) => {
       // Store changes to UI state
-      dispatch(setUiState({ state: uiState.toJSON() }));
+      dispatch(setUIStateState(uiState.toJSON()));
     });
   }, [dispatch, uiState]);
 
@@ -151,6 +150,6 @@ export const WorkspaceUI = () => {
   );
 };
 
-// The app uses EuiResizableContainer that triggers a rerender for ever mouseover action.
+// The app uses EuiResizableContainer that triggers a rerender for every mouseover action.
 // To prevent this child component from unnecessarily rerendering in that instance, it needs to be memoized
 export const Workspace = React.memo(WorkspaceUI);

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

@@ -22,7 +22,6 @@ export interface MetadataState {
     state: EditorState;
   };
   originatingApp?: string;
-  uiState: any;
 }
 
 const initialState: MetadataState = {
@@ -31,7 +30,6 @@ const initialState: MetadataState = {
     state: 'loading',
   },
   originatingApp: undefined,
-  uiState: {},
 };
 
 export const getPreloadedState = async ({
@@ -63,15 +61,11 @@ export const slice = createSlice({
     setOriginatingApp: (state, action: PayloadAction<{ state?: string }>) => {
       state.originatingApp = action.payload.state;
     },
-    setUiState: (state, action: PayloadAction<{ state: any; editorState?: EditorState }>) => {
-      state.uiState = action.payload.state;
-      state.editor.state = action.payload.editorState || 'dirty';
-    },
     setState: (_state, action: PayloadAction<MetadataState>) => {
       return action.payload;
     },
   },
 });
 
 export const { reducer } = slice;
-export const { setError, setEditorState, setOriginatingApp, setState, setUiState } = slice.actions;
+export const { setError, setEditorState, setOriginatingApp, setState } = slice.actions;

src/plugins/vis_builder/public/application/utils/state_management/redux_persistence.test.tsx

@@ -18,6 +18,7 @@ describe('test redux state persistence', () => {
       style: 'style',
       visualization: 'visualization',
       metadata: 'metadata',
+      ui: 'ui',
     };
   });
 
@@ -32,8 +33,8 @@ describe('test redux state persistence', () => {
       metadata: {
         editor: { errors: {}, state: 'loading' },
         originatingApp: undefined,
-        uiState: {},
       },
+      ui: {},
     };
 
     const returnStates = await loadReduxState(mockServices);

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

@@ -21,13 +21,13 @@ export const loadReduxState = async (services: VisBuilderServices) => {
 };
 
 export const persistReduxState = (
-  { style, visualization, metadata },
+  { style, visualization, metadata, ui }: RootState,
   services: VisBuilderServices
 ) => {
   try {
     services.osdUrlStateStorage.set<RootState>(
       '_a',
-      { style, visualization, metadata },
+      { style, visualization, metadata, ui },
       {
         replace: true,
       }

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

@@ -8,12 +8,14 @@ import { isEqual } from 'lodash';
 import { reducer as styleReducer } from './style_slice';
 import { reducer as visualizationReducer } from './visualization_slice';
 import { reducer as metadataReducer } from './metadata_slice';
+import { reducer as uiStateReducer } from './ui_state_slice';
 import { VisBuilderServices } from '../../..';
 import { loadReduxState, persistReduxState } from './redux_persistence';
 import { handlerEditorState } from './handlers/editor_state';
 import { handlerParentAggs } from './handlers/parent_aggs';
 
 const rootReducer = combineReducers({
+  ui: uiStateReducer,
   style: styleReducer,
   visualization: visualizationReducer,
   metadata: metadataReducer,
@@ -61,3 +63,4 @@ export type AppDispatch = Store['dispatch'];
 export { setState as setStyleState, StyleState } from './style_slice';
 export { setState as setVisualizationState, VisualizationState } from './visualization_slice';
 export { MetadataState } from './metadata_slice';
+export { setState as setUIStateState, UIStateState } from './ui_state_slice';

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

@@ -0,0 +1,34 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { createSlice, PayloadAction } from '@reduxjs/toolkit';
+
+export type UIStateState<T = any> = T;
+
+const initialState = {} as UIStateState;
+
+export const uiStateSlice = createSlice({
+  name: 'ui',
+  initialState,
+  reducers: {
+    setState<T>(state: T, action: PayloadAction<UIStateState<T>>) {
+      return action.payload;
+    },
+    updateState<T>(state: T, action: PayloadAction<Partial<UIStateState<T>>>) {
+      state = {
+        ...state,
+        ...action.payload,
+      };
+    },
+  },
+});
+
+// Exposing the state functions as generics
+export const setState = uiStateSlice.actions.setState as <T>(payload: T) => PayloadAction<T>;
+export const updateState = uiStateSlice.actions.updateState as <T>(
+  payload: Partial<T>
+) => PayloadAction<Partial<T>>;
+
+export const { reducer } = uiStateSlice;

src/plugins/vis_builder/public/application/utils/use/use_saved_vis_builder_vis.ts

@@ -14,9 +14,14 @@ import {
 import { EDIT_PATH, PLUGIN_ID } from '../../../../common';
 import { VisBuilderServices } from '../../../types';
 import { getCreateBreadcrumbs, getEditBreadcrumbs } from '../breadcrumbs';
-import { useTypedDispatch, setStyleState, setVisualizationState } from '../state_management';
+import {
+  useTypedDispatch,
+  setStyleState,
+  setVisualizationState,
+  setUIStateState,
+} from '../state_management';
 import { useOpenSearchDashboards } from '../../../../../opensearch_dashboards_react/public';
-import { setEditorState, setUiState } from '../state_management/metadata_slice';
+import { setEditorState } from '../state_management/metadata_slice';
 import { getStateFromSavedObject } from '../../../saved_visualizations/transforms';
 
 // This function can be used when instantiating a saved vis or creating a new one
@@ -53,13 +58,14 @@ export const useSavedVisBuilderVis = (visualizationIdFromUrl: string | undefined
         );
 
         if (savedVisBuilderVis.id) {
-          const { title, state, uiState } = getStateFromSavedObject(savedVisBuilderVis);
+          const { title, state } = getStateFromSavedObject(savedVisBuilderVis);
           chrome.setBreadcrumbs(getEditBreadcrumbs(title, navigateToApp));
           chrome.docTitle.change(title);
 
-          dispatch(setUiState({ state: uiState, editorState: 'loaded' }));
+          dispatch(setUIStateState(state.ui));
           dispatch(setStyleState(state.style));
           dispatch(setVisualizationState(state.visualization));
+          dispatch(setEditorState({ state: 'loaded' }));
         } else {
           chrome.setBreadcrumbs(getCreateBreadcrumbs(navigateToApp));
         }

src/plugins/vis_builder/public/embeddable/vis_builder_embeddable.tsx

@@ -116,7 +116,7 @@ export class VisBuilderEmbeddable extends Embeddable<VisBuilderInput, VisBuilder
 
     this.deps = deps;
     this.savedVis = savedVis;
-    this.uiState = new PersistedState(savedVis.uiState);
+    this.uiState = new PersistedState(savedVis.state.ui);
     this.uiState.on('change', this.uiStateChangeHandler);
     this.uiState.on('reload', this.reload);
 

src/plugins/vis_builder/public/saved_visualizations/transforms.test.ts

@@ -24,7 +24,7 @@ describe('transforms', () => {
       TEST_INDEX_PATTERN_ID = 'test-pattern';
       savedObject = {} as VisBuilderSavedObject;
       rootState = {
-        metadata: { editor: { state: 'loading', errors: {} }, uiState: {} },
+        metadata: { editor: { state: 'loading', errors: {} } },
         style: {},
         visualization: {
           searchField: '',
@@ -34,6 +34,7 @@ describe('transforms', () => {
             aggConfigParams: [],
           },
         },
+        ui: {},
       };
       indexPattern = getStubIndexPattern(
         TEST_INDEX_PATTERN_ID,
@@ -49,7 +50,7 @@ describe('transforms', () => {
 
       expect(savedObject.visualizationState).not.toContain(TEST_INDEX_PATTERN_ID);
       expect(savedObject.styleState).toEqual(JSON.stringify(rootState.style));
-      expect(savedObject.uiState).toEqual(JSON.stringify(rootState.metadata.uiState));
+      expect(savedObject.uiState).toEqual(JSON.stringify(rootState.ui));
       expect(savedObject.searchSourceFields?.index?.id).toEqual(TEST_INDEX_PATTERN_ID);
     });
 
@@ -77,7 +78,7 @@ describe('transforms', () => {
     } as VisBuilderSavedObjectAttributes;
 
     test('should return saved object with state', () => {
-      const { state, uiState } = getStateFromSavedObject(defaultVBSaveObj);
+      const { state } = getStateFromSavedObject(defaultVBSaveObj);
 
       expect(state).toMatchInlineSnapshot(`
         Object {
@@ -86,10 +87,9 @@ describe('transforms', () => {
             "indexPattern": "test-index",
             "searchField": "",
           },
+          "ui": Object {},
         }
       `);
-
-      expect(uiState).toMatchInlineSnapshot(`"{}"`);
     });
 
     test('should throw error if state is invalid', () => {

src/plugins/vis_builder/public/saved_visualizations/transforms.ts

@@ -27,15 +27,14 @@ export const saveStateToSavedObject = (
   );
   obj.styleState = JSON.stringify(state.style);
   obj.searchSourceFields = { index: indexPattern };
-  obj.uiState = JSON.stringify(state.metadata.uiState);
+  obj.uiState = JSON.stringify(state.ui);
 
   return obj;
 };
 
 export interface VisBuilderSavedVis
   extends Pick<VisBuilderSavedObjectAttributes, 'id' | 'title' | 'description'> {
   state: RenderState;
-  uiState: any;
 }
 
 export const getStateFromSavedObject = (
@@ -51,7 +50,7 @@ export const getStateFromSavedObject = (
     indexPattern: obj.searchSourceFields?.index,
   };
 
-  const validateResult = validateVisBuilderState({ styleState, visualizationState });
+  const validateResult = validateVisBuilderState({ styleState, visualizationState, uiState });
 
   if (!validateResult.valid) {
     throw new InvalidJSONProperty(
@@ -78,7 +77,7 @@ export const getStateFromSavedObject = (
     state: {
       visualization: visualizationState,
       style: styleState,
+      ui: uiState,
     },
-    uiState,
   };
 };