Rework column selection UI

src/Explorer/Panes/TableColumnSelectionPane/TableColumnSelectionPane.less

@@ -0,0 +1,3 @@
+.tableColumnSelectionCheckbox label {
+  padding: 4px 8px;
+}

src/Explorer/Panes/TableColumnSelectionPane/TableColumnSelectionPane.tsx

@@ -0,0 +1,108 @@
+import { Checkbox } from "@fluentui/react";
+import {
+  Button,
+  FluentProvider,
+  InputOnChangeData,
+  SearchBox,
+  SearchBoxChangeEvent,
+  Text,
+} from "@fluentui/react-components";
+import { configContext } from "ConfigContext";
+import { ColumnDefinition } from "Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent";
+import { getPlatformTheme } from "Explorer/Theme/ThemeUtil";
+import React from "react";
+import { useSidePanel } from "../../../hooks/useSidePanel";
+import "./TableColumnSelectionPane.less";
+
+export interface TableColumnSelectionPaneProps {
+  columnDefinitions: ColumnDefinition[];
+  selectedColumnIds: string[];
+  onSelectionChange: (newSelectedColumnIds: string[]) => void;
+}
+
+export const TableColumnSelectionPane: React.FC<TableColumnSelectionPaneProps> = ({
+  columnDefinitions,
+  selectedColumnIds,
+  onSelectionChange,
+}: TableColumnSelectionPaneProps): JSX.Element => {
+  const closeSidePanel = useSidePanel((state) => state.closeSidePanel);
+  const originalSelectedColumnIds = React.useMemo(() => selectedColumnIds, []);
+  const [columnSearchText, setColumnSearchText] = React.useState<string>("");
+  const [newSelectedColumnIds, setNewSelectedColumnIds] = React.useState<string[]>(originalSelectedColumnIds);
+
+  const selectedColumnIdsSet = new Set(newSelectedColumnIds);
+  const onCheckedValueChange = (id: string, checked?: boolean): void => {
+    if (checked) {
+      selectedColumnIdsSet.add(id);
+    } else {
+      selectedColumnIdsSet.delete(id);
+    }
+    setNewSelectedColumnIds([...selectedColumnIdsSet]);
+  };
+
+  const onSave = (): void => {
+    onSelectionChange(newSelectedColumnIds);
+    closeSidePanel();
+  };
+
+  const onSearchChange: (event: SearchBoxChangeEvent, data: InputOnChangeData) => void = (_, data) =>
+    // eslint-disable-next-line react/prop-types
+    setColumnSearchText(data.value);
+
+  const theme = getPlatformTheme(configContext.platform);
+
+  // Filter and move partition keys to the top
+  const columnDefinitionList = columnDefinitions
+    .filter((def) => !columnSearchText || def.label.toLowerCase().includes(columnSearchText.toLowerCase()))
+    .sort((a, b) => {
+      const ID = "id";
+      // "id" always at the top, then partition keys, then everything else sorted
+      if (a.id === ID) {
+        return b.id === ID ? 0 : -1;
+      } else if (b.id === ID) {
+        return a.id === ID ? 0 : 1;
+      } else if (a.isPartitionKey && !b.isPartitionKey) {
+        return -1;
+      } else if (b.isPartitionKey && !a.isPartitionKey) {
+        return 1;
+      } else {
+        return a.label.localeCompare(b.label);
+      }
+    });
+
+  return (
+    <FluentProvider theme={theme} style={{ height: "100%" }}>
+      <div className="panelFormWrapper">
+        <div className="panelMainContent" style={{ display: "flex", flexDirection: "column" }}>
+          <Text>Select which columns to display in your view of items in your container.</Text>
+          <div /* Wrap <SearchBox> to avoid margin-bottom set by panelMainContent css */>
+            <SearchBox
+              value={columnSearchText}
+              onChange={onSearchChange}
+              style={{ width: "100%" }}
+              placeholder="Search fields"
+            />
+          </div>
+
+          {columnDefinitionList.map((columnDefinition) => (
+            <Checkbox
+              className="tableColumnSelectionCheckbox"
+              key={columnDefinition.id}
+              label={columnDefinition.label}
+              checked={selectedColumnIdsSet.has(columnDefinition.id)}
+              onChange={(_, checked) => onCheckedValueChange(columnDefinition.id, checked)}
+            />
+          ))}
+        </div>
+        <div className="panelFooter" style={{ display: "flex", gap: theme.spacingHorizontalS }}>
+          <Button appearance="primary" onClick={onSave}>
+            Save
+          </Button>
+          <Button appearance="secondary" onClick={closeSidePanel}>
+            Cancel
+          </Button>
+        </div>
+      </div>
+    </FluentProvider>
+  );
+};

