[Fix] VisType switching persistence and selectively show warning

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

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import React, { useState } from 'react';
+import React, { useCallback, useMemo, useState } from 'react';
 import {
   EuiSuperSelect,
   EuiSuperSelectOption,
@@ -18,28 +18,48 @@ import './side_nav.scss';
 import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react/public';
 import { VisBuilderServices } from '../../types';
 import {
+  ActiveVisPayload,
   setActiveVisualization,
   useTypedDispatch,
   useTypedSelector,
 } from '../utils/state_management';
-import { usePersistedAggParams } from '../utils/use/use_persisted_agg_params';
+import { getPersistedAggParams } from '../utils/get_persisted_agg_params';
 
 export const RightNav = () => {
-  const [newVisType, setNewVisType] = useState<string>();
+  const { ui, name: activeVisName } = useVisualizationType();
+  const [confirmAggs, setConfirmAggs] = useState<ActiveVisPayload | undefined>();
   const {
     services: { types },
   } = useOpenSearchDashboards<VisBuilderServices>();
-  const { ui, name: activeVisName } = useVisualizationType();
   const dispatch = useTypedDispatch();
   const StyleSection = ui.containerConfig.style.render;
 
   const { activeVisualization } = useTypedSelector((state) => state.visualization);
-  const aggConfigParams = activeVisualization?.aggConfigParams ?? [];
-  const persistedAggParams = usePersistedAggParams(
-    types,
-    aggConfigParams,
-    activeVisName,
-    newVisType
+  const aggConfigParams = useMemo(() => activeVisualization?.aggConfigParams ?? [], [
+    activeVisualization,
+  ]);
+
+  const handleVisTypeChange = useCallback(
+    (newVisName) => {
+      const currentVisSchemas = types.get(activeVisName)?.ui.containerConfig.data.schemas.all ?? [];
+      const newVisSchemas = types.get(newVisName)?.ui.containerConfig.data.schemas.all ?? [];
+      const persistedAggParams = getPersistedAggParams(
+        aggConfigParams,
+        currentVisSchemas,
+        newVisSchemas
+      );
+
+      const newVis = {
+        name: newVisName,
+        aggConfigParams: persistedAggParams,
+        style: types.get(newVisName)?.ui.containerConfig.style.defaults,
+      };
+
+      if (persistedAggParams.length < aggConfigParams.length) return setConfirmAggs(newVis);
+
+      dispatch(setActiveVisualization(newVis));
+    },
+    [activeVisName, aggConfigParams, dispatch, types]
   );
 
   const options: Array<EuiSuperSelectOption<string>> = types.all().map(({ name, icon, title }) => ({
@@ -55,17 +75,15 @@ export const RightNav = () => {
         <EuiSuperSelect
           options={options}
           valueOfSelected={activeVisName}
-          onChange={(name) => {
-            setNewVisType(name);
-          }}
+          onChange={handleVisTypeChange}
           fullWidth
           data-test-subj="chartPicker"
         />
       </div>
       <div className="vbSidenav__style">
         <StyleSection />
       </div>
-      {newVisType && (
+      {confirmAggs && (
         <EuiConfirmModal
           title={i18n.translate('visBuilder.rightNav.changeVisType.modalTitle', {
             defaultMessage: 'Change visualization type',
@@ -76,25 +94,19 @@ export const RightNav = () => {
           cancelButtonText={i18n.translate('visBuilder.rightNav.changeVisType.cancelText', {
             defaultMessage: 'Cancel',
           })}
-          onCancel={() => setNewVisType(undefined)}
+          onCancel={() => setConfirmAggs(undefined)}
           onConfirm={() => {
-            dispatch(
-              setActiveVisualization({
-                name: newVisType,
-                style: types.get(newVisType)?.ui.containerConfig.style.defaults,
-                aggConfigParams: persistedAggParams,
-              })
-            );
+            dispatch(setActiveVisualization(confirmAggs));
 
-            setNewVisType(undefined);
+            setConfirmAggs(undefined);
           }}
           maxWidth="300px"
           data-test-subj="confirmVisChangeModal"
         >
           <p>
             <FormattedMessage
               id="visBuilder.rightNav.changeVisType.modalDescription"
-              defaultMessage="Changing the visualization type will reset all field selections. Do you want to continue?"
+              defaultMessage="Certain field configurations may be lost when changing visualization types and you may need to reconfigure those fields. Do you want to continue?"
             />
           </p>
         </EuiConfirmModal>

src/plugins/vis_builder/public/application/utils/get_persisted_agg_params.test.tsx

@@ -0,0 +1,145 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { AggGroupName, AggGroupNames, CreateAggConfigParams } from '../../../../data/common';
+import { Schema } from '../../../../vis_default_editor/public';
+import { getPersistedAggParams } from './get_persisted_agg_params';
+
+describe('getPersistedAggParams', () => {
+  const getSchema = (
+    name: string,
+    group: AggGroupName,
+    aggFilter: string[] = ['*'],
+    max: number = Infinity
+  ): Schema => ({
+    name,
+    group,
+    max,
+    min: 0,
+    aggFilter,
+    defaults: [],
+    editor: true,
+    params: [],
+    title: name,
+  });
+
+  test('Should return the same aggConfigParams when the new vis type schemas have the same schemas as the existing schemas', () => {
+    const aggConfigParams: CreateAggConfigParams[] = [
+      { type: 'avg', schema: 'm1' },
+      { type: 'avg', schema: 'm2' },
+      { type: 'avg', schema: 'b2' },
+      { type: 'avg', schema: 'b2' },
+    ];
+
+    const schemas = [
+      getSchema('m1', AggGroupNames.Metrics),
+      getSchema('m2', AggGroupNames.Metrics),
+      getSchema('b2', AggGroupNames.Buckets),
+    ];
+
+    const persistResult = getPersistedAggParams(aggConfigParams, schemas, schemas);
+
+    expect(persistResult).toEqual(aggConfigParams);
+  });
+
+  test('Should select the next compatible schema when aggConfigParam schema exists but exceeds the max count of the schema', () => {
+    const aggConfigParams: CreateAggConfigParams[] = [
+      { type: 'avg', schema: 'm1' },
+      { type: 'avg', schema: 'm1' },
+      { type: 'avg', schema: 'b2' },
+      { type: 'avg', schema: 'b2' },
+    ];
+
+    const schemas = [
+      getSchema('m1', AggGroupNames.Metrics, ['avg'], 1),
+      getSchema('m2', AggGroupNames.Metrics),
+      getSchema('b2', AggGroupNames.Buckets),
+    ];
+
+    const persistResult = getPersistedAggParams(aggConfigParams, schemas, schemas);
+
+    const expected: CreateAggConfigParams[] = [...aggConfigParams];
+    expected[1].schema = 'm2';
+    expect(persistResult).toEqual(expected);
+  });
+
+  test('Should select the next compatible schema when aggConfigParam schema exists but does not match the group in the new schema', () => {
+    const aggConfigParams: CreateAggConfigParams[] = [
+      { type: 'avg', schema: 'm1' },
+      { type: 'avg', schema: 'm1' },
+      { type: 'avg', schema: 'b1' },
+      { type: 'avg', schema: 'b1' },
+    ];
+
+    const oldSchemas = [
+      getSchema('m1', AggGroupNames.Metrics),
+      getSchema('b1', AggGroupNames.Buckets),
+    ];
+
+    const newSchemas = [
+      getSchema('m1', AggGroupNames.Buckets),
+      getSchema('m2', AggGroupNames.Metrics),
+      getSchema('b1', AggGroupNames.Buckets),
+    ];
+
+    const persistResult = getPersistedAggParams(aggConfigParams, oldSchemas, newSchemas);
+
+    const expected: CreateAggConfigParams[] = [...aggConfigParams];
+    expected[0].schema = 'm2';
+    expected[1].schema = 'm2';
+    expect(persistResult).toEqual(expected);
+  });
+
+  test('Should select the next compatible schema with the correct aggfilters', () => {
+    const aggConfigParams: CreateAggConfigParams[] = [
+      { type: 'count', schema: 'm1' },
+      { type: 'avg', schema: 'm1' },
+      { type: 'avg', schema: 'b1' },
+      { type: 'avg', schema: 'b1' },
+    ];
+
+    const oldSchemas = [
+      getSchema('m1', AggGroupNames.Metrics),
+      getSchema('b1', AggGroupNames.Buckets),
+    ];
+
+    const newSchemas = [
+      getSchema('m2', AggGroupNames.Metrics, ['count']),
+      getSchema('m3', AggGroupNames.Metrics, ['!count']),
+      getSchema('b1', AggGroupNames.Buckets),
+    ];
+
+    const persistResult = getPersistedAggParams(aggConfigParams, oldSchemas, newSchemas);
+
+    const expected: CreateAggConfigParams[] = [...aggConfigParams];
+    expected[0].schema = 'm2';
+    expected[1].schema = 'm3';
+    expect(persistResult).toEqual(expected);
+  });
+
+  test('Should drop aggConfigParam when no compatible schema is found', () => {
+    const aggConfigParams: CreateAggConfigParams[] = [
+      { type: 'avg', schema: 'm1' },
+      { type: 'avg', schema: 'm1' },
+      { type: 'avg', schema: 'b1' },
+      { type: 'avg', schema: 'b1' },
+    ];
+
+    const oldSchemas = [
+      getSchema('m1', AggGroupNames.Metrics),
+      getSchema('b1', AggGroupNames.Buckets),
+    ];
+
+    const newSchemas = [
+      getSchema('m2', AggGroupNames.Metrics, ['count']),
+      getSchema('m3', AggGroupNames.Metrics, ['!count'], 1),
+      getSchema('b1', AggGroupNames.Buckets),
+    ];
+
+    const persistResult = getPersistedAggParams(aggConfigParams, oldSchemas, newSchemas);
+
+    expect(persistResult.length).toBe(3);
+  });
+});

src/plugins/vis_builder/public/application/utils/get_persisted_agg_params.ts

@@ -0,0 +1,78 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { CreateAggConfigParams, propFilter } from '../../../../data/common';
+import { Schema } from '../../../../vis_default_editor/public';
+
+const filterByType = propFilter('type');
+
+export const getPersistedAggParams = (
+  aggConfigParams: CreateAggConfigParams[],
+  oldVisSchemas: Schema[] = [],
+  newVisSchemas: Schema[] = []
+): CreateAggConfigParams[] => {
+  const updatedAggConfigParams: CreateAggConfigParams[] = [];
+  const newVisSchemaCounts: Record<string, number> = newVisSchemas.reduce((acc, schema: Schema) => {
+    acc[schema.name] = schema.max;
+    return acc;
+  }, {});
+
+  // For each aggConfigParam, check if a compatible schema exists in the new visualization type
+  aggConfigParams.forEach((aggConfigParam) => {
+    const currentSchema = oldVisSchemas.find((schema: Schema) => {
+      return schema.name === aggConfigParam.schema;
+    });
+
+    // see if a matching schma exists in the new visualization type
+    const matchingSchema = newVisSchemas.find((schema: Schema) => {
+      return schema.name === aggConfigParam.schema;
+    });
+
+    // if the matching schema is same as the current schema, add the aggConfigParam to the updatedAggConfigParams
+    if (
+      isSchemaEqual(matchingSchema, currentSchema) &&
+      newVisSchemaCounts[matchingSchema!.name] > 0
+    ) {
+      updatedAggConfigParams.push(aggConfigParam);
+      newVisSchemaCounts[matchingSchema!.name] -= 1;
+      return;
+    }
+
+    // if a matching schema does not exist, check if a compatible schema exists
+    for (const schema of newVisSchemas) {
+      // Check if the schema group is the same
+      if (schema.group !== currentSchema!.group) continue;
+
+      const compatibleSchema = filterByType([aggConfigParam], schema.aggFilter).length !== 0;
+
+      if (compatibleSchema && newVisSchemaCounts[schema.name] > 0) {
+        updatedAggConfigParams.push({
+          ...aggConfigParam,
+          schema: schema.name,
+        });
+        newVisSchemaCounts[schema.name] -= 1;
+        break;
+      }
+    }
+  });
+
+  return updatedAggConfigParams;
+};
+
+function isSchemaEqual(schema1?: Schema, schema2?: Schema) {
+  // Check if schema1 and schema2 exist
+  if (!schema1 || !schema2) return false;
+
+  if (schema1.name !== schema2.name) return false;
+  if (schema1.group !== schema2.group) return false;
+
+  // Check if aggFilter is the same
+  if (schema1.aggFilter.length !== schema2.aggFilter.length) return false;
+  for (let i = 0; i < schema1.aggFilter.length; i++) {
+    if (schema1.aggFilter[i] !== schema2.aggFilter[i]) return false;
+  }
+
+  return true;
+}

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

@@ -7,7 +7,7 @@ import { createAction } from '@reduxjs/toolkit';
 import { CreateAggConfigParams } from '../../../../../data/common';
 import { VisualizationType } from '../../../services/type_service/visualization_type';
 
-interface ActiveVisPayload {
+export interface ActiveVisPayload {
   name: VisualizationType['name'];
   style: VisualizationType['ui']['containerConfig']['style']['defaults'];
   aggConfigParams: CreateAggConfigParams[];