Fix row selection issue in DocumentsTab when sorting rows (#1997)

* Fix bug clicking on item highlights wrong row. Remove unused prop.

* Fix clicking on table row on sorted rows and multi-select using ctrl

* Update test snaphosts

* Remove unnecessary setTimeout

src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx

@@ -610,7 +610,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
   // Table user clicked on this row
   const [clickedRowIndex, setClickedRowIndex] = useState<number>(RESET_INDEX);
   // Table multiple selection
-  const [selectedRows, setSelectedRows] = React.useState<Set<TableRowId>>(() => new Set<TableRowId>([0]));
+  const [selectedRows, setSelectedRows] = React.useState<Set<TableRowId>>(() => new Set<TableRowId>());
 
   // Command buttons
   const [editorState, setEditorState] = useState<ViewModels.DocumentExplorerState>(
@@ -663,23 +663,6 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     }
   }, [isFilterFocused]);
 
-  // Clicked row must be defined
-  useEffect(() => {
-    if (documentIds.length > 0) {
-      let currentClickedRowIndex = clickedRowIndex;
-      if (
-        (currentClickedRowIndex === RESET_INDEX &&
-          editorState === ViewModels.DocumentExplorerState.noDocumentSelected) ||
-        currentClickedRowIndex > documentIds.length - 1
-      ) {
-        // reset clicked row or the current clicked row is out of bounds
-        currentClickedRowIndex = INITIAL_SELECTED_ROW_INDEX;
-        setSelectedRows(new Set([INITIAL_SELECTED_ROW_INDEX]));
-        onDocumentClicked(currentClickedRowIndex, documentIds);
-      }
-    }
-  }, [documentIds, clickedRowIndex, editorState]);
-
   /**
    * Recursively delete all documents by retrying throttled requests (429).
    * This only works for NoSQL, because the bulk response includes status for each delete document request.
@@ -2232,7 +2215,6 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
                     <DocumentsTableComponent
                       onRefreshTable={() => refreshDocumentsGrid(false)}
                       items={tableItems}
-                      onItemClicked={(index) => onDocumentClicked(index, documentIds)}
                       onSelectedRowsChange={onSelectedRowsChange}
                       selectedRows={selectedRows}
                       size={tableContainerSizePx}

src/Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent.test.tsx

@@ -14,7 +14,6 @@ describe("DocumentsTableComponent", () => {
       { [ID_HEADER]: "2", [PARTITION_KEY_HEADER]: "pk2" },
       { [ID_HEADER]: "3", [PARTITION_KEY_HEADER]: "pk3" },
     ],
-    onItemClicked: (): void => {},
     onSelectedRowsChange: (): void => {},
     selectedRows: new Set<TableRowId>(),
     size: {

src/Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent.tsx

@@ -68,7 +68,6 @@ export type ColumnDefinition = {
 export interface IDocumentsTableComponentProps {
   onRefreshTable: () => void;
   items: DocumentsTableComponentItem[];
-  onItemClicked: (index: number) => void;
   onSelectedRowsChange: (selectedItemsIndices: Set<TableRowId>) => void;
   selectedRows: Set<TableRowId>;
   size: { height: number; width: number };
@@ -98,6 +97,7 @@ const defaultSize = {
   idealWidth: 200,
   minWidth: 50,
 };
+
 export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> = ({
   onRefreshTable,
   items,
@@ -115,6 +115,8 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
 }: IDocumentsTableComponentProps) => {
   const styles = useDocumentsTabStyles();
 
+  const sortedRowsRef = React.useRef(null);
+
   const [columnSizingOptions, setColumnSizingOptions] = React.useState<TableColumnSizingOptions>(() => {
     const columnSizesMap: ColumnSizesMap = readSubComponentState(SubComponentName.ColumnSizes, collection, {});
     const columnSizesPx: TableColumnSizingOptions = {};
@@ -291,22 +293,42 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
 
   const [selectionStartIndex, setSelectionStartIndex] = React.useState<number>(INITIAL_SELECTED_ROW_INDEX);
   const onTableCellClicked = useCallback(
-    (e: React.MouseEvent, index: number) => {
+    (e: React.MouseEvent | undefined, index: number, rowId: TableRowId) => {
       if (isSelectionDisabled) {
         // Only allow click
-        onSelectedRowsChange(new Set<TableRowId>([index]));
+        onSelectedRowsChange(new Set<TableRowId>([rowId]));
         setSelectionStartIndex(index);
         return;
       }
 
+      // The selection helper computes in the index space (what's visible to the user in the table, ie the sorted array).
+      // selectedRows is in the rowId space (the index of the original unsorted array), so it must be converted to the index space.
+      const selectedRowsIndex = new Set<number>();
+      selectedRows.forEach((rowId) => {
+        const index = sortedRowsRef.current.findIndex((row: TableRowData) => row.rowId === rowId);
+        if (index !== -1) {
+          selectedRowsIndex.add(index);
+        } else {
+          // This should never happen
+          console.error(`Row with rowId ${rowId} not found in sorted rows`);
+        }
+      });
+
       const result = selectionHelper(
-        selectedRows as Set<number>,
+        selectedRowsIndex,
         index,
-        isEnvironmentShiftPressed(e),
+        e && isEnvironmentShiftPressed(e),
-        isEnvironmentCtrlPressed(e),
+        e && isEnvironmentCtrlPressed(e),
         selectionStartIndex,
       );
-      onSelectedRowsChange(result.selection);
+
+      // Convert selectionHelper result from index space back to rowId space
+      const selectedRowIds = new Set<TableRowId>();
+      result.selection.forEach((index) => {
+        selectedRowIds.add(sortedRowsRef.current[index].rowId);
+      });
+      onSelectedRowsChange(selectedRowIds);
+
       if (result.selectionStartIndex !== undefined) {
         setSelectionStartIndex(result.selectionStartIndex);
       }
@@ -320,16 +342,20 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
    * - a key is down and the cell is clicked by the mouse
    */
   const onIdClicked = useCallback(
-    (e: React.KeyboardEvent<Element>, index: number) => {
+    (e: React.KeyboardEvent<Element>, rowId: TableRowId) => {
       if (e.key === NormalizedEventKey.Enter || e.key === NormalizedEventKey.Space) {
-        onSelectedRowsChange(new Set<TableRowId>([index]));
+        onSelectedRowsChange(new Set<TableRowId>([rowId]));
       }
     },
     [onSelectedRowsChange],
   );
 
   const RenderRow = ({ index, style, data }: ReactWindowRenderFnProps) => {
-    const { item, selected, appearance, onClick, onKeyDown } = data[index];
+    // WARNING: because the table sorts the data, 'index' is not the same as 'rowId'
+    // The rowId is the index of the item in the original array,
+    // while the index is the index of the item in the sorted array
+    const { item, selected, appearance, onClick, onKeyDown, rowId } = data[index];
+
     return (
       <TableRow
         aria-rowindex={index + 2}
@@ -359,8 +385,8 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
             key={column.columnId}
             className={styles.tableCell}
             // When clicking on a cell with shift/ctrl key, onKeyDown is called instead of onClick.
-            onClick={(e: React.MouseEvent<Element, MouseEvent>) => onTableCellClicked(e, index)}
+            onClick={(e: React.MouseEvent<Element, MouseEvent>) => onTableCellClicked(e, index, rowId)}
-            onKeyPress={(e: React.KeyboardEvent<Element>) => onIdClicked(e, index)}
+            onKeyPress={(e: React.KeyboardEvent<Element>) => onIdClicked(e, rowId)}
             {...columnSizing.getTableCellProps(column.columnId)}
             tabIndex={column.columnId === "id" ? 0 : -1}
           >
@@ -420,6 +446,19 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
     }),
   );
 
