[Lens] Improve layer suggestions and fix layer crash in metric (#49389)

* Add loading indicator to Lens workspace panel

* [Expressions] [Lens] Handle loading and errors in ExpressionRenderer

* Using loading$ observable and improve tests

* [Lens] Fix layer crash and improve layer suggestions

* Using CSS and to handle layout of expression renderer

Added TODO for using chart loader when area is completely empty

* Improve error handling and simplify code

* Fix cleanup behavior

* Fix double render and prevent error cases in xy chart

* Fix context for use in dashboards

* Remove className from expression rendere component

* Improve handling of additional interpreter args

* More layout fixes

- Hide chart if Empty not Loading
- Fix relative positioning for progress bar since className is no longer passed (super hacky)

* Build suggestions that remove layers

* Update tests and add keptLayerIds everywhere

* Fix bug where datatable would accept multi-layer suggestions

* Build more suggestions that work with metric/datatable

* Fix issue with chart switching from empty

* Fix datatable multiple layer issue

x-pack/legacy/plugins/lens/public/datatable_visualization_plugin/visualization.test.tsx

@@ -12,7 +12,7 @@ import {
   DataTableLayer,
 } from './visualization';
 import { mount } from 'enzyme';
-import { Operation, DataType, FramePublicAPI } from '../types';
+import { Operation, DataType, FramePublicAPI, TableSuggestionColumn } from '../types';
 import { generateId } from '../id_generator';
 
 jest.mock('../id_generator');
@@ -72,6 +72,112 @@ describe('Datatable Visualization', () => {
     });
   });
 
