[Lens] Fixes the types of the DatasourceLayers (elastic#138371)

* [Lens] Fixes the types of the DatasourceLayers

* Remove unecessary if

x-pack/plugins/lens/public/editor_frame_service/editor_frame/config_panel/buttons/draggable_dimension_button.tsx

@@ -47,7 +47,7 @@ export function DraggableDimensionButton({
   group: VisualizationDimensionGroupConfig;
   label: string;
   children: ReactElement;
-  layerDatasource: Datasource<unknown, unknown>;
+  layerDatasource?: Datasource<unknown, unknown>;
   datasourceLayers: DatasourceLayers;
   state: unknown;
   accessorIndex: number;

x-pack/plugins/lens/public/editor_frame_service/editor_frame/config_panel/buttons/empty_dimension_button.tsx

@@ -118,7 +118,7 @@ export function EmptyDimensionButton({
   onDrop: (source: DragDropIdentifier, dropTarget: DragDropIdentifier, dropType?: DropType) => void;
   onClick: (id: string) => void;
   group: VisualizationDimensionGroupConfig;
-  layerDatasource: Datasource<unknown, unknown>;
+  layerDatasource?: Datasource<unknown, unknown>;
   datasourceLayers: DatasourceLayers;
   state: unknown;
 }) {

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

@@ -75,20 +75,22 @@ export function LayerPanels(
     [activeVisualization, dispatchLens]
   );
   const updateDatasource = useMemo(
-    () => (datasourceId: string, newState: unknown) => {
-      dispatchLens(
-        updateDatasourceState({
-          updater: (prevState: unknown) =>
-            typeof newState === 'function' ? newState(prevState) : newState,
-          datasourceId,
-          clearStagedPreview: false,
-        })
-      );
+    () => (datasourceId: string | undefined, newState: unknown) => {
+      if (datasourceId) {
+        dispatchLens(
+          updateDatasourceState({
+            updater: (prevState: unknown) =>
+              typeof newState === 'function' ? newState(prevState) : newState,
+            datasourceId,
+            clearStagedPreview: false,
+          })
+        );
+      }
     },
     [dispatchLens]
   );
   const updateDatasourceAsync = useMemo(
-    () => (datasourceId: string, newState: unknown) => {
+    () => (datasourceId: string | undefined, newState: unknown) => {
       // React will synchronously update if this is triggered from a third party component,
       // which we don't want. The timeout lets user interaction have priority, then React updates.
       setTimeout(() => {
@@ -99,43 +101,49 @@ export function LayerPanels(
   );
 
   const updateAll = useMemo(
-    () => (datasourceId: string, newDatasourceState: unknown, newVisualizationState: unknown) => {
-      // React will synchronously update if this is triggered from a third party component,
-      // which we don't want. The timeout lets user interaction have priority, then React updates.
-
-      setTimeout(() => {
-        dispatchLens(
-          updateState({
-            updater: (prevState) => {
-              const updatedDatasourceState =
-                typeof newDatasourceState === 'function'
-                  ? newDatasourceState(prevState.datasourceStates[datasourceId].state)
-                  : newDatasourceState;
-
-              const updatedVisualizationState =
-                typeof newVisualizationState === 'function'
-                  ? newVisualizationState(prevState.visualization.state)
-                  : newVisualizationState;
-
-              return {
-                ...prevState,
-                datasourceStates: {
-                  ...prevState.datasourceStates,
-                  [datasourceId]: {
-                    state: updatedDatasourceState,
-                    isLoading: false,
+    () =>
+      (
+        datasourceId: string | undefined,
+        newDatasourceState: unknown,
+        newVisualizationState: unknown
+      ) => {
+        if (!datasourceId) return;
+        // React will synchronously update if this is triggered from a third party component,
+        // which we don't want. The timeout lets user interaction have priority, then React updates.
+
+        setTimeout(() => {
+          dispatchLens(
+            updateState({
+              updater: (prevState) => {
+                const updatedDatasourceState =
+                  typeof newDatasourceState === 'function'
+                    ? newDatasourceState(prevState.datasourceStates[datasourceId].state)
+                    : newDatasourceState;
+
+                const updatedVisualizationState =
+                  typeof newVisualizationState === 'function'
+                    ? newVisualizationState(prevState.visualization.state)
+                    : newVisualizationState;
+
+                return {
+                  ...prevState,
+                  datasourceStates: {
+                    ...prevState.datasourceStates,
+                    [datasourceId]: {
+                      state: updatedDatasourceState,
+                      isLoading: false,
+                    },
                   },
-                },
-                visualization: {
-                  ...prevState.visualization,
-                  state: updatedVisualizationState,
-                },
-              };
-            },
-          })
-        );
-      }, 0);
-    },
+                  visualization: {
+                    ...prevState.visualization,
+                    state: updatedVisualizationState,
+                  },
+                };
+              },
+            })
+          );
+        }, 0);
+      },
     [dispatchLens]
   );
 
@@ -187,10 +195,10 @@ export function LayerPanels(
           onRemoveLayer={() => {
             const datasourcePublicAPI = props.framePublicAPI.datasourceLayers?.[layerId];
             const datasourceId = datasourcePublicAPI?.datasourceId;
-            const layerDatasource = datasourceMap[datasourceId];
-            const layerDatasourceState = datasourceStates?.[datasourceId]?.state;
 
             if (datasourceId) {
+              const layerDatasource = datasourceMap[datasourceId];
+              const layerDatasourceState = datasourceStates?.[datasourceId]?.state;
               const trigger = props.uiActions.getTrigger(UPDATE_FILTER_REFERENCES_TRIGGER);
               const action = props.uiActions.getAction(UPDATE_FILTER_REFERENCES_ACTION);
 

x-pack/plugins/lens/public/editor_frame_service/editor_frame/config_panel/layer_panel.tsx

@@ -54,10 +54,10 @@ export function LayerPanel(
     layerIndex: number;
     isOnlyLayer: boolean;
     updateVisualization: StateSetter<unknown>;
-    updateDatasource: (datasourceId: string, newState: unknown) => void;
-    updateDatasourceAsync: (datasourceId: string, newState: unknown) => void;
+    updateDatasource: (datasourceId: string | undefined, newState: unknown) => void;
+    updateDatasourceAsync: (datasourceId: string | undefined, newState: unknown) => void;
     updateAll: (
-      datasourceId: string,
+      datasourceId: string | undefined,
       newDatasourcestate: unknown,
       newVisualizationState: unknown
     ) => void;
@@ -112,8 +112,8 @@ export function LayerPanel(
 
   const datasourcePublicAPI = framePublicAPI.datasourceLayers?.[layerId];
   const datasourceId = datasourcePublicAPI?.datasourceId;
-  const layerDatasourceState = datasourceStates?.[datasourceId]?.state;
-  const layerDatasource = props.datasourceMap[datasourceId];
+  const layerDatasourceState = datasourceId ? datasourceStates?.[datasourceId]?.state : undefined;
+  const layerDatasource = datasourceId ? props.datasourceMap[datasourceId] : undefined;
 
   const layerDatasourceConfigProps = {
     state: layerDatasourceState,
@@ -339,11 +339,11 @@ export function LayerPanel(
                         nextPublicAPI.getTableSpec().map(({ columnId }) => columnId)
                       );
                       const removed = datasourcePublicAPI
-                        .getTableSpec()
+                        ?.getTableSpec()
                         .map(({ columnId }) => columnId)
                         .filter((columnId) => !nextTable.has(columnId));
                       let nextVisState = props.visualizationState;
-                      removed.forEach((columnId) => {
+                      removed?.forEach((columnId) => {
                         nextVisState = activeVisualization.removeDimension({
                           layerId,
                           columnId,
@@ -433,7 +433,7 @@ export function LayerPanel(
                             groupIndex={groupIndex}
                             key={columnId}
                             state={layerDatasourceState}
-                            label={columnLabelMap?.[columnId]}
+                            label={columnLabelMap?.[columnId] ?? ''}
                             layerDatasource={layerDatasource}
                             datasourceLayers={framePublicAPI.datasourceLayers}
                             layerIndex={layerIndex}
@@ -445,7 +445,7 @@ export function LayerPanel(
                             <div className="lnsLayerPanel__dimension">
                               <DimensionButton
                                 accessorConfig={accessorConfig}
-                                label={columnLabelMap?.[accessorConfig.columnId]}
+                                label={columnLabelMap?.[accessorConfig.columnId] ?? ''}
                                 group={group}
                                 onClick={(id: string) => {
                                   setActiveDimension({
@@ -508,7 +508,7 @@ export function LayerPanel(
                                   <>
                                     {activeVisualization?.renderDimensionTrigger?.({
                                       columnId,
-                                      label: columnLabelMap[columnId],
+                                      label: columnLabelMap?.[columnId] ?? '',
                                       hideTooltip,
                                       invalid: group.invalid,
                                       invalidMessage: group.invalidMessage,

x-pack/plugins/lens/public/editor_frame_service/editor_frame/suggestion_helpers.ts

@@ -269,7 +269,7 @@ export function getTopSuggestionForField(
   field: DragDropIdentifier
 ) {
   const hasData = Object.values(datasourceLayers).some(
-    (datasourceLayer) => datasourceLayer.getTableSpec().length > 0
+    (datasourceLayer) => datasourceLayer && datasourceLayer.getTableSpec().length > 0
   );
 
   const activeVisualization = visualization.activeId

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

@@ -197,7 +197,7 @@ describe('chart_switch', () => {
     const visualizationMap = mockVisualizationMap();
     visualizationMap.visB.getSuggestions.mockReturnValueOnce([]);
     const frame = mockFrame(['a']);
-    (frame.datasourceLayers.a.getTableSpec as jest.Mock).mockReturnValue([]);
+    (frame.datasourceLayers.a?.getTableSpec as jest.Mock).mockReturnValue([]);
     const datasourceMap = mockDatasourceMap();
     const datasourceStates = mockDatasourceStates();
     const { instance, lensStore } = await mountWithProvider(
@@ -407,7 +407,7 @@ describe('chart_switch', () => {
     const visualizationMap = mockVisualizationMap();
     visualizationMap.visB.getSuggestions.mockReturnValueOnce([]);
     const frame = mockFrame(['a']);
-    (frame.datasourceLayers.a.getTableSpec as jest.Mock).mockReturnValue([]);
+    (frame.datasourceLayers.a?.getTableSpec as jest.Mock).mockReturnValue([]);
 
     const { instance } = await mountWithProvider(
       <ChartSwitch
@@ -599,7 +599,8 @@ describe('chart_switch', () => {
 
   it('should not remove layers and initialize with existing state when switching between subtypes without data', async () => {
     const frame = mockFrame(['a']);
-    frame.datasourceLayers.a.getTableSpec = jest.fn().mockReturnValue([]);
+    const datasourceLayers = frame.datasourceLayers as Record<string, DatasourcePublicAPI>;
+    datasourceLayers.a.getTableSpec = jest.fn().mockReturnValue([]);
     const visualizationMap = mockVisualizationMap();
     visualizationMap.visC.getSuggestions = jest.fn().mockReturnValue([]);
     visualizationMap.visC.switchVisualizationType = jest.fn(() => 'switched');

x-pack/plugins/lens/public/editor_frame_service/editor_frame/workspace_panel/chart_switch.tsx

@@ -155,7 +155,7 @@ export const ChartSwitch = memo(function ChartSwitch(props: Props) {
       ((_type: string, initialState: unknown) => initialState);
     const layers = Object.entries(props.framePublicAPI.datasourceLayers);
     const containsData = layers.some(
-      ([_layerId, datasource]) => datasource.getTableSpec().length > 0
+      ([_layerId, datasource]) => datasource && datasource.getTableSpec().length > 0
     );
     // Always show the active visualization as a valid selection
     if (
@@ -190,7 +190,7 @@ export const ChartSwitch = memo(function ChartSwitch(props: Props) {
       dataLoss = 'everything';
     } else if (layers.length > 1 && layers.length !== topSuggestion.keptLayerIds.length) {
       dataLoss = 'layers';
-    } else if (topSuggestion.columns !== layers[0][1].getTableSpec().length) {
+    } else if (topSuggestion.columns !== layers[0][1]?.getTableSpec().length) {
       dataLoss = 'columns';
     } else {
       dataLoss = 'nothing';

x-pack/plugins/lens/public/shared_components/axis_extent/helpers.test.ts

@@ -7,7 +7,7 @@
 
 import { Datatable } from '@kbn/expressions-plugin/common';
 import { createMockDatasource } from '../../mocks';
-import { FramePublicAPI, OperationDescriptor } from '../../types';
+import { OperationDescriptor, DatasourcePublicAPI } from '../../types';
 import {
   hasNumericHistogramDimension,
   validateAxisDomain,
@@ -58,7 +58,7 @@ describe('validateZeroInclusivityExtent', () => {
 });
 
 describe('hasNumericHistogramDimension', () => {
-  const datasourceLayers: FramePublicAPI['datasourceLayers'] = {
+  const datasourceLayers: Record<string, DatasourcePublicAPI> = {
     first: createMockDatasource('test').publicAPIMock,
   };
   it('should return true if a numeric histogram is present', () => {

x-pack/plugins/lens/public/shared_components/axis_extent/helpers.ts

@@ -45,7 +45,7 @@ export function validateAxisDomain(extents?: { lowerBound?: number; upperBound?:
  * @returns boolean
  */
 export function hasNumericHistogramDimension(
-  datasourceLayer: DatasourcePublicAPI,
+  datasourceLayer: DatasourcePublicAPI | undefined,
   columnId?: string
 ) {
   if (!columnId) {

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

@@ -718,7 +718,7 @@ export interface VisualizationSuggestion<T = unknown> {
   previewIcon: IconType;
 }
 
-export type DatasourceLayers = Record<string, DatasourcePublicAPI>;
+export type DatasourceLayers = Partial<Record<string, DatasourcePublicAPI>>;
 
 export interface FramePublicAPI {
   datasourceLayers: DatasourceLayers;

x-pack/plugins/lens/public/visualizations/datatable/components/dimension_editor.test.tsx

@@ -12,6 +12,7 @@ import {
   FramePublicAPI,
   OperationDescriptor,
   VisualizationDimensionEditorProps,
+  DatasourcePublicAPI,
 } from '../../../types';
 import { DatatableVisualizationState } from '../visualization';
 import { createMockDatasource, createMockFramePublicAPI } from '../../../mocks';
@@ -221,7 +222,8 @@ describe('data table dimension editor', () => {
 
   it('should not show the dynamic coloring option for a bucketed operation', () => {
     frame.activeData!.first.columns[0].meta.type = 'number';
-    frame.datasourceLayers.first.getOperationForColumnId = jest.fn(
+    const datasourceLayers = frame.datasourceLayers as Record<string, DatasourcePublicAPI>;
+    datasourceLayers.first.getOperationForColumnId = jest.fn(
       () => ({ isBucketed: true } as OperationDescriptor)
     );
     state.columns[0].colorMode = 'cell';

x-pack/plugins/lens/public/visualizations/datatable/visualization.tsx

@@ -587,10 +587,10 @@ function getDataSourceAndSortedColumns(
   layerId: string
 ) {
   const datasource = datasourceLayers[state.layerId];
-  const originalOrder = datasource.getTableSpec().map(({ columnId }) => columnId);
+  const originalOrder = datasource?.getTableSpec().map(({ columnId }) => columnId);
   // When we add a column it could be empty, and therefore have no order
   const sortedColumns = Array.from(
-    new Set(originalOrder.concat(state.columns.map(({ columnId }) => columnId)))
+    new Set(originalOrder?.concat(state.columns.map(({ columnId }) => columnId)))
   );
   return { datasource, sortedColumns };
 }

x-pack/plugins/lens/public/visualizations/gauge/visualization.tsx

@@ -122,7 +122,7 @@ const toExpression = (
   const datasource = datasourceLayers[state.layerId];
   const datasourceExpression = datasourceExpressionsByLayers[state.layerId];
 
-  const originalOrder = datasource.getTableSpec().map(({ columnId }) => columnId);
+  const originalOrder = datasource?.getTableSpec().map(({ columnId }) => columnId);
   if (!originalOrder || !state.metricAccessor) {
     return null;
   }

x-pack/plugins/lens/public/visualizations/heatmap/visualization.tsx

@@ -160,7 +160,7 @@ export const getHeatmapVisualization = ({
   getConfiguration({ state, frame, layerId }) {
     const datasourceLayer = frame.datasourceLayers[layerId];
 
-    const originalOrder = datasourceLayer.getTableSpec().map(({ columnId }) => columnId);
+    const originalOrder = datasourceLayer?.getTableSpec().map(({ columnId }) => columnId);
     if (!originalOrder) {
       return { groups: [] };
     }
@@ -312,7 +312,7 @@ export const getHeatmapVisualization = ({
     const datasource = datasourceLayers[state.layerId];
     const datasourceExpression = datasourceExpressionsByLayers[state.layerId];
 
-    const originalOrder = datasource.getTableSpec().map(({ columnId }) => columnId);
+    const originalOrder = datasource?.getTableSpec().map(({ columnId }) => columnId);
     // When we add a column it could be empty, and therefore have no order
 
     if (!originalOrder || !state.valueAccessor) {
@@ -402,7 +402,7 @@ export const getHeatmapVisualization = ({
     const datasource = datasourceLayers[state.layerId];
     const datasourceExpression = datasourceExpressionsByLayers[state.layerId];
 
-    const originalOrder = datasource.getTableSpec().map(({ columnId }) => columnId);
+    const originalOrder = datasource?.getTableSpec().map(({ columnId }) => columnId);
     // When we add a column it could be empty, and therefore have no order
 
     if (!originalOrder) {
@@ -513,7 +513,7 @@ export const getHeatmapVisualization = ({
     const hasArrayValues = rows.some((row) => Array.isArray(row[state.valueAccessor!]));
 
     const datasource = frame.datasourceLayers[state.layerId];
-    const operation = datasource.getOperationForColumnId(state.valueAccessor);
+    const operation = datasource?.getOperationForColumnId(state.valueAccessor);
 
     return hasArrayValues
       ? [