[8.x] [ES|QL] [Discover] Keeps the preferred chart configuration when possible (#197453) (#199069)

# Backport This will backport the following commits from `main` to `8.x`: - [[ES|QL] [Discover] Keeps the preferred chart configuration when possible (#197453)](https://github.com/elastic/kibana/pull/197453)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Stratoula Kalafateli <efstratia.kalafateli@elastic.co>
2025-04-23 09:19:04 -04:00 · 2024-11-06 19:04:19 +11:00 · 2024-11-06 19:04:19 +11:00 · 1263c97491
commit 1263c97491
parent 113055b89d
16 changed files with 858 additions and 81 deletions
--- a/packages/kbn-visualization-utils/index.ts
+++ b/packages/kbn-visualization-utils/index.ts
@ -11,3 +11,6 @@ export { getTimeZone } from './src/get_timezone';
 export { getLensAttributesFromSuggestion } from './src/get_lens_attributes';
 export { TooltipWrapper } from './src/tooltip_wrapper';
 export { useDebouncedValue } from './src/debounced_value';
+export { ChartType } from './src/types';
+export { getDatasourceId } from './src/get_datasource_id';
+export { mapVisToChartType } from './src/map_vis_to_chart_type';
--- a/packages/kbn-visualization-utils/src/get_datasource_id.ts
+++ b/packages/kbn-visualization-utils/src/get_datasource_id.ts
@ -0,0 +1,17 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+export const getDatasourceId = (datasourceStates: Record<string, unknown>) => {
+  const datasourceId: 'formBased' | 'textBased' | undefined = [
+    'formBased' as const,
+    'textBased' as const,
+  ].find((key) => Boolean(datasourceStates[key]));
+
+  return datasourceId;
+};
--- a/packages/kbn-visualization-utils/src/map_vis_to_chart_type.ts
+++ b/packages/kbn-visualization-utils/src/map_vis_to_chart_type.ts
@ -0,0 +1,32 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { ChartType, LensVisualizationType } from './types';
+
+type ValueOf<T> = T[keyof T];
+type LensToChartMap = {
+  [K in ValueOf<typeof LensVisualizationType>]: ChartType;
+};
+const lensTypesToChartTypes: LensToChartMap = {
+  [LensVisualizationType.XY]: ChartType.XY,
+  [LensVisualizationType.Metric]: ChartType.Metric,
+  [LensVisualizationType.LegacyMetric]: ChartType.Metric,
+  [LensVisualizationType.Pie]: ChartType.Pie,
+  [LensVisualizationType.Heatmap]: ChartType.Heatmap,
+  [LensVisualizationType.Gauge]: ChartType.Gauge,
+  [LensVisualizationType.Datatable]: ChartType.Table,
+};
+function isLensVisualizationType(value: string): value is LensVisualizationType {
+  return Object.values(LensVisualizationType).includes(value as LensVisualizationType);
+}
+export const mapVisToChartType = (visualizationType: string) => {
+  if (isLensVisualizationType(visualizationType)) {
+    return lensTypesToChartTypes[visualizationType];
+  }
+};
--- a/packages/kbn-visualization-utils/src/types.ts
+++ b/packages/kbn-visualization-utils/src/types.ts
@ -43,3 +43,30 @@ export interface Suggestion<T = unknown, V = unknown> {
  changeType: TableChangeType;
  keptLayerIds: string[];
 }
+
+export enum ChartType {
+  XY = 'XY',
+  Gauge = 'Gauge',
+  Bar = 'Bar',
+  Line = 'Line',
+  Area = 'Area',
+  Donut = 'Donut',
+  Heatmap = 'Heatmap',
+  Metric = 'Metric',
+  Treemap = 'Treemap',
+  Tagcloud = 'Tagcloud',
+  Waffle = 'Waffle',
+  Pie = 'Pie',
+  Mosaic = 'Mosaic',
+  Table = 'Table',
+}
+
+export enum LensVisualizationType {
+  XY = 'lnsXY',
+  Metric = 'lnsMetric',
+  Pie = 'lnsPie',
+  Heatmap = 'lnsHeatmap',
+  Gauge = 'lnsGauge',
+  Datatable = 'lnsDatatable',
+  LegacyMetric = 'lnsLegacyMetric',
+}
--- a/src/plugins/unified_histogram/public/services/lens_vis_service.ts
+++ b/src/plugins/unified_histogram/public/services/lens_vis_service.ts
@ -27,7 +27,11 @@ import type {
 import type { AggregateQuery, TimeRange } from '@kbn/es-query';
 import { getAggregateQueryMode, isOfAggregateQueryType } from '@kbn/es-query';
 import { i18n } from '@kbn/i18n';
-import { getLensAttributesFromSuggestion } from '@kbn/visualization-utils';
+import {
+  getLensAttributesFromSuggestion,
+  ChartType,
+  mapVisToChartType,
+} from '@kbn/visualization-utils';
 import { LegendSize } from '@kbn/visualizations-plugin/public';
 import { XYConfiguration } from '@kbn/visualizations-plugin/common';
 import type { Datatable, DatatableColumn } from '@kbn/expressions-plugin/common';
@ -42,6 +46,7 @@ import {
  isSuggestionShapeAndVisContextCompatible,
  deriveLensSuggestionFromLensAttributes,
  type QueryParams,
+  injectESQLQueryIntoLensLayers,
 } from '../utils/external_vis_context';
 import { computeInterval } from '../utils/compute_interval';
 import { fieldSupportsBreakdown } from '../utils/field_supports_breakdown';
@ -147,7 +152,10 @@ export class LensVisService {
      externalVisContextStatus: UnifiedHistogramExternalVisContextStatus
    ) => void;
  }) => {
-    const allSuggestions = this.getAllSuggestions({ queryParams });
+    const allSuggestions = this.getAllSuggestions({
+      queryParams,
+      preferredVisAttributes: externalVisContext?.attributes,
+    });

    const suggestionState = this.getCurrentSuggestionState({
      externalVisContext,
@ -252,6 +260,7 @@ export class LensVisService {
      const histogramSuggestionForESQL = this.getHistogramSuggestionForESQL({
        queryParams,
        breakdownField,
+        preferredVisAttributes: externalVisContext?.attributes,
      });
      if (histogramSuggestionForESQL) {
        // In case if histogram suggestion, we want to empty the array and push the new suggestion
@ -463,9 +472,11 @@ export class LensVisService {
  private getHistogramSuggestionForESQL = ({
    queryParams,
    breakdownField,
+    preferredVisAttributes,
  }: {
    queryParams: QueryParams;
    breakdownField?: DataViewField;
+    preferredVisAttributes?: UnifiedHistogramVisContext['attributes'];
  }): Suggestion | undefined => {
    const { dataView, query, timeRange, columns } = queryParams;
    const breakdownColumn = breakdownField?.name
@ -510,7 +521,22 @@ export class LensVisService {
      if (breakdownColumn) {
        context.textBasedColumns.push(breakdownColumn);
      }
-      const suggestions = this.lensSuggestionsApi(context, dataView, ['lnsDatatable']) ?? [];
+
+      // here the attributes contain the main query and not the histogram one
+      const updatedAttributesWithQuery = preferredVisAttributes
+        ? injectESQLQueryIntoLensLayers(preferredVisAttributes, {
+            esql: esqlQuery,
+          })
+        : undefined;
+
+      const suggestions =
+        this.lensSuggestionsApi(
+          context,
+          dataView,
+          ['lnsDatatable'],
+          ChartType.XY,
+          updatedAttributesWithQuery
+        ) ?? [];
      if (suggestions.length) {
        const suggestion = suggestions[0];
        const suggestionVisualizationState = Object.assign({}, suggestion?.visualizationState);
@ -574,9 +600,25 @@ export class LensVisService {
    );
  };

-  private getAllSuggestions = ({ queryParams }: { queryParams: QueryParams }): Suggestion[] => {
+  private getAllSuggestions = ({
+    queryParams,
+    preferredVisAttributes,
+  }: {
+    queryParams: QueryParams;
+    preferredVisAttributes?: UnifiedHistogramVisContext['attributes'];
+  }): Suggestion[] => {
    const { dataView, columns, query, isPlainRecord } = queryParams;

+    const preferredChartType = preferredVisAttributes
+      ? mapVisToChartType(preferredVisAttributes.visualizationType)
+      : undefined;
+
+    let visAttributes = preferredVisAttributes;
+
+    if (query && isOfAggregateQueryType(query) && preferredVisAttributes) {
+      visAttributes = injectESQLQueryIntoLensLayers(preferredVisAttributes, query);
+    }
+
    const context = {
      dataViewSpec: dataView?.toSpec(),
      fieldName: '',
@ -584,7 +626,13 @@ export class LensVisService {
      query: query && isOfAggregateQueryType(query) ? query : undefined,
    };
    const allSuggestions = isPlainRecord
-      ? this.lensSuggestionsApi(context, dataView, ['lnsDatatable']) ?? []
+      ? this.lensSuggestionsApi(
+          context,
+          dataView,
+          ['lnsDatatable'],
+          preferredChartType,
+          visAttributes
+        ) ?? []
      : [];

    return allSuggestions;
--- a/src/plugins/unified_histogram/public/utils/external_vis_context.test.ts
+++ b/src/plugins/unified_histogram/public/utils/external_vis_context.test.ts
@ -13,6 +13,7 @@ import {
  canImportVisContext,
  exportVisContext,
  isSuggestionShapeAndVisContextCompatible,
+  injectESQLQueryIntoLensLayers,
 } from './external_vis_context';
 import { getLensVisMock } from '../__mocks__/lens_vis';
 import { dataViewWithTimefieldMock } from '../__mocks__/data_view_with_timefield';
@ -162,4 +163,63 @@ describe('external_vis_context', () => {
      ).toBe(true);
    });
  });
+
+  describe('injectESQLQueryIntoLensLayers', () => {
+    it('should return the Lens attributes as they are for unknown datasourceId', async () => {
+      const attributes = {
+        visualizationType: 'lnsXY',
+        state: {
+          visualization: { preferredSeriesType: 'line' },
+          datasourceStates: { unknownId: { layers: {} } },
+        },
+      } as unknown as UnifiedHistogramVisContext['attributes'];
+      expect(injectESQLQueryIntoLensLayers(attributes, { esql: 'from foo' })).toStrictEqual(
+        attributes
+      );
+    });
+
+    it('should return the Lens attributes as they are for DSL config (formbased)', async () => {
+      const attributes = {
+        visualizationType: 'lnsXY',
+        state: {
+          visualization: { preferredSeriesType: 'line' },
+          datasourceStates: { formBased: { layers: {} } },
+        },
+      } as UnifiedHistogramVisContext['attributes'];
+      expect(injectESQLQueryIntoLensLayers(attributes, { esql: 'from foo' })).toStrictEqual(
+        attributes
+      );
+    });
+
+    it('should inject the query to the Lens attributes for ES|QL config (textbased)', async () => {
+      const attributes = {
+        visualizationType: 'lnsXY',
+        state: {
+          visualization: { preferredSeriesType: 'line' },
+          datasourceStates: { textBased: { layers: { layer1: { query: { esql: 'from foo' } } } } },
+        },
+      } as unknown as UnifiedHistogramVisContext['attributes'];
+
+      const expectedAttributes = {
+        ...attributes,
+        state: {
+          ...attributes.state,
+          datasourceStates: {
+            ...attributes.state.datasourceStates,
+            textBased: {
+              ...attributes.state.datasourceStates.textBased,
+              layers: {
+                layer1: {
+                  query: { esql: 'from foo | stats count(*)' },
+                },
+              },
+            },
+          },
+        },
+      } as unknown as UnifiedHistogramVisContext['attributes'];
+      expect(
+        injectESQLQueryIntoLensLayers(attributes, { esql: 'from foo | stats count(*)' })
+      ).toStrictEqual(expectedAttributes);
+    });
+  });
 });
--- a/src/plugins/unified_histogram/public/utils/external_vis_context.ts
+++ b/src/plugins/unified_histogram/public/utils/external_vis_context.ts
@ -7,9 +7,10 @@
 * License v3.0 only", or the "Server Side Public License, v 1".
 */

-import { isEqual } from 'lodash';
+import { isEqual, cloneDeep } from 'lodash';
 import type { DataView } from '@kbn/data-views-plugin/common';
 import type { AggregateQuery, Filter, Query, TimeRange } from '@kbn/es-query';
+import { getDatasourceId } from '@kbn/visualization-utils';
 import type { DatatableColumn } from '@kbn/expressions-plugin/common';
 import type { PieVisualizationState, Suggestion, XYState } from '@kbn/lens-plugin/public';
 import { UnifiedHistogramSuggestionType, UnifiedHistogramVisContext } from '../types';
@ -103,6 +104,42 @@ export const isSuggestionShapeAndVisContextCompatible = (
  );
 };

+export const injectESQLQueryIntoLensLayers = (
+  visAttributes: UnifiedHistogramVisContext['attributes'],
+  query: AggregateQuery
+) => {
+  const datasourceId = getDatasourceId(visAttributes.state.datasourceStates);
+
+  // if the datasource is formBased, we should not fix the query
+  if (!datasourceId || datasourceId === 'formBased') {
+    return visAttributes;
+  }
+
+  if (!visAttributes.state.datasourceStates[datasourceId]) {
+    return visAttributes;
+  }
+
+  const datasourceState = cloneDeep(visAttributes.state.datasourceStates[datasourceId]);
+
+  if (datasourceState && datasourceState.layers) {
+    Object.values(datasourceState.layers).forEach((layer) => {
+      if (!isEqual(layer.query, query)) {
+        layer.query = query;
+      }
+    });
+  }
+  return {
+    ...visAttributes,
+    state: {
+      ...visAttributes.state,
+      datasourceStates: {
+        ...visAttributes.state.datasourceStates,
+        [datasourceId]: datasourceState,
+      },
+    },
+  };
+};
+
 export function deriveLensSuggestionFromLensAttributes({
  externalVisContext,
  queryParams,
@ -122,10 +159,7 @@ export function deriveLensSuggestionFromLensAttributes({
      }

      // it should be one of 'formBased'/'textBased' and have value
-      const datasourceId: 'formBased' | 'textBased' | undefined = [
-        'formBased' as const,
-        'textBased' as const,
-      ].find((key) => Boolean(externalVisContext.attributes.state.datasourceStates[key]));
+      const datasourceId = getDatasourceId(externalVisContext.attributes.state.datasourceStates);

      if (!datasourceId) {
        return undefined;
--- a/test/functional/apps/discover/esql/_esql_view.ts
+++ b/test/functional/apps/discover/esql/_esql_view.ts
@ -676,6 +676,94 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
          `from logstash-* | sort @timestamp desc | limit 10000 | stats countB = count(bytes) by geo.dest | sort countB | where countB > 0\nAND \`geo.dest\`=="BT"`
        );
      });
+
+      it('should append a where clause by clicking the table without changing the chart type', async () => {
+        await discover.selectTextBaseLang();
+        const testQuery = `from logstash-* | sort @timestamp desc | limit 10000 | stats countB = count(bytes) by geo.dest | sort countB`;
+        await monacoEditor.setCodeEditorValue(testQuery);
+
+        await testSubjects.click('querySubmitButton');
+        await header.waitUntilLoadingHasFinished();
+        await discover.waitUntilSearchingHasFinished();
+        await unifiedFieldList.waitUntilSidebarHasLoaded();
+
+        // change the type to line
+        await testSubjects.click('unifiedHistogramEditFlyoutVisualization');
+        await header.waitUntilLoadingHasFinished();
+        await testSubjects.click('lnsChartSwitchPopover');
+        await testSubjects.click('lnsChartSwitchPopover_line');
+        await header.waitUntilLoadingHasFinished();
+        await testSubjects.click('applyFlyoutButton');
+
+        await dataGrid.clickCellFilterForButtonExcludingControlColumns(0, 1);
+        await header.waitUntilLoadingHasFinished();
+        await discover.waitUntilSearchingHasFinished();
+        await unifiedFieldList.waitUntilSidebarHasLoaded();
+
+        const editorValue = await monacoEditor.getCodeEditorValue();
+        expect(editorValue).to.eql(
+          `from logstash-* | sort @timestamp desc | limit 10000 | stats countB = count(bytes) by geo.dest | sort countB\n| WHERE \`geo.dest\`=="BT"`
+        );
+
+        // check that the type is still line
+        await testSubjects.click('unifiedHistogramEditFlyoutVisualization');
+        await header.waitUntilLoadingHasFinished();
+        const chartSwitcher = await testSubjects.find('lnsChartSwitchPopover');
+        const type = await chartSwitcher.getVisibleText();
+        expect(type).to.be('Line');
+      });
+
+      it('should append a where clause by clicking the table without changing the chart type nor the visualization state', async () => {
+        await discover.selectTextBaseLang();
+        const testQuery = `from logstash-* | sort @timestamp desc | limit 10000 | stats countB = count(bytes) by geo.dest | sort countB`;
+        await monacoEditor.setCodeEditorValue(testQuery);
+
+        await testSubjects.click('querySubmitButton');
+        await header.waitUntilLoadingHasFinished();
+        await discover.waitUntilSearchingHasFinished();
+        await unifiedFieldList.waitUntilSidebarHasLoaded();
+
+        // change the type to line
+        await testSubjects.click('unifiedHistogramEditFlyoutVisualization');
+        await header.waitUntilLoadingHasFinished();
+        await testSubjects.click('lnsChartSwitchPopover');
+        await testSubjects.click('lnsChartSwitchPopover_line');
+
+        // change the color to red
+        await testSubjects.click('lnsXY_yDimensionPanel');
+        const colorPickerInput = await testSubjects.find('~indexPattern-dimension-colorPicker');
+        await colorPickerInput.clearValueWithKeyboard();
+        await colorPickerInput.type('#ff0000');
+        await common.sleep(1000); // give time for debounced components to rerender
+
+        await header.waitUntilLoadingHasFinished();
+        await testSubjects.click('lns-indexPattern-dimensionContainerClose');
+        await testSubjects.click('applyFlyoutButton');
+
+        await dataGrid.clickCellFilterForButtonExcludingControlColumns(0, 1);
+        await header.waitUntilLoadingHasFinished();
+        await discover.waitUntilSearchingHasFinished();
+        await unifiedFieldList.waitUntilSidebarHasLoaded();
+
+        const editorValue = await monacoEditor.getCodeEditorValue();
+        expect(editorValue).to.eql(
+          `from logstash-* | sort @timestamp desc | limit 10000 | stats countB = count(bytes) by geo.dest | sort countB\n| WHERE \`geo.dest\`=="BT"`
+        );
+
+        // check that the type is still line
+        await testSubjects.click('unifiedHistogramEditFlyoutVisualization');
+        await header.waitUntilLoadingHasFinished();
+        const chartSwitcher = await testSubjects.find('lnsChartSwitchPopover');
+        const type = await chartSwitcher.getVisibleText();
+        expect(type).to.be('Line');
+
+        // check that the color is still red
+        await testSubjects.click('lnsXY_yDimensionPanel');
+        const colorPickerInputAfterFilter = await testSubjects.find(
+          '~indexPattern-dimension-colorPicker'
+        );
+        expect(await colorPickerInputAfterFilter.getAttribute('value')).to.be('#FF0000');
+      });
    });

    describe('histogram breakdown', () => {
--- a/test/functional/apps/discover/group3/_lens_vis.ts
+++ b/test/functional/apps/discover/group3/_lens_vis.ts
@ -288,7 +288,8 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
      await header.waitUntilLoadingHasFinished();
      await discover.waitUntilSearchingHasFinished();

-      expect(await getCurrentVisTitle()).to.be('Bar');
+      // Line has been retained although the query changed!
+      expect(await getCurrentVisTitle()).to.be('Line');

      await checkESQLHistogramVis(defaultTimespanESQL, '100');

@ -567,15 +568,14 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
      await testSubjects.existOrFail('partitionVisChart');
      expect(await discover.getVisContextSuggestionType()).to.be('lensSuggestion');

-      await monacoEditor.setCodeEditorValue(
-        'from logstash-* | stats averageB = avg(bytes) by extension.raw'
-      );
+      // reset to histogram
+      await monacoEditor.setCodeEditorValue('from logstash-*');
      await testSubjects.click('querySubmitButton');
      await header.waitUntilLoadingHasFinished();
      await discover.waitUntilSearchingHasFinished();

      expect(await getCurrentVisTitle()).to.be('Bar');
-      expect(await discover.getVisContextSuggestionType()).to.be('lensSuggestion');
+      expect(await discover.getVisContextSuggestionType()).to.be('histogramForESQL');

      await testSubjects.existOrFail('unsavedChangesBadge');

--- a/x-pack/plugins/lens/public/lens_suggestions_api/helpers.test.ts
+++ b/x-pack/plugins/lens/public/lens_suggestions_api/helpers.test.ts
@ -0,0 +1,206 @@
+/*
+ * 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 { DatatableColumn } from '@kbn/expressions-plugin/common';
+import { mergeSuggestionWithVisContext } from './helpers';
+import { mockAllSuggestions } from '../mocks';
+import type { TypedLensByValueInput } from '../embeddable/embeddable_component';
+
+const context = {
+  dataViewSpec: {
+    id: 'index1',
+    title: 'index1',
+    name: 'DataView',
+  },
+  fieldName: '',
+  textBasedColumns: [
+    {
+      id: 'field1',
+      name: 'field1',
+      meta: {
+        type: 'number',
+      },
+    },
+    {
+      id: 'field2',
+      name: 'field2',
+      meta: {
+        type: 'string',
+      },
+    },
+  ] as DatatableColumn[],
+  query: {
+    esql: 'FROM index1 | keep field1, field2',
+  },
+};
+
+describe('lens suggestions api helpers', () => {
+  describe('mergeSuggestionWithVisContext', () => {
+    it('should return the suggestion as it is if the visualization types do not match', async () => {
+      const suggestion = mockAllSuggestions[0];
+      const visAttributes = {
+        visualizationType: 'lnsXY',
+        state: {
+          visualization: {
+            preferredSeriesType: 'bar_stacked',
+          },
+          datasourceStates: { textBased: { layers: {} } },
+        },
+      } as unknown as TypedLensByValueInput['attributes'];
+      expect(mergeSuggestionWithVisContext({ suggestion, visAttributes, context })).toStrictEqual(
+        suggestion
+      );
+    });
+
+    it('should return the suggestion as it is if the context is not from ES|QL', async () => {
+      const nonESQLContext = {
+        dataViewSpec: {
+          id: 'index1',
+          title: 'index1',
+          name: 'DataView',
+        },
+        fieldName: 'field1',
+      };
+      const suggestion = mockAllSuggestions[0];
+      const visAttributes = {
+        visualizationType: 'lnsHeatmap',
+        state: {
+          visualization: {
+            preferredSeriesType: 'bar_stacked',
+          },
+          datasourceStates: { textBased: { layers: {} } },
+        },
+      } as unknown as TypedLensByValueInput['attributes'];
+      expect(
+        mergeSuggestionWithVisContext({ suggestion, visAttributes, context: nonESQLContext })
+      ).toStrictEqual(suggestion);
+    });
+
+    it('should return the suggestion as it is for DSL config (formbased)', async () => {
+      const suggestion = mockAllSuggestions[0];
+      const visAttributes = {
+        visualizationType: 'lnsHeatmap',
+        state: {
+          visualization: {
+            preferredSeriesType: 'bar_stacked',
+          },
+          datasourceStates: { formBased: { layers: {} } },
+        },
+      } as unknown as TypedLensByValueInput['attributes'];
+      expect(mergeSuggestionWithVisContext({ suggestion, visAttributes, context })).toStrictEqual(
+        suggestion
+      );
+    });
+
+    it('should return the suggestion as it is for columns that dont match the context', async () => {
+      const suggestion = mockAllSuggestions[0];
+      const visAttributes = {
+        visualizationType: 'lnsHeatmap',
+        state: {
+          visualization: {
+            shape: 'heatmap',
+          },
+          datasourceStates: {
+            textBased: {
+              layers: {
+                layer1: {
+                  index: 'layer1',
+                  query: {
+                    esql: 'FROM kibana_sample_data_flights | keep Dest, AvgTicketPrice',
+                  },
+                  columns: [
+                    {
+                      columnId: 'colA',
+                      fieldName: 'Dest',
+                      meta: {
+                        type: 'string',
+                      },
+                    },
+                    {
+                      columnId: 'colB',
+                      fieldName: 'AvgTicketPrice',
+                      meta: {
+                        type: 'number',
+                      },
+                    },
+                  ],
+                  timeField: 'timestamp',
+                },
+              },
+            },
+          },
+        },
+      } as unknown as TypedLensByValueInput['attributes'];
+      expect(mergeSuggestionWithVisContext({ suggestion, visAttributes, context })).toStrictEqual(
+        suggestion
+      );
+    });
+
+    it('should return the suggestion updated with the attributes if the visualization types and the context columns match', async () => {
+      const suggestion = mockAllSuggestions[0];
+      const visAttributes = {
+        visualizationType: 'lnsHeatmap',
+        state: {
+          visualization: {
+            shape: 'heatmap',
+            layerId: 'layer1',
+            layerType: 'data',
+            legend: {
+              isVisible: false,
+              position: 'left',
+              type: 'heatmap_legend',
+            },
+            gridConfig: {
+              type: 'heatmap_grid',
+              isCellLabelVisible: true,
+              isYAxisLabelVisible: false,
+              isXAxisLabelVisible: false,
+              isYAxisTitleVisible: false,
+              isXAxisTitleVisible: false,
+            },
+            valueAccessor: 'acc1',
+            xAccessor: 'acc2',
+          },
+          datasourceStates: {
+            textBased: {
+              layers: {
+                layer1: {
+                  index: 'layer1',
+                  query: {
+                    esql: 'FROM index1 | keep field1, field2',
+                  },
+                  columns: [
+                    {
+                      columnId: 'field2',
+                      fieldName: 'field2',
+                      meta: {
+                        type: 'string',
+                      },
+                    },
+                    {
+                      columnId: 'field1',
+                      fieldName: 'field1',
+                      meta: {
+                        type: 'number',
+                      },
+                    },
+                  ],
+                  timeField: 'timestamp',
+                },
+              },
+            },
+          },
+        },
+      } as unknown as TypedLensByValueInput['attributes'];
+      const updatedSuggestion = mergeSuggestionWithVisContext({
+        suggestion,
+        visAttributes,
+        context,
+      });
+      expect(updatedSuggestion.visualizationState).toStrictEqual(visAttributes.state.visualization);
+    });
+  });
+});
--- a/x-pack/plugins/lens/public/lens_suggestions_api/helpers.ts
+++ b/x-pack/plugins/lens/public/lens_suggestions_api/helpers.ts
@ -0,0 +1,76 @@
+/*
+ * 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 { VisualizeFieldContext } from '@kbn/ui-actions-plugin/public';
+import { getDatasourceId } from '@kbn/visualization-utils';
+import type { VisualizeEditorContext, Suggestion } from '../types';
+import type { TypedLensByValueInput } from '../embeddable/embeddable_component';
+
+/**
+ * Returns the suggestion updated with external visualization state for ES|QL charts
+ * The visualization state is merged with the suggestion if the datasource is textBased, the columns match the context and the visualization type matches
+ * @param suggestion the suggestion to be updated
+ * @param visAttributes the preferred visualization attributes
+ * @param context the lens suggestions api context as being set by the consumers
+ * @returns updated suggestion
+ */
+
+export function mergeSuggestionWithVisContext({
+  suggestion,
+  visAttributes,
+  context,
+}: {
+  suggestion: Suggestion;
+  visAttributes: TypedLensByValueInput['attributes'];
+  context: VisualizeFieldContext | VisualizeEditorContext;
+}): Suggestion {
+  if (
+    visAttributes.visualizationType !== suggestion.visualizationId ||
+    !('textBasedColumns' in context)
+  ) {
+    return suggestion;
+  }
+
+  // it should be one of 'formBased'/'textBased' and have value
+  const datasourceId = getDatasourceId(visAttributes.state.datasourceStates);
+
+  // if the datasource is formBased, we should not merge
+  if (!datasourceId || datasourceId === 'formBased') {
+    return suggestion;
+  }
+  const datasourceState = Object.assign({}, visAttributes.state.datasourceStates[datasourceId]);
+
+  // should be based on same columns
+  if (
+    !datasourceState?.layers ||
+    Object.values(datasourceState?.layers).some(
+      (layer) =>
+        layer.columns?.some(
+          (c: { fieldName: string }) =>
+            !context?.textBasedColumns?.find((col) => col.name === c.fieldName)
+        ) || layer.columns?.length !== context?.textBasedColumns?.length
+    )
+  ) {
+    return suggestion;
+  }
+  const layerIds = Object.keys(datasourceState.layers);
+  try {
+    return {
+      title: visAttributes.title,
+      visualizationId: visAttributes.visualizationType,
+      visualizationState: visAttributes.state.visualization,
+      keptLayerIds: layerIds,
+      datasourceState,
+      datasourceId,
+      columns: suggestion.columns,
+      changeType: suggestion.changeType,
+      score: suggestion.score,
+      previewIcon: suggestion.previewIcon,
+    };
+  } catch {
+    return suggestion;
+  }
+}
--- a/x-pack/plugins/lens/public/lens_suggestions_api/index.ts
+++ b/x-pack/plugins/lens/public/lens_suggestions_api/index.ts
@ -6,22 +6,12 @@
 */
 import type { VisualizeFieldContext } from '@kbn/ui-actions-plugin/public';
 import type { DataView } from '@kbn/data-views-plugin/public';
-import { getSuggestions } from './editor_frame_service/editor_frame/suggestion_helpers';
-import type { DatasourceMap, VisualizationMap, VisualizeEditorContext } from './types';
-import type { DataViewsState } from './state_management';
-
-export enum ChartType {
-  XY = 'XY',
-  Bar = 'Bar',
-  Line = 'Line',
-  Area = 'Area',
-  Donut = 'Donut',
-  Heatmap = 'Heat map',
-  Treemap = 'Treemap',
-  Tagcloud = 'Tag cloud',
-  Waffle = 'Waffle',
-  Table = 'Table',
-}
+import type { ChartType } from '@kbn/visualization-utils';
+import { getSuggestions } from '../editor_frame_service/editor_frame/suggestion_helpers';
+import type { DatasourceMap, VisualizationMap, VisualizeEditorContext } from '../types';
+import type { DataViewsState } from '../state_management';
+import type { TypedLensByValueInput } from '../embeddable/embeddable_component';
+import { mergeSuggestionWithVisContext } from './helpers';

 interface SuggestionsApiProps {
  context: VisualizeFieldContext | VisualizeEditorContext;
@ -30,6 +20,7 @@ interface SuggestionsApiProps {
  datasourceMap?: DatasourceMap;
  excludedVisualizations?: string[];
  preferredChartType?: ChartType;
+  preferredVisAttributes?: TypedLensByValueInput['attributes'];
 }

 export const suggestionsApi = ({
@ -39,6 +30,7 @@ export const suggestionsApi = ({
  visualizationMap,
  excludedVisualizations,
  preferredChartType,
+  preferredVisAttributes,
 }: SuggestionsApiProps) => {
  const initialContext = context;
  if (!datasourceMap || !visualizationMap || !dataView.id) return undefined;
@ -79,32 +71,7 @@ export const suggestionsApi = ({
    dataViews,
  });
  if (!suggestions.length) return [];
-  // check if there is an XY chart suggested
-  // if user has requested for a line or area, we want to sligthly change the state
-  // to return line / area instead of a bar chart
-  const chartType = preferredChartType?.toLowerCase();
-  const XYSuggestion = suggestions.find((sug) => sug.visualizationId === 'lnsXY');
-  if (XYSuggestion && chartType && ['area', 'line'].includes(chartType)) {
-    const visualizationState = visualizationMap[
-      XYSuggestion.visualizationId
-    ]?.switchVisualizationType?.(chartType, XYSuggestion?.visualizationState);
-    return [
-      {
-        ...XYSuggestion,
-        visualizationState,
-      },
-    ];
-  }
-  // in case the user asks for another type (except from area, line) check if it exists
-  // in suggestions and return this instead
-  if (suggestions.length > 1 && preferredChartType) {
-    const suggestionFromModel = suggestions.find(
-      (s) => s.title.includes(preferredChartType) || s.visualizationId.includes(preferredChartType)
-    );
-    if (suggestionFromModel) {
-      return [suggestionFromModel];
-    }
-  }
+
  const activeVisualization = suggestions[0];
  if (
    activeVisualization.incomplete ||
@ -126,7 +93,46 @@ export const suggestionsApi = ({
    visualizationState: activeVisualization.visualizationState,
    dataViews,
  }).filter((sug) => !sug.hide && sug.visualizationId !== 'lnsLegacyMetric');
+
+  // check if there is an XY chart suggested
+  // if user has requested for a line or area, we want to sligthly change the state
+  // to return line / area instead of a bar chart
+  const chartType = preferredChartType?.toLowerCase();
+  const XYSuggestion = newSuggestions.find((s) => s.visualizationId === 'lnsXY');
+  // a type can be area, line, area_stacked, area_percentage etc
+  const isAreaOrLine = ['area', 'line'].some((type) => chartType?.includes(type));
+  if (XYSuggestion && chartType && isAreaOrLine) {
+    const visualizationState = visualizationMap[
+      XYSuggestion.visualizationId
+    ]?.switchVisualizationType?.(chartType, XYSuggestion?.visualizationState);
+
+    return [
+      {
+        ...XYSuggestion,
+        visualizationState,
+      },
+    ];
+  }
+  // in case the user asks for another type (except from area, line) check if it exists
+  // in suggestions and return this instead
  const suggestionsList = [activeVisualization, ...newSuggestions];
+  if (suggestionsList.length > 1 && preferredChartType) {
+    const compatibleSuggestion = suggestionsList.find(
+      (s) => s.title.includes(preferredChartType) || s.visualizationId.includes(preferredChartType)
+    );
+
+    if (compatibleSuggestion) {
+      const suggestion = preferredVisAttributes
+        ? mergeSuggestionWithVisContext({
+            suggestion: compatibleSuggestion,
+            visAttributes: preferredVisAttributes,
+            context,
+          })
+        : compatibleSuggestion;
+
+      return [suggestion];
+    }
+  }

  // if there is no preference from the user, send everything
  // until we separate the text based suggestions logic from the dataview one,
--- a/x-pack/plugins/lens/public/lens_suggestions_api/lens_suggestions_api.test.ts
+++ b/x-pack/plugins/lens/public/lens_suggestions_api/lens_suggestions_api.test.ts
@ -6,9 +6,11 @@
 */
 import type { DataView } from '@kbn/data-views-plugin/public';
 import type { DatatableColumn } from '@kbn/expressions-plugin/common';
-import { createMockVisualization, DatasourceMock, createMockDatasource } from './mocks';
-import { DatasourceSuggestion } from './types';
-import { suggestionsApi, ChartType } from './lens_suggestions_api';
+import { ChartType } from '@kbn/visualization-utils';
+import { createMockVisualization, DatasourceMock, createMockDatasource } from '../mocks';
+import { DatasourceSuggestion } from '../types';
+import { suggestionsApi } from '.';
+import type { TypedLensByValueInput } from '../embeddable/embeddable_component';

 const generateSuggestion = (state = {}, layerId: string = 'first'): DatasourceSuggestion => ({
  state,
@ -264,6 +266,9 @@ describe('suggestionsApi', () => {
    datasourceMap.textBased.getDatasourceSuggestionsForVisualizeField.mockReturnValue([
      generateSuggestion(),
    ]);
+    datasourceMap.textBased.getDatasourceSuggestionsFromCurrentState.mockReturnValue([
+      generateSuggestion(),
+    ]);
    const context = {
      dataViewSpec: {
        id: 'index1',
@ -284,8 +289,7 @@ describe('suggestionsApi', () => {
      preferredChartType: ChartType.Line,
    });
    expect(suggestions?.length).toEqual(1);
-    expect(suggestions?.[0]).toMatchInlineSnapshot(
-      `
+    expect(suggestions?.[0]).toMatchInlineSnapshot(`
      Object {
        "changeType": "unchanged",
        "columns": 0,
@ -302,8 +306,111 @@ describe('suggestionsApi', () => {
          "preferredSeriesType": "line",
        },
      }
-    `
-    );
+    `);
+  });
+
+  test('returns the suggestion with the preferred attributes ', async () => {
+    const dataView = { id: 'index1' } as unknown as DataView;
+    const visualizationMap = {
+      lnsXY: {
+        ...mockVis,
+        switchVisualizationType(seriesType: string, state: unknown) {
+          return {
+            ...(state as Record<string, unknown>),
+            preferredSeriesType: seriesType,
+          };
+        },
+        getSuggestions: () => [
+          {
+            score: 0.8,
+            title: 'bar',
+            state: {
+              preferredSeriesType: 'bar_stacked',
+              legend: {
+                isVisible: true,
+                position: 'right',
+              },
+            },
+            previewIcon: 'empty',
+            visualizationId: 'lnsXY',
+          },
+          {
+            score: 0.8,
+            title: 'Test2',
+            state: {},
+            previewIcon: 'empty',
+          },
+          {
+            score: 0.8,
+            title: 'Test2',
+            state: {},
+            previewIcon: 'empty',
+            incomplete: true,
+          },
+        ],
+      },
+    };
+    datasourceMap.textBased.getDatasourceSuggestionsForVisualizeField.mockReturnValue([
+      generateSuggestion(),
+    ]);
+    datasourceMap.textBased.getDatasourceSuggestionsFromCurrentState.mockReturnValue([
+      generateSuggestion(),
+    ]);
+    const context = {
+      dataViewSpec: {
+        id: 'index1',
+        title: 'index1',
+        name: 'DataView',
+      },
+      fieldName: '',
+      textBasedColumns: textBasedQueryColumns,
+      query: {
+        esql: 'FROM "index1" | keep field1, field2',
+      },
+    };
+    const suggestions = suggestionsApi({
+      context,
+      dataView,
+      datasourceMap,
+      visualizationMap,
+      preferredChartType: ChartType.XY,
+      preferredVisAttributes: {
+        visualizationType: 'lnsXY',
+        state: {
+          visualization: {
+            preferredSeriesType: 'bar_stacked',
+            legend: {
+              isVisible: false,
+              position: 'left',
+            },
+          },
+          datasourceStates: { textBased: { layers: {} } },
+        },
+      } as unknown as TypedLensByValueInput['attributes'],
+    });
+    expect(suggestions?.length).toEqual(1);
+    expect(suggestions?.[0]).toMatchInlineSnapshot(`
+      Object {
+        "changeType": "unchanged",
+        "columns": 0,
+        "datasourceId": "textBased",
+        "datasourceState": Object {
+          "layers": Object {},
+        },
+        "keptLayerIds": Array [],
+        "previewIcon": "empty",
+        "score": 0.8,
+        "title": undefined,
+        "visualizationId": "lnsXY",
+        "visualizationState": Object {
+          "legend": Object {
+            "isVisible": false,
+            "position": "left",
+          },
+          "preferredSeriesType": "bar_stacked",
+        },
+      }
+    `);
  });

  test('filters out the suggestion if exists on excludedVisualizations', async () => {
--- a/x-pack/plugins/lens/public/lens_suggestions_api/readme.md
+++ b/x-pack/plugins/lens/public/lens_suggestions_api/readme.md
@ -0,0 +1,77 @@
+# Lens Suggestions API
+
+This document provides an overview of the Lens Suggestions API. It is used mostly for suggesting ES|QL charts based on an ES|QL query. It is used by the observability assistant, Discover and Dashboards ES|QL charts.
+
+## Overview
+
+The Lens Suggestions API is designed to provide suggestions for visualizations based on a given ES|QL query. It helps users to quickly find the most relevant visualizations for their data.
+
+## Getting Started
+
+To use the Lens Suggestions API, you need to import it from the Lens plugin:
+
+```typescript
+import useAsync from 'react-use/lib/useAsync';
+
+const lensHelpersAsync = useAsync(() => {
+    return lensService?.stateHelperApi() ?? Promise.resolve(null);
+  }, [lensService]);
+
+  if (lensHelpersAsync.value) {
+    const suggestionsApi = lensHelpersAsync.value.suggestions;
+  }
+```
+
+## The api
+
+The api returns an array of suggestions.
+
+#### Parameters
+
+  dataView: DataView;
+  visualizationMap?: VisualizationMap;
+  datasourceMap?: DatasourceMap;
+  excludedVisualizations?: string[];
+  preferredChartType?: ChartType;
+  preferredVisAttributes?: TypedLensByValueInput['attributes'];
+
+- `context`: The context as descibed by the VisualizeFieldContext. 
+- `dataView`: The dataView, can be an adhoc one too. For ES|QL you can create a dataview like this
+
+```typescript
+const indexName = (await getIndexForESQLQuery({ dataViews })) ?? '*';
+const dataView = await getESQLAdHocDataview(`from ${indexName}`, dataViews);
+```
+Optional parameters:
+- `preferredChartType`: Use this if you want the suggestions api to prioritize a specific suggestion type.
+- `preferredVisAttributes`: Use this with the preferredChartType if you want to prioritize a specific suggestion type with a non-default visualization state.
+
+#### Returns
+
+An array of suggestion objects
+
+## Example Usage
+
+```typescript
+const abc = new AbortController();
+
+const columns = await getESQLQueryColumns({
+  esqlQuery,
+  search: dataService.search.search,
+  signal: abc.signal,
+  timeRange: dataService.query.timefilter.timefilter.getAbsoluteTime(),
+});
+
+const context = {
+  dataViewSpec: dataView?.toSpec(false),
+  fieldName: '',
+  textBasedColumns: columns,
+  query: { esql: esqlQuery },
+};
+
+const chartSuggestions = lensHelpersAsync.value.suggestions(context, dataView);
+
+suggestions.forEach(suggestion => {
+  console.log(`Suggestion: ${suggestion.title}, Score: ${suggestion.score}`);
+});
+```
--- a/x-pack/plugins/lens/public/plugin.ts
+++ b/x-pack/plugins/lens/public/plugin.ts
@ -62,6 +62,7 @@ import {
  ContentManagementPublicStart,
 } from '@kbn/content-management-plugin/public';
 import { i18n } from '@kbn/i18n';
+import type { ChartType } from '@kbn/visualization-utils';
 import type { ServerlessPluginStart } from '@kbn/serverless/public';
 import { LicensingPluginStart } from '@kbn/licensing-plugin/public';
 import type { EditorFrameService as EditorFrameServiceType } from './editor_frame_service';
@ -137,7 +138,7 @@ import {
 } from '../common/content_management';
 import type { EditLensConfigurationProps } from './app_plugin/shared/edit_on_the_fly/get_edit_lens_configuration';
 import { savedObjectToEmbeddableAttributes } from './lens_attribute_service';
-import { ChartType } from './lens_suggestions_api';
+import type { TypedLensByValueInput } from './embeddable/embeddable_component';

 export type { SaveProps } from './app_plugin';

@ -281,7 +282,8 @@ export type LensSuggestionsApi = (
  context: VisualizeFieldContext | VisualizeEditorContext,
  dataViews: DataView,
  excludedVisualizations?: string[],
-  preferredChartType?: ChartType
+  preferredChartType?: ChartType,
+  preferredVisAttributes?: TypedLensByValueInput['attributes']
 ) => Suggestion[] | undefined;

 export class LensPlugin {
@ -713,7 +715,13 @@ export class LensPlugin {
        return {
          formula: createFormulaPublicApi(),
          chartInfo: createChartInfoApi(startDependencies.dataViews, this.editorFrameService),
-          suggestions: (context, dataView, excludedVisualizations, preferredChartType) => {
+          suggestions: (
+            context,
+            dataView,
+            excludedVisualizations,
+            preferredChartType,
+            preferredVisAttributes
+          ) => {
            return suggestionsApi({
              datasourceMap,
              visualizationMap,
@ -721,6 +729,7 @@ export class LensPlugin {
              dataView,
              excludedVisualizations,
              preferredChartType,
+              preferredVisAttributes,
            });
          },
        };
--- a/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/visualize_esql.tsx
+++ b/x-pack/plugins/observability_solution/observability_ai_assistant_app/public/functions/visualize_esql.tsx
@ -37,7 +37,7 @@ import {
  VisualizeESQLUserIntention,
 } from '@kbn/observability-ai-assistant-plugin/public';
 import type { UiActionsStart } from '@kbn/ui-actions-plugin/public';
-import { getLensAttributesFromSuggestion } from '@kbn/visualization-utils';
+import { getLensAttributesFromSuggestion, ChartType } from '@kbn/visualization-utils';
 import React, { useCallback, useContext, useEffect, useMemo, useState } from 'react';
 import ReactDOM from 'react-dom';
 import useAsync from 'react-use/lib/useAsync';
@ -48,19 +48,6 @@ import type {
 } from '../../common/functions/visualize_esql';
 import { ObservabilityAIAssistantAppPluginStartDependencies } from '../types';

-enum ChartType {
-  XY = 'XY',
-  Bar = 'Bar',
-  Line = 'Line',
-  Area = 'Area',
-  Donut = 'Donut',
-  Heatmap = 'Heat map',
-  Treemap = 'Treemap',
-  Tagcloud = 'Tag cloud',
-  Waffle = 'Waffle',
-  Table = 'Table',
-}
-
 interface VisualizeESQLProps {
  /** Lens start contract, get the ES|QL charts suggestions api */
  lens: LensPublicStart;