src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx

@@ -1,11 +1,10 @@
 import { Item, ItemDefinition, PartitionKey, PartitionKeyDefinition, QueryIterator, Resource } from "@azure/cosmos";
 import { Button, FluentProvider, Input, TableRowId } from "@fluentui/react-components";
-import { ArrowClockwise16Filled, Dismiss16Filled } from "@fluentui/react-icons";
+import { Dismiss16Filled } from "@fluentui/react-icons";
 import Split from "@uiw/react-split";
 import { KeyCodes, QueryCopilotSampleContainerId, QueryCopilotSampleDatabaseId } from "Common/Constants";
 import { getErrorMessage, getErrorStack } from "Common/ErrorHandlingUtils";
 import MongoUtility from "Common/MongoUtility";
-import { StyleConstants } from "Common/StyleConstants";
 import { createDocument } from "Common/dataAccess/createDocument";
 import { deleteDocuments as deleteNoSqlDocuments } from "Common/dataAccess/deleteDocument";
 import { queryDocuments } from "Common/dataAccess/queryDocuments";
@@ -1184,16 +1183,6 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     documentsIterator, // loadNextPage: disabled as it will trigger a circular dependency and infinite loop
   ]);
 
-  const onRefreshKeyInput: KeyboardEventHandler<HTMLButtonElement> = (event) => {
-    if (event.key === " " || event.key === "Enter") {
-      const focusElement = event.target as HTMLElement;
-      refreshDocumentsGrid(false);
-      focusElement && focusElement.focus();
-      event.stopPropagation();
-      event.preventDefault();
-    }
-  };
-
   const onLoadMoreKeyInput: KeyboardEventHandler<HTMLAnchorElement> = (event) => {
     if (event.key === " " || event.key === "Enter") {
       const focusElement = event.target as HTMLElement;
@@ -1242,13 +1231,13 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
       .filter((key) => typeof (document as any)[key] === "string" || typeof (document as any)[key] === "number") // Only allow safe types for displayable React children
       .map((key) =>
         key === "id"
-          ? { id: key, label: isPreferredApiMongoDB ? "_id" : "id", group: undefined }
-          : { id: key, label: key, group: undefined },
+          ? { id: key, label: isPreferredApiMongoDB ? "_id" : "id", isPartitionKey: false }
+          : { id: key, label: key, isPartitionKey: false },
       );
 
     if (showPartitionKey(_collection, isPreferredApiMongoDB)) {
       columnDefinitions.push(
-        ...partitionKeyPropertyHeaders.map((key) => ({ id: key, label: key, group: "Partition Key" })),
+        ...partitionKeyPropertyHeaders.map((key) => ({ id: key, label: key, isPartitionKey: true })),
       );
 
       // Remove properties that are the partition keys, since they are already included
@@ -1908,23 +1897,6 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
               }}
               ref={tableContainerRef}
             >
-              <Button
-                appearance="transparent"
-                aria-label="Refresh"
-                size="small"
-                icon={<ArrowClockwise16Filled />}
-                style={{
-                  position: "absolute",
-                  top: 6,
-                  right: 0,
-                  float: "right",
-                  backgroundColor: "white",
-                  zIndex: 1,
-                  color: StyleConstants.AccentMedium,
-                }}
-                onClick={() => refreshDocumentsGrid(false)}
-                onKeyDown={onRefreshKeyInput}
-              />
               <div
                 style={
                   {
@@ -1934,6 +1906,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
                 }
               >
                 <DocumentsTableComponent
+                  onRefreshTable={() => refreshDocumentsGrid(false)}
                   items={tableItems}
                   onItemClicked={(index) => onDocumentClicked(index, documentIds)}
                   onSelectedRowsChange={onSelectedRowsChange}

src/Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent.tsx

@@ -1,22 +1,12 @@
 import {
   Button,
-  InputOnChangeData,
   Menu,
-  MenuCheckedValueChangeData,
-  MenuCheckedValueChangeEvent,
   MenuDivider,
-  MenuGroup,
-  MenuGroupHeader,
   MenuItem,
-  MenuItemCheckbox,
   MenuList,
   MenuPopover,
-  MenuProps,
   MenuTrigger,
-  PositioningImperativeRef,
   TableRowData as RowStateBase,
-  SearchBox,
-  SearchBoxChangeEvent,
   Table,
   TableBody,
   TableCell,
@@ -29,16 +19,23 @@ import {
   TableRowId,
   TableSelectionCell,
   useArrowNavigationGroup,
-  useRestoreFocusTarget,
   useTableColumnSizing_unstable,
   useTableFeatures,
   useTableSelection,
 } from "@fluentui/react-components";
-import { Add16Regular, Subtract12Regular } from "@fluentui/react-icons";
+import {
+  ArrowClockwise16Regular,
+  DeleteRegular,
+  EditRegular,
+  MoreHorizontalRegular,
+  TableResizeColumnRegular,
+} from "@fluentui/react-icons";
 import { NormalizedEventKey } from "Common/Constants";
+import { TableColumnSelectionPane } from "Explorer/Panes/TableColumnSelectionPane/TableColumnSelectionPane";
 import { selectionHelper } from "Explorer/Tabs/DocumentsTabV2/SelectionHelper";
 import { isEnvironmentCtrlPressed, isEnvironmentShiftPressed } from "Utils/KeyboardUtils";
-import React, { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import { useSidePanel } from "hooks/useSidePanel";
+import React, { useCallback, useMemo } from "react";
 import { FixedSizeList as List, ListChildComponentProps } from "react-window";
 
 export type DocumentsTableComponentItem = {
@@ -48,10 +45,11 @@ export type DocumentsTableComponentItem = {
 export type ColumnDefinition = {
   id: string;
   label: string;
+  isPartitionKey: boolean;
   defaultWidthPx?: number;
-  group: string | undefined;
 };
 export interface IDocumentsTableComponentProps {
+  onRefreshTable: () => void;
   items: DocumentsTableComponentItem[];
   onItemClicked: (index: number) => void;
   onSelectedRowsChange: (selectedItemsIndices: Set<TableRowId>) => void;
@@ -80,6 +78,7 @@ const MIN_COLUMN_WIDTH_PX = 20;
 const COLUMNS_MENU_NAME = "columnsMenu";
 
 export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> = ({
+  onRefreshTable,
   items,
   onSelectedRowsChange,
   selectedRows,
@@ -102,31 +101,9 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
   const [columnSizingOptions, setColumnSizingOptions] = React.useState<TableColumnSizingOptions>(initialSizingOptions);
 
   // This is for the menu to select columns
-  const [columnSearchText, setColumnSearchText] = useState<string>("");
-  const [isColumnSelectionMenuOpen, setIsColumnSelectionMenuOpen] = useState<boolean>(false);
-  const columnSelectionMenuButtonRef = useRef<HTMLButtonElement>(null);
-  const columnSelectionMenuPositionRef = useRef<HTMLDivElement>(null);
-  const positioningRef = React.useRef<PositioningImperativeRef>(null);
-  const onColumnSelectionMenuOpenChange: MenuProps["onOpenChange"] = (e, data) => {
-    // do not close menu as an outside click if clicking on the custom trigger/target
-    // this prevents it from closing & immediately re-opening when clicking custom triggers
-    if (
-      data.type === "clickOutside" &&
-      (e.target === columnSelectionMenuButtonRef.current || e.target === columnSelectionMenuPositionRef.current)
-    ) {
-      return;
-    }
+  // const [columnSearchText, setColumnSearchText] = useState<string>("");
 
-    setIsColumnSelectionMenuOpen(data.open);
-  };
-
-  useEffect(() => {
-    if (columnSelectionMenuPositionRef.current) {
-      positioningRef.current?.setTarget(columnSelectionMenuPositionRef.current);
-    }
-  }, [columnSelectionMenuPositionRef, positioningRef]);
-
-  const restoreFocusTargetAttribute = useRestoreFocusTarget();
+  // const restoreFocusTargetAttribute = useRestoreFocusTarget();
 
   const onColumnResize = React.useCallback(
     (_, { columnId, width }) => {
@@ -162,12 +139,36 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
           renderHeaderCell: () => (
             <>
               <span title={column.label}>{column.label}</span>
+              <Menu>
+                <MenuTrigger disableButtonEnhancement>
                   <Button
+                    // {...restoreFocusTargetAttribute}
                     appearance="transparent"
-                aria-label="De-select column"
+                    aria-label="Select column"
                     size="small"
-                icon={<Subtract12Regular />}
-                style={{ position: "absolute", right: -8 }}
+                    icon={<MoreHorizontalRegular />}
+                    style={{ position: "absolute", right: -6 }}
+                  />
+                </MenuTrigger>
+                <MenuPopover>
+                  <MenuList>
+                    <MenuItem key="refresh" icon={<ArrowClockwise16Regular />} onClick={onRefreshTable}>
+                      Refresh
+                    </MenuItem>
+                    <MenuItem key="editcolumns" icon={<EditRegular />} onClick={openColumnSelectionPane}>
+                      Edit columns
+                    </MenuItem>
+                    <MenuDivider />
+                    <MenuItem
+                      key="keyboardresize"
+                      icon={<TableResizeColumnRegular />}
+                      onClick={columnSizing.enableKeyboardMode(column.id)}
+                    >
+                      Resize with left/right arrow keys
+                    </MenuItem>
+                    <MenuItem
+                      key="remove"
+                      icon={<DeleteRegular />}
                       onClick={() => {
                         // Remove column id from selectedColumnIds
                         const index = selectedColumnIds.indexOf(column.id);
@@ -178,7 +179,12 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
                         newSelectedColumnIds.splice(index, 1);
                         onColumnSelectionChange(newSelectedColumnIds);
                       }}
-              />
+                    >
+                      Remove column
+                    </MenuItem>
+                  </MenuList>
+                </MenuPopover>
+              </Menu>
             </>
           ),
           renderCell: (item) => (
@@ -187,7 +193,7 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
             </TableCellLayout>
           ),
         })),
-    [columnDefinitions, selectedColumnIds],
+    [columnDefinitions, onColumnSelectionChange, selectedColumnIds],
   );
 
   const [selectionStartIndex, setSelectionStartIndex] = React.useState<number>(undefined);
@@ -337,58 +343,17 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
       selectedColumnIds.includes(columnDefinition.id) && checkedValues[COLUMNS_MENU_NAME].push(columnDefinition.id),
   );
 
-  const onCheckedValueChange = (_: MenuCheckedValueChangeEvent, data: MenuCheckedValueChangeData) => {
-    // eslint-disable-next-line react/prop-types
-    onColumnSelectionChange(data.checkedItems);
-  };
-
-  const onSearchChange: (event: SearchBoxChangeEvent, data: InputOnChangeData) => void = (_, data) =>
-    // eslint-disable-next-line react/prop-types
-    setColumnSearchText(data.value);
-
-  const getMenuList = (columnDefinitions: ColumnDefinition[]): JSX.Element => {
-    // Group by group. Unnamed group first
-    const unnamedGroup: ColumnDefinition[] = [];
-    const groupMap = new Map<string, ColumnDefinition[]>();
-    columnDefinitions.forEach((column) => {
-      if (column.group) {
-        if (!groupMap.has(column.group)) {
-          groupMap.set(column.group, []);
-        }
-        groupMap.get(column.group).push(column);
-      } else {
-        unnamedGroup.push(column);
-      }
-    });
-
-    const menuList: JSX.Element[] = [];
-    menuList.push(<SearchBox key="search" size="small" value={columnSearchText} onChange={onSearchChange} />);
-    if (unnamedGroup.length > 0) {
-      menuList.push(
-        ...unnamedGroup
-          .filter((def) => !columnSearchText || def.label.toLowerCase().includes(columnSearchText.toLowerCase()))
-          .map((column) => (
-            <MenuItemCheckbox key={column.id} name={COLUMNS_MENU_NAME} value={column.id}>
-              {column.label}
-            </MenuItemCheckbox>
-          )),
+  const openColumnSelectionPane = (): void => {
+    useSidePanel
+      .getState()
+      .openSidePanel(
+        "Save Query",
+        <TableColumnSelectionPane
+          selectedColumnIds={selectedColumnIds}
+          columnDefinitions={columnDefinitions}
+          onSelectionChange={onColumnSelectionChange}
+        />,
       );
-    }
-    groupMap.forEach((columns, group) => {
-      menuList.push(<MenuDivider key={`divider${group}`} />);
-      menuList.push(
-        <MenuGroup key={group}>
-          <MenuGroupHeader>{group}</MenuGroupHeader>
-          {...columns.map((column) => (
-            <MenuItemCheckbox key={column.id} name={COLUMNS_MENU_NAME} value={column.id}>
-              {column.label}
-            </MenuItemCheckbox>
-          ))}
-        </MenuGroup>,
-      );
-    });
-
-    return <>{menuList}</>;
   };
 
   return (
@@ -397,6 +362,7 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
         <TableRow style={{ width: size ? size.width - 15 : "100%" }}>
           {!isSelectionDisabled && (
             <TableSelectionCell
+              key="selectcell"
               checked={allRowsSelected ? true : someRowsSelected ? "mixed" : false}
               onClick={toggleAllRows}
               onKeyDown={toggleAllKeydown}
@@ -404,8 +370,6 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
             />
           )}
           {columns.map((column) => (
-            <Menu openOnContext key={column.columnId}>
-              <MenuTrigger>
             <TableHeaderCell
               className="documentsTableCell"
               key={column.columnId}
@@ -413,43 +377,7 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
             >
               {column.renderHeaderCell()}
             </TableHeaderCell>
-              </MenuTrigger>
-              <MenuPopover>
-                <MenuList>
-                  <MenuItem onClick={columnSizing.enableKeyboardMode(column.columnId)}>
-                    Enable Left/Right Arrow keys to resize
-                  </MenuItem>
-                </MenuList>
-              </MenuPopover>
-            </Menu>
           ))}
-          <Button
-            {...restoreFocusTargetAttribute}
-            appearance="transparent"
-            aria-label="Select column"
-            size="small"
-            icon={<Add16Regular />}
-            style={{ position: "absolute", right: 25 }}
-            onClick={() => setIsColumnSelectionMenuOpen((s) => !s)}
-          />
-          <div
-            {...restoreFocusTargetAttribute}
-            ref={columnSelectionMenuPositionRef}
-            style={{ height: 0, position: "absolute", right: 20, top: 0 }}
-          />
-          <Menu
-            checkedValues={checkedValues}
-            onCheckedValueChange={onCheckedValueChange}
-            open={isColumnSelectionMenuOpen}
-            onOpenChange={onColumnSelectionMenuOpenChange}
-            positioning={{ positioningRef }}
-          >
-            <MenuPopover>
-              <MenuList style={{ maxHeight: size?.height, overflowY: "auto", overflowX: "hidden" }}>
-                {getMenuList(columnDefinitions)}
-              </MenuList>
-            </MenuPopover>
-          </Menu>
         </TableRow>
       </TableHeader>
       <TableBody>