[Lens] Fix empty dimensions after upgrading from pre-7.7

x-pack/legacy/plugins/lens/public/editor_frame_service/editor_frame/editor_frame.test.tsx

@@ -116,7 +116,7 @@ describe('editor_frame', () => {
       expect(mockDatasource.initialize).toHaveBeenCalled();
     });
 
-    it('should not initialize datasource and visualization if no initial one is specificed', () => {
+    it('should not initialize datasource and visualization if no initial one is specified', () => {
       act(() => {
         mount(
           <EditorFrame
@@ -226,13 +226,17 @@ describe('editor_frame', () => {
       expect(mockVisualization.initialize).toHaveBeenCalled();
     });
 
-    it('should pass the public frame api into visualization initialize', async () => {
+    it('should use initialize the visualization and update with private state', async () => {
+      const initializeMock = {
+        ...mockVisualization,
+        initialize: jest.fn().mockReturnValue('initialized'),
+      };
       await act(async () => {
         mount(
           <EditorFrame
             {...getDefaultProps()}
             visualizationMap={{
-              testVis: mockVisualization,
+              testVis: initializeMock,
             }}
             datasourceMap={{
               testDatasource: mockDatasource,
@@ -243,17 +247,33 @@ describe('editor_frame', () => {
             dateRange={{ fromDate: 'now-7d', toDate: 'now' }}
           />
         );
-        expect(mockVisualization.initialize).not.toHaveBeenCalled();
+        expect(initializeMock.initialize).not.toHaveBeenCalled();
       });
 
-      expect(mockVisualization.initialize).toHaveBeenCalledWith({
-        datasourceLayers: {},
-        addNewLayer: expect.any(Function),
-        removeLayers: expect.any(Function),
-        query: { query: '', language: 'lucene' },
-        filters: [],
-        dateRange: { fromDate: 'now-7d', toDate: 'now' },
-      });
+      expect(initializeMock.initialize).toHaveBeenCalledTimes(2);
+      expect(initializeMock.initialize).toHaveBeenCalledWith(
+        {
+          datasourceLayers: {},
+          addNewLayer: expect.any(Function),
+          removeLayers: expect.any(Function),
+          query: { query: '', language: 'lucene' },
+          filters: [],
+          dateRange: { fromDate: 'now-7d', toDate: 'now' },
+        },
+        null
+      );
+
+      expect(initializeMock.initialize).toHaveBeenCalledWith(
+        {
+          datasourceLayers: {},
+          addNewLayer: expect.any(Function),
+          removeLayers: expect.any(Function),
+          query: { query: '', language: 'lucene' },
+          filters: [],
+          dateRange: { fromDate: 'now-7d', toDate: 'now' },
+        },
+        'initialized'
+      );
     });
 
     it('should add new layer on active datasource on frame api call', async () => {

x-pack/legacy/plugins/lens/public/editor_frame_service/editor_frame/editor_frame.tsx

@@ -4,6 +4,7 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
+import { isEqual } from 'lodash';
 import React, { useEffect, useReducer } from 'react';
 import { CoreSetup, CoreStart } from 'src/core/public';
 import { ReactExpressionRendererType } from '../../../../../../../src/plugins/expressions/public';
@@ -163,13 +164,21 @@ export function EditorFrame(props: EditorFrameProps) {
 
   // Initialize visualization as soon as all datasources are ready
   useEffect(() => {
-    if (allLoaded && state.visualization.state === null && activeVisualization) {
-      const initialVisualizationState = activeVisualization.initialize(framePublicAPI);
-      dispatch({
-        type: 'UPDATE_VISUALIZATION_STATE',
-        visualizationId: activeVisualization.id,
-        newState: initialVisualizationState,
-      });
+    if (allLoaded && activeVisualization) {
+      // Initialize from either empty state or saved doc
+      const initializedState = activeVisualization.initialize(
+        framePublicAPI,
+        state.visualization.state ?? null
+      );
+
+      // Visualizations can have private state that is not saved
+      if (!isEqual(state.visualization.state, initializedState)) {
+        dispatch({
+          type: 'UPDATE_VISUALIZATION_STATE',
+          visualizationId: activeVisualization.id,
+          newState: initializedState,
+        });
+      }
     }
   }, [allLoaded, activeVisualization, state.visualization.state]);
 

x-pack/legacy/plugins/lens/public/types.ts

@@ -398,7 +398,16 @@ export interface Visualization<T = unknown, P = unknown> {
 
   switchVisualizationType?: (visualizationTypeId: string, state: T) => T;
 
-  // For initializing from saved object
+  /**
+   * Initialize is called when opening the visualization, such as in the following cases:
+   *
+   * - Loading Lens in an empty state
+   * - Loading a saved Lens visualization which uses this visualization
+   * - When switching to this visualization in the chart switcher
+   *
+   * The response of the initialize call will be used to update the editor state, such as
+   * if the visualization tracks UI state that is not saved
+   */
   initialize: (frame: FramePublicAPI, state?: P) => T;
 
   getPersistableState: (state: T) => P;

x-pack/legacy/plugins/lens/public/xy_visualization/xy_visualization.test.ts

@@ -111,10 +111,27 @@ describe('xy_visualization', () => {
       `);
     });
 
-    it('loads from persisted state', () => {
-      expect(xyVisualization.initialize(createMockFramePublicAPI(), exampleState())).toEqual(
-        exampleState()
-      );
+    it('removes dimensions from saved state if they are not on datasource', () => {
+      const state = exampleState();
+      expect(xyVisualization.initialize(createMockFramePublicAPI(), state)).toEqual({
+        ...state,
+        layers: [
+          {
+            layerId: 'first',
+            seriesType: 'area',
+            splitAccessor: undefined,
+            xAccessor: undefined,
+            accessors: [],
+          },
+        ],
+      });
+    });
+
+    it('keeps dimensions from saved state if they are in datasource', () => {
+      const state = exampleState();
+      const frame = createMockFramePublicAPI();
+      (frame.datasourceLayers.first.getOperationForColumnId as jest.Mock).mockReturnValue(true);
+      expect(xyVisualization.initialize(frame, state)).toEqual(state);
     });
   });
 

x-pack/legacy/plugins/lens/public/xy_visualization/xy_visualization.tsx

@@ -126,28 +126,50 @@ export const xyVisualization: Visualization<State, PersistableState> = {
   getSuggestions,
 
   initialize(frame, state) {
-    return (
-      state || {
-        title: 'Empty XY chart',
-        legend: { isVisible: true, position: Position.Right },
-        preferredSeriesType: defaultSeriesType,
-        layers: [
-          {
-            layerId: frame.addNewLayer(),
-            accessors: [],
-            position: Position.Top,
-            seriesType: defaultSeriesType,
-            showGridlines: false,
-          },
-        ],
-      }
-    );
+    if (state) {
+      // Before 7.7, the XY visualization would generate IDs for missing dimensions
+      // Instead of migrating, we use the initialization step to hide these
+      return {
+        ...state,
+        layers: state.layers.map(layer => {
+          const datasource = frame.datasourceLayers[layer.layerId];
+          const hasXAccessor =
+            layer.xAccessor && datasource.getOperationForColumnId(layer.xAccessor);
+          const hasSplitAccessor =
+            layer.splitAccessor && datasource.getOperationForColumnId(layer.splitAccessor);
+          const yAccessors = layer.accessors.filter(accessor =>
+            datasource.getOperationForColumnId(accessor)
+          );
+          return {
+            ...layer,
+            accessors: yAccessors,
+            xAccessor: hasXAccessor ? layer.xAccessor : undefined,
+            splitAccessor: hasSplitAccessor ? layer.splitAccessor : undefined,
+          };
+        }),
+      };
+    }
+    return {
+      title: 'Empty XY chart',
+      legend: { isVisible: true, position: Position.Right },
+      preferredSeriesType: defaultSeriesType,
+      layers: [
+        {
+          layerId: frame.addNewLayer(),
+          accessors: [],
+          position: Position.Top,
+          seriesType: defaultSeriesType,
+          showGridlines: false,
+        },
+      ],
+    };
   },
 
   getPersistableState: state => state,
 
   getConfiguration(props) {
     const layer = props.state.layers.find(l => l.layerId === props.layerId)!;
+
     return {
       groups: [
         {