Add column selection from right-click

package.json

@@ -13,7 +13,7 @@
     "@babel/plugin-proposal-class-properties": "7.12.1",
     "@babel/plugin-proposal-decorators": "7.12.12",
     "@fluentui/react": "8.112.1",
-    "@fluentui/react-components": "9.34.0",
+    "@fluentui/react-components": "9.54.2",
     "@jupyterlab/services": "6.0.2",
     "@jupyterlab/terminal": "3.0.3",
     "@microsoft/applicationinsights-web": "2.6.1",
@@ -245,4 +245,4 @@
     "printWidth": 120,
     "endOfLine": "auto"
   }
-}
+}

src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx

@@ -18,6 +18,7 @@ import { EditorReact } from "Explorer/Controls/Editor/EditorReact";
 import Explorer from "Explorer/Explorer";
 import { useCommandBar } from "Explorer/Menus/CommandBar/CommandBarComponentAdapter";
 import { querySampleDocuments, readSampleDocument } from "Explorer/QueryCopilot/QueryCopilotUtilities";
+import { usePrevious } from "Explorer/Tabs/DocumentsTabV2/SelectionHelper";
 import {
   DocumentsTabPrefs,
   readDocumentsTabPrefs,
@@ -48,11 +49,11 @@ import * as DataModels from "../../../Contracts/DataModels";
 import * as ViewModels from "../../../Contracts/ViewModels";
 import * as TelemetryProcessor from "../../../Shared/Telemetry/TelemetryProcessor";
 import * as QueryUtils from "../../../Utils/QueryUtils";
-import { extractPartitionKeyValues } from "../../../Utils/QueryUtils";
+import { defaultQueryFields, extractPartitionKeyValues } from "../../../Utils/QueryUtils";
 import DocumentId from "../../Tree/DocumentId";
 import ObjectId from "../../Tree/ObjectId";
 import TabsBase from "../TabsBase";
-import { ColumnsDefinition, DocumentsTableComponent, DocumentsTableComponentItem } from "./DocumentsTableComponent";
+import { ColumnDefinition, DocumentsTableComponent, DocumentsTableComponentItem } from "./DocumentsTableComponent";
 
 export class DocumentsTabV2 extends TabsBase {
   public partitionKey: DataModels.PartitionKey;
@@ -415,12 +416,19 @@ export const buildQuery = (
   filter: string,
   partitionKeyProperties?: string[],
   partitionKey?: DataModels.PartitionKey,
+  additionalField?: string[],
 ): string => {
   if (isMongo) {
     return filter || "{}";
   }
 
-  return QueryUtils.buildDocumentsQuery(filter, partitionKeyProperties, partitionKey);
+  // Filter out fields starting with "/" (partition keys)
+  return QueryUtils.buildDocumentsQuery(
+    filter,
+    partitionKeyProperties,
+    partitionKey,
+    additionalField.filter((f) => !f.startsWith("/")),
+  );
 };
 
 // Export to expose to unit tests
@@ -542,10 +550,19 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     [partitionKeyPropertyHeaders],
   );
 
+  const [selectedColumnIds, setSelectedColumnIds] = useState<string[]>(() => {
+    const columnsIds = ["id"];
+    if (showPartitionKey(_collection, isPreferredApiMongoDB)) {
+      columnsIds.push(...partitionKeyPropertyHeaders);
+    }
+    return columnsIds;
+  });
+
   // new DocumentId() requires a DocumentTab which we mock with only the required properties
   const newDocumentId = useCallback(
-    (rawDocument: DataModels.DocumentId, partitionKeyProperties: string[], partitionKeyValue: string[]) =>
+    (rawDocument: DataModels.DocumentId, partitionKeyProperties: string[], partitionKeyValue: string[]) => ({
-      new DocumentId(
+      ...rawDocument,
+      ...new DocumentId(
         {
           partitionKey,
           partitionKeyProperties,
@@ -556,6 +573,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
         rawDocument,
         partitionKeyValue,
       ),
+    }),
     [partitionKey],
   );
 
@@ -975,7 +993,13 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     const _queryAbortController = new AbortController();
     setQueryAbortController(_queryAbortController);
     const filter: string = filterContent.trim();
-    const query: string = buildQuery(isPreferredApiMongoDB, filter, partitionKeyProperties, partitionKey);
+    const query: string = buildQuery(
+      isPreferredApiMongoDB,
+      filter,
+      partitionKeyProperties,
+      partitionKey,
+      selectedColumnIds,
+    );
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     const options: any = {};
     // TODO: Property 'enableCrossPartitionQuery' does not exist on type 'FeedOptions'.
@@ -998,6 +1022,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     resourceTokenPartitionKey,
     isQueryCopilotSampleContainer,
     _collection,
+    selectedColumnIds,
   ]);
 
   const onHideFilterClick = (): void => {
@@ -1030,6 +1055,8 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
       );
       setOnLoadStartKey(undefined);
     }
