From 4778183e5050e015b440b14a0b7185766e370fd9 Mon Sep 17 00:00:00 2001
From: Laurent Nguyen <laurent.nguyen@microsoft.com>
Date: Fri, 23 Aug 2024 09:23:38 +0200
Subject: [PATCH] Save columns definition (schema) along with selected columns.

---
 .../DocumentsTabV2/DocumentsTabStateUtil.ts   |  3 +-
 .../Tabs/DocumentsTabV2/DocumentsTabV2.tsx    | 41 +++++++++++++++----
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabStateUtil.ts b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabStateUtil.ts
index b580cff92..0189179de 100644
--- a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabStateUtil.ts
+++ b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabStateUtil.ts
@@ -1,5 +1,6 @@
 // Definitions of State data
 
+import { ColumnDefinition } from "Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent";
 import { deleteState, loadState, saveState, saveStateDebounced } from "Shared/AppStatePersistenceUtility";
 import { userContext } from "UserContext";
 import * as ViewModels from "../../../Contracts/ViewModels";
@@ -19,7 +20,7 @@ export type ColumnSizesMap = { [columnId: string]: WidthDefinition };
 export type FilterHistory = string[];
 export type WidthDefinition = { idealWidth?: number; minWidth?: number };
 export type TabDivider = { leftPaneWidthPercent: number };
-export type ColumnsSelection = string[];
+export type ColumnsSelection = { selectedColumnIds: string[]; columnDefinitions: ColumnDefinition[] };
 export type ColumnSort = { columnId: string; direction: "ascending" | "descending" };
 
 /**
diff --git a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
index 922cf0d1f..8cafbdfec 100644
--- a/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
+++ b/src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx
@@ -663,10 +663,22 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
       defaultColumnsIds.push(...partitionKeyPropertyHeaders);
     }
 
-    return readSubComponentState<ColumnsSelection>(SubComponentName.ColumnsSelection, _collection, defaultColumnsIds);
+    return defaultColumnsIds;
   };
 
-  const [selectedColumnIds, setSelectedColumnIds] = useState<string[]>(getInitialColumnSelection);
+  const [selectedColumnIds, setSelectedColumnIds] = useState<string[]>(() => {
+    const persistedColumnsSelection = readSubComponentState<ColumnsSelection>(
+      SubComponentName.ColumnsSelection,
+      _collection,
+      undefined,
+    );
+
+    if (!persistedColumnsSelection) {
+      return getInitialColumnSelection();
+    }
+
+    return persistedColumnsSelection.selectedColumnIds;
+  });
 
   // new DocumentId() requires a DocumentTab which we mock with only the required properties
   const newDocumentId = useCallback(
@@ -1515,11 +1527,21 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
   }, []);
 
   // Column definition is a map<id, ColumnDefinition> to garantee uniqueness
-  const [columnDefinitions, setColumnDefinitions] = useState<ColumnDefinition[]>(() =>
-    extractColumnDefinitionsFromDocument({
-      id: "id",
-    }),
-  );
+  const [columnDefinitions, setColumnDefinitions] = useState<ColumnDefinition[]>(() => {
+    const persistedColumnsSelection = readSubComponentState<ColumnsSelection>(
+      SubComponentName.ColumnsSelection,
+      _collection,
+      undefined,
+    );
+
+    if (!persistedColumnsSelection) {
+      return extractColumnDefinitionsFromDocument({
+        id: "id",
+      });
+    }
+
+    return persistedColumnsSelection.columnDefinitions;
+  });
 
   const onSelectedRowsChange = (selectedRows: Set<TableRowId>) => {
     confirmDiscardingChange(() => {
@@ -1858,7 +1880,10 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
 
     setSelectedColumnIds(newSelectedColumnIds);
 
-    saveSubComponentState<ColumnsSelection>(SubComponentName.ColumnsSelection, _collection, newSelectedColumnIds);
+    saveSubComponentState<ColumnsSelection>(SubComponentName.ColumnsSelection, _collection, {
+      selectedColumnIds: newSelectedColumnIds,
+      columnDefinitions,
+    });
   };
 
   const prevSelectedColumnIds = usePrevious({ selectedColumnIds, setSelectedColumnIds });