[Unified Data Table] Prevent undefined errors when row accessed via rows[rowIndex]

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

@@ -15,18 +15,14 @@ import { servicesMock } from './services';
 import { DataTableContext } from '../src/table_context';
 import { convertValueToString } from '../src/utils/convert_value_to_string';
 import { buildDataTableRecord } from '@kbn/discover-utils';
-import type { DataTableRecord, EsHitRecord } from '@kbn/discover-utils/types';
+import type { DataTableRecord } 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) => {
-    return buildDataTableRecord(row, dataView);
-  });
-
+const buildTableContext = (dataView: DataView, rows: DataTableRecord[]): DataTableContext => {
   return {
     expanded: undefined,
     setExpanded: jest.fn(),
-    rows: usedRows,
+    getRowByIndex: jest.fn((index) => rows[index]),
     onFilter: jest.fn(),
     dataView,
     isDarkMode: false,
@@ -38,16 +34,27 @@ const buildTableContext = (dataView: DataView, rows: EsHitRecord[]): DataTableCo
         rowIndex,
         columnId,
         fieldFormats: servicesMock.fieldFormats,
-        rows: usedRows,
+        rows,
         dataView,
         options,
       }),
   };
 };
 
-export const dataTableContextMock = buildTableContext(dataViewMock, esHitsMock);
+export const dataTableContextRowsMock = esHitsMock.map((row) =>
+  buildDataTableRecord(row, dataViewMock)
+);
+
+export const dataTableContextMock = buildTableContext(dataViewMock, dataTableContextRowsMock);
+
+export const dataTableContextComplexRowsMock = esHitsComplex.map((row) =>
+  buildDataTableRecord(row, dataViewComplexMock)
+);
 
-export const dataTableContextComplexMock = buildTableContext(dataViewComplexMock, esHitsComplex);
+export const dataTableContextComplexMock = buildTableContext(
+  dataViewComplexMock,
+  dataTableContextComplexRowsMock
+);
 
 export function buildSelectedDocsState(selectedDocIds: string[]): UseSelectedDocsState {
   const selectedDocsSet = new Set(selectedDocIds);

packages/kbn-unified-data-table/src/components/custom_control_columns/additional_row_control/row_control_column.test.tsx

@@ -54,6 +54,6 @@ describe('getRowControlColumn', () => {
 
     button.click();
 
-    expect(mockClick).toHaveBeenCalledWith({ record: contextMock.rows[1], rowIndex: 1 });
+    expect(mockClick).toHaveBeenCalledWith({ record: contextMock.getRowByIndex(1), rowIndex: 1 });
   });
 });

packages/kbn-unified-data-table/src/components/custom_control_columns/additional_row_control/row_control_column.tsx

@@ -26,7 +26,7 @@ export const RowControlCell = ({
 }: EuiDataGridCellValueElementProps & {
   renderControl: RowControlColumn['renderControl'];
 }) => {
-  const rowProps = useControlColumn(props);
+  const { record, rowIndex } = useControlColumn(props);
 
   const Control: React.FC<RowControlProps> = useMemo(
     () =>
@@ -50,17 +50,19 @@ export const RowControlCell = ({
                 color={color ?? 'text'}
                 aria-label={label}
                 onClick={() => {
-                  onClick?.(rowProps);
+                  if (record) {
+                    onClick?.({ record, rowIndex });
+                  }
                 }}
               />
             </EuiToolTip>
           </DataTableRowControl>
         );
       },
-    [props.columnId, rowProps]
+    [props.columnId, record, rowIndex]
   );
 
-  return renderControl(Control, rowProps);
+  return record ? renderControl(Control, { record, rowIndex }) : null;
 };
 
 export const getRowControlColumn = (

packages/kbn-unified-data-table/src/components/custom_control_columns/additional_row_control/row_menu_control_column.test.tsx

@@ -66,6 +66,6 @@ describe('getRowMenuControlColumn', () => {
     expect(button).toBeInTheDocument();
 
     button.click();
-    expect(mockClick).toHaveBeenCalledWith({ record: contextMock.rows[1], rowIndex: 1 });
+    expect(mockClick).toHaveBeenCalledWith({ record: contextMock.getRowByIndex(1), rowIndex: 1 });
   });
 });

