From 88f38d65223cd035a2f609f0521fb25c46d42290 Mon Sep 17 00:00:00 2001
From: Laurent Nguyen <laurent.nguyen@microsoft.com>
Date: Mon, 24 Jun 2024 13:38:32 +0200
Subject: [PATCH] Move table values under its own property

---
 .../Tabs/DocumentsTabV2/DocumentsTabV2.tsx    | 26 ++++++++++++-------
 .../DocumentsTableComponent.tsx               | 16 +++++++++---
 2 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
index 3d8cc2f9a..0badf9415 100644
--- a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
+++ b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
@@ -445,6 +445,9 @@ export interface IDocumentsTabComponentProps {
   isTabActive: boolean;
 }
 
+// Extend DocumentId to include fields displayed in the table
+type ExtendedDocumentId = DocumentId & { tableFields?: DocumentsTableComponentItem };
+
 // Export to expose to unit tests
 export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabComponentProps> = ({
   isPreferredApiMongoDB,
@@ -463,7 +466,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
   const [isFilterFocused, setIsFilterFocused] = useState<boolean>(false);
   const [appliedFilter, setAppliedFilter] = useState<string>("");
   const [filterContent, setFilterContent] = useState<string>("");
-  const [documentIds, setDocumentIds] = useState<DocumentId[]>([]);
+  const [documentIds, setDocumentIds] = useState<ExtendedDocumentId[]>([]);
   const [isExecuting, setIsExecuting] = useState<boolean>(false);
   const filterInput = useRef<HTMLInputElement>(null);
 
@@ -560,9 +563,12 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
 
   // new DocumentId() requires a DocumentTab which we mock with only the required properties
   const newDocumentId = useCallback(
-    (rawDocument: DataModels.DocumentId, partitionKeyProperties: string[], partitionKeyValue: string[]) => ({
-      ...rawDocument,
-      ...new DocumentId(
+    (
+      rawDocument: DataModels.DocumentId,
+      partitionKeyProperties: string[],
+      partitionKeyValue: string[],
+    ): ExtendedDocumentId => {
+      const extendedDocumentId = new DocumentId(
         {
           partitionKey,
           partitionKeyProperties,
@@ -572,8 +578,10 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
         },
         rawDocument,
         partitionKeyValue,
-      ),
-    }),
+      ) as ExtendedDocumentId;
+      extendedDocumentId.tableFields = { ...rawDocument };
+      return extendedDocumentId;
+    },
     [partitionKey],
   );
 
@@ -1209,9 +1217,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
 
   // Table config here
   const tableItems: DocumentsTableComponentItem[] = documentIds.map((documentId) => {
-    const item: Record<string, string> & { id: string } = {
-      id: documentId.id(),
-    };
+    const item: DocumentsTableComponentItem = documentId.tableFields || { id: documentId.id() };
 
     if (partitionKeyPropertyHeaders && documentId.stringPartitionKeyValues) {
       for (let i = 0; i < partitionKeyPropertyHeaders.length; i++) {
@@ -1219,7 +1225,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
       }
     }
 
-    return { ...documentId, ...item };
+    return item;
   });
 
   const extractColumnDefinitionsFromDocument = (document: unknown): ColumnDefinition[] => {
diff --git a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent.tsx b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent.tsx
index c60736ecd..1b840796e 100644
--- a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent.tsx
+++ b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent.tsx
@@ -43,7 +43,7 @@ import { FixedSizeList as List, ListChildComponentProps } from "react-window";
 
 export type DocumentsTableComponentItem = {
   id: string;
-} & Record<string, string>;
+} & Record<string, string | number>;
 
 export type ColumnDefinition = {
   id: string;
@@ -149,7 +149,16 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
         .filter((column) => selectedColumnIds.includes(column.id))
         .map((column) => ({
           columnId: column.id,
-          compare: (a, b) => a[column.id].localeCompare(b[column.id]),
+          compare: (a, b) => {
+            if (typeof a[column.id] === "string") {
+              return (a[column.id] as string).localeCompare(b[column.id] as string);
+            } else if (typeof a[column.id] === "number") {
+              return (a[column.id] as number) - (b[column.id] as number);
+            } else {
+              // Should not happen
+              return 0;
+            }
+          },
           renderHeaderCell: () => (
             <>
               <span title={column.label}>{column.label}</span>
@@ -173,7 +182,7 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
             </>
           ),
           renderCell: (item) => (
-            <TableCellLayout truncate title={item[column.id]}>
+            <TableCellLayout truncate title={`${item[column.id]}`}>
               {item[column.id]}
             </TableCellLayout>
           ),
@@ -334,6 +343,7 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
   };
 
   const onSearchChange: (event: SearchBoxChangeEvent, data: InputOnChangeData) => void = (_, data) =>
+    // eslint-disable-next-line react/prop-types
     setColumnSearchText(data.value);
 
   const getMenuList = (columnDefinitions: ColumnDefinition[]): JSX.Element => {