+  describe('#getSuggestions', () => {
+    function numCol(columnId: string): TableSuggestionColumn {
+      return {
+        columnId,
+        operation: {
+          dataType: 'number',
+          label: `Avg ${columnId}`,
+          isBucketed: false,
+        },
+      };
+    }
+
+    function strCol(columnId: string): TableSuggestionColumn {
+      return {
+        columnId,
+        operation: {
+          dataType: 'string',
+          label: `Top 5 ${columnId}`,
+          isBucketed: true,
+        },
+      };
+    }
+
+    it('should accept a single-layer suggestion', () => {
+      const suggestions = datatableVisualization.getSuggestions({
+        state: {
+          layers: [{ layerId: 'first', columns: ['col1'] }],
+        },
+        table: {
+          isMultiRow: true,
+          layerId: 'first',
+          changeType: 'initial',
+          columns: [numCol('col1'), strCol('col2')],
+        },
+        keptLayerIds: [],
+      });
+
+      expect(suggestions.length).toBeGreaterThan(0);
+    });
+
+    it('should not make suggestions when the table is unchanged', () => {
+      const suggestions = datatableVisualization.getSuggestions({
+        state: {
+          layers: [{ layerId: 'first', columns: ['col1'] }],
+        },
+        table: {
+          isMultiRow: true,
+          layerId: 'first',
+          changeType: 'unchanged',
+          columns: [numCol('col1')],
+        },
+        keptLayerIds: ['first'],
+      });
+
+      expect(suggestions).toEqual([]);
+    });
+
+    it('should not make suggestions when multiple layers are involved', () => {
+      const suggestions = datatableVisualization.getSuggestions({
+        state: {
+          layers: [{ layerId: 'first', columns: ['col1'] }],
+        },
+        table: {
+          isMultiRow: true,
+          layerId: 'first',
+          changeType: 'unchanged',
+          columns: [numCol('col1')],
+        },
+        keptLayerIds: ['first', 'second'],
+      });
+
+      expect(suggestions).toEqual([]);
+    });
+
+    it('should not make suggestions when the suggestion keeps a different layer', () => {
+      const suggestions = datatableVisualization.getSuggestions({
+        state: {
+          layers: [{ layerId: 'older', columns: ['col1'] }],
+        },
+        table: {
+          isMultiRow: true,
+          layerId: 'newer',
+          changeType: 'initial',
+          columns: [numCol('col1'), strCol('col2')],
+        },
+        keptLayerIds: ['older'],
+      });
+
+      expect(suggestions).toEqual([]);
+    });
+
+    it('should suggest unchanged tables when the state is not passed in', () => {
+      const suggestions = datatableVisualization.getSuggestions({
+        table: {
+          isMultiRow: true,
+          layerId: 'first',
+          changeType: 'unchanged',
+          columns: [numCol('col1')],
+        },
+        keptLayerIds: ['first'],
+      });
+
+      expect(suggestions.length).toBeGreaterThan(0);
+    });
+  });
+
   describe('DataTableLayer', () => {
     it('allows all kinds of operations', () => {
       const setState = jest.fn();

x-pack/legacy/plugins/lens/public/datatable_visualization_plugin/visualization.tsx

@@ -134,10 +134,15 @@ export const datatableVisualization: Visualization<
   getSuggestions({
     table,
     state,
+    keptLayerIds,
   }: SuggestionRequest<DatatableVisualizationState>): Array<
     VisualizationSuggestion<DatatableVisualizationState>
   > {
-    if (state && table.changeType === 'unchanged') {
+    if (
+      keptLayerIds.length > 1 ||
+      (keptLayerIds.length && table.layerId !== keptLayerIds[0]) ||
+      (state && table.changeType === 'unchanged')
+    ) {
       return [];
     }
     const title =

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

@@ -94,6 +94,7 @@ describe('chart_switch', () => {
           layerId: 'a',
           changeType: 'unchanged',
         },
+        keptLayerIds: ['a'],
       },
     ]);
     return {
@@ -180,6 +181,8 @@ describe('chart_switch', () => {
 
     switchTo('subvisB', component);
 
+    expect(frame.removeLayers).toHaveBeenCalledWith(['a']);
+
     expect(dispatch).toHaveBeenCalledWith({
       initialState: 'visB initial state',
       newVisualizationId: 'visB',
@@ -219,6 +222,7 @@ describe('chart_switch', () => {
           isMultiRow: true,
           changeType: 'unchanged',
         },
+        keptLayerIds: [],
       },
     ]);
     datasourceMap.testDatasource.publicAPIMock.getTableSpec.mockReturnValue([
@@ -307,11 +311,7 @@ describe('chart_switch', () => {
 
   it('should not indicate data loss if visualization is not changed', () => {
     const dispatch = jest.fn();
-    const removeLayers = jest.fn();
-    const frame = {
-      ...mockFrame(['a', 'b', 'c']),
-      removeLayers,
-    };
+    const frame = mockFrame(['a', 'b', 'c']);
     const visualizations = mockVisualizations();
     const switchVisualizationType = jest.fn(() => 'therebedragons');
 
@@ -332,30 +332,6 @@ describe('chart_switch', () => {
     expect(getMenuItem('subvisC2', component).prop('betaBadgeIconType')).toBeUndefined();
   });
 
-  it('should remove unused layers', () => {
-    const removeLayers = jest.fn();
-    const frame = {
-      ...mockFrame(['a', 'b', 'c']),
-      removeLayers,
-    };
-    const component = mount(
-      <ChartSwitch
-        visualizationId="visA"
-        visualizationState={{}}
-        visualizationMap={mockVisualizations()}
-        dispatch={jest.fn()}
-        framePublicAPI={frame}
-        datasourceMap={mockDatasourceMap()}
-        datasourceStates={mockDatasourceStates()}
-      />
-    );
-
-    switchTo('subvisB', component);
-
-    expect(removeLayers).toHaveBeenCalledTimes(1);
-    expect(removeLayers).toHaveBeenCalledWith(['b', 'c']);
-  });
-
   it('should remove all layers if there is no suggestion', () => {
     const dispatch = jest.fn();
     const visualizations = mockVisualizations();
@@ -378,15 +354,24 @@ describe('chart_switch', () => {
 
     expect(frame.removeLayers).toHaveBeenCalledTimes(1);
     expect(frame.removeLayers).toHaveBeenCalledWith(['a', 'b', 'c']);
+
+    expect(visualizations.visB.getSuggestions).toHaveBeenCalledWith(
+      expect.objectContaining({
+        keptLayerIds: ['a'],
+      })
+    );
+
+    expect(dispatch).toHaveBeenCalledWith(
+      expect.objectContaining({
+        type: 'SWITCH_VISUALIZATION',
+        initialState: 'visB initial state',
+      })
+    );
   });
 
   it('should not remove layers if the visualization is not changing', () => {
     const dispatch = jest.fn();
-    const removeLayers = jest.fn();
-    const frame = {
-      ...mockFrame(['a', 'b', 'c']),
-      removeLayers,
-    };
+    const frame = mockFrame(['a', 'b', 'c']);
     const visualizations = mockVisualizations();
     const switchVisualizationType = jest.fn(() => 'therebedragons');
 
@@ -405,7 +390,6 @@ describe('chart_switch', () => {
     );
 
     switchTo('subvisC2', component);
-    expect(removeLayers).not.toHaveBeenCalled();
     expect(switchVisualizationType).toHaveBeenCalledWith('subvisC2', 'therebegriffins');
     expect(dispatch).toHaveBeenCalledWith(
       expect.objectContaining({
@@ -447,6 +431,7 @@ describe('chart_switch', () => {
           isMultiRow: true,
           changeType: 'unchanged',
         },
+        keptLayerIds: [],
       },
     ]);
 

x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/chart_switch.tsx

@@ -88,6 +88,10 @@ export function ChartSwitch(props: Props) {
       },
       'SWITCH_VISUALIZATION'
     );
+
+    if (!selection.datasourceId || selection.dataLoss === 'everything') {
+      props.framePublicAPI.removeLayers(Object.keys(props.framePublicAPI.datasourceLayers));
+    }
   };
 
   function getSelection(

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

@@ -35,6 +35,7 @@ function generateSuggestion(state = {}): DatasourceSuggestion {
       layerId: 'first',
       changeType: 'unchanged',
     },
+    keptLayerIds: ['first'],
   };
 }
 
@@ -928,6 +929,7 @@ describe('editor_frame', () => {
             layerId: 'first',
             changeType: 'unchanged',
           },
+          keptLayerIds: [],
         },
       ]);
 
@@ -1073,6 +1075,7 @@ describe('editor_frame', () => {
             isMultiRow: true,
             layerId: 'first',
           },
+          keptLayerIds: [],
         },
       ]);
       mount(

x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/suggestion_helpers.test.ts

@@ -16,6 +16,7 @@ const generateSuggestion = (state = {}, layerId: string = 'first'): DatasourceSu
     layerId,
     changeType: 'unchanged',
   },
+  keptLayerIds: [layerId],
 });
 
 let datasourceMap: Record<string, DatasourceMock>;
@@ -235,8 +236,8 @@ describe('suggestion helpers', () => {
       changeType: 'unchanged',
     };
     datasourceMap.mock.getDatasourceSuggestionsFromCurrentState.mockReturnValue([
-      { state: {}, table: table1 },
-      { state: {}, table: table2 },
+      { state: {}, table: table1, keptLayerIds: ['first'] },
+      { state: {}, table: table2, keptLayerIds: ['first'] },
     ]);
     getSuggestions({
       visualizationMap: {
@@ -343,45 +344,4 @@ describe('suggestion helpers', () => {
       })
     );
   });
-
-  it('should drop other layers only on visualization switch', () => {
-    const mockVisualization1 = createMockVisualization();
-    const mockVisualization2 = createMockVisualization();
-    datasourceMap.mock.getDatasourceSuggestionsFromCurrentState.mockReturnValue([
-      generateSuggestion(),
-    ]);
-    datasourceMap.mock.getLayers.mockReturnValue(['first', 'second']);
-    const suggestions = getSuggestions({
-      visualizationMap: {
-        vis1: {
-          ...mockVisualization1,
-          getSuggestions: () => [
-            {
-              score: 0.8,
-              title: 'Test2',
-              state: {},
-              previewIcon: 'empty',
-            },
-          ],
-        },
-        vis2: {
-          ...mockVisualization2,
-          getSuggestions: () => [
-            {
-              score: 0.6,
-              title: 'Test3',
-              state: {},
-              previewIcon: 'empty',
-            },
-          ],
-        },
-      },
-      activeVisualizationId: 'vis1',
-      visualizationState: {},
-      datasourceMap,
-      datasourceStates,
-    });
-    expect(suggestions[0].keptLayerIds).toEqual(['first', 'second']);
-    expect(suggestions[1].keptLayerIds).toEqual(['first']);
-  });
 });