[OnWeek][Discover] Allow to change current sample size and save it with a saved search

packages/kbn-unified-data-table/src/components/data_table.test.tsx

@@ -10,6 +10,7 @@ import { ReactWrapper } from 'enzyme';
 import {
   EuiButton,
   EuiCopy,
+  EuiDataGrid,
   EuiDataGridCellValueElementProps,
   EuiDataGridCustomBodyProps,
 } from '@elastic/eui';
@@ -52,7 +53,7 @@ function getProps(): UnifiedDataTableProps {
     onSetColumns: jest.fn(),
     onSort: jest.fn(),
     rows: esHitsMock.map((hit) => buildDataTableRecord(hit, dataViewMock)),
-    sampleSize: 30,
+    sampleSizeState: 30,
     searchDescription: '',
     searchTitle: '',
     setExpandedDoc: jest.fn(),
@@ -301,6 +302,74 @@ describe('UnifiedDataTable', () => {
     });
   });
 
+  describe('display settings', () => {
+    it('should include additional display settings if onUpdateSampleSize is provided', async () => {
+      const component = await getComponent({
+        ...getProps(),
+        sampleSizeState: 150,
+        onUpdateSampleSize: jest.fn(),
+        onUpdateRowHeight: jest.fn(),
+      });
+
+      expect(component.find(EuiDataGrid).prop('toolbarVisibility')).toMatchInlineSnapshot(`
+        Object {
+          "additionalControls": <React.Fragment />,
+          "showColumnSelector": false,
+          "showDisplaySelector": Object {
+            "additionalDisplaySettings": <UnifiedDataTableAdditionalDisplaySettings
+              onChangeSampleSize={[MockFunction]}
+              sampleSize={150}
+            />,
+            "allowDensity": false,
+            "allowResetButton": false,
+            "allowRowHeight": true,
+          },
+          "showFullScreenSelector": true,
+          "showSortSelector": true,
+        }
+      `);
+    });
+
+    it('should not include additional display settings if onUpdateSampleSize is not provided', async () => {
+      const component = await getComponent({
+        ...getProps(),
+        sampleSizeState: 200,
+        onUpdateRowHeight: jest.fn(),
+      });
+
+      expect(component.find(EuiDataGrid).prop('toolbarVisibility')).toMatchInlineSnapshot(`
+        Object {
+          "additionalControls": <React.Fragment />,
+          "showColumnSelector": false,
+          "showDisplaySelector": Object {
+            "allowDensity": false,
+            "allowRowHeight": true,
+          },
+          "showFullScreenSelector": true,
+          "showSortSelector": true,
+        }
+      `);
+    });
+
+    it('should hide display settings if no handlers provided', async () => {
+      const component = await getComponent({
+        ...getProps(),
+        onUpdateRowHeight: undefined,
+        onUpdateSampleSize: undefined,
+      });
+
+      expect(component.find(EuiDataGrid).prop('toolbarVisibility')).toMatchInlineSnapshot(`
+        Object {
+          "additionalControls": <React.Fragment />,
+          "showColumnSelector": false,
+          "showDisplaySelector": undefined,
+          "showFullScreenSelector": true,
+          "showSortSelector": true,
+        }
+      `);
+    });
+  });
+
   describe('externalControlColumns', () => {
     it('should render external leading control columns', async () => {
       const component = await getComponent({

packages/kbn-unified-data-table/src/components/data_table.tsx

@@ -27,6 +27,7 @@ import {
   EuiDataGridControlColumn,
   EuiDataGridCustomBodyProps,
   EuiDataGridCellValueElementProps,
+  EuiDataGridToolBarVisibilityDisplaySelectorOptions,
   EuiDataGridStyle,
 } from '@elastic/eui';
 import type { DataView } from '@kbn/data-views-plugin/public';
@@ -63,6 +64,7 @@ import {
   toolbarVisibility as toolbarVisibilityDefaults,
 } from '../constants';
 import { UnifiedDataTableFooter } from './data_table_footer';
+import { UnifiedDataTableAdditionalDisplaySettings } from './data_table_additional_display_settings';
 
 export type SortOrder = [string, string];
 
@@ -137,10 +139,6 @@ export interface UnifiedDataTableProps {
    * Array of documents provided by Elasticsearch
    */
   rows?: DataTableRecord[];
-  /**
-   * The max size of the documents returned by Elasticsearch
-   */
-  sampleSize: number;
   /**
    * Function to set the expanded document, which is displayed in a flyout
    */
@@ -205,6 +203,18 @@ export interface UnifiedDataTableProps {
    * Update rows per page state
    */
   onUpdateRowsPerPage?: (rowsPerPage: number) => void;
+  /**
+   * Configuration option to limit sample size slider
+   */
+  maxAllowedSampleSize?: number;
+  /**
+   * The max size of the documents returned by Elasticsearch
+   */
+  sampleSizeState: number;
+  /**
+   * Update rows per page state
+   */
+  onUpdateSampleSize?: (sampleSize: number) => void;
   /**
    * Callback to execute on edit runtime field
    */
@@ -328,7 +338,6 @@ export const UnifiedDataTable = ({
   onSetColumns,
   onSort,
   rows,
-  sampleSize,
   searchDescription,
   searchTitle,
   settings,
@@ -342,6 +351,9 @@ export const UnifiedDataTable = ({
   className,
   rowHeightState,
   onUpdateRowHeight,
+  maxAllowedSampleSize,
+  sampleSizeState,
+  onUpdateSampleSize,
   isPlainRecord = false,
   rowsPerPageState,
   onUpdateRowsPerPage,
@@ -715,16 +727,27 @@ export const UnifiedDataTable = ({
     [usedSelectedDocs, isFilterActive, rows, externalAdditionalControls]
   );
 
-  const showDisplaySelector = useMemo(
-    () =>
-      !!onUpdateRowHeight
-        ? {
-            allowDensity: false,
-            allowRowHeight: true,
-          }
-        : undefined,
-    [onUpdateRowHeight]
-  );
+  const showDisplaySelector = useMemo(() => {
+    const options: EuiDataGridToolBarVisibilityDisplaySelectorOptions = {};
+
+    if (onUpdateRowHeight) {
+      options.allowDensity = false;
+      options.allowRowHeight = true;
+    }
+
+    if (onUpdateSampleSize) {
+      options.allowResetButton = false;
+      options.additionalDisplaySettings = (
+        <UnifiedDataTableAdditionalDisplaySettings
+          maxAllowedSampleSize={maxAllowedSampleSize}
+          sampleSize={sampleSizeState}
+          onChangeSampleSize={onUpdateSampleSize}
+        />
+      );
+    }
+
+    return Object.keys(options).length ? options : undefined;
+  }, [maxAllowedSampleSize, sampleSizeState, onUpdateRowHeight, onUpdateSampleSize]);
 
   const inMemory = useMemo(() => {
     return isPlainRecord && columns.length
@@ -837,7 +860,7 @@ export const UnifiedDataTable = ({
             <UnifiedDataTableFooter
               isLoadingMore={loadingState === DataLoadingState.loadingMore}
               rowCount={rowCount}
-              sampleSize={sampleSize}
+              sampleSize={sampleSizeState}
               pageCount={pageCount}
               pageIndex={paginationObj?.pageIndex}
               totalHits={totalHits}

packages/kbn-unified-data-table/src/components/data_table_additional_display_settings.test.tsx

@@ -0,0 +1,111 @@
+/*
+ * 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 React from 'react';
+import { mountWithIntl } from '@kbn/test-jest-helpers';
+import { act } from 'react-dom/test-utils';
+import { findTestSubject } from '@elastic/eui/lib/test';
+import { UnifiedDataTableAdditionalDisplaySettings } from './data_table_additional_display_settings';
+import lodash from 'lodash';
+
+jest.spyOn(lodash, 'debounce').mockImplementation((fn: any) => fn);
+
+describe('UnifiedDataTableAdditionalDisplaySettings', function () {
+  describe('sampleSize', function () {
+    it('should work correctly', async () => {
+      const onChangeSampleSizeMock = jest.fn();
+
+      const component = mountWithIntl(
+        <UnifiedDataTableAdditionalDisplaySettings
+          sampleSize={10}
+          onChangeSampleSize={onChangeSampleSizeMock}
+        />
+      );
+      const input = findTestSubject(component, 'unifiedDataTableSampleSizeInput').last();
+      expect(input.prop('value')).toBe(10);
+
+      await act(async () => {
+        input.simulate('change', {
+          target: {
+            value: 100,
+          },
+        });
+      });
+
+      expect(onChangeSampleSizeMock).toHaveBeenCalledWith(100);
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      component.update();
+
+      expect(
+        findTestSubject(component, 'unifiedDataTableSampleSizeInput').last().prop('value')
+      ).toBe(100);
+    });
+
+    it('should not execute the callback for an invalid input', async () => {
+      const invalidValue = 600;
+      const onChangeSampleSizeMock = jest.fn();
+
+      const component = mountWithIntl(
+        <UnifiedDataTableAdditionalDisplaySettings
+          maxAllowedSampleSize={500}
+          sampleSize={50}
+          onChangeSampleSize={onChangeSampleSizeMock}
+        />
+      );
+      const input = findTestSubject(component, 'unifiedDataTableSampleSizeInput').last();
+      expect(input.prop('value')).toBe(50);
+
+      await act(async () => {
+        input.simulate('change', {
+          target: {
+            value: invalidValue,
+          },
+        });
+      });
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      component.update();
+
+      expect(
+        findTestSubject(component, 'unifiedDataTableSampleSizeInput').last().prop('value')
+      ).toBe(invalidValue);
+
+      expect(onChangeSampleSizeMock).not.toHaveBeenCalled();
+    });
+
+    it('should render value changes correctly', async () => {
+      const onChangeSampleSizeMock = jest.fn();
+
+      const component = mountWithIntl(
+        <UnifiedDataTableAdditionalDisplaySettings
+          sampleSize={200}
+          onChangeSampleSize={onChangeSampleSizeMock}
+        />
+      );
+
+      expect(
+        findTestSubject(component, 'unifiedDataTableSampleSizeInput').last().prop('value')
+      ).toBe(200);
+
+      component.setProps({
+        sampleSize: 500,
+        onChangeSampleSize: onChangeSampleSizeMock,
+      });
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      component.update();
+
+      expect(
+        findTestSubject(component, 'unifiedDataTableSampleSizeInput').last().prop('value')
+      ).toBe(500);
+
+      expect(onChangeSampleSizeMock).not.toHaveBeenCalled();
+    });
+  });
+});