optimize code

Signed-off-by: Qxisylolo <qianxisy@amazon.com>

Author: Qxisylolo

src/plugins/explore/public/application/in_context_vis_editor/component/save_vis_button.test.tsx

@@ -32,6 +32,26 @@ jest.mock('./save_vis_modal', () => ({
   SaveVisModal: () => <div data-test-subj="save-vis-modal">Modal</div>,
 }));
 
+jest.mock('../../../components/data_transformations', () => {
+  const mockObs = {
+    subscribe: jest.fn().mockReturnValue({ unsubscribe: jest.fn() }),
+    getValue: jest.fn().mockReturnValue([]),
+  };
+  const mockMapObs = {
+    subscribe: jest.fn().mockReturnValue({ unsubscribe: jest.fn() }),
+    getValue: jest.fn().mockReturnValue(new Map()),
+  };
+  return {
+    useTransformationService: jest.fn().mockReturnValue({
+      clearPipeline: jest.fn(),
+      pipeline$: mockObs,
+      getPipeline$: () => mockObs,
+      stageSchemas$: mockMapObs,
+    }),
+    TransformPanel: () => null,
+  };
+});
+
 jest.mock('@osd/i18n', () => ({
   i18n: {
     translate: jest.fn((key, options) => options.defaultMessage),
@@ -67,6 +87,7 @@ const buildVisualizationBuilder = () => ({
   visualizationBuilderForEditor: {
     visConfig$: { value: { type: 'bar', styles: {}, axesMapping: { x: 'field' } } },
     isVisDirty$: new BehaviorSubject(false),
+    getTransformationService: jest.fn().mockReturnValue({ pipeline$: new BehaviorSubject([]) }),
   },
 });
 

src/plugins/explore/public/application/in_context_vis_editor/component/save_vis_button.tsx

@@ -25,7 +25,6 @@ import { useCurrentExploreId } from '../hooks/use_explore_id';
 import { useVisualizationBuilder } from '../hooks/use_visualization_builder';
 import { EditorMode } from '../../utils/state_management/types';
 import { ContainerState, CONTAINER_URL_KEY } from '../types';
-import { useTransformationService } from '../hooks/use_transformation_service';
 
 export interface OnSaveProps {
   savedExplore: SavedExplore;

src/plugins/explore/public/application/in_context_vis_editor/component/visualization_editor_bottom_left_container.test.tsx

@@ -3,7 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import React from 'react';
 import { render, screen } from '@testing-library/react';
 import { ResizableQueryPanelAndVisualization } from './visualization_editor_bottom_left_container';
 import { QueryExecutionStatus } from '../../utils/state_management/types';
@@ -12,6 +11,26 @@ import { useVisualizationBuilder } from '../hooks/use_visualization_builder';
 
 jest.mock('../hooks/use_query_builder_state', () => ({ useQueryBuilderState: jest.fn() }));
 jest.mock('../hooks/use_visualization_builder', () => ({ useVisualizationBuilder: jest.fn() }));
+
+jest.mock('../../../components/data_transformations', () => {
+  const mockObs = {
+    subscribe: jest.fn().mockReturnValue({ unsubscribe: jest.fn() }),
+    getValue: jest.fn().mockReturnValue([]),
+  };
+  const mockMapObs = {
+    subscribe: jest.fn().mockReturnValue({ unsubscribe: jest.fn() }),
+    getValue: jest.fn().mockReturnValue(new Map()),
+  };
+  return {
+    useTransformationService: jest.fn().mockReturnValue({
+      clearPipeline: jest.fn(),
+      pipeline$: mockObs,
+      getPipeline$: () => mockObs,
+      stageSchemas$: mockMapObs,
+    }),
+    TransformPanel: () => null,
+  };
+});
 jest.mock('../../../components/query_panel/utils/use_search_context', () => ({
   useSearchContext: jest.fn().mockReturnValue({}),
 }));
@@ -44,8 +63,12 @@ jest.mock('@osd/i18n', () => ({
 }));
 
 const mockQueryEditorState$ = { getValue: jest.fn() };
+const mockResultState$ = { subscribe: jest.fn().mockReturnValue({ unsubscribe: jest.fn() }) };
 const mockQueryBuilder = {
   queryEditorState$: mockQueryEditorState$,
+  resultState$: mockResultState$,
+  updateQueryEditorState: jest.fn(),
+  setOnDatasetChanged: jest.fn(),
 };
 
 const buildState = (options: Record<string, any>) => ({

src/plugins/explore/public/application/in_context_vis_editor/component/visualization_editor_bottom_left_container.tsx

@@ -31,7 +31,7 @@ import { useVisualizationBuilder } from '../hooks/use_visualization_builder';
 
 import '../visualization_editor.scss';
 import {
-  TransformationService,
+  ITransformationService,
   useTransformationService,
   TransformPanel,
 } from '../../../components/data_transformations';
@@ -69,6 +69,12 @@ export const ResizableQueryPanelAndVisualization = () => {
     }, [queryBuilder]),
   });
 
+  useEffect(() => {
+    queryBuilder.setOnDatasetChanged(() => {
+      transformServices.clearPipeline();
+    });
+  }, [queryBuilder, transformServices]);
+
   useEffect(() => {
     queryBuilder.updateQueryEditorState({ activeBottomPanelTab: activeTab });
   }, [activeTab, queryBuilder]);
@@ -164,10 +170,7 @@ export const ResizableQueryPanelAndVisualization = () => {
                 {activeTab === 'QUERY_TAB' ? (
                   <QueryPanel queryEditorState$={queryBuilder.queryEditorState$} />
                 ) : (
-                  <TransformPanel
-                    transformationService={transformServices}
-                    visualizationBuilder={visualizationBuilderForEditor}
-                  />
+                  <TransformPanel transformationService={transformServices} />
                 )}
               </EuiPanel>
             </EuiResizablePanel>
@@ -179,7 +182,7 @@ export const ResizableQueryPanelAndVisualization = () => {
 };
 
 export const VisualizationContainer = React.memo(
-  ({ transformationService }: { transformationService: TransformationService }) => {
+  ({ transformationService }: { transformationService: ITransformationService }) => {
     const searchContext = useSearchContext();
 
     const { visualizationBuilderForEditor: visualizationBuilder } = useVisualizationBuilder();

src/plugins/explore/public/application/in_context_vis_editor/hooks/use_initial_save_explore.ts

@@ -24,17 +24,14 @@ export const useInitialSaveExplore = () => {
   }, [savedExplore]);
 
   // parse saved vis state and transformation from saved explore
-  const { savedVisConfig } = useMemo(() => {
-    if (!savedExplore?.visualization)
-      return { savedVisConfig: undefined, savedTransformationPipeline: undefined };
+  const savedVisConfig = useMemo(() => {
+    if (!savedExplore?.visualization) return undefined;
     try {
       const parsedVisualization = JSON.parse(savedExplore.visualization);
 
-      return {
-        savedVisConfig: parsedVisualization,
-      };
+      return parsedVisualization;
     } catch {
-      return { savedVisConfig: undefined, savedTransformationPipeline: undefined };
+      return undefined;
     }
   }, [savedExplore]);
 

src/plugins/explore/public/application/in_context_vis_editor/hooks/use_transformation_service.ts

@@ -136,6 +136,7 @@ export class QueryBuilder {
   private getServices: () => ExploreServices;
   private interpolationService?: IVariableInterpolationService;
   public lastExecutedInterpolatedQuery?: string;
+  private onDatasetChangedCallback?: () => void;
 
   constructor(getServices: () => ExploreServices) {
     this.getServices = getServices;
@@ -296,6 +297,9 @@ export class QueryBuilder {
           // sync dataset change
           // check isLanguageChanged and isInitialized for the initial sync
           if (isDatasetChanged || isLanguageChanged || !this.isInitialized) {
+            if (isDatasetChanged && this.isInitialized) {
+              this.onDatasetChangedCallback?.();
+            }
             this.datasetView$.next({ ...this.datasetView$.getValue(), isLoading: true });
             return from(this.handleDatasetChange(newQuery.dataset));
           }
@@ -610,6 +614,12 @@ export class QueryBuilder {
     return this.editorRef;
   }
 
+  // register a callback that fires when the dataset changes,
+  // for example, used to clear the transformation pipeline on dataset switch.
+  setOnDatasetChanged(callback: () => void) {
+    this.onDatasetChangedCallback = callback;
+  }
+
   getQueryEditorState() {
     return this.queryEditorState$.value;
   }

src/plugins/explore/public/application/in_context_vis_editor/query_builder/query_builder.ts

@@ -30,7 +30,6 @@ import { getBreadcrumbs } from './utils';
 import { useInitialContainerContext } from './hooks/use_initial_container_context';
 import { useVariables } from './hooks/use_variables';
 import { DashboardSelectionModal } from './component/dashboard_selection_modal';
-import { cleanupGlobalTransformationService } from './hooks/use_transformation_service';
 
 export const VisualizationEditorPage = ({
   setHeaderActionMenu,
@@ -156,7 +155,6 @@ export const VisualizationEditorPage = ({
     return () => {
       queryBuilder.reset();
       visualizationBuilderForEditor.reset();
-      cleanupGlobalTransformationService();
     };
   }, [queryBuilder, visualizationBuilderForEditor]);
 

src/plugins/explore/public/application/in_context_vis_editor/visualization_editor_page.tsx

@@ -27,40 +27,42 @@ import { TransformationInstance, FieldSchema } from './types';
 import { TransformationService } from './transformation_service';
 import { FIELD_TYPE_MAP } from '../visualizations/constants';
 import { VisFieldType } from '../visualizations/types';
-import { VisualizationBuilder } from '../visualizations/visualization_builder';
 
 const DROPPABLE_ID = 'transformationPipelineDroppable';
 
 export interface TransformPanelProps {
   transformationService: TransformationService;
-  visualizationBuilder: VisualizationBuilder;
 }
 
-export const TransformPanel = ({
-  transformationService,
-  visualizationBuilder,
-}: TransformPanelProps) => {
+export const TransformPanel = ({ transformationService }: TransformPanelProps) => {
   const pipeline = useObservable(
     transformationService.pipeline$,
     transformationService.pipeline$.getValue()
   );
 
-  const stageSchemas = useObservable(
-    visualizationBuilder.stageSchemas$,
-    visualizationBuilder.stageSchemas$.getValue()
+  const stageSchemasMap = useObservable(
+    transformationService.stageSchemas$,
+    transformationService.stageSchemas$.getValue()
   );
 
-  const availableFieldsByStage = useMemo(() => {
-    return stageSchemas.map((raw) =>
-      raw.map((field) => ({
-        name: field.name || '',
-        visFieldType: FIELD_TYPE_MAP[field.type || ''] || VisFieldType.Unknown,
-      }))
-    );
-  }, [stageSchemas]);
-
-  const getAvailableFieldsForIndex = (index: number): FieldSchema[] => {
-    return availableFieldsByStage[index] ?? [];
+  // stageSchemas$ is updated asynchronously via handleData
+  // among this gap, each editor still use its last own stage schemas
+  const availableFieldsForInstanceId = useMemo(() => {
+    const result = new Map<string, FieldSchema[]>();
+    stageSchemasMap.forEach((raw, instanceId) => {
+      result.set(
+        instanceId,
+        raw.map((field) => ({
+          name: field.name || '',
+          visFieldType: FIELD_TYPE_MAP[field.type || ''] || VisFieldType.Unknown,
+        }))
+      );
+    });
+    return result;
+  }, [stageSchemasMap]);
+
+  const getAvailableFieldsForInstance = (instanceId: string): FieldSchema[] => {
+    return availableFieldsForInstanceId.get(instanceId) ?? [];
   };
 
   const onSelectTransformation = (id: string) => {
@@ -142,7 +144,7 @@ export const TransformPanel = ({
                       onConfigChange={onConfigChange}
                       onToggleHide={onToggleHide}
                       dragHandleProps={provided.dragHandleProps}
-                      availableFields={getAvailableFieldsForIndex(index)}
+                      availableFields={getAvailableFieldsForInstance(instance.instance_id)}
                     />
                   )}
                 </EuiDraggable>
@@ -204,7 +206,7 @@ const TransformationCard = ({
           </EuiFlexGroup>
         }
         extraAction={
-          <EuiFlexGroup gutterSize="xs" alignItems="center" responsive={false}>
+          <EuiFlexGroup gutterSize="xs" alignItems="center">
             <EuiFlexItem grow={false}>
               <EuiButtonIcon
                 iconType={instance.hide ? 'eyeClosed' : 'eye'}
@@ -226,8 +228,6 @@ const TransformationCard = ({
                 {...dragHandleProps}
                 style={{
                   cursor: 'grab',
-                  display: 'flex',
-                  alignItems: 'center',
                   marginLeft: '10px',
                 }}
               >