+
+    // Update column definitions
   };
 
   let loadNextPage = useCallback(
@@ -1194,9 +1221,35 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
       }
     }
 
-    return item;
+    return { ...documentId, ...item };
   });
 
+  const extractColumnDefinitionsFromDocument = (document: unknown): ColumnDefinition[] => {
+    let columnDefinitions: ColumnDefinition[] = Object.keys(document).map((key) =>
+      key === "id"
+        ? { id: key, label: isPreferredApiMongoDB ? "_id" : "id", group: undefined }
+        : { id: key, label: key, group: undefined },
+    );
+
+    if (showPartitionKey(_collection, isPreferredApiMongoDB)) {
+      columnDefinitions.push(
+        ...partitionKeyPropertyHeaders.map((key) => ({ id: key, label: key, group: "Partition Key" })),
+      );
+
+      // Remove properties that are the partition keys, since they are already included
+      columnDefinitions = columnDefinitions.filter(
+        (columnDefinition) => !partitionKeyProperties.includes(columnDefinition.id),
+      );
+    }
+
+    return columnDefinitions;
+  };
+
+  const setColumnDefinitionsFromDocument = (document: unknown): void => {
+    // TODO Add fields rather than replace
+    setColumnDefinitions(extractColumnDefinitionsFromDocument(document));
+  };
+
   /**
    * replicate logic of selectedDocument.click();
    * Document has been clicked on in table
@@ -1212,6 +1265,9 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     (_isQueryCopilotSampleContainer ? readSampleDocument(documentId) : readDocument(_collection, documentId)).then(
       (content) => {
         initDocumentEditor(documentId, content);
+
+        // Update columns
+        setColumnDefinitionsFromDocument(content);
       },
     );
 
@@ -1302,23 +1358,11 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     return () => resizeObserver.disconnect(); // clean up
   }, []);
 
-  const columnsDefinition: ColumnsDefinition = [
+  const [columnDefinitions, setColumnDefinitions] = useState<ColumnDefinition[]>(() =>
-    {
+    extractColumnDefinitionsFromDocument({
       id: "id",
-      label: isPreferredApiMongoDB ? "_id" : "id",
+    }),
-      defaultWidthPx: prefs.columnWidths ? prefs.columnWidths["id"] : undefined,
+  );
-    },
-  ];
-
-  if (showPartitionKey(_collection, isPreferredApiMongoDB)) {
-    partitionKeyPropertyHeaders.forEach((header) => {
-      columnsDefinition.push({
-        id: header,
-        label: header,
-        defaultWidthPx: prefs.columnWidths ? prefs.columnWidths[header] : undefined,
-      });
-    });
-  }
 
   const onSelectedRowsChange = (selectedRows: Set<TableRowId>) => {
     confirmDiscardingChange(() => {
@@ -1607,7 +1651,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
       setIsExecuting(true);
       onExecutionErrorChange(false);
       const filter: string = filterContent.trim();
-      const query: string = buildQuery(isPreferredApiMongoDB, filter);
+      const query: string = buildQuery(isPreferredApiMongoDB, filter, selectedColumnIds);
 
       return MongoProxyClient.queryDocuments(
         _collection.databaseId,
@@ -1697,6 +1741,31 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     setPrefs({ ...prefs });
   };
 
+  const onColumnSelectionChange = (newSelectedColumnIds: string[]): void => {
+    setSelectedColumnIds(newSelectedColumnIds);
+  };
+
+  const prevSelectedColumnIds = usePrevious({ selectedColumnIds, setSelectedColumnIds });
+
+  useEffect(() => {
+    // If we are adding a field, let's refresh to include the field in the query
+    let addedField = false;
+    for (const field of selectedColumnIds) {
+      if (
+        !defaultQueryFields.includes(field) &&
+        prevSelectedColumnIds &&
+        !prevSelectedColumnIds.selectedColumnIds.includes(field)
+      ) {
+        addedField = true;
+        break;
+      }
+    }
+
+    if (addedField) {
+      refreshDocumentsGrid(false);
+    }
+  }, [prevSelectedColumnIds, refreshDocumentsGrid, selectedColumnIds]);
+
   return (
     <FluentProvider theme={getPlatformTheme(configContext.platform)} style={{ height: "100%" }}>
       <div className="tab-pane active documentsTab" role="tabpanel" style={{ display: "flex" }}>
@@ -1840,11 +1909,13 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
                   onSelectedRowsChange={onSelectedRowsChange}
                   selectedRows={selectedRows}
                   size={tableContainerSizePx}
-                  columnsDefinition={columnsDefinition}
+                  selectedColumnIds={selectedColumnIds}
+                  columnDefinitions={columnDefinitions}
                   isSelectionDisabled={
                     configContext.platform === Platform.Fabric && userContext.fabricContext?.isReadOnly
                   }
                   onColumnResize={onTableColumnResize}
+                  onColumnSelectionChange={onColumnSelectionChange}
                 />
                 {tableItems.length > 0 && (
                   <a

src/Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent.tsx

@@ -1,6 +1,13 @@
+import { SearchBox } from "@fluentui/react";
 import {
   Menu,
+  MenuCheckedValueChangeData,
+  MenuCheckedValueChangeEvent,
+  MenuDivider,
+  MenuGroup,
+  MenuGroupHeader,
   MenuItem,
+  MenuItemCheckbox,
   MenuList,
   MenuPopover,
   MenuTrigger,
@@ -19,33 +26,36 @@ import {
   useArrowNavigationGroup,
   useTableColumnSizing_unstable,
   useTableFeatures,
-  useTableSelection,
+  useTableSelection
 } from "@fluentui/react-components";
 import { NormalizedEventKey } from "Common/Constants";
 import { selectionHelper } from "Explorer/Tabs/DocumentsTabV2/SelectionHelper";
 import { isEnvironmentCtrlPressed, isEnvironmentShiftPressed } from "Utils/KeyboardUtils";
-import React, { useCallback, useMemo } from "react";
+import React, { ChangeEvent, useCallback, useMemo } from "react";
 import { FixedSizeList as List, ListChildComponentProps } from "react-window";
 
 export type DocumentsTableComponentItem = {
   id: string;
 } & Record<string, string>;
 
-export type ColumnsDefinition = {
+export type ColumnDefinition = {
   id: string;
   label: string;
   defaultWidthPx?: number;
-}[];
+  group: string | undefined;
+};
 export interface IDocumentsTableComponentProps {
   items: DocumentsTableComponentItem[];
   onItemClicked: (index: number) => void;
   onSelectedRowsChange: (selectedItemsIndices: Set<TableRowId>) => void;
   selectedRows: Set<TableRowId>;
   size: { height: number; width: number };
-  columnsDefinition: ColumnsDefinition;
+  selectedColumnIds: string[];
+  columnDefinitions: ColumnDefinition[];
   style?: React.CSSProperties;
   isSelectionDisabled?: boolean;
   onColumnResize?: (columnId: string, width: number) => void;
+  onColumnSelectionChange?: (newSelectedColumnIds: string[]) => void;
 }
 
 interface TableRowData extends RowStateBase<DocumentsTableComponentItem> {
@@ -60,6 +70,7 @@ interface ReactWindowRenderFnProps extends ListChildComponentProps {
 
 const DEFAULT_COLUMN_WIDTH_PX = 200;
 const MIN_COLUMN_WIDTH_PX = 20;
+const COLUMNS_MENU_NAME = "columnsMenu";
 
 export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> = ({
   items,
@@ -67,12 +78,14 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
   selectedRows,
   style,
   size,
-  columnsDefinition,
+  selectedColumnIds,
+  columnDefinitions,
   isSelectionDisabled,
   onColumnResize: _onColumnResize,
+  onColumnSelectionChange,
 }: IDocumentsTableComponentProps) => {
   const initialSizingOptions: TableColumnSizingOptions = {};
-  columnsDefinition.forEach((column) => {
+  columnDefinitions.forEach((column) => {
     initialSizingOptions[column.id] = {
       idealWidth: column.defaultWidthPx || DEFAULT_COLUMN_WIDTH_PX, // 0 is not a valid width
       minWidth: MIN_COLUMN_WIDTH_PX,
@@ -80,6 +93,7 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
   });
 
   const [columnSizingOptions, setColumnSizingOptions] = React.useState<TableColumnSizingOptions>(initialSizingOptions);
+  const [columnSearchText, setColumnSearchText] = React.useState<string>("");
 
   const onColumnResize = React.useCallback(
     (_, { columnId, width }) => {
@@ -98,17 +112,19 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
   // Columns must be a static object and cannot change on re-renders otherwise React will complain about too many refreshes
   const columns: TableColumnDefinition<DocumentsTableComponentItem>[] = useMemo(
     () =>
-      columnsDefinition.map((column) => ({
+      columnDefinitions
-        columnId: column.id,
+        .filter((column) => selectedColumnIds.includes(column.id))
-        compare: (a, b) => a[column.id].localeCompare(b[column.id]),
+        .map((column) => ({
-        renderHeaderCell: () => <span title={column.label}>{column.label}</span>,
+          columnId: column.id,
-        renderCell: (item) => (
+          compare: (a, b) => a[column.id].localeCompare(b[column.id]),
-          <TableCellLayout truncate title={item[column.id]}>
+          renderHeaderCell: () => <span title={column.label}>{column.label}</span>,
-            {item[column.id]}
+          renderCell: (item) => (
-          </TableCellLayout>
+            <TableCellLayout truncate title={item[column.id]}>
-        ),
+              {item[column.id]}
-      })),
+            </TableCellLayout>
-    [columnsDefinition],
+          ),
+        })),
+    [columnDefinitions, selectedColumnIds],
   );
 
   const [selectionStartIndex, setSelectionStartIndex] = React.useState<number>(undefined);
@@ -250,6 +266,66 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
     ...style,
   };
 
+  const checkedValues: { [COLUMNS_MENU_NAME]: string[] } = {
+    [COLUMNS_MENU_NAME]: [],
+  };
+  columnDefinitions.forEach(
+    (columnDefinition) =>
+      selectedColumnIds.includes(columnDefinition.id) && checkedValues[COLUMNS_MENU_NAME].push(columnDefinition.id),
+  );
+
+  const onCheckedValueChange = (_: MenuCheckedValueChangeEvent, data: MenuCheckedValueChangeData) => {
+    // TODO this is expensive
+    // eslint-disable-next-line react/prop-types
+    onColumnSelectionChange(data.checkedItems);
+  };
+
+  const onSearchChange: (event?: ChangeEvent<HTMLInputElement>, newValue?: string) => void = (_, newValue) =>
+    setColumnSearchText(newValue);
+
+  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" value={columnSearchText} onChange={onSearchChange} />)
+    if (unnamedGroup.length > 0) {
+      menuList.push(
+        ...unnamedGroup.filter(def => !columnSearchText || def.label.startsWith(columnSearchText)).map((column) => (
+          <MenuItemCheckbox key={column.id} name={COLUMNS_MENU_NAME} value={column.id}>
+            {column.label}
+          </MenuItemCheckbox>
+        )),
+      );
+    }
+    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 (
     <Table className="documentsTable" noNativeElements {...tableProps}>
       <TableHeader className="documentsTableHeader">
@@ -263,7 +339,12 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
             />
           )}
           {columns.map((column /* index */) => (
-            <Menu openOnContext key={column.columnId}>
+            <Menu
+              openOnContext
+              key={column.columnId}
+              checkedValues={checkedValues}
+              onCheckedValueChange={onCheckedValueChange}
+            >
               <MenuTrigger>
                 <TableHeaderCell
                   className="documentsTableCell"
@@ -274,9 +355,11 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
                 </TableHeaderCell>
               </MenuTrigger>
               <MenuPopover>
-                <MenuList>
+                <MenuList style={{ maxHeight: size?.height, overflowY: "auto", overflowX: "hidden" }}>
+                  {getMenuList(columnDefinitions)}
+                  <MenuDivider />
                   <MenuItem onClick={columnSizing.enableKeyboardMode(column.columnId)}>
-                    Keyboard Column Resizing
+                    Use Left/Right Arrow keys to resize
                   </MenuItem>
                 </MenuList>
               </MenuPopover>

src/Explorer/Tabs/DocumentsTabV2/SelectionHelper.ts

@@ -1,3 +1,5 @@
+import { useEffect, useRef } from "react";
+
 /**
  * Utility class to help with selection.
  * This emulates File Explorer selection behavior.
@@ -90,3 +92,12 @@ export const selectionHelper = (
     }
   }
 };
+
+// To get previous values of a state in useEffect
+export const usePrevious = <T>(value: T): T | undefined => {
+  const ref = useRef<T>();
+  useEffect(() => {
+    ref.current = value;
+  });
+  return ref.current;
+};

src/Explorer/Tabs/DocumentsTabV2/documentsTabPrefs.ts

@@ -4,7 +4,7 @@ import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
 
 export interface DocumentsTabPrefs {
   leftPaneWidthPercent: number;
-  columnWidths?: { [columnId: string]: number };
+  columnWidths?: { [columnId: string]: number }; // TODO save per database/collection
 }
 
 const defaultPrefs: DocumentsTabPrefs = {

src/Utils/QueryUtils.ts

@@ -2,18 +2,28 @@ import { PartitionKey, PartitionKeyDefinition } from "@azure/cosmos";
 import * as DataModels from "../Contracts/DataModels";
 import * as ViewModels from "../Contracts/ViewModels";
 
+export const defaultQueryFields = ["id", "_self", "_rid", "_ts"];
+
 export function buildDocumentsQuery(
   filter: string,
   partitionKeyProperties: string[],
   partitionKey: DataModels.PartitionKey,
+  additionalField: string[] = [],
 ): string {
+  const fieldSet = new Set<string>(defaultQueryFields);
+  additionalField.forEach((prop) => fieldSet.add(prop));
+
+  const objectListSpec = [...fieldSet]
+    .filter((f) => !partitionKeyProperties.includes(f))
+    .map((prop) => `c.${prop}`)
+    .join(",");
   let query =
     partitionKeyProperties && partitionKeyProperties.length > 0
-      ? `select c.id, c._self, c._rid, c._ts, [${buildDocumentsQueryPartitionProjections(
+      ? `select ${objectListSpec}, [${buildDocumentsQueryPartitionProjections(
           "c",
           partitionKey,
         )}] as _partitionKeyValue from c`
-      : `select c.id, c._self, c._rid, c._ts from c`;
+      : `select ${objectListSpec} from c`;
 
   if (filter) {
     query += " " + filter;