packages/kbn-unified-data-table/src/components/custom_control_columns/additional_row_control/row_menu_control_column.tsx

@@ -34,7 +34,7 @@ export const RowMenuControlCell = ({
 }: EuiDataGridCellValueElementProps & {
   rowControlColumns: RowControlColumn[];
 }) => {
-  const rowProps = useControlColumn(props);
+  const { record, rowIndex } = useControlColumn(props);
   const [isMoreActionsPopoverOpen, setIsMoreActionsPopoverOpen] = useState<boolean>(false);
 
   const buttonLabel = i18n.translate('unifiedDataTable.grid.additionalRowActions', {
@@ -51,15 +51,17 @@ export const RowMenuControlCell = ({
             icon={iconType}
             color={color}
             onClick={() => {
-              onClick?.(rowProps);
+              if (record) {
+                onClick?.({ record, rowIndex });
+              }
               setIsMoreActionsPopoverOpen(false);
             }}
           >
             {label}
           </EuiContextMenuItem>
         );
       },
-    [rowProps, setIsMoreActionsPopoverOpen]
+    [record, rowIndex]
   );
 
   const popoverMenuItems = useMemo(
@@ -68,11 +70,11 @@ export const RowMenuControlCell = ({
         const Control = getControlComponent(rowControlColumn.id);
         return (
           <Fragment key={rowControlColumn.id}>
-            {rowControlColumn.renderControl(Control, rowProps)}
+            {record ? rowControlColumn.renderControl(Control, { record, rowIndex }) : null}
           </Fragment>
         );
       }),
-    [rowControlColumns, rowProps, getControlComponent]
+    [rowControlColumns, getControlComponent, record, rowIndex]
   );
 
   return (

packages/kbn-unified-data-table/src/components/custom_control_columns/color_indicator/color_indicator_control_column.test.tsx

@@ -39,6 +39,6 @@ describe('ColorIndicatorControlColumn', () => {
         />
       </UnifiedDataTableContext.Provider>
     );
-    expect(getRowIndicator).toHaveBeenCalledWith(contextMock.rows[1], expect.any(Object));
+    expect(getRowIndicator).toHaveBeenCalledWith(contextMock.getRowByIndex(1), expect.any(Object));
   });
 });

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

@@ -69,7 +69,7 @@ import {
   SELECT_ROW,
   OPEN_DETAILS,
 } from './data_table_columns';
-import { UnifiedDataTableContext } from '../table_context';
+import { DataTableContext, UnifiedDataTableContext } from '../table_context';
 import { getSchemaDetectors } from './data_table_schema';
 import { DataTableDocumentToolbarBtn } from './data_table_document_selection';
 import { useRowHeightsOptions } from '../hooks/use_row_heights_options';
@@ -629,11 +629,11 @@ export const UnifiedDataTable = ({
     );
   }, [currentPageSize, setPagination]);
 
