[Textbased languages] Render Lens suggestions in Discover

src/plugins/discover/public/__mocks__/data_view.ts

@@ -109,6 +109,7 @@ export const buildDataViewMock = ({
     getFormatterForField: jest.fn(() => ({ convert: (value: unknown) => value })),
     isTimeNanosBased: () => false,
     isPersisted: () => true,
+    toSpec: () => ({}),
     getTimeField: () => {
       return dataViewFields.find((field) => field.name === timeFieldName);
     },

src/plugins/discover/public/__mocks__/services.ts

@@ -152,7 +152,14 @@ export function createDiscoverServicesMock(): DiscoverServices {
     savedObjectsTagging: {},
     dataViews: dataPlugin.dataViews,
     timefilter: dataPlugin.query.timefilter.timefilter,
-    lens: { EmbeddableComponent: jest.fn(() => null) },
+    lens: {
+      EmbeddableComponent: jest.fn(() => null),
+      stateHelperApi: jest.fn(() => {
+        return {
+          suggestionsApi: jest.fn(),
+        };
+      }),
+    },
     locator: {
       useUrl: jest.fn(() => ''),
       navigate: jest.fn(),

src/plugins/discover/public/application/main/components/layout/discover_histogram_layout.tsx

@@ -46,6 +46,7 @@ export const DiscoverHistogramLayout = ({
     inspectorAdapters,
     savedSearchFetch$: stateContainer.dataState.fetch$,
     searchSessionId,
+    isPlainRecord,
     ...commonProps,
   });
 

src/plugins/discover/public/application/main/components/layout/use_discover_histogram.test.tsx

@@ -7,6 +7,7 @@
  */
 
 import React, { ReactElement } from 'react';
+import { AggregateQuery, Query } from '@kbn/es-query';
 import { buildDataTableRecord } from '../../../../utils/build_data_record';
 import { esHits } from '../../../../__mocks__/es_hits';
 import { act, renderHook, WrapperComponent } from '@testing-library/react-hooks';
@@ -38,11 +39,11 @@ import { checkHitCount, sendErrorTo } from '../../hooks/use_saved_search_message
 import type { InspectorAdapters } from '../../hooks/use_inspector';
 
 const mockData = dataPluginMock.createStartContract();
-const mockQueryState = {
+let mockQueryState = {
   query: {
     query: 'query',
     language: 'kuery',
-  },
+  } as Query | AggregateQuery,
   filters: [],
   time: {
     from: 'now-15m',
@@ -112,19 +113,21 @@ describe('useDiscoverHistogram', () => {
       foundDocuments: true,
     }) as DataMain$,
     savedSearchFetch$ = new Subject() as DataFetch$,
+    documents$ = new BehaviorSubject({
+      fetchStatus: FetchStatus.COMPLETE,
+      result: esHits.map((esHit) => buildDataTableRecord(esHit, dataViewWithTimefieldMock)),
+    }) as DataDocuments$,
+    isPlainRecord = false,
   }: {
     stateContainer?: DiscoverStateContainer;
     searchSessionId?: string;
     inspectorAdapters?: InspectorAdapters;
     totalHits$?: DataTotalHits$;
     main$?: DataMain$;
     savedSearchFetch$?: DataFetch$;
+    documents$?: DataDocuments$;
+    isPlainRecord?: boolean;
   } = {}) => {
-    const documents$ = new BehaviorSubject({
-      fetchStatus: FetchStatus.COMPLETE,
-      result: esHits.map((esHit) => buildDataTableRecord(esHit, dataViewWithTimefieldMock)),
-    }) as DataDocuments$;
-
     const availableFields$ = new BehaviorSubject({
       fetchStatus: FetchStatus.COMPLETE,
       fields: [] as string[],
@@ -144,6 +147,7 @@ describe('useDiscoverHistogram', () => {
       dataView: dataViewWithTimefieldMock,
       inspectorAdapters,
       searchSessionId,
+      isPlainRecord,
     };
 
     const Wrapper: WrapperComponent<UseDiscoverHistogramProps> = ({ children }) => (
@@ -186,6 +190,7 @@ describe('useDiscoverHistogram', () => {
         'timeRange',
         'chartHidden',
         'timeInterval',
+        'columns',
         'breakdownField',
         'searchSessionId',
         'totalHitsStatus',
@@ -196,6 +201,9 @@ describe('useDiscoverHistogram', () => {
   });
 
   describe('state', () => {
+    beforeEach(() => {
+      mockCheckHitCount.mockClear();
+    });
     it('should subscribe to state changes', async () => {
       const { hook } = await renderUseDiscoverHistogram();
       const api = createMockUnifiedHistogramApi({ initialized: true });
@@ -347,7 +355,6 @@ describe('useDiscoverHistogram', () => {
     });
 
     it('should update total hits when the total hits state changes', async () => {
-      mockCheckHitCount.mockClear();
       const totalHits$ = new BehaviorSubject({
         fetchStatus: FetchStatus.LOADING,
         result: undefined,
@@ -383,8 +390,66 @@ describe('useDiscoverHistogram', () => {
       expect(mockCheckHitCount).toHaveBeenCalledWith(main$, 100);
     });
 
+    it('should fetch total hits for text based', async () => {
+      mockQueryState = {
+        query: { sql: 'select * from index' },
+        filters: [],
+        time: {
+          from: 'now-15m',
+          to: 'now',
+        },
+      };
+
+      mockData.query.getState = () => mockQueryState;
+      const documents$ = new BehaviorSubject({
+        fetchStatus: FetchStatus.COMPLETE,
+        result: [{ raw: {}, flattened: {}, id: '1' }],
+      }) as unknown as DataDocuments$;
+      const main$ = new BehaviorSubject({
+        fetchStatus: FetchStatus.COMPLETE,
+        recordRawType: RecordRawType.DOCUMENT,
+        foundDocuments: true,
+      }) as DataMain$;
+      const stateContainer = getStateContainer();
+      const { hook } = await renderUseDiscoverHistogram({
+        stateContainer,
+        documents$,
+        main$,
+        isPlainRecord: true,
+      });
+      const api = createMockUnifiedHistogramApi({ initialized: true });
+      let params: Partial<UnifiedHistogramState> = {};
+      api.setTotalHits = jest.fn((p) => {
+        params = { ...params, ...p };
+      });
+      act(() => {
+        hook.result.current.setUnifiedHistogramApi(api);
+      });
+      expect(api.setTotalHits).toHaveBeenCalledWith({
+        totalHitsResult: 1,
+        totalHitsStatus: FetchStatus.COMPLETE,
+      });
+      expect(documents$.value).toEqual({
+        fetchStatus: FetchStatus.COMPLETE,
+        result: [{ raw: {}, flattened: {}, id: '1' }],
+      });
+      expect(mockCheckHitCount).not.toHaveBeenCalled();
+    });
+
     it('should not update total hits when the total hits state changes to an error', async () => {
-      mockCheckHitCount.mockClear();
+      mockQueryState = {
+        query: {
+          query: 'query',
+          language: 'kuery',
+        } as Query | AggregateQuery,
+        filters: [],
+        time: {
+          from: 'now-15m',
+          to: 'now',
+        },
+      };
+
+      mockData.query.getState = () => mockQueryState;
       const totalHits$ = new BehaviorSubject({
         fetchStatus: FetchStatus.UNINITIALIZED,
         result: undefined,

src/plugins/discover/public/application/main/components/layout/use_discover_histogram.ts

@@ -15,7 +15,8 @@ import {
   UnifiedHistogramState,
 } from '@kbn/unified-histogram-plugin/public';
 import { isEqual } from 'lodash';
-import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
+import { AggregateQuery, Query, isOfAggregateQueryType } from '@kbn/es-query';
+import { useCallback, useEffect, useRef, useState } from 'react';
 import { distinctUntilChanged, map, Observable } from 'rxjs';
 import { useDiscoverServices } from '../../../../hooks/use_discover_services';
 import { getUiActions } from '../../../../kibana_services';
@@ -34,6 +35,7 @@ export interface UseDiscoverHistogramProps {
   inspectorAdapters: InspectorAdapters;
   savedSearchFetch$: DataFetch$;
   searchSessionId: string | undefined;
+  isPlainRecord: boolean;
 }
 
 export const useDiscoverHistogram = ({
@@ -43,6 +45,7 @@ export const useDiscoverHistogram = ({
   inspectorAdapters,
   savedSearchFetch$,
   searchSessionId,
+  isPlainRecord,
 }: UseDiscoverHistogramProps) => {
   const services = useDiscoverServices();
   const timefilter = services.data.query.timefilter.timefilter;
@@ -52,6 +55,12 @@ export const useDiscoverHistogram = ({
    */
 
   const [unifiedHistogram, setUnifiedHistogram] = useState<UnifiedHistogramInitializedApi>();
+  const prev = useRef<{ query: AggregateQuery | Query | undefined; columns: string[] | undefined }>(
+    {
+      columns: [],
+      query: undefined,
+    }
+  );
 
   const setUnifiedHistogramApi = useCallback(
     (api: UnifiedHistogramApi | null) => {
@@ -66,12 +75,15 @@ export const useDiscoverHistogram = ({
           hideChart: chartHidden,
           interval: timeInterval,
           breakdownField,
+          columns,
         } = stateContainer.appState.getState();
 
         const { fetchStatus: totalHitsStatus, result: totalHitsResult } =
           savedSearchData$.totalHits$.getValue();
 
         const { query, filters, time: timeRange } = services.data.query.getState();
+        prev.current.query = query;
+        prev.current.columns = columns;
 
         api.initialize({
           services: { ...services, uiActions: getUiActions() },
@@ -85,6 +97,7 @@ export const useDiscoverHistogram = ({
             timeRange,
             chartHidden,
             timeInterval,
+            columns,
             breakdownField,
             searchSessionId,
             totalHitsStatus: totalHitsStatus.toString() as UnifiedHistogramFetchStatus,
@@ -134,18 +147,8 @@ export const useDiscoverHistogram = ({
   /**
    * Update Unified Histgoram request params
    */
-
-  const {
-    query,
-    filters,
-    fromDate: from,
-    toDate: to,
-  } = useQuerySubscriber({ data: services.data });
-
-  const timeRange = useMemo(
-    () => (from && to ? { from, to } : timefilter.getTimeDefaults()),
-    [timefilter, from, to]
-  );
+  const { query, filters } = useQuerySubscriber({ data: services.data });
+  const timeRange = timefilter.getAbsoluteTime();
 
   useEffect(() => {
     unifiedHistogram?.setRequestParams({
@@ -180,13 +183,13 @@ export const useDiscoverHistogram = ({
     // We only want to show the partial results on the first load,
     // or there will be a flickering effect as the loading spinner
     // is quickly shown and hidden again on fetches
-    if (!firstLoadComplete.current) {
+    if (!firstLoadComplete.current && !isPlainRecord) {
       unifiedHistogram?.setTotalHits({
         totalHitsStatus: totalHitsStatus.toString() as UnifiedHistogramFetchStatus,
         totalHitsResult,
       });
     }
-  }, [totalHitsResult, totalHitsStatus, unifiedHistogram]);
+  }, [isPlainRecord, totalHitsResult, totalHitsStatus, unifiedHistogram]);
 
   /**
    * Sync URL query params with Unified Histogram
@@ -214,6 +217,35 @@ export const useDiscoverHistogram = ({
     unifiedHistogram?.setBreakdownField(breakdownField);
   }, [breakdownField, unifiedHistogram]);
 
+  const columns = useAppStateSelector((state) => state.columns);
+  const documentState = useDataState(savedSearchData$.documents$);
+
+  useEffect(() => {
+    const currentQueryIsTextBased = query && isOfAggregateQueryType(query);
+    const initialQueryIsTextBased =
+      prev.current.query && isOfAggregateQueryType(prev.current.query);
+    // Update the columns only when the query changes
+    if (!isEqual(columns, prev.current.columns) && !isEqual(query, prev.current.query)) {
+      // transitioning from dataview mode to text based mode should set the columns to []
+      if (currentQueryIsTextBased && !initialQueryIsTextBased) {
+        unifiedHistogram?.setColumns([]);
+      } else {
+        unifiedHistogram?.setColumns(columns);
+      }
+      prev.current.query = query;
+      prev.current.columns = columns;
+    }
+  }, [columns, query, unifiedHistogram]);
+
+  useEffect(() => {
+    if (isPlainRecord) {
+      unifiedHistogram?.setTotalHits({
+        totalHitsResult: documentState.result?.length,
+        totalHitsStatus: documentState.fetchStatus.toString() as UnifiedHistogramFetchStatus,
+      });
+    }
+  }, [documentState.fetchStatus, documentState.result, isPlainRecord, unifiedHistogram]);
+
   /**
    * Total hits
    */
@@ -227,6 +259,10 @@ export const useDiscoverHistogram = ({
           return;
         }
 
+        if (isPlainRecord) {
+          return;
+        }
+
         const { recordRawType } = savedSearchData$.totalHits$.getValue();
 
         // Sync the totalHits$ observable with the unified histogram state
@@ -252,7 +288,13 @@ export const useDiscoverHistogram = ({
     return () => {
       subscription?.unsubscribe();
     };
-  }, [savedSearchData$.main$, savedSearchData$.totalHits$, services.data, unifiedHistogram]);
+  }, [
+    isPlainRecord,
+    savedSearchData$.main$,
+    savedSearchData$.totalHits$,
+    services.data,
+    unifiedHistogram,
+  ]);
 
   /**
    * Data fetching

src/plugins/discover/public/application/main/components/sidebar/discover_sidebar.test.tsx

@@ -262,12 +262,4 @@ describe('discover sidebar', function () {
     const createDataViewButton = findTestSubject(compWithPickerInViewerMode, 'dataview-create-new');
     expect(createDataViewButton.length).toBe(0);
   });
-
-  it('should render the Visualize in Lens button in text based languages mode', async () => {
-    const compInViewerMode = await mountComponent(getCompProps(), {
-      query: { sql: 'SELECT * FROM test' },
-    });
-    const visualizeField = findTestSubject(compInViewerMode, 'textBased-visualize');
-    expect(visualizeField.length).toBe(1);
-  });
 });