[lens] Datasources and visualizations should coordinate when switching visualizations

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

@@ -19,6 +19,30 @@ interface ConfigPanelWrapperProps {
   datasourcePublicAPI: DatasourcePublicAPI;
 }
 
+function getSuggestedVisualizationState(
+  visualization: Visualization,
+  datasource: DatasourcePublicAPI
+) {
+  const suggestions = visualization.getSuggestions({
+    tables: [
+      {
+        datasourceSuggestionId: 0,
+        isMultiRow: true,
+        columns: datasource.getTableSpec().map(col => ({
+          ...col,
+          operation: datasource.getOperationForColumnId(col.columnId)!,
+        })),
+      },
+    ],
+  });
+
+  if (!suggestions.length) {
+    return visualization.initialize(datasource);
+  }
+
+  return visualization.initialize(datasource, suggestions[0].state);
+}
+
 export function ConfigPanelWrapper(props: ConfigPanelWrapperProps) {
   const context = useContext(DragContext);
   const setVisualizationState = useMemo(
@@ -41,14 +65,14 @@ export function ConfigPanelWrapper(props: ConfigPanelWrapperProps) {
         }))}
         value={props.activeVisualizationId || undefined}
         onChange={e => {
+          const newState = getSuggestedVisualizationState(
+            props.visualizationMap[e.target.value],
+            props.datasourcePublicAPI
+          );
           props.dispatch({
             type: 'SWITCH_VISUALIZATION',
             newVisualizationId: e.target.value,
-            // TODO we probably want to have a separate API to "force" a visualization switch
-            // which isn't a result of a picked suggestion
-            initialState: props.visualizationMap[e.target.value].initialize(
-              props.datasourcePublicAPI
-            ),
+            initialState: newState,
           });
         }}
       />

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

@@ -37,7 +37,7 @@ describe('editor_frame', () => {
   let mockVisualization: Visualization;
   let mockDatasource: DatasourceMock;
 
-  let mockVisualization2: Visualization;
+  let mockVisualization2: jest.Mocked<Visualization>;
   let mockDatasource2: DatasourceMock;
 
   let expressionRendererMock: ExpressionRenderer;
@@ -520,19 +520,51 @@ Object {
       expect(mockVisualization2.initialize).toHaveBeenCalled();
     });
 
-    it('should call visualization render with new state on switch', async () => {
-      const initialState = {};
-      mockVisualization2.initialize = () => initialState;
+    it('should use suggestions to switch to new visualization', async () => {
+      const initialState = { suggested: true };
+      mockVisualization2.initialize.mockReturnValueOnce({ initial: true });
+      mockVisualization2.getSuggestions.mockReturnValueOnce([
+        {
+          title: 'Suggested vis',
+          score: 1,
+          datasourceSuggestionId: 0,
+          state: initialState,
+        },
+      ]);
 
       act(() => {
         instance
           .find('select[data-test-subj="visualization-switch"]')
           .simulate('change', { target: { value: 'testVis2' } });
       });
 
+      expect(mockDatasource.publicAPIMock.getTableSpec).toHaveBeenCalled();
+      expect(mockVisualization2.getSuggestions).toHaveBeenCalled();
+      expect(mockVisualization2.initialize).toHaveBeenCalledWith(
+        mockDatasource.publicAPIMock,
+        initialState
+      );
       expect(mockVisualization2.renderConfigPanel).toHaveBeenCalledWith(
         expect.any(Element),
-        expect.objectContaining({ state: initialState })
+        expect.objectContaining({ state: { initial: true } })
+      );
+    });
+
+    it('should fall back when switching visualizations if the visualization has no suggested use', async () => {
+      mockVisualization2.initialize.mockReturnValueOnce({ initial: true });
+
+      act(() => {
+        instance
+          .find('select[data-test-subj="visualization-switch"]')
+          .simulate('change', { target: { value: 'testVis2' } });
+      });
+
+      expect(mockDatasource.publicAPIMock.getTableSpec).toHaveBeenCalled();
+      expect(mockVisualization2.getSuggestions).toHaveBeenCalled();
+      expect(mockVisualization2.initialize).toHaveBeenCalledWith(mockDatasource.publicAPIMock);
+      expect(mockVisualization2.renderConfigPanel).toHaveBeenCalledWith(
+        expect.any(Element),
+        expect.objectContaining({ state: { initial: true } })
       );
     });
   });

x-pack/legacy/plugins/lens/public/indexpattern_plugin/indexpattern.test.tsx

@@ -450,6 +450,46 @@ describe('IndexPattern Data Source', () => {
     });
   });
 