-  const unifiedDataTableContextValue = useMemo(
+  const unifiedDataTableContextValue = useMemo<DataTableContext>(
     () => ({
       expanded: expandedDoc,
       setExpanded: setExpandedDoc,
-      rows: displayedRows,
+      getRowByIndex: (index: number) => displayedRows[index],
       onFilter,
       dataView,
       isDarkMode: darkMode,
@@ -876,7 +876,7 @@ export const UnifiedDataTable = ({
   const canSetExpandedDoc = Boolean(setExpandedDoc && !!renderDocumentView);
 
   const leadingControlColumns: EuiDataGridControlColumn[] = useMemo(() => {
-    const defaultControlColumns = getLeadControlColumns(canSetExpandedDoc);
+    const defaultControlColumns = getLeadControlColumns({ rows: displayedRows, canSetExpandedDoc });
     const internalControlColumns = controlColumnIds
       ? // reorder the default controls as per controlColumnIds
         controlColumnIds.reduce((acc, id) => {
@@ -907,6 +907,7 @@ export const UnifiedDataTable = ({
   }, [
     canSetExpandedDoc,
     controlColumnIds,
+    displayedRows,
     externalControlColumns,
     getRowIndicator,
     rowAdditionalLeadingControls,

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

@@ -18,12 +18,13 @@ import {
 import { type DataView, DataViewField } from '@kbn/data-views-plugin/public';
 import { ToastsStart, IUiSettingsClient } from '@kbn/core/public';
 import { DocViewFilterFn } from '@kbn/unified-doc-viewer/types';
+import type { DataTableRecord } from '@kbn/discover-utils';
 import { ExpandButton } from './data_table_expand_button';
 import { CustomGridColumnsConfiguration, UnifiedDataTableSettings } from '../types';
 import type { ValueToStringConverter, DataTableColumnsMeta } from '../types';
 import { buildCellActions } from './default_cell_actions';
 import { getSchemaByKbnType } from './data_table_schema';
-import { SelectButton, SelectAllButton } from './data_table_document_selection';
+import { SelectButton, getSelectAllButton } from './data_table_document_selection';
 import {
   defaultTimeColumnWidth,
   ROWS_HEIGHT_OPTIONS,
@@ -73,18 +74,24 @@ const openDetails = {
   rowCellRender: ExpandButton,
 };
 
-const select = {
+const getSelect = (rows: DataTableRecord[]) => ({
   id: SELECT_ROW,
   width: DEFAULT_CONTROL_COLUMN_WIDTH,
   rowCellRender: SelectButton,
-  headerCellRender: SelectAllButton,
-};
+  headerCellRender: getSelectAllButton(rows),
+});
 
-export function getLeadControlColumns(canSetExpandedDoc: boolean) {
+export function getLeadControlColumns({
+  rows,
+  canSetExpandedDoc,
+}: {
+  rows: DataTableRecord[];
+  canSetExpandedDoc: boolean;
+}) {
   if (!canSetExpandedDoc) {
-    return [select];
+    return [getSelect(rows)];
   }
-  return [openDetails, select];
+  return [openDetails, getSelect(rows)];
 }
 
 function buildEuiGridColumn({

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

@@ -11,20 +11,23 @@ import React, { useContext, useState } from 'react';
 import { FormattedMessage } from '@kbn/i18n-react';
 import { EuiContextMenuItem } from '@elastic/eui';
 import type { ToastsStart } from '@kbn/core/public';
+import type { DataTableRecord } from '@kbn/discover-utils';
 import { copyRowsAsJsonToClipboard } from '../utils/copy_value_to_clipboard';
 import { UnifiedDataTableContext } from '../table_context';
 
 interface DataTableCopyRowsAsJsonProps {
+  rows: DataTableRecord[];
   toastNotifications: ToastsStart;
   onCompleted: () => void;
 }
 
 export const DataTableCopyRowsAsJson: React.FC<DataTableCopyRowsAsJsonProps> = ({
+  rows,
   toastNotifications,
   onCompleted,
 }) => {
   const [isProcessing, setIsProcessing] = useState<boolean>(false);
-  const { rows, selectedDocsState, isPlainRecord } = useContext(UnifiedDataTableContext);
+  const { selectedDocsState, isPlainRecord } = useContext(UnifiedDataTableContext);
   const { getSelectedDocsOrderedByRows } = selectedDocsState;
 
   return (

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

@@ -12,23 +12,25 @@ import { uniq } from 'lodash';
 import { FormattedMessage } from '@kbn/i18n-react';
 import { EuiContextMenuItem } from '@elastic/eui';
 import type { ToastsStart } from '@kbn/core/public';
-import { calcFieldCounts } from '@kbn/discover-utils';
+import { DataTableRecord, calcFieldCounts } from '@kbn/discover-utils';
 import { copyRowsAsTextToClipboard } from '../utils/copy_value_to_clipboard';
 import { UnifiedDataTableContext } from '../table_context';
 
 interface DataTableCopyRowsAsTextProps {
+  rows: DataTableRecord[];
   toastNotifications: ToastsStart;
   columns: string[];
   onCompleted: () => void;
 }
 
 export const DataTableCopyRowsAsText: React.FC<DataTableCopyRowsAsTextProps> = ({
+  rows,
   toastNotifications,
   columns,
   onCompleted,
 }) => {
   const [isProcessing, setIsProcessing] = useState<boolean>(false);
-  const { valueToStringConverter, dataView, rows, selectedDocsState } =
+  const { valueToStringConverter, dataView, selectedDocsState } =
     useContext(UnifiedDataTableContext);
   const { isDocSelected } = selectedDocsState;
 

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

@@ -15,9 +15,13 @@ import {
   DataTableCompareToolbarBtn,
   DataTableDocumentToolbarBtn,
   SelectButton,
-  SelectAllButton,
+  getSelectAllButton,
 } from './data_table_document_selection';
-import { buildSelectedDocsState, dataTableContextMock } from '../../__mocks__/table_context';
+import {
+  buildSelectedDocsState,
+  dataTableContextMock,
+  dataTableContextRowsMock,
+} from '../../__mocks__/table_context';
 import { UnifiedDataTableContext } from '../table_context';
 import { getDocId } from '@kbn/discover-utils';
 import { render, screen } from '@testing-library/react';
@@ -49,6 +53,7 @@ describe('document selection', () => {
       const contextMock = {
         ...dataTableContextMock,
       };
+      const SelectAllButton = getSelectAllButton(dataTableContextRowsMock);
 
       const component = mountWithIntl(
         <UnifiedDataTableContext.Provider value={contextMock}>
@@ -65,6 +70,7 @@ describe('document selection', () => {
         ...dataTableContextMock,
         selectedDocsState: buildSelectedDocsState(['i::1::']),
       };
+      const SelectAllButton = getSelectAllButton(dataTableContextRowsMock);
 
       const component = mountWithIntl(
         <UnifiedDataTableContext.Provider value={contextMock}>
@@ -197,7 +203,7 @@ describe('document selection', () => {
       const props = {
         isPlainRecord: false,
         isFilterActive: false,
-        rows: dataTableContextMock.rows,
+        rows: dataTableContextRowsMock,
         selectedDocsState: buildSelectedDocsState(['i::1::', 'i::2::']),
         setIsFilterActive: jest.fn(),
         enableComparisonMode: true,
@@ -242,7 +248,7 @@ describe('document selection', () => {
       const props = {
         isPlainRecord: false,
         isFilterActive: false,
-        rows: dataTableContextMock.rows,
+        rows: dataTableContextRowsMock,
         selectedDocsState: buildSelectedDocsState(['i::1::']),
         setIsFilterActive: jest.fn(),
         enableComparisonMode: true,
@@ -271,7 +277,7 @@ describe('document selection', () => {
       const props = {
         isPlainRecord: false,
         isFilterActive: false,
-        rows: dataTableContextMock.rows,
+        rows: dataTableContextRowsMock,
         selectedDocsState: buildSelectedDocsState(['i::1::', 'i::2::']),
         setIsFilterActive: jest.fn(),
         enableComparisonMode: true,
@@ -307,7 +313,7 @@ describe('document selection', () => {
       const props = {
         isPlainRecord: false,
         isFilterActive: false,
-        rows: dataTableContextMock.rows,
+        rows: dataTableContextRowsMock,
         selectedDocsState: buildSelectedDocsState(['i::1::', 'i::2::']),
         setIsFilterActive: jest.fn(),
         enableComparisonMode: true,
@@ -336,8 +342,8 @@ describe('document selection', () => {
       const props = {
         isPlainRecord: false,
         isFilterActive: false,
-        rows: dataTableContextMock.rows,
-        selectedDocsState: buildSelectedDocsState(dataTableContextMock.rows.map((row) => row.id)),
+        rows: dataTableContextRowsMock,
+        selectedDocsState: buildSelectedDocsState(dataTableContextRowsMock.map((row) => row.id)),
         setIsFilterActive: jest.fn(),
         enableComparisonMode: true,
         setIsCompareActive: jest.fn(),
@@ -357,7 +363,7 @@ describe('document selection', () => {
         </UnifiedDataTableContext.Provider>
       );
       expect(findTestSubject(component, 'unifiedDataTableSelectionBtn').text()).toBe(
-        `Selected${dataTableContextMock.rows.length}`
+        `Selected${dataTableContextRowsMock.length}`
       );
 
       expect(findTestSubject(component, 'dscGridSelectAllDocs').exists()).toBe(false);
@@ -368,7 +374,7 @@ describe('document selection', () => {
     const props = {
       isPlainRecord: false,
       isFilterActive: false,
-      rows: dataTableContextMock.rows,
+      rows: dataTableContextRowsMock,
       selectedDocsState: buildSelectedDocsState([]),
       setIsFilterActive: jest.fn(),
       enableComparisonMode: true,