[Discover] Allow to select/deselect all rows in the grid at once

packages/kbn-unified-data-table/__mocks__/table_context.ts

@@ -15,6 +15,7 @@ import { DataTableContext } from '../src/table_context';
 import { convertValueToString } from '../src/utils/convert_value_to_string';
 import { buildDataTableRecord } from '@kbn/discover-utils';
 import type { EsHitRecord } from '@kbn/discover-utils/types';
+import type { UseSelectedDocsState } from '../src/hooks/use_selected_docs';
 
 const buildTableContext = (dataView: DataView, rows: EsHitRecord[]): DataTableContext => {
   const usedRows = rows.map((row) => {
@@ -28,8 +29,9 @@ const buildTableContext = (dataView: DataView, rows: EsHitRecord[]): DataTableCo
     onFilter: jest.fn(),
     dataView,
     isDarkMode: false,
-    selectedDocs: [],
-    setSelectedDocs: jest.fn(),
+    selectedDocsState: buildSelectedDocsState([]),
+    pageIndex: 0,
+    pageSize: 10,
     valueToStringConverter: (rowIndex, columnId, options) =>
       convertValueToString({
         rowIndex,
@@ -45,3 +47,21 @@ const buildTableContext = (dataView: DataView, rows: EsHitRecord[]): DataTableCo
 export const dataTableContextMock = buildTableContext(dataViewMock, esHitsMock);
 
 export const dataTableContextComplexMock = buildTableContext(dataViewComplexMock, esHitsComplex);
+
+export function buildSelectedDocsState(selectedDocIds: string[]): UseSelectedDocsState {
+  const selectedDocsSet = new Set(selectedDocIds);
+
+  return {
+    isDocSelected: (docId: string) => selectedDocsSet.has(docId),
+    getCountOfSelectedDocs: (docIds: string[]) =>
+      docIds.reduce((acc, docId) => (selectedDocsSet.has(docId) ? acc + 1 : acc), 0),
+    hasSelectedDocs: selectedDocsSet.size > 0,
+    selectedDocIds,
+    toggleDocSelection: jest.fn(),
+    selectAllDocs: jest.fn(),
+    selectMoreDocs: jest.fn(),
+    deselectSomeDocs: jest.fn(),
+    replaceSelectedDocs: jest.fn(),
+    clearAllSelectedDocs: jest.fn(),
+  };
+}

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

@@ -52,7 +52,7 @@ const getDocById = (id: string) => docs.find((doc) => doc.raw._id === id);
 const renderCompareDocuments = ({
   forceShowAllFields = false,
 }: { forceShowAllFields?: boolean } = {}) => {
-  const setSelectedDocs = jest.fn();
+  const replaceSelectedDocs = jest.fn();
   const getCompareDocuments = (props?: Partial<CompareDocumentsProps>) => (
     <CompareDocuments
       id="test"
@@ -63,20 +63,20 @@ const renderCompareDocuments = ({
       dataView={dataViewWithTimefieldMock}
       isPlainRecord={false}
       selectedFieldNames={['message', 'extension', 'bytes']}
-      selectedDocs={['0', '1', '2']}
+      selectedDocIds={['0', '1', '2']}
       schemaDetectors={[]}
       forceShowAllFields={forceShowAllFields}
       showFullScreenButton={true}
       fieldFormats={{} as any}
       getDocById={getDocById}
-      setSelectedDocs={setSelectedDocs}
+      replaceSelectedDocs={replaceSelectedDocs}
       setIsCompareActive={jest.fn()}
       {...props}
     />
   );
   const { rerender } = render(getCompareDocuments());
   return {
-    setSelectedDocs,
+    replaceSelectedDocs,
     rerender: (props?: Partial<CompareDocumentsProps>) => rerender(getCompareDocuments(props)),
   };
 };
@@ -146,10 +146,10 @@ describe('CompareDocuments', () => {
   });
 
   it('should set selected docs when columns change', () => {
-    const { setSelectedDocs } = renderCompareDocuments();
+    const { replaceSelectedDocs } = renderCompareDocuments();
     const visibleColumns = ['fields_generated-id', '0', '1', '2'];
     mockDataGridProps?.columnVisibility.setVisibleColumns(visibleColumns);
-    expect(setSelectedDocs).toHaveBeenCalledWith(visibleColumns.slice(1));
+    expect(replaceSelectedDocs).toHaveBeenCalledWith(visibleColumns.slice(1));
   });
 
   it('should force show all fields when prop is true', () => {

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

@@ -42,13 +42,13 @@ export interface CompareDocumentsProps {
   dataView: DataView;
   isPlainRecord: boolean;
   selectedFieldNames: string[];
-  selectedDocs: string[];
+  selectedDocIds: string[];
   schemaDetectors: EuiDataGridSchemaDetector[];
   forceShowAllFields: boolean;
   showFullScreenButton?: boolean;
   fieldFormats: FieldFormatsStart;
   getDocById: (id: string) => DataTableRecord | undefined;
-  setSelectedDocs: (selectedDocs: string[]) => void;
+  replaceSelectedDocs: (docIds: string[]) => void;
   setIsCompareActive: (isCompareActive: boolean) => void;
   additionalFieldGroups?: AdditionalFieldGroups;
 }
@@ -69,13 +69,13 @@ const CompareDocuments = ({
   isPlainRecord,
   selectedFieldNames,
   additionalFieldGroups,
-  selectedDocs,
+  selectedDocIds,
   schemaDetectors,
   forceShowAllFields,
   showFullScreenButton,
   fieldFormats,
   getDocById,
-  setSelectedDocs,
+  replaceSelectedDocs,
   setIsCompareActive,
 }: CompareDocumentsProps) => {
   // Memoize getDocById to ensure we don't lose access to the comparison docs if, for example,
@@ -104,7 +104,7 @@ const CompareDocuments = ({
     dataView,
     selectedFieldNames,
     additionalFieldGroups,
-    selectedDocs,
+    selectedDocIds,
     showAllFields: Boolean(forceShowAllFields || showAllFields),
     showMatchingValues: Boolean(showMatchingValues),
     getDocById: memoizedGetDocById,
@@ -113,25 +113,25 @@ const CompareDocuments = ({
     wrapper,
     isPlainRecord,
     fieldColumnId,
-    selectedDocs,
+    selectedDocIds,
     getDocById: memoizedGetDocById,
-    setSelectedDocs,
+    replaceSelectedDocs,
   });
   const comparisonColumnVisibility = useMemo<EuiDataGridColumnVisibility>(
     () => ({
       visibleColumns: comparisonColumns.map(({ id: columnId }) => columnId),
       setVisibleColumns: (visibleColumns) => {
         const [_fieldColumnId, ...newSelectedDocs] = visibleColumns;
-        setSelectedDocs(newSelectedDocs);
+        replaceSelectedDocs(newSelectedDocs);
       },
     }),
-    [comparisonColumns, setSelectedDocs]
+    [comparisonColumns, replaceSelectedDocs]
   );
   const additionalControls = useMemo(
     () => (
       <ComparisonControls
         isPlainRecord={isPlainRecord}
-        selectedDocs={selectedDocs}
+        selectedDocIds={selectedDocIds}
         showDiff={showDiff}
         diffMode={diffMode}
         showDiffDecorations={showDiffDecorations}
@@ -150,7 +150,7 @@ const CompareDocuments = ({
       diffMode,
       forceShowAllFields,
       isPlainRecord,
-      selectedDocs,
+      selectedDocIds,
       setDiffMode,
       setIsCompareActive,
       setShowAllFields,
@@ -184,7 +184,7 @@ const CompareDocuments = ({
     dataView,
     comparisonFields,
     fieldColumnId,
-    selectedDocs,
+    selectedDocIds,
     diffMode: showDiff ? diffMode : undefined,
     fieldFormats,
     getDocById: memoizedGetDocById,

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

@@ -20,7 +20,7 @@ const renderComparisonControls = ({
   isPlainRecord?: ComparisonControlsProps['isPlainRecord'];
   forceShowAllFields?: ComparisonControlsProps['forceShowAllFields'];
 } = {}) => {
-  const selectedDocs = ['0', '1', '2'];
+  const selectedDocIds = ['0', '1', '2'];
   const Wrapper = () => {
     const [showDiff, setShowDiff] = useState(true);
     const [diffMode, setDiffMode] = useState<DocumentDiffMode>('basic');
@@ -34,7 +34,7 @@ const renderComparisonControls = ({
         <IntlProvider locale="en">
           <ComparisonControls
             isPlainRecord={isPlainRecord}
-            selectedDocs={selectedDocs}
+            selectedDocIds={selectedDocIds}
             showDiff={showDiff}
             diffMode={diffMode}
             showDiffDecorations={showDiffDecorations}
@@ -67,7 +67,7 @@ const renderComparisonControls = ({
   return {
     getComparisonCountDisplay: () =>
       screen.getByText(
-        `Comparing ${selectedDocs.length} ${isPlainRecord ? 'results' : 'documents'}`
+        `Comparing ${selectedDocIds.length} ${isPlainRecord ? 'results' : 'documents'}`
       ),
     getComparisonSettingsButton,
     clickComparisonSettingsButton: () => userEvent.click(getComparisonSettingsButton()),

packages/kbn-unified-data-table/src/components/compare_documents/comparison_controls.tsx

@@ -30,7 +30,7 @@ import type { DocumentDiffMode } from './types';
 
 export interface ComparisonControlsProps {
   isPlainRecord?: boolean;
-  selectedDocs: string[];
+  selectedDocIds: string[];
   showDiff: boolean | undefined;
   diffMode: DocumentDiffMode | undefined;
   showDiffDecorations: boolean | undefined;
@@ -47,7 +47,7 @@ export interface ComparisonControlsProps {
 
 export const ComparisonControls = ({
   isPlainRecord,
-  selectedDocs,
+  selectedDocIds,
   showDiff,
   diffMode,
   showDiffDecorations,
@@ -72,13 +72,13 @@ export const ComparisonControls = ({
               <FormattedMessage
                 id="unifiedDataTable.comparingResults"
                 defaultMessage="Comparing {documentCount} results"
-                values={{ documentCount: selectedDocs.length }}
+                values={{ documentCount: selectedDocIds.length }}
               />
             ) : (
               <FormattedMessage
                 id="unifiedDataTable.comparingDocuments"
                 defaultMessage="Comparing {documentCount} documents"
-                values={{ documentCount: selectedDocs.length }}
+                values={{ documentCount: selectedDocIds.length }}
               />
             )}
           </strong>

packages/kbn-unified-data-table/src/components/compare_documents/hooks/use_comparison_cell_value.test.tsx

@@ -55,7 +55,7 @@ const renderComparisonCellValue = (props: Partial<UseComparisonCellValueProps> =
     dataView: dataViewWithTimefieldMock,
     comparisonFields: ['message', 'extension', 'bytes'],
     fieldColumnId,
-    selectedDocs: ['0', '1', '2'],
+    selectedDocIds: ['0', '1', '2'],
     diffMode: undefined,
     fieldFormats: fieldFormatsMock,
     getDocById,
@@ -408,7 +408,7 @@ describe('useComparisonCellValue', () => {
     expect(calculateDiff).toHaveBeenCalledTimes(2);
     renderComparisonCell(cellProps2);
     expect(calculateDiff).toHaveBeenCalledTimes(2);
-    rerender({ diffMode: 'words', selectedDocs: ['1', '2', '0'] });
+    rerender({ diffMode: 'words', selectedDocIds: ['1', '2', '0'] });
     const cellProps3 = {
       ...cellProps1,
       columnId: '2',
@@ -425,7 +425,7 @@ describe('useComparisonCellValue', () => {
     expect(calculateDiff).toHaveBeenCalledTimes(4);
     renderComparisonCell(cellProps4);
     expect(calculateDiff).toHaveBeenCalledTimes(4);
-    rerender({ diffMode: 'lines', selectedDocs: ['2', '0', '1'] });
+    rerender({ diffMode: 'lines', selectedDocIds: ['2', '0', '1'] });
     const cellProps5 = {
       ...cellProps1,
       columnId: '0',

packages/kbn-unified-data-table/src/components/compare_documents/hooks/use_comparison_cell_value.tsx

@@ -39,7 +39,7 @@ export interface UseComparisonCellValueProps {
   dataView: DataView;
   comparisonFields: string[];
   fieldColumnId: string;
-  selectedDocs: string[];
+  selectedDocIds: string[];
   diffMode: DocumentDiffMode | undefined;
   fieldFormats: FieldFormatsStart;
   getDocById: (id: string) => DataTableRecord | undefined;
@@ -50,13 +50,13 @@ export const useComparisonCellValue = ({
   dataView,
   comparisonFields,
   fieldColumnId,
-  selectedDocs,
+  selectedDocIds,
   diffMode,
   fieldFormats,
   getDocById,
   additionalFieldGroups,
 }: UseComparisonCellValueProps) => {
-  const baseDocId = selectedDocs[0];
+  const baseDocId = selectedDocIds[0];
   const baseDoc = useMemo(() => getDocById(baseDocId)?.flattened, [baseDocId, getDocById]);
   const [calculateDiffMemoized] = useState(() => createCalculateDiffMemoized());
 
@@ -92,7 +92,7 @@ export const useComparisonCellValue = ({
   );
 };
 
-type CellValueProps = Omit<UseComparisonCellValueProps, 'selectedDocs'> &
+type CellValueProps = Omit<UseComparisonCellValueProps, 'selectedDocIds'> &
   EuiDataGridCellValueElementProps & {
     baseDocId: string;
     baseDoc: DataTableRecord['flattened'] | undefined;