+  describe('#getDatasourceSuggestionsFromCurrentState', () => {
+    it('returns no suggestions if there are no columns', () => {
+      expect(
+        indexPatternDatasource.getDatasourceSuggestionsFromCurrentState({
+          indexPatterns: expectedIndexPatterns,
+          columnOrder: [],
+          columns: {},
+          currentIndexPatternId: '1',
+        })
+      ).toEqual([]);
+    });
+
+    it('returns a single suggestion containing the current columns', async () => {
+      const state = await indexPatternDatasource.initialize(persistedState);
+      expect(indexPatternDatasource.getDatasourceSuggestionsFromCurrentState(state)).toEqual([
+        {
+          state: {
+            ...persistedState,
+            indexPatterns: expectedIndexPatterns,
+          },
+          table: {
+            datasourceSuggestionId: 0,
+            isMultiRow: true,
+            columns: [
+              {
+                columnId: 'col1',
+                operation: {
+                  id: 'op1',
+                  label: 'My Op',
+                  dataType: 'string',
+                  isBucketed: true,
+                },
+              },
+            ],
+          },
+        },
+      ]);
+    });
+  });
+
   describe('#getPublicAPI', () => {
     let publicAPI: DatasourcePublicAPI;
 
@@ -477,6 +517,10 @@ describe('IndexPattern Data Source', () => {
           isBucketed: true,
         } as Operation);
       });
+
+      it('should return null for non-existant columns', () => {
+        expect(publicAPI.getOperationForColumnId('col2')).toBe(null);
+      });
     });
   });
 });

x-pack/legacy/plugins/lens/public/indexpattern_plugin/indexpattern.tsx

@@ -11,7 +11,7 @@ import { Chrome } from 'ui/chrome';
 import { ToastNotifications } from 'ui/notify/toasts/toast_notifications';
 import { EuiComboBox } from '@elastic/eui';
 import uuid from 'uuid';
-import { Datasource, DataType } from '../../public';
+import { Datasource, DataType } from '..';
 import {
   DatasourceDimensionPanelProps,
   DatasourceDataPanelProps,
@@ -257,6 +257,9 @@ export function getIndexPatternDatasource(chrome: Chrome, toastNotifications: To
           return state.columnOrder.map(colId => ({ columnId: colId }));
         },
         getOperationForColumnId: (columnId: string) => {
+          if (!state.columns[columnId]) {
+            return null;
+          }
           return columnToOperation(state.columns[columnId]);
         },
         generateColumnId: () => {
@@ -369,8 +372,26 @@ export function getIndexPatternDatasource(chrome: Chrome, toastNotifications: To
       return [];
     },
 
-    getDatasourceSuggestionsFromCurrentState(state) {
-      return [];
+    getDatasourceSuggestionsFromCurrentState(
+      state
+    ): Array<DatasourceSuggestion<IndexPatternPrivateState>> {
+      if (!state.columnOrder.length) {
+        return [];
+      }
+      return [
+        {
+          state,
+
+          table: {
+            columns: state.columnOrder.map(id => ({
+              columnId: id,
+              operation: columnToOperation(state.columns[id]),
+            })),
+            isMultiRow: true,
+            datasourceSuggestionId: 0,
+          },
+        },
+      ];
     },
   };