[EuiInMemoryTable] Replace basic usages of deprecated ref method with controlled selection.selected API

x-pack/packages/kbn-elastic-assistant/impl/data_anonymization_editor/context_editor/index.test.tsx

@@ -12,9 +12,12 @@ import React from 'react';
 import { ContextEditor } from '.';
 
 describe('ContextEditor', () => {
-  const allow = ['field1', 'field2'];
+  const allow = Array.from({ length: 20 }, (_, i) => `field${i + 1}`);
   const allowReplacement = ['field1'];
-  const rawData = { field1: ['value1'], field2: ['value2'] };
+  const rawData = allow.reduce(
+    (acc, field, index) => ({ ...acc, [field]: [`value${index + 1}`] }),
+    {}
+  );
 
   const onListUpdated = jest.fn();
 
@@ -36,13 +39,17 @@ describe('ContextEditor', () => {
   });
 
   it('renders the select all fields button with the expected count', () => {
-    expect(screen.getByTestId('selectAllFields')).toHaveTextContent('Select all 2 fields');
+    expect(screen.getByTestId('selectAllFields')).toHaveTextContent('Select all 20 fields');
   });
 
   it('updates the table selection when "Select all n fields" is clicked', () => {
-    userEvent.click(screen.getByTestId('selectAllFields'));
+    // The table select all checkbox should only select the number of rows visible on the page
+    userEvent.click(screen.getByTestId('checkboxSelectAll'));
+    expect(screen.getByTestId('selectedFields')).toHaveTextContent('Selected 10 fields');
 
-    expect(screen.getByTestId('selectedFields')).toHaveTextContent('Selected 2 fields');
+    // The select all button should select all rows regardless of visibility
+    userEvent.click(screen.getByTestId('selectAllFields'));
+    expect(screen.getByTestId('selectedFields')).toHaveTextContent('Selected 20 fields');
   });
 
   it('calls onListUpdated with the expected values when the update button is clicked', () => {

x-pack/packages/kbn-elastic-assistant/impl/data_anonymization_editor/context_editor/index.tsx

@@ -7,7 +7,7 @@
 
 import { EuiInMemoryTable } from '@elastic/eui';
 import type { EuiSearchBarProps, EuiTableSelectionType } from '@elastic/eui';
-import React, { useCallback, useMemo, useRef, useState } from 'react';
+import React, { useCallback, useMemo, useState, useRef } from 'react';
 
 import { getColumns } from './get_columns';
 import { getRows } from './get_rows';
@@ -59,16 +59,25 @@ const ContextEditorComponent: React.FC<Props> = ({
   rawData,
   pageSize = DEFAULT_PAGE_SIZE,
 }) => {
+  const isAllSelected = useRef(false); // Must be a ref and not state in order not to re-render `selectionValue`, which fires `onSelectionChange` twice
   const [selected, setSelection] = useState<ContextEditorRow[]>([]);
   const selectionValue: EuiTableSelectionType<ContextEditorRow> = useMemo(
     () => ({
       selectable: () => true,
-      onSelectionChange: (newSelection) => setSelection(newSelection),
-      initialSelected: [],
+      onSelectionChange: (newSelection) => {
+        if (isAllSelected.current === true) {
+          // If passed every possible row (including non-visible ones), EuiInMemoryTable
+          // will fire `onSelectionChange` with only the visible rows - we need to
+          // ignore this call when that happens and continue to pass all rows
+          isAllSelected.current = false;
+        } else {
+          setSelection(newSelection);
+        }
+      },
+      selected,
     }),
-    []
+    [selected]
   );
-  const tableRef = useRef<EuiInMemoryTable<ContextEditorRow> | null>(null);
 
   const columns = useMemo(() => getColumns({ onListUpdated, rawData }), [onListUpdated, rawData]);
 
@@ -83,9 +92,8 @@ const ContextEditorComponent: React.FC<Props> = ({
   );
 
   const onSelectAll = useCallback(() => {
-    tableRef.current?.setSelection(rows); // updates selection in the EuiInMemoryTable
-
-    setTimeout(() => setSelection(rows), 0); // updates selection in the component state
+    isAllSelected.current = true;
+    setSelection(rows);
   }, [rows]);
 
   const pagination = useMemo(() => {
@@ -106,7 +114,7 @@ const ContextEditorComponent: React.FC<Props> = ({
         totalFields={rows.length}
       />
     ),
-    [onListUpdated, onReset, onSelectAll, rawData, rows.length, selected]
+    [onListUpdated, onReset, onSelectAll, rawData, rows, selected]
   );
 
   return (
@@ -120,7 +128,6 @@ const ContextEditorComponent: React.FC<Props> = ({
       itemId={FIELDS.FIELD}
       items={rows}
       pagination={pagination}
-      ref={tableRef}
       search={search}
       selection={selectionValue}
       sorting={defaultSort}

x-pack/plugins/cases/public/components/all_cases/all_cases_list.tsx

@@ -5,8 +5,8 @@
  * 2.0.
  */
 
-import React, { useCallback, useMemo, useRef, useState } from 'react';
-import type { EuiBasicTable, EuiTableSelectionType } from '@elastic/eui';
+import React, { useCallback, useMemo, useState } from 'react';
+import type { EuiTableSelectionType } from '@elastic/eui';
 import { EuiProgress } from '@elastic/eui';
 import { difference, head, isEmpty } from 'lodash/fp';
 import styled, { css } from 'styled-components';
@@ -112,11 +112,8 @@ export const AllCasesList = React.memo<AllCasesListProps>(
       [queryParams.sortField, queryParams.sortOrder]
     );
 
-    const tableRef = useRef<EuiBasicTable | null>(null);
-
     const deselectCases = useCallback(() => {
       setSelectedCases([]);
-      tableRef.current?.setSelection([]);
     }, [setSelectedCases]);
 
     const tableOnChangeCallback = useCallback(
@@ -175,7 +172,7 @@ export const AllCasesList = React.memo<AllCasesListProps>(
     const euiBasicTableSelectionProps = useMemo<EuiTableSelectionType<CaseUI>>(
       () => ({
         onSelectionChange: setSelectedCases,
-        initialSelected: selectedCases,
+        selected: selectedCases,
         selectable: () => !isReadOnlyPermissions(permissions),
       }),
       [permissions, selectedCases]
@@ -229,7 +226,6 @@ export const AllCasesList = React.memo<AllCasesListProps>(
           selectedCases={selectedCases}
           selection={euiBasicTableSelectionProps}
           sorting={sorting}
-          tableRef={tableRef}
           tableRowProps={tableRowProps}
           deselectCases={deselectCases}
           selectedColumns={selectedColumns}

x-pack/plugins/cases/public/components/all_cases/table.tsx

@@ -35,7 +35,7 @@ interface CasesTableProps {
   selectedCases: CasesUI;
   selection: EuiTableSelectionType<CaseUI>;
   sorting: EuiBasicTableProps<CaseUI>['sorting'];
-  tableRef: MutableRefObject<EuiBasicTable | null>;
+  tableRef?: MutableRefObject<EuiBasicTable | null>;
   tableRowProps: EuiBasicTableProps<CaseUI>['rowProps'];
   deselectCases: () => void;
   selectedColumns: CasesColumnSelection[];

x-pack/plugins/infra/public/pages/metrics/hosts/components/hosts_table.tsx

@@ -33,7 +33,6 @@ export const HostsTable = () => {
     selection,
     selectedItemsCount,
     filterSelectedHosts,
-    refs,
   } = useHostsTableContext();
 
   return (
@@ -43,7 +42,6 @@ export const HostsTable = () => {
         filterSelectedHosts={filterSelectedHosts}
       />
       <EuiBasicTable
-        ref={refs.tableRef}
         data-test-subj={`hostsView-table-${loading ? 'loading' : 'loaded'}`}
         itemId="id"
         isSelectable

x-pack/plugins/infra/public/pages/metrics/hosts/hooks/use_hosts_table.tsx

@@ -5,13 +5,8 @@
  * 2.0.
  */
 
-import React, { useCallback, useMemo, useRef, useState } from 'react';
-import {
-  EuiBasicTableColumn,
-  CriteriaWithPagination,
-  EuiTableSelectionType,
-  type EuiBasicTable,
-} from '@elastic/eui';
+import React, { useCallback, useMemo, useState } from 'react';
+import { EuiBasicTableColumn, CriteriaWithPagination, EuiTableSelectionType } from '@elastic/eui';
 import createContainer from 'constate';
 import useAsync from 'react-use/lib/useAsync';
 import { isEqual } from 'lodash';
@@ -142,8 +137,6 @@ export const useHostsTable = () => {
   } = useKibanaContextForPlugin();
   const { dataView } = useMetricsDataViewContext();
 
-  const tableRef = useRef<EuiBasicTable | null>(null);
-
   const closeFlyout = useCallback(() => setProperties({ detailsItemId: null }), [setProperties]);
 
   const onSelectionChange = (newSelectedItems: HostNodeRow[]) => {
@@ -163,7 +156,6 @@ export const useHostsTable = () => {
 
     filterManagerService.addFilters(newFilter);
     setSelectedItems([]);
-    tableRef.current?.setSelection([]);
   }, [dataView, filterManagerService, selectedItems]);
 
   const reportHostEntryClick = useCallback(
@@ -358,6 +350,7 @@ export const useHostsTable = () => {
   const selection: EuiTableSelectionType<HostNodeRow> = {
     onSelectionChange,
     selectable: (item: HostNodeRow) => !!item.name,
+    selected: selectedItems,
   };
 
   return {
@@ -373,9 +366,6 @@ export const useHostsTable = () => {
     selection,
     selectedItemsCount: selectedItems.length,
     filterSelectedHosts,
-    refs: {
-      tableRef,
-    },
   };
 };
 

x-pack/plugins/saved_objects_tagging/public/management/components/table.tsx

@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import React, { useRef, useEffect, FC, ReactNode } from 'react';
+import React, { FC, ReactNode } from 'react';
 import { EuiInMemoryTable, EuiBasicTableColumn, EuiLink, Query } from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
 import { FormattedMessage } from '@kbn/i18n-react';
@@ -57,14 +57,6 @@ export const TagTable: FC<TagTableProps> = ({
   actionBar,
   actions,
 }) => {
-  const tableRef = useRef<EuiInMemoryTable<TagWithRelations>>(null);
-
-  useEffect(() => {
-    if (tableRef.current) {
-      tableRef.current.setSelection(selectedTags);
-    }
-  }, [selectedTags]);
-
   const columns: Array<EuiBasicTableColumn<TagWithRelations>> = [
     {
       field: 'name',
@@ -144,7 +136,6 @@ export const TagTable: FC<TagTableProps> = ({
   return (
     <EuiInMemoryTable
       data-test-subj={`tagsManagementTable ${testSubjectState}`}
-      ref={tableRef}
       childrenBetween={actionBar}
       loading={loading}
       itemId={'id'}
@@ -159,7 +150,7 @@ export const TagTable: FC<TagTableProps> = ({
       selection={
         allowSelection
           ? {
-              initialSelected: selectedTags,
+              selected: selectedTags,
               onSelectionChange,
             }
           : undefined