[Lens] Duplicated datatable available as inspector data for Heatmap chart (#126786)

* [Lens] Duplicated datatable available as inspector data for Heatmap chart

Close: #123176

* update workspace_panel

* revert allowCsvExport

* fix CI

* fix PR comments

* apply pr comments

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

src/plugins/expressions/common/util/tables_adapter.ts

@@ -7,7 +7,7 @@
  */
 
 import { EventEmitter } from 'events';
-import { Datatable } from '../expression_types/specs';
+import type { Datatable } from '../expression_types/specs';
 
 export class TablesAdapter extends EventEmitter {
   private _tables: { [key: string]: Datatable } = {};
@@ -17,6 +17,11 @@ export class TablesAdapter extends EventEmitter {
     this.emit('change', this.tables);
   }
 
+  public reset() {
+    this._tables = {};
+    this.emit('change', this.tables);
+  }
+
   public get tables() {
     return this._tables;
   }

x-pack/plugins/lens/common/expressions/datatable/datatable_fn.ts

@@ -15,6 +15,7 @@ import type {
   ExecutionContext,
 } from '../../../../../../src/plugins/expressions';
 import type { DatatableExpressionFunction } from './types';
+import { logDataTable } from '../expressions_utils';
 
 function isRange(meta: { params?: { id?: string } } | undefined) {
   return meta?.params?.id === 'range';
@@ -25,6 +26,10 @@ export const datatableFn =
     getFormatFactory: (context: ExecutionContext) => FormatFactory | Promise<FormatFactory>
   ): DatatableExpressionFunction['fn'] =>
   async (data, args, context) => {
+    if (context?.inspectorAdapters?.tables) {
+      logDataTable(context.inspectorAdapters.tables, data.tables);
+    }
+
     let untransposedData: LensMultiTable | undefined;
     // do the sorting at this level to propagate it also at CSV download
     const [firstTable] = Object.values(data.tables);

x-pack/plugins/lens/common/expressions/expressions_utils.ts

@@ -0,0 +1,16 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { TablesAdapter } from '../../../../../src/plugins/expressions';
+import type { Datatable } from '../../../../../src/plugins/expressions';
+
+export const logDataTable = (
+  tableAdapter: TablesAdapter,
+  datatables: Record<string, Datatable> = {}
+) => {
+  Object.entries(datatables).forEach(([key, table]) => tableAdapter.logDatatable(key, table));
+};

x-pack/plugins/lens/common/expressions/merge_tables/index.ts

@@ -50,14 +50,16 @@ export const mergeTables: ExpressionFunctionDefinition<
   inputTypes: ['kibana_context', 'null'],
   fn(input, { layerIds, tables }, context) {
     const resultTables: Record<string, Datatable> = {};
+
+    if (context.inspectorAdapters?.tables) {
+      context.inspectorAdapters.tables.reset();
+      context.inspectorAdapters.tables.allowCsvExport = true;
+    }
+
     tables.forEach((table, index) => {
       resultTables[layerIds[index]] = table;
-      // adapter is always defined at that point because we make sure by the beginning of the function
-      if (context?.inspectorAdapters?.tables) {
-        context.inspectorAdapters.tables.allowCsvExport = true;
-        context.inspectorAdapters.tables.logDatatable(layerIds[index], table);
-      }
     });
+
     return {
       type: 'lens_multitable',
       tables: resultTables,

x-pack/plugins/lens/common/expressions/merge_tables/merge_tables.test.ts

@@ -54,17 +54,17 @@ describe('lens_merge_tables', () => {
     });
   });
 
-  it('should store the current tables in the tables inspector', () => {
-    const adapters: DefaultInspectorAdapters = {
+  it('should reset the current tables in the tables inspector', () => {
+    const adapters = {
       tables: new TablesAdapter(),
-      requests: {} as never,
-      expression: {} as never,
-    };
+    } as DefaultInspectorAdapters;
+
+    const resetSpy = jest.spyOn(adapters.tables, 'reset');
+
     mergeTables.fn(null, { layerIds: ['first', 'second'], tables: [sampleTable1, sampleTable2] }, {
       inspectorAdapters: adapters,
     } as ExecutionContext<DefaultInspectorAdapters, ExpressionValueSearchContext>);
-    expect(adapters.tables!.tables.first).toBe(sampleTable1);
-    expect(adapters.tables!.tables.second).toBe(sampleTable2);
+    expect(resetSpy).toHaveBeenCalled();
   });
 
   it('should pass the date range along', () => {

x-pack/plugins/lens/common/expressions/xy_chart/xy_chart.ts

@@ -10,6 +10,7 @@ import type { ExpressionValueSearchContext } from '../../../../../../src/plugins
 import type { LensMultiTable } from '../../types';
 import type { XYArgs } from './xy_args';
 import { fittingFunctionDefinitions } from './fitting_function';
+import { logDataTable } from '../expressions_utils';
 
 export interface XYChartProps {
   data: LensMultiTable;
@@ -157,6 +158,9 @@ export const xyChart: ExpressionFunctionDefinition<
     },
   },
   fn(data: LensMultiTable, args: XYArgs, handlers) {
+    if (handlers?.inspectorAdapters?.tables) {
+      logDataTable(handlers.inspectorAdapters.tables, data.tables);
+    }
     return {
       type: 'render',
       as: 'lens_xy_chart_renderer',

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

@@ -410,7 +410,7 @@ describe('workspace_panel', () => {
       first: mockDatasource.publicAPIMock,
     };
     mockDatasource.toExpression.mockReturnValue('datasource');
-    mockDatasource.getLayers.mockReturnValue(['first']);
+    mockDatasource.getLayers.mockReturnValue(['table1']);
 
     const mounted = await mountWithProvider(
       <WorkspacePanel

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

@@ -52,7 +52,7 @@ import {
   getUnknownVisualizationTypeError,
 } from '../../error_helper';
 import { getMissingIndexPattern, validateDatasourceAndVisualization } from '../state_helpers';
-import { DefaultInspectorAdapters } from '../../../../../../../src/plugins/expressions/common';
+import type { DefaultInspectorAdapters } from '../../../../../../../src/plugins/expressions/common';
 import {
   onActiveDataChange,
   useLensDispatch,
@@ -68,10 +68,12 @@ import {
   selectSearchSessionId,
   selectAutoApplyEnabled,
   selectTriggerApplyChanges,
+  selectDatasourceLayers,
 } from '../../../state_management';
 import type { LensInspector } from '../../../lens_inspector_service';
 import { inferTimeField } from '../../../utils';
 import { setChangesApplied } from '../../../state_management/lens_slice';
+import type { Datatable } from '../../../../../../../src/plugins/expressions/public';
 
 export interface WorkspacePanelProps {
   visualizationMap: VisualizationMap;
@@ -381,6 +383,7 @@ export const InnerWorkspacePanel = React.memo(function InnerWorkspacePanel({
     if (localState.expressionToRender === null) {
       return renderEmptyWorkspace();
     }
+
     return (
       <VisualizationWrapper
         expression={localState.expressionToRender}
@@ -391,6 +394,7 @@ export const InnerWorkspacePanel = React.memo(function InnerWorkspacePanel({
         localState={{ ...localState, configurationValidationError, missingRefsErrors }}
         ExpressionRendererComponent={ExpressionRendererComponent}
         application={core.application}
+        datasourceMap={datasourceMap}
         activeDatasourceId={activeDatasourceId}
       />
     );
@@ -455,6 +459,7 @@ export const VisualizationWrapper = ({
   ExpressionRendererComponent,
   application,
   activeDatasourceId,
+  datasourceMap,
 }: {
   expression: string | null | undefined;
   framePublicAPI: FramePublicAPI;
@@ -473,6 +478,7 @@ export const VisualizationWrapper = ({
   ExpressionRendererComponent: ReactExpressionRendererType;
   application: ApplicationStart;
   activeDatasourceId: string | null;
+  datasourceMap: DatasourceMap;
 }) => {
   const context = useLensSelector(selectExecutionContext);
   const searchContext: ExecutionContextSearch = useMemo(
@@ -487,16 +493,26 @@ export const VisualizationWrapper = ({
     [context]
   );
   const searchSessionId = useLensSelector(selectSearchSessionId);
-
+  const datasourceLayers = useLensSelector((state) => selectDatasourceLayers(state, datasourceMap));
   const dispatchLens = useLensDispatch();
+  const [defaultLayerId] = Object.keys(datasourceLayers);
 
   const onData$ = useCallback(
     (data: unknown, adapters?: Partial<DefaultInspectorAdapters>) => {
       if (adapters && adapters.tables) {
-        dispatchLens(onActiveDataChange({ ...adapters.tables.tables }));
+        dispatchLens(
+          onActiveDataChange(
+            Object.entries(adapters.tables?.tables).reduce<Record<string, Datatable>>(
+              (acc, [key, value], index, tables) => ({
+                [tables.length === 1 ? defaultLayerId : key]: value,
+              }),
+              {}
+            )
+          )
+        );
       }
     },
-    [dispatchLens]
+    [defaultLayerId, dispatchLens]
   );
 
   function renderFixAction(