+  // Store the sorted rows in a ref which won't trigger a re-render (as opposed to a state)
+  sortedRowsRef.current = rows;
+
+  // If there are no selected rows, auto select the first row
+  const [autoSelectFirstDoc, setAutoSelectFirstDoc] = React.useState<boolean>(true);
+  React.useEffect(() => {
+    if (autoSelectFirstDoc && sortedRowsRef.current?.length > 0 && selectedRows.size === 0) {
+      setAutoSelectFirstDoc(false);
+      const DOC_INDEX_TO_SELECT = 0;
+      onTableCellClicked(undefined, DOC_INDEX_TO_SELECT, sortedRowsRef.current[DOC_INDEX_TO_SELECT].rowId);
+    }
+  }, [selectedRows, onTableCellClicked, autoSelectFirstDoc]);
+
   const toggleAllKeydown = React.useCallback(
     (e: React.KeyboardEvent<HTMLDivElement>) => {
       if (e.key === " ") {

src/Explorer/Tabs/DocumentsTabV2/__snapshots__/DocumentsTabV2.test.tsx.snap

@@ -90,7 +90,6 @@ exports[`Documents tab (noSql API) when rendered should render the page 1`] = `
                   isRowSelectionDisabled={true}
                   items={[]}
                   onColumnSelectionChange={[Function]}
-                  onItemClicked={[Function]}
                   onRefreshTable={[Function]}
                   onSelectedRowsChange={[Function]}
                   selectedColumnIds={
@@ -98,11 +97,7 @@ exports[`Documents tab (noSql API) when rendered should render the page 1`] = `
                       "id",
                     ]
                   }
-                  selectedRows={
+                  selectedRows={Set {}}
-                    Set {
-                      0,
-                    }
-                  }
                 />
               </div>
             </div>

src/Explorer/Tabs/DocumentsTabV2/__snapshots__/DocumentsTableComponent.test.tsx.snap

@@ -39,7 +39,6 @@ exports[`DocumentsTableComponent should not render selection column when isSelec
       },
     ]
   }
-  onItemClicked={[Function]}
   onRefreshTable={[Function]}
   onSelectedRowsChange={[Function]}
   selectedColumnIds={[]}
@@ -504,7 +503,6 @@ exports[`DocumentsTableComponent should render documents and partition keys in h
       },
     ]
   }
-  onItemClicked={[Function]}
   onRefreshTable={[Function]}
   onSelectedRowsChange={[Function]}
   selectedColumnIds={[]}