[Lens] Initialize visualization with datasource api

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

@@ -44,11 +44,14 @@ export function ConfigPanelWrapper(props: ConfigPanelWrapperProps) {
             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(),
+            initialState: props.visualizationMap[e.target.value].initialize(
+              undefined,
+              props.datasourcePublicAPI
+            ),
           });
         }}
       />
-      {props.activeVisualizationId && (
+      {props.activeVisualizationId && props.visualizationState !== null && (
         <NativeRenderer
           render={props.visualizationMap[props.activeVisualizationId].renderConfigPanel}
           nativeProps={{

x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/data_panel_wrapper.tsx

@@ -4,7 +4,7 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-import React, { useMemo } from 'react';
+import React, { useMemo, memo } from 'react';
 import { EuiSelect } from '@elastic/eui';
 import { DatasourceDataPanelProps, Datasource } from '../..';
 import { NativeRenderer } from '../../native_renderer';
@@ -18,7 +18,7 @@ interface DataPanelWrapperProps {
   dispatch: (action: Action) => void;
 }
 
-export function DataPanelWrapper(props: DataPanelWrapperProps) {
+export const DataPanelWrapper = memo((props: DataPanelWrapperProps) => {
   const setDatasourceState = useMemo(
     () => (newState: unknown) => {
       props.dispatch({
@@ -55,4 +55,4 @@ export function DataPanelWrapper(props: DataPanelWrapperProps) {
       )}
     </>
   );
-}
+});

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

@@ -8,12 +8,13 @@ import React from 'react';
 import { ReactWrapper } from 'enzyme';
 import { mountWithIntl as mount } from 'test_utils/enzyme_helpers';
 import { EditorFrame } from './editor_frame';
-import { Visualization, Datasource, DatasourcePublicAPI, DatasourceSuggestion } from '../../types';
+import { Visualization, DatasourcePublicAPI, DatasourceSuggestion } from '../../types';
 import { act } from 'react-dom/test-utils';
 import {
   createMockVisualization,
   createMockDatasource,
   createExpressionRendererMock,
+  DatasourceMock,
 } from '../mocks';
 import { ExpressionRenderer } from 'src/legacy/core_plugins/data/public';
 
@@ -34,10 +35,10 @@ function generateSuggestion(datasourceSuggestionId = 1, state = {}): DatasourceS
 
 describe('editor_frame', () => {
   let mockVisualization: Visualization;
-  let mockDatasource: Datasource;
+  let mockDatasource: DatasourceMock;
 
   let mockVisualization2: Visualization;
-  let mockDatasource2: Datasource;
+  let mockDatasource2: DatasourceMock;
 
   let expressionRendererMock: ExpressionRenderer;
 
@@ -52,7 +53,7 @@ describe('editor_frame', () => {
   });
 
   describe('initialization', () => {
-    it('should initialize initial datasource and visualization if present', () => {
+    it('should initialize initial datasource', () => {
       act(() => {
         mount(
           <EditorFrame
@@ -69,7 +70,6 @@ describe('editor_frame', () => {
         );
       });
 
-      expect(mockVisualization.initialize).toHaveBeenCalled();
       expect(mockDatasource.initialize).toHaveBeenCalled();
     });
 
@@ -115,6 +115,57 @@ describe('editor_frame', () => {
       expect(mockDatasource.renderDataPanel).not.toHaveBeenCalled();
     });
 
+    it('should not initialize visualization before datasource is initialized', async () => {
+      act(() => {
+        mount(
+          <EditorFrame
+            visualizationMap={{
+              testVis: mockVisualization,
+            }}
+            datasourceMap={{
+              testDatasource: mockDatasource,
+            }}
+            initialDatasourceId="testDatasource"
+            initialVisualizationId="testVis"
+            ExpressionRenderer={expressionRendererMock}
+          />
+        );
+      });
+
+      expect(mockVisualization.initialize).not.toHaveBeenCalled();
+
+      await waitForPromises();
+
+      expect(mockVisualization.initialize).toHaveBeenCalled();
+    });
+
+    it('should pass the public datasource api into visualization initialize', async () => {
+      act(() => {
+        mount(
+          <EditorFrame
+            visualizationMap={{
+              testVis: mockVisualization,
+            }}
+            datasourceMap={{
+              testDatasource: mockDatasource,
+            }}
+            initialDatasourceId="testDatasource"
+            initialVisualizationId="testVis"
+            ExpressionRenderer={expressionRendererMock}
+          />
+        );
+      });
+
+      expect(mockVisualization.initialize).not.toHaveBeenCalled();
+
+      await waitForPromises();
+
+      expect(mockVisualization.initialize).toHaveBeenCalledWith(
+        undefined,
+        mockDatasource.publicAPIMock
+      );
+    });
+
     it('should render data panel after initialization is complete', async () => {
       const initialState = {};
       let databaseInitialized: ({}) => void;
@@ -252,9 +303,6 @@ Object {
           state: updatedState,
         })
       );
-
-      // don't re-render datasource when visulization changes
-      expect(mockDatasource.renderDataPanel).toHaveBeenCalledTimes(1);
     });
 
     it('should re-render data panel after state update', async () => {
@@ -309,7 +357,7 @@ Object {
 
       const updatedPublicAPI = {};
       mockDatasource.getPublicAPI = jest.fn(
-        () => (updatedPublicAPI as unknown) as DatasourcePublicAPI
+        _ => (updatedPublicAPI as unknown) as DatasourcePublicAPI
       );
 
       const setDatasourceState = (mockDatasource.renderDataPanel as jest.Mock).mock.calls[0][1]
@@ -330,10 +378,6 @@ Object {
 
   describe('datasource public api communication', () => {
     it('should pass the datasource api to the visualization', async () => {
-      const publicAPI = ({} as unknown) as DatasourcePublicAPI;
-
-      mockDatasource.getPublicAPI = () => publicAPI;
-
       mount(
         <EditorFrame
           visualizationMap={{
@@ -352,13 +396,13 @@ Object {
 
       expect(mockVisualization.renderConfigPanel).toHaveBeenCalledWith(
         expect.any(Element),
-        expect.objectContaining({ datasource: publicAPI })
+        expect.objectContaining({ datasource: mockDatasource.publicAPIMock })
       );
     });
 
     it('should give access to the datasource state in the datasource factory function', async () => {
       const datasourceState = {};
-      mockDatasource.initialize = () => Promise.resolve(datasourceState);
+      mockDatasource.initialize.mockResolvedValue(datasourceState);
 
       mount(
         <EditorFrame
@@ -456,7 +500,7 @@ Object {
 
     it('should call datasource render with new state on switch', async () => {
       const initialState = {};
-      mockDatasource2.initialize = () => Promise.resolve(initialState);
+      mockDatasource2.initialize.mockResolvedValue(initialState);
 
       instance
         .find('select[data-test-subj="datasource-switch"]')

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

@@ -27,6 +27,29 @@ export interface EditorFrameProps {
 export function EditorFrame(props: EditorFrameProps) {
   const [state, dispatch] = useReducer(reducer, props, getInitialState);
 
+  // create public datasource api for current state
+  // as soon as datasource is available and memoize it
+  const datasourcePublicAPI = useMemo(
+    () =>
+      state.datasource.activeId && !state.datasource.isLoading
+        ? props.datasourceMap[state.datasource.activeId].getPublicAPI(
+            state.datasource.state,
+            (newState: unknown) => {
+              dispatch({
+                type: 'UPDATE_DATASOURCE_STATE',
+                newState,
+              });
+            }
+          )
+        : undefined,
+    [
+      props.datasourceMap,
+      state.datasource.isLoading,
+      state.datasource.activeId,
+      state.datasource.state,
+    ]
+  );
+
   // Initialize current datasource
   useEffect(
     () => {
@@ -49,27 +72,24 @@ export function EditorFrame(props: EditorFrameProps) {
     [state.datasource.activeId, state.datasource.isLoading]
   );
 
-  // create public datasource api for current state
-  // as soon as datasource is available and memoize it
-  const datasourcePublicAPI = useMemo(
-    () =>
-      state.datasource.activeId && !state.datasource.isLoading
-        ? props.datasourceMap[state.datasource.activeId].getPublicAPI(
-            state.datasource.state,
-            (newState: unknown) => {
-              dispatch({
-                type: 'UPDATE_DATASOURCE_STATE',
-                newState,
-              });
-            }
-          )
-        : undefined,
-    [
-      props.datasourceMap,
-      state.datasource.isLoading,
-      state.datasource.activeId,
-      state.datasource.state,
-    ]
+  // Initialize visualization as soon as datasource is ready
+  useEffect(
+    () => {
+      if (
+        datasourcePublicAPI &&
+        state.visualization.state === null &&
+        state.visualization.activeId !== null
+      ) {
+        const initialVisualizationState = props.visualizationMap[
+          state.visualization.activeId
+        ].initialize(undefined, datasourcePublicAPI);
+        dispatch({
+          type: 'UPDATE_VISUALIZATION_STATE',
+          newState: initialVisualizationState,
+        });
+      }
+    },
+    [datasourcePublicAPI, state.visualization.activeId, state.visualization.state]
   );
 
   if (state.datasource.activeId && !state.datasource.isLoading) {

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

@@ -29,20 +29,19 @@ describe('editor_frame state management', () => {
       expect(initialState.visualization.activeId).toEqual('testVis');
     });
 
-    it('should initialize visualization', () => {
-      const initialVisState = {};
-      props.visualizationMap.testVis.initialize = jest.fn(() => initialVisState);
-
+    it('should not initialize visualization but set active id', () => {
       const initialState = getInitialState(props);
 
-      expect(initialState.visualization.state).toBe(initialVisState);
-      expect(props.visualizationMap.testVis.initialize).toHaveBeenCalled();
+      expect(initialState.visualization.state).toBe(null);
+      expect(initialState.visualization.activeId).toBe('testVis');
+      expect(props.visualizationMap.testVis.initialize).not.toHaveBeenCalled();
     });
 
-    it('should not initialize visualization if no initial visualization is passed in', () => {
+    it('should not set active id if no initial visualization is passed in', () => {
       const initialState = getInitialState({ ...props, initialVisualizationId: null });
 
       expect(initialState.visualization.state).toEqual(null);
+      expect(initialState.visualization.activeId).toEqual(null);
       expect(props.visualizationMap.testVis.initialize).not.toHaveBeenCalled();
     });
   });

x-pack/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts

@@ -46,9 +46,7 @@ export const getInitialState = (props: EditorFrameProps): EditorFrameState => {
       activeId: props.initialDatasourceId,
     },
     visualization: {
-      state: props.initialVisualizationId
-        ? props.visualizationMap[props.initialVisualizationId].initialize()
-        : null,
+      state: null,
       activeId: props.initialVisualizationId,
     },
   };

x-pack/plugins/lens/public/editor_frame_plugin/mocks.tsx

@@ -27,6 +27,7 @@ export function createMockDatasource(): DatasourceMock {
   const publicAPIMock: jest.Mocked<DatasourcePublicAPI> = {
     getTableSpec: jest.fn(() => []),
     getOperationForColumnId: jest.fn(),
+    generateColumnId: jest.fn(),
     renderDimensionPanel: jest.fn(),
     removeColumnInTableSpec: jest.fn(),
     moveColumnTo: jest.fn(),

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

@@ -10,6 +10,7 @@ import { Chrome } from 'ui/chrome';
 import { ToastNotifications } from 'ui/notify/toasts/toast_notifications';
 import { EuiComboBox } from '@elastic/eui';
 import { Datasource, DataType } from '..';
+import uuid from 'uuid';
 import { DatasourceDimensionPanelProps, DatasourceDataPanelProps } from '../types';
 import { getIndexPatterns } from './loader';
 
@@ -232,6 +233,10 @@ export function getIndexPatternDatasource(chrome: Chrome, toastNotifications: To
             isBucketed,
           };
         },
+        generateColumnId: () => {
+          // TODO: Come up with a more compact form of generating unique column ids
+          return uuid.v4();
+        },
 
         renderDimensionPanel: (domElement: Element, props: DatasourceDimensionPanelProps) => {
           render(