[ES|QL] [Discover] Keeps the preferred chart configuration when possi…

…ble (elastic#197453) ## Summary Closes elastic#184631 It keeps the chart configuration when the user is doing actions compatible with the current query such as: - Adding a where filter (by clicking the table, the sidebar, the chart) - Changes the breakdown field and the field type is compatible with the current chart - Changing to a compatible chart type (from example from bar to line or pie to treemap) - Changing the query that doesnt affect the generated columns mapped to a chart. For example adding a limit or creating a runtime field etc. The logic depends on the suggestions. If the suggestions return the preferred chart type, then we are going to use this. So it really depends on the api and the type / number of columns. It is as smarter as it can in order to not create bugs. I am quite happy with the result. It is much better than what we have so far. ![meow](https://github.com/user-attachments/assets/c4249e5e-e785-4e57-8651-d1f660f5a61a) ### Next steps I would love to do the same on the dahsboard too, needs more time though. But the changes made here will def work in favor ### Checklist - [x] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios - [x] [Flaky Test Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was used on any tests changed --------- Co-authored-by: Marta Bondyra <[email protected]>
mgadewoll · Nov 7, 2024 · 080e352 · 080e352
1 parent 30d48f0
commit 080e352
Show file tree

Hide file tree

Showing 16 changed files with 858 additions and 81 deletions.
diff --git 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';
diff --git 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;
+};
diff --git 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];
+  }
+};
diff --git 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',
+}
diff --git 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,17 +600,39 @@ 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: '',
       textBasedColumns: columns,
       query: query && isOfAggregateQueryType(query) ? query : undefined,
     };
     const allSuggestions = isPlainRecord
-      ? this.lensSuggestionsApi(context, dataView, ['lnsDatatable']) ?? []
+      ? this.lensSuggestionsApi(
+          context,
+          dataView,
+          ['lnsDatatable'],
+          preferredChartType,
+          visAttributes
+        ) ?? []
       : [];
 
     return allSuggestions;

diff --git 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);
+    });
+  });
 });
diff --git 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;