[Lens] value based ES|QL data source (#171081)

## Summary Closes #167631 Enhances the ES|QL lens datasource to receive a DataTable as an input (so no datafetching happens in lens) --------- Co-authored-by: kibanamachine <[email protected]> Co-authored-by: Stratoula Kalafateli <[email protected]> Co-authored-by: Stratoula Kalafateli <[email protected]> Co-authored-by: Julia Rechkunova <[email protected]>
elastic · Jan 26, 2024 · 0e3513a · 0e3513a
1 parent 3f92578
commit 0e3513a
Show file tree

Hide file tree

Showing 15 changed files with 250 additions and 71 deletions.
diff --git a/src/plugins/discover/public/application/main/components/layout/discover_histogram_layout.tsx b/src/plugins/discover/public/application/main/components/layout/discover_histogram_layout.tsx
@@ -6,10 +6,11 @@
  * Side Public License, v 1.
  */
 
-import React, { useCallback } from 'react';
+import React, { useCallback, useMemo } from 'react';
 import { UnifiedHistogramContainer } from '@kbn/unified-histogram-plugin/public';
 import { css } from '@emotion/react';
 import useObservable from 'react-use/lib/useObservable';
+import { Datatable } from '@kbn/expressions-plugin/common';
 import { useDiscoverHistogram } from './use_discover_histogram';
 import { type DiscoverMainContentProps, DiscoverMainContent } from './discover_main_content';
 import { useAppStateSelector } from '../../services/discover_app_state_container';
@@ -40,6 +41,7 @@ export const DiscoverHistogramLayout = ({
     isPlainRecord,
   });
 
+  const datatable = useObservable(dataState.data$.documents$);
   const renderCustomChartToggleActions = useCallback(
     () =>
       React.isValidElement(panelsToggle)
@@ -48,6 +50,21 @@ export const DiscoverHistogramLayout = ({
     [panelsToggle]
   );
 
+  const table: Datatable | undefined = useMemo(() => {
+    if (
+      isPlainRecord &&
+      datatable &&
+      datatable &&
+      ['partial', 'complete'].includes(datatable.fetchStatus)
+    ) {
+      return {
+        type: 'datatable' as 'datatable',
+        rows: datatable.result!.map((r) => r.raw),
+        columns: datatable.textBasedQueryColumns || [],
+      };
+    }
+  }, [datatable, isPlainRecord]);
+
   // Initialized when the first search has been requested or
   // when in text-based mode since search sessions are not supported
   if (!searchSessionId && !isPlainRecord) {
@@ -59,6 +76,7 @@ export const DiscoverHistogramLayout = ({
       {...unifiedHistogramProps}
       searchSessionId={searchSessionId}
       requestAdapter={dataState.inspectorAdapters.requests}
+      table={table}
       container={container}
       css={histogramLayoutCss}
       renderCustomChartToggleActions={renderCustomChartToggleActions}

diff --git a/src/plugins/unified_histogram/public/chart/chart.tsx b/src/plugins/unified_histogram/public/chart/chart.tsx
@@ -20,6 +20,8 @@ import type { DataView, DataViewField } from '@kbn/data-views-plugin/public';
 import type { LensEmbeddableInput } from '@kbn/lens-plugin/public';
 import type { AggregateQuery, Filter, Query, TimeRange } from '@kbn/es-query';
 import { Subject } from 'rxjs';
+import type { LensAttributes } from '@kbn/lens-embeddable-utils';
+import type { TextBasedPersistedState } from '@kbn/lens-plugin/public/datasources/text_based/types';
 import { Histogram } from './histogram';
 import type {
   UnifiedHistogramBreakdownContext,
@@ -267,6 +269,37 @@ export function Chart({
     });
   }
 
+  const removeTables = (attributes: LensAttributes) => {
+    if (!attributes.state.datasourceStates.textBased) {
+      return attributes;
+    }
+    const layers = attributes.state.datasourceStates.textBased?.layers;
+
+    const newState = {
+      ...attributes,
+      state: {
+        ...attributes.state,
+        datasourceStates: {
+          ...attributes.state.datasourceStates,
+          textBased: {
+            ...(attributes.state.datasourceStates.textBased || {}),
+            layers: {} as TextBasedPersistedState['layers'],
+          },
+        },
+      },
+    };
+
+    if (layers) {
+      for (const key of Object.keys(layers)) {
+        const newLayer = { ...layers[key] };
+        delete newLayer.table;
+        newState.state.datasourceStates.textBased!.layers[key] = newLayer;
+      }
+    }
+
+    return newState;
+  };
+
   return (
     <EuiFlexGroup
       {...a11yCommonProps}
@@ -397,7 +430,9 @@ export function Chart({
       )}
       {canSaveVisualization && isSaveModalVisible && lensAttributesContext.attributes && (
         <LensSaveModalComponent
-          initialInput={lensAttributesContext.attributes as unknown as LensEmbeddableInput}
+          initialInput={
+            removeTables(lensAttributesContext.attributes) as unknown as LensEmbeddableInput
+          }
           onSave={() => {}}
           onClose={() => setIsSaveModalVisible(false)}
           isSaveable={false}

diff --git a/src/plugins/unified_histogram/public/chart/histogram.tsx b/src/plugins/unified_histogram/public/chart/histogram.tsx
@@ -225,6 +225,7 @@ export function Histogram({
           {...lensProps}
           disableTriggers={disableTriggers}
           disabledActions={disabledActions}
+          shouldUseSizeTransitionVeil={false}
           onFilter={onFilter}
           onBrushEnd={onBrushEnd}
           withDefaultActions={withDefaultActions}

diff --git a/src/plugins/unified_histogram/public/container/container.tsx b/src/plugins/unified_histogram/public/container/container.tsx
@@ -11,6 +11,7 @@ import { Subject } from 'rxjs';
 import { pick } from 'lodash';
 import useMount from 'react-use/lib/useMount';
 import { LensSuggestionsApi } from '@kbn/lens-plugin/public';
+import type { Datatable } from '@kbn/expressions-plugin/common';
 import { UnifiedHistogramLayout, UnifiedHistogramLayoutProps } from '../layout';
 import type { UnifiedHistogramInputMessage, UnifiedHistogramRequestContext } from '../types';
 import {
@@ -43,6 +44,7 @@ export type UnifiedHistogramContainerProps = {
   searchSessionId?: UnifiedHistogramRequestContext['searchSessionId'];
   requestAdapter?: UnifiedHistogramRequestContext['adapter'];
   isChartLoading?: boolean;
+  table?: Datatable;
 } & Pick<
   UnifiedHistogramLayoutProps,
   | 'services'

diff --git a/src/plugins/unified_histogram/public/layout/helpers.ts b/src/plugins/unified_histogram/public/layout/helpers.ts
@@ -0,0 +1,25 @@
+/*
+ * 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 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 or the Server
+ * Side Public License, v 1.
+ */
+
+import { AggregateQuery } from '@kbn/es-query';
+
+const TRANSFORMATIONAL_COMMANDS = ['stats', 'project', 'keep'];
+
+export const shouldDisplayHistogram = (query: AggregateQuery) => {
+  let queryHasTransformationalCommands = false;
+  if ('esql' in query) {
+    TRANSFORMATIONAL_COMMANDS.forEach((command: string) => {
+      if (query.esql.toLowerCase().includes(command)) {
+        queryHasTransformationalCommands = true;
+        return;
+      }
+    });
+  }
+
+  return !queryHasTransformationalCommands;
+};
diff --git a/src/plugins/unified_histogram/public/layout/hooks/use_lens_suggestions.ts b/src/plugins/unified_histogram/public/layout/hooks/use_lens_suggestions.ts
@@ -15,12 +15,12 @@ import {
   Query,
   TimeRange,
 } from '@kbn/es-query';
-import type { DatatableColumn } from '@kbn/expressions-plugin/common';
+import type { Datatable, DatatableColumn } from '@kbn/expressions-plugin/common';
 import { LensSuggestionsApi, Suggestion } from '@kbn/lens-plugin/public';
 import { isEqual } from 'lodash';
 import { useEffect, useMemo, useRef, useState } from 'react';
 import { computeInterval } from './compute_interval';
-const TRANSFORMATIONAL_COMMANDS = ['stats', 'project', 'keep'];
+import { shouldDisplayHistogram } from '../helpers';
 
 export const useLensSuggestions = ({
   dataView,
@@ -42,6 +42,7 @@ export const useLensSuggestions = ({
   timeRange?: TimeRange;
   lensSuggestionsApi: LensSuggestionsApi;
   onSuggestionChange?: (suggestion: Suggestion | undefined) => void;
+  table?: Datatable;
 }) => {
   const suggestions = useMemo(() => {
     const context = {
@@ -57,10 +58,11 @@ export const useLensSuggestions = ({
     const [firstSuggestion] = allSuggestions;
 
     return { firstSuggestion, allSuggestions };
-  }, [dataView, isPlainRecord, lensSuggestionsApi, query, columns]);
+  }, [dataView, columns, query, isPlainRecord, lensSuggestionsApi]);
 
   const [allSuggestions, setAllSuggestions] = useState(suggestions.allSuggestions);
-  const currentSuggestion = originalSuggestion ?? suggestions.firstSuggestion;
+  const currentSuggestion = originalSuggestion || suggestions.firstSuggestion;
+
   const suggestionDeps = useRef(getSuggestionDeps({ dataView, query, columns }));
   const histogramQuery = useRef<AggregateQuery | undefined>();
   const histogramSuggestion = useMemo(() => {
@@ -72,17 +74,8 @@ export const useLensSuggestions = ({
       getAggregateQueryMode(query) === 'esql' &&
       timeRange
     ) {
-      let queryHasTransformationalCommands = false;
-      if ('esql' in query) {
-        TRANSFORMATIONAL_COMMANDS.forEach((command: string) => {
-          if (query.esql.toLowerCase().includes(command)) {
-            queryHasTransformationalCommands = true;
-            return;
-          }
-        });
-      }
-
-      if (queryHasTransformationalCommands) return undefined;
+      const isOnHistogramMode = shouldDisplayHistogram(query);
+      if (!isOnHistogramMode) return undefined;
 
       const interval = computeInterval(timeRange, data);
       const language = getAggregateQueryMode(query);

diff --git a/src/plugins/unified_histogram/public/layout/layout.tsx b/src/plugins/unified_histogram/public/layout/layout.tsx
@@ -7,11 +7,11 @@
  */
 
 import { EuiSpacer, useEuiTheme, useIsWithinBreakpoints } from '@elastic/eui';
-import React, { PropsWithChildren, ReactElement, useState } from 'react';
+import React, { PropsWithChildren, ReactElement, useMemo, useState } from 'react';
 import { Observable } from 'rxjs';
 import { createHtmlPortalNode, InPortal, OutPortal } from 'react-reverse-portal';
 import { css } from '@emotion/css';
-import type { DatatableColumn } from '@kbn/expressions-plugin/common';
+import type { Datatable, DatatableColumn } from '@kbn/expressions-plugin/common';
 import type { DataView, DataViewField } from '@kbn/data-views-plugin/public';
 import type {
   EmbeddableComponentProps,
@@ -20,12 +20,13 @@ import type {
   LensSuggestionsApi,
   Suggestion,
 } from '@kbn/lens-plugin/public';
-import type { AggregateQuery, Filter, Query, TimeRange } from '@kbn/es-query';
+import { AggregateQuery, Filter, isOfAggregateQueryType, Query, TimeRange } from '@kbn/es-query';
 import {
   ResizableLayout,
   ResizableLayoutMode,
   ResizableLayoutDirection,
 } from '@kbn/resizable-layout';
+import { TextBasedPersistedState } from '@kbn/lens-plugin/public/datasources/text_based/types';
 import { Chart, checkChartAvailability } from '../chart';
 import type {
   UnifiedHistogramChartContext,
@@ -38,6 +39,7 @@ import type {
   UnifiedHistogramInput$,
 } from '../types';
 import { useLensSuggestions } from './hooks/use_lens_suggestions';
+import { shouldDisplayHistogram } from './helpers';
 
 const ChartMemoized = React.memo(Chart);
 
@@ -179,6 +181,8 @@ export interface UnifiedHistogramLayoutProps extends PropsWithChildren<unknown>
    * Allows users to enable/disable default actions
    */
   withDefaultActions?: EmbeddableComponentProps['withDefaultActions'];
+
+  table?: Datatable;
 }
 
 export const UnifiedHistogramLayout = ({
@@ -207,6 +211,7 @@ export const UnifiedHistogramLayout = ({
   disabledActions,
   lensSuggestionsApi,
   input$,
+  table,
   onTopPanelHeightChange,
   onChartHiddenChange,
   onTimeIntervalChange,
@@ -237,6 +242,36 @@ export const UnifiedHistogramLayout = ({
     onSuggestionChange,
   });
 
+  // apply table to current suggestion
+  const usedSuggestion = useMemo(() => {
+    if (
+      currentSuggestion &&
+      table &&
+      query &&
+      isOfAggregateQueryType(query) &&
+      !shouldDisplayHistogram(query)
+    ) {
+      const { layers } = currentSuggestion.datasourceState as TextBasedPersistedState;
+
+      const newState = {
+        ...currentSuggestion,
+        datasourceState: {
+          ...(currentSuggestion.datasourceState as TextBasedPersistedState),
+          layers: {} as Record<string, unknown>,
+        },
+      };
+
+      for (const key of Object.keys(layers)) {
+        const newLayer = { ...layers[key], table };
+        newState.datasourceState.layers[key] = newLayer;
+      }
+
+      return newState;
+    } else {
+      return currentSuggestion;
+    }
+  }, [currentSuggestion, query, table]);
+
   const chart = suggestionUnsupported ? undefined : originalChart;
   const isChartAvailable = checkChartAvailability({ chart, dataView, isPlainRecord });
 
@@ -283,7 +318,7 @@ export const UnifiedHistogramLayout = ({
           relativeTimeRange={relativeTimeRange}
           request={request}
           hits={hits}
-          currentSuggestion={currentSuggestion}
+          currentSuggestion={usedSuggestion}
           isChartLoading={isChartLoading}
           allSuggestions={allSuggestions}
           isPlainRecord={isPlainRecord}

diff --git a/src/plugins/unified_histogram/tsconfig.json b/src/plugins/unified_histogram/tsconfig.json
@@ -27,6 +27,7 @@
     "@kbn/resizable-layout",
     "@kbn/shared-ux-button-toolbar",
     "@kbn/calculate-width-from-char-count",
+    "@kbn/lens-embeddable-utils",
     "@kbn/i18n-react",
     "@kbn/field-utils",
   ],