[Discover] Add "Shift + Select" functionality to Discover grid

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

@@ -69,6 +69,7 @@ export function buildSelectedDocsState(selectedDocIds: string[]): UseSelectedDoc
     selectedDocsCount: selectedDocsSet.size,
     docIdsInSelectionOrder: selectedDocIds,
     toggleDocSelection: jest.fn(),
+    toggleMultipleDocsSelection: jest.fn(),
     selectAllDocs: jest.fn(),
     selectMoreDocs: jest.fn(),
     deselectSomeDocs: jest.fn(),

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

@@ -497,8 +497,14 @@ export const UnifiedDataTable = ({
   const [isCompareActive, setIsCompareActive] = useState(false);
   const displayedColumns = getDisplayedColumns(columns, dataView);
   const defaultColumns = displayedColumns.includes('_source');
-  const docMap = useMemo(() => new Map(rows?.map((row) => [row.id, row]) ?? []), [rows]);
-  const getDocById = useCallback((id: string) => docMap.get(id), [docMap]);
+  const docMap = useMemo(
+    () =>
+      new Map<string, { doc: DataTableRecord; docIndex: number }>(
+        rows?.map((row, docIndex) => [row.id, { doc: row, docIndex }]) ?? []
+      ),
+    [rows]
+  );
+  const getDocById = useCallback((id: string) => docMap.get(id)?.doc, [docMap]);
   const selectedDocsState = useSelectedDocs(docMap);
   const {
     isDocSelected,

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

@@ -38,7 +38,7 @@ export const SelectButton = (props: EuiDataGridCellValueElementProps) => {
   const { record, rowIndex } = useControlColumn(props);
   const { euiTheme } = useEuiTheme();
   const { selectedDocsState } = useContext(UnifiedDataTableContext);
-  const { isDocSelected, toggleDocSelection } = selectedDocsState;
+  const { isDocSelected, toggleDocSelection, toggleMultipleDocsSelection } = selectedDocsState;
 
   const toggleDocumentSelectionLabel = i18n.translate('unifiedDataTable.grid.selectDoc', {
     defaultMessage: `Select document ''{rowNumber}''`,
@@ -66,8 +66,12 @@ export const SelectButton = (props: EuiDataGridCellValueElementProps) => {
           aria-label={toggleDocumentSelectionLabel}
           checked={isDocSelected(record.id)}
           data-test-subj={`dscGridSelectDoc-${record.id}`}
-          onChange={() => {
-            toggleDocSelection(record.id);
+          onChange={(event) => {
+            if ((event.nativeEvent as MouseEvent)?.shiftKey) {
+              toggleMultipleDocsSelection(record.id);
+            } else {
+              toggleDocSelection(record.id);
+            }
           }}
         />
       </EuiFlexItem>

packages/kbn-unified-data-table/src/hooks/use_selected_docs.test.ts

@@ -17,7 +17,7 @@ describe('useSelectedDocs', () => {
   const docs = generateEsHits(dataViewWithTimefieldMock, 5).map((hit) =>
     buildDataTableRecord(hit, dataViewWithTimefieldMock)
   );
-  const docsMap = new Map(docs.map((doc) => [doc.id, doc]));
+  const docsMap = new Map(docs.map((doc, docIndex) => [doc.id, { doc, docIndex }]));
 
   test('should have a correct default state', () => {
     const { result } = renderHook(() => useSelectedDocs(docsMap));
@@ -223,4 +223,30 @@ describe('useSelectedDocs', () => {
     expect(result.current.getCountOfFilteredSelectedDocs([docs[0].id])).toBe(0);
     expect(result.current.getCountOfFilteredSelectedDocs([docs[2].id, docs[3].id])).toBe(0);
   });
+
+  test('should toggleMultipleDocsSelection correctly', () => {
+    const { result } = renderHook(() => useSelectedDocs(docsMap));
+    const docIds = docs.map((doc) => doc.id);
+
+    // select `0`
+    act(() => {
+      result.current.toggleDocSelection(docs[0].id);
+    });
+
+    expect(result.current.getCountOfFilteredSelectedDocs(docIds)).toBe(1);
+
+    // select from `0` to `4`
+    act(() => {
+      result.current.toggleMultipleDocsSelection(docs[4].id);
+    });
+
+    expect(result.current.getCountOfFilteredSelectedDocs(docIds)).toBe(5);
+
+    // deselect from `2` to `4`
+    act(() => {
+      result.current.toggleMultipleDocsSelection(docs[2].id);
+    });
+
+    expect(result.current.getCountOfFilteredSelectedDocs(docIds)).toBe(2);
+  });
 });

packages/kbn-unified-data-table/src/hooks/use_selected_docs.ts

@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { useCallback, useMemo, useState } from 'react';
+import { useCallback, useMemo, useRef, useState } from 'react';
 import type { DataTableRecord } from '@kbn/discover-utils';
 
 export interface UseSelectedDocsState {
@@ -17,6 +17,7 @@ export interface UseSelectedDocsState {
   selectedDocsCount: number;
   docIdsInSelectionOrder: string[];
   toggleDocSelection: (docId: string) => void;
+  toggleMultipleDocsSelection: (toDocId: string) => void;
   selectAllDocs: () => void;
   selectMoreDocs: (docIds: string[]) => void;
   deselectSomeDocs: (docIds: string[]) => void;
@@ -25,8 +26,11 @@ export interface UseSelectedDocsState {
   getSelectedDocsOrderedByRows: (rows: DataTableRecord[]) => DataTableRecord[];
 }
 
-export const useSelectedDocs = (docMap: Map<string, DataTableRecord>): UseSelectedDocsState => {
+export const useSelectedDocs = (
+  docMap: Map<string, { doc: DataTableRecord; docIndex: number }>
+): UseSelectedDocsState => {
   const [selectedDocsSet, setSelectedDocsSet] = useState<Set<string>>(new Set());
+  const lastCheckboxToggledDocId = useRef<string | undefined>();
 
   const toggleDocSelection = useCallback((docId: string) => {
     setSelectedDocsSet((prevSelectedRowsSet) => {
@@ -38,6 +42,7 @@ export const useSelectedDocs = (docMap: Map<string, DataTableRecord>): UseSelect
       }
       return newSelectedRowsSet;
     });
+    lastCheckboxToggledDocId.current = docId;
   }, []);
 
   const replaceSelectedDocs = useCallback((docIds: string[]) => {
@@ -73,6 +78,42 @@ export const useSelectedDocs = (docMap: Map<string, DataTableRecord>): UseSelect
     [selectedDocsSet, docMap]
   );
 
+  const toggleMultipleDocsSelection = useCallback(
+    (toDocId: string) => {
+      const shouldSelect = !isDocSelected(toDocId);
+
+      const lastToggledDocIdIndex = docMap.get(
+        lastCheckboxToggledDocId.current ?? toDocId
+      )?.docIndex;
+      const currentToggledDocIdIndex = docMap.get(toDocId)?.docIndex;
+      const docIds: string[] = [];
+
+      if (
+        typeof lastToggledDocIdIndex === 'number' &&
+        typeof currentToggledDocIdIndex === 'number' &&
+        lastToggledDocIdIndex !== currentToggledDocIdIndex
+      ) {
+        const startIndex = Math.min(lastToggledDocIdIndex, currentToggledDocIdIndex);
+        const endIndex = Math.max(lastToggledDocIdIndex, currentToggledDocIdIndex);
+
+        docMap.forEach(({ doc, docIndex }) => {
+          if (docIndex >= startIndex && docIndex <= endIndex) {
+            docIds.push(doc.id);
+          }
+        });
+      }
+
+      if (shouldSelect) {
+        selectMoreDocs(docIds);
+      } else {
+        deselectSomeDocs(docIds);
+      }
+
+      lastCheckboxToggledDocId.current = toDocId;
+    },
+    [selectMoreDocs, deselectSomeDocs, docMap, isDocSelected]
+  );
+
   const getSelectedDocsOrderedByRows = useCallback(
     (rows: DataTableRecord[]) => {
       return rows.filter((row) => isDocSelected(row.id));
@@ -101,6 +142,7 @@ export const useSelectedDocs = (docMap: Map<string, DataTableRecord>): UseSelect
       docIdsInSelectionOrder: selectedDocIds,
       getCountOfFilteredSelectedDocs,
       toggleDocSelection,
+      toggleMultipleDocsSelection,
       selectAllDocs,
       selectMoreDocs,
       deselectSomeDocs,
@@ -112,6 +154,7 @@ export const useSelectedDocs = (docMap: Map<string, DataTableRecord>): UseSelect
       isDocSelected,
       getCountOfFilteredSelectedDocs,
       toggleDocSelection,
+      toggleMultipleDocsSelection,
       selectAllDocs,
       selectMoreDocs,
       deselectSomeDocs,

test/functional/apps/discover/group2_data_grid3/_data_grid_row_selection.ts

@@ -84,6 +84,60 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
       });
     });
 
+    it('should be able to select multiple rows holding Shift key', async () => {
+      expect(await dataGrid.isSelectedRowsMenuVisible()).to.be(false);
+
+      // select 1 row
+      await dataGrid.selectRow(1);
+
+      await retry.try(async () => {
+        expect(await dataGrid.isSelectedRowsMenuVisible()).to.be(true);
+        expect(await dataGrid.getNumberOfSelectedRowsOnCurrentPage()).to.be(1);
+        expect(await dataGrid.getNumberOfSelectedRows()).to.be(1);
+      });
+
+      // select 3 more
+      await dataGrid.selectRow(4, { pressShiftKey: true });
+
+      await retry.try(async () => {
+        expect(await dataGrid.isSelectedRowsMenuVisible()).to.be(true);
+        expect(await dataGrid.getNumberOfSelectedRowsOnCurrentPage()).to.be(4);
+        expect(await dataGrid.getNumberOfSelectedRows()).to.be(4);
+      });
+
+      // deselect index 3 and 4
+      await dataGrid.selectRow(3, { pressShiftKey: true });
+
+      await retry.try(async () => {
+        expect(await dataGrid.isSelectedRowsMenuVisible()).to.be(true);
+        expect(await dataGrid.getNumberOfSelectedRowsOnCurrentPage()).to.be(2);
+        expect(await dataGrid.getNumberOfSelectedRows()).to.be(2);
+      });
+
+      // select from index 3 to 0
+      await dataGrid.selectRow(0, { pressShiftKey: true });
+
+      await retry.try(async () => {
+        expect(await dataGrid.isSelectedRowsMenuVisible()).to.be(true);
+        expect(await dataGrid.getNumberOfSelectedRowsOnCurrentPage()).to.be(4);
+        expect(await dataGrid.getNumberOfSelectedRows()).to.be(4);
+      });
+
+      // select from both pages
+      await testSubjects.click('pagination-button-1');
+      await retry.try(async () => {
+        expect(await dataGrid.getNumberOfSelectedRowsOnCurrentPage()).to.be(0);
+      });
+
+      await dataGrid.selectRow(2, { pressShiftKey: true });
+
+      await retry.try(async () => {
+        expect(await dataGrid.isSelectedRowsMenuVisible()).to.be(true);
+        expect(await dataGrid.getNumberOfSelectedRowsOnCurrentPage()).to.be(3);
+        expect(await dataGrid.getNumberOfSelectedRows()).to.be(8);
+      });
+    });
+
     it('should be able to bulk select rows', async () => {
       expect(await dataGrid.isSelectedRowsMenuVisible()).to.be(false);
       expect(await testSubjects.getAttribute('selectAllDocsOnPageToggle', 'title')).to.be(

test/functional/services/data_grid.ts

@@ -700,12 +700,21 @@ export class DataGridService extends FtrService {
     await this.checkCurrentRowsPerPageToBe(newValue);
   }
 
-  public async selectRow(rowIndex: number) {
+  public async selectRow(rowIndex: number, { pressShiftKey }: { pressShiftKey?: boolean } = {}) {
     const checkbox = await this.find.byCssSelector(
       `.euiDataGridRow[data-grid-visible-row-index="${rowIndex}"] [data-gridcell-column-id="select"] .euiCheckbox__input`
     );
 
-    await checkbox.click();
+    if (pressShiftKey) {
+      await this.browser
+        .getActions()
+        .keyDown(Key.SHIFT)
+        .click(checkbox._webElement)
+        .keyUp(Key.SHIFT)
+        .perform();
+    } else {
+      await checkbox.click();
+    }
   }
 
   public async getNumberOfSelectedRows() {