[VisBuilder] Adds UIState to vis, adds index patterns to embeddable, …

…bug fixes (#3751)

* adds uiActions to visBuilder
* prevents multiple errors on load
* fixes visbuilder type errors
* fixes save
* adds ui state to vis builder
* fixes tests
* Adds support in embeddables for multiple indices
* Moves ui state to separate slice
* Adds changelog

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

---------

Signed-off-by: Ashwin P Chandran <ashwinpc@amazon.com>
Co-authored-by: Josh Romero <rmerqg@amazon.com>
(cherry picked from commit ee32d20)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

# Conflicts:
#	CHANGELOG.md

src/plugins/vis_builder/common/vis_builder_saved_object_attributes.ts

@@ -13,6 +13,7 @@ export interface VisBuilderSavedObjectAttributes extends SavedObjectAttributes {
   visualizationState?: string;
   updated_at?: string;
   styleState?: string;
+  uiState?: string;
   version: number;
   searchSourceFields?: {
     index?: string;

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 { useTypedSelector } from '../utils/state_management';
+import { useTypedDispatch, useTypedSelector, setUIStateState } from '../utils/state_management';
 import { useAggs, useVisualizationType } from '../utils/use';
 import { PersistedState } from '../../../../visualizations/public';
 
@@ -39,8 +39,25 @@ export const WorkspaceUI = () => {
     timeRange: data.query.timefilter.timefilter.getTime(),
   });
   const rootState = useTypedSelector((state) => state);
-  // Visualizations require the uiState to persist even when the expression changes
-  const uiState = useMemo(() => new PersistedState(), []);
+  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.ui), []);
+
+  useEffect(() => {
+    if (rootState.metadata.editor.state === 'loaded') {
+      uiState.setSilent(rootState.ui);
+    }
+    // To update uiState once saved object data is loaded
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [rootState.metadata.editor.state, uiState]);
+
+  useEffect(() => {
+    uiState.on('change', (args) => {
+      // Store changes to UI state
+      dispatch(setUIStateState(uiState.toJSON()));
+    });
+  }, [dispatch, uiState]);
 
   useEffect(() => {
     async function loadExpression() {
@@ -133,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/schema.json

@@ -1,28 +1,47 @@
 {
-    "type": "object",
-    "properties": {
-      "styleState": {
-        "type": "object"
-      },
-      "visualizationState": {
-        "type": "object",
-        "properties": {
-          "activeVisualization": {
-            "type": "object",
-            "properties": {
-              "name": { "type": "string" },
-              "aggConfigParams": { "type": "array" }
+  "type": "object",
+  "properties": {
+    "styleState": {
+      "type": "object"
+    },
+    "visualizationState": {
+      "type": "object",
+      "properties": {
+        "activeVisualization": {
+          "type": "object",
+          "properties": {
+            "name": {
+              "type": "string"
             },
-            "required": ["name", "aggConfigParams"],
-            "additionalProperties": false
+            "aggConfigParams": {
+              "type": "array"
+            }
           },
-          "indexPattern": { "type": "string" },
-          "searchField": { "type": "string" }
+          "required": [
+            "name",
+            "aggConfigParams"
+          ],
+          "additionalProperties": false
         },
-        "required": ["searchField"],
-        "additionalProperties": false
-      }
+        "indexPattern": {
+          "type": "string"
+        },
+        "searchField": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "searchField"
+      ],
+      "additionalProperties": false
     },
-    "required": ["styleState", "visualizationState"],
-    "additionalProperties": false
+    "uiState": {
+      "type": "object"
+    }
+  },
+  "required": [
+    "styleState",
+    "visualizationState"
+  ],
+  "additionalProperties": false
 }

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

@@ -11,7 +11,7 @@ import { VisBuilderServices } from '../../../types';
  * Clean state: when viz finished loading and ready to be edited
  * Dirty state: when there are changes applied to the viz after it finished loading
  */
-type EditorState = 'loading' | 'clean' | 'dirty';
+type EditorState = 'loading' | 'loaded' | 'clean' | 'dirty';
 
 export interface MetadataState {
   editor: {

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

@@ -8,6 +8,7 @@ import { VisBuilderServices } from '../../..';
 import { getPreloadedState as getPreloadedStyleState } from './style_slice';
 import { getPreloadedState as getPreloadedVisualizationState } from './visualization_slice';
 import { getPreloadedState as getPreloadedMetadataState } from './metadata_slice';
+import { getPreloadedState as getPreloadedUIState } from './ui_state_slice';
 import { RootState } from './store';
 
 export const getPreloadedState = async (
@@ -16,10 +17,12 @@ export const getPreloadedState = async (
   const styleState = await getPreloadedStyleState(services);
   const visualizationState = await getPreloadedVisualizationState(services);
   const metadataState = await getPreloadedMetadataState(services);
+  const uiState = await getPreloadedUIState(services);
 
   return {
     style: styleState,
     visualization: visualizationState,
     metadata: metadataState,
+    ui: uiState,
   };
 };

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',
     };
   });
 
@@ -33,6 +34,7 @@ describe('test redux state persistence', () => {
         editor: { errors: {}, state: 'loading' },
         originatingApp: undefined,
       },
+      ui: {},
     };
 
     const returnStates = await loadReduxState(mockServices);

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

@@ -12,22 +12,21 @@ export const loadReduxState = async (services: VisBuilderServices) => {
     const serializedState = services.osdUrlStateStorage.get<RootState>('_a');
     if (serializedState !== null) return serializedState;
   } catch (err) {
-    /* eslint-disable no-console */
+    // eslint-disable-next-line no-console
     console.error(err);
-    /* eslint-enable no-console */
   }
 
   return await getPreloadedState(services);
 };
 
 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,
@@ -60,3 +62,5 @@ 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,42 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { createSlice, PayloadAction } from '@reduxjs/toolkit';
+import { VisBuilderServices } from '../../../types';
+
+export type UIStateState<T = any> = T;
+
+const initialState = {} as UIStateState;
+
+export const getPreloadedState = async ({
+  types,
+  data,
+}: VisBuilderServices): Promise<UIStateState> => {
+  return initialState;
+};
+
+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,7 +14,12 @@ 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 } from '../state_management/metadata_slice';
 import { getStateFromSavedObject } from '../../../saved_visualizations/transforms';
@@ -46,6 +51,7 @@ export const useSavedVisBuilderVis = (visualizationIdFromUrl: string | undefined
 
     const loadSavedVisBuilderVis = async () => {
       try {
+        dispatch(setEditorState({ state: 'loading' }));
         const savedVisBuilderVis = await getSavedVisBuilderVis(
           savedVisBuilderLoader,
           visualizationIdFromUrl
@@ -56,8 +62,10 @@ export const useSavedVisBuilderVis = (visualizationIdFromUrl: string | undefined
           chrome.setBreadcrumbs(getEditBreadcrumbs(title, navigateToApp));
           chrome.docTitle.change(title);
 
+          dispatch(setUIStateState(state.ui));
           dispatch(setStyleState(state.style));
           dispatch(setVisualizationState(state.visualization));
+          dispatch(setEditorState({ state: 'loaded' }));
         } else {
           chrome.setBreadcrumbs(getCreateBreadcrumbs(navigateToApp));
         }