Add changes for ARM/AAD Token renewal

src/Common/CosmosClient.ts

@@ -8,11 +8,12 @@ import { AuthType } from "../AuthType";
 import { BackendApi, PriorityLevel } from "../Common/Constants";
 import * as Logger from "../Common/Logger";
 import { Platform, configContext } from "../ConfigContext";
-import { userContext } from "../UserContext";
+import { updateUserContext, userContext } from "../UserContext";
 import { logConsoleError } from "../Utils/NotificationConsoleUtils";
 import * as PriorityBasedExecutionUtils from "../Utils/PriorityBasedExecutionUtils";
 import { EmulatorMasterKey, HttpHeaders } from "./Constants";
 import { getErrorMessage } from "./ErrorHandlingUtils";
+import { runCommand } from "hooks/useDatabaseAccounts";
 
 const _global = typeof self === "undefined" ? window : self;
 
@@ -32,6 +33,7 @@ export const tokenProvider = async (requestInfo: Cosmos.RequestInfo) => {
       return null;
     }
     const AUTH_PREFIX = `type=aad&ver=1.0&sig=`;
+    console.log("AAD Token ", userContext.aadToken);
     const authorizationToken = `${AUTH_PREFIX}${userContext.aadToken}`;
     return authorizationToken;
   }
@@ -209,10 +211,14 @@ export function client(): Cosmos.CosmosClient {
     _defaultHeaders["x-ms-cosmos-priority-level"] = PriorityLevel.Default;
   }
 
+  const wrappedTokenProvider = async (requestInfo: Cosmos.RequestInfo) => {
+    return await runCommand(tokenProvider, requestInfo);
+  };
+
   const options: Cosmos.CosmosClientOptions = {
     endpoint: endpoint() || "https://cosmos.azure.com", // CosmosClient gets upset if we pass a bad URL. This should never actually get called
     key: userContext.dataPlaneRbacEnabled ? "" : userContext.masterKey,
-    tokenProvider,
+    tokenProvider: wrappedTokenProvider,
     userAgentSuffix: "Azure Portal",
     defaultHeaders: _defaultHeaders,
     connectionPolicy: {

src/Contracts/ActionContracts.ts

@@ -9,7 +9,6 @@ export enum TabKind {
   Graph,
   SQLQuery,
   ScaleSettings,
-  MongoQuery,
 }
 
 /**
@@ -52,8 +51,6 @@ export interface OpenCollectionTab extends OpenTab {
  */
 export interface OpenQueryTab extends OpenCollectionTab {
   query: QueryInfo;
-  splitterDirection?: "vertical" | "horizontal";
-  queryViewSizePercent?: number;
 }
 
 /**

src/Contracts/ViewModels.ts

@@ -115,13 +115,7 @@ export interface CollectionBase extends TreeNode {
   isSampleCollection?: boolean;
 
   onDocumentDBDocumentsClick(): void;
-  onNewQueryClick(
-    source: any,
-    event?: MouseEvent,
-    queryText?: string,
-    splitterDirection?: "horizontal" | "vertical",
-    queryViewSizePercent?: number,
-  ): void;
+  onNewQueryClick(source: any, event?: MouseEvent, queryText?: string): void;
   expandCollection(): void;
   collapseCollection(): void;
   getDatabase(): Database;
@@ -157,13 +151,7 @@ export interface Collection extends CollectionBase {
   onSettingsClick: () => Promise<void>;
 
   onNewGraphClick(): void;
-  onNewMongoQueryClick(
-    source: any,
-    event?: MouseEvent,
-    queryText?: string,
-    splitterDirection?: "horizontal" | "vertical",
-    queryViewSizePercent?: number,
-  ): void;
+  onNewMongoQueryClick(source: any, event?: MouseEvent, queryText?: string): void;
   onNewMongoShellClick(): void;
   onNewStoredProcedureClick(source: Collection, event?: MouseEvent): void;
   onNewUserDefinedFunctionClick(source: Collection, event?: MouseEvent): void;
@@ -323,8 +311,6 @@ export interface QueryTabOptions extends TabOptions {
   partitionKey?: DataModels.PartitionKey;
   queryText?: string;
   resourceTokenPartitionKey?: string;
-  splitterDirection?: "horizontal" | "vertical";
-  queryViewSizePercent?: number;
 }
 
 export interface ScriptTabOption extends TabOptions {

src/Explorer/Controls/Settings/SettingsSubComponents/PartitionKeyComponent.tsx

@@ -14,7 +14,6 @@ import * as ViewModels from "../../../../Contracts/ViewModels";
 
 import { handleError } from "Common/ErrorHandlingUtils";
 import { cancelDataTransferJob, pollDataTransferJob } from "Common/dataAccess/dataTransfers";
-import { Platform, configContext } from "ConfigContext";
 import Explorer from "Explorer/Explorer";
 import { ChangePartitionKeyPane } from "Explorer/Panes/ChangePartitionKeyPane/ChangePartitionKeyPane";
 import {
@@ -149,71 +148,69 @@ export const PartitionKeyComponent: React.FC<PartitionKeyComponentProps> = ({ da
   };
 
   return (
-    configContext.platform !== Platform.Emulator && (
-      <Stack tokens={{ childrenGap: 20 }} styles={{ root: { maxWidth: 600 } }}>
-        <Stack tokens={{ childrenGap: 10 }}>
-          <Text styles={textHeadingStyle}>Change {partitionKeyName.toLowerCase()}</Text>
-          <Stack horizontal tokens={{ childrenGap: 20 }}>
-            <Stack tokens={{ childrenGap: 5 }}>
-              <Text styles={textSubHeadingStyle}>Current {partitionKeyName.toLowerCase()}</Text>
-              <Text styles={textSubHeadingStyle}>Partitioning</Text>
-            </Stack>
-            <Stack tokens={{ childrenGap: 5 }}>
-              <Text>{partitionKeyValue}</Text>
-              <Text>{isHierarchicalPartitionedContainer() ? "Hierarchical" : "Non-hierarchical"}</Text>
-            </Stack>
+    <Stack tokens={{ childrenGap: 20 }} styles={{ root: { maxWidth: 600 } }}>
+      <Stack tokens={{ childrenGap: 10 }}>
+        <Text styles={textHeadingStyle}>Change {partitionKeyName.toLowerCase()}</Text>
+        <Stack horizontal tokens={{ childrenGap: 20 }}>
+          <Stack tokens={{ childrenGap: 5 }}>
+            <Text styles={textSubHeadingStyle}>Current {partitionKeyName.toLowerCase()}</Text>
+            <Text styles={textSubHeadingStyle}>Partitioning</Text>
+          </Stack>
+          <Stack tokens={{ childrenGap: 5 }}>
+            <Text>{partitionKeyValue}</Text>
+            <Text>{isHierarchicalPartitionedContainer() ? "Hierarchical" : "Non-hierarchical"}</Text>
           </Stack>
         </Stack>
-        <MessageBar messageBarType={MessageBarType.warning}>
-          To safeguard the integrity of the data being copied to the new container, ensure that no updates are made to
-          the source container for the entire duration of the partition key change process.
-          <Link
-            href="https://learn.microsoft.com/azure/cosmos-db/container-copy#how-does-container-copy-work"
-            target="_blank"
-            underline
+      </Stack>
+      <MessageBar messageBarType={MessageBarType.warning}>
+        To safeguard the integrity of the data being copied to the new container, ensure that no updates are made to the
+        source container for the entire duration of the partition key change process.
+        <Link
+          href="https://learn.microsoft.com/azure/cosmos-db/container-copy#how-does-container-copy-work"
+          target="_blank"
+          underline
+        >
+          Learn more
+        </Link>
+      </MessageBar>
+      <Text>
+        To change the partition key, a new destination container must be created or an existing destination container
+        selected. Data will then be copied to the destination container.
+      </Text>
+      <PrimaryButton
+        styles={{ root: { width: "fit-content" } }}
+        text="Change"
+        onClick={startPartitionkeyChangeWorkflow}
+        disabled={isCurrentJobInProgress(portalDataTransferJob)}
+      />
+      {portalDataTransferJob && (
+        <Stack>
+          <Text styles={textHeadingStyle}>{partitionKeyName} change job</Text>
+          <Stack
+            horizontal
+            tokens={{ childrenGap: 20 }}
+            styles={{
+              root: {
+                alignItems: "center",
+              },
+            }}
           >
-            Learn more
-          </Link>
-        </MessageBar>
-        <Text>
-          To change the partition key, a new destination container must be created or an existing destination container
-          selected. Data will then be copied to the destination container.
-        </Text>
-        <PrimaryButton
-          styles={{ root: { width: "fit-content" } }}
-          text="Change"
-          onClick={startPartitionkeyChangeWorkflow}
-          disabled={isCurrentJobInProgress(portalDataTransferJob)}
-        />
-        {portalDataTransferJob && (
-          <Stack>
-            <Text styles={textHeadingStyle}>{partitionKeyName} change job</Text>
-            <Stack
-              horizontal
-              tokens={{ childrenGap: 20 }}
+            <ProgressIndicator
+              label={portalDataTransferJob?.properties?.jobName}
+              description={getProgressDescription()}
+              percentComplete={getPercentageComplete()}
               styles={{
                 root: {
-                  alignItems: "center",
+                  width: "85%",
                 },
               }}
-            >
-              <ProgressIndicator
-                label={portalDataTransferJob?.properties?.jobName}
-                description={getProgressDescription()}
-                percentComplete={getPercentageComplete()}
-                styles={{
-                  root: {
-                    width: "85%",
-                  },
-                }}
-              ></ProgressIndicator>
-              {isCurrentJobInProgress(portalDataTransferJob) && (
-                <DefaultButton text="Cancel" onClick={() => cancelRunningDataTransferJob(portalDataTransferJob)} />
-              )}
-            </Stack>
+            ></ProgressIndicator>
+            {isCurrentJobInProgress(portalDataTransferJob) && (
+              <DefaultButton text="Cancel" onClick={() => cancelRunningDataTransferJob(portalDataTransferJob)} />
+            )}
           </Stack>
-        )}
-      </Stack>
-    )
+        </Stack>
+      )}
+    </Stack>
   );
 };

src/Explorer/Explorer.tsx

@@ -1134,7 +1134,7 @@ export default class Explorer {
     if (userContext.apiType !== "Postgres" && userContext.apiType !== "VCoreMongo") {
       userContext.authType === AuthType.ResourceToken
         ? this.refreshDatabaseForResourceToken()
-        : await this.refreshAllDatabases(); // await: we rely on the databases to be loaded before restoring the tabs further in the flow
+        : this.refreshAllDatabases();
     }
     await useNotebook.getState().refreshNotebooksEnabledStateForAccount();
 

src/Explorer/OpenActions/OpenActions.tsx

@@ -1,5 +1,4 @@
 // TODO convert this file to an action registry in order to have actions and their handlers be more tightly coupled.
-import { configContext, Platform } from "ConfigContext";
 import { useDatabases } from "Explorer/useDatabases";
 import React from "react";
 import { ActionContracts } from "../../Contracts/ExplorerContracts";
@@ -57,19 +56,6 @@ function openCollectionTab(
         continue;
       }
 
-      if (
-        configContext.platform === Platform.Fabric &&
-        !(
-          // whitelist the tab kinds that are allowed to be opened in Fabric
-          (
-            action.tabKind === ActionContracts.TabKind.SQLDocuments ||
-            action.tabKind === ActionContracts.TabKind.SQLQuery
-          )
-        )
-      ) {
-        continue;
-      }
-
       //expand database first if not expanded to load the collections
       if (!database.isDatabaseExpanded?.()) {
         database.expandDatabase?.();
@@ -135,28 +121,10 @@ function openCollectionTab(
             action.tabKind === ActionContracts.TabKind.SQLQuery ||
             action.tabKind === ActionContracts.TabKind[ActionContracts.TabKind.SQLQuery]
           ) {
-            const openQueryTabAction = action as ActionContracts.OpenQueryTab;
             collection.onNewQueryClick(
               collection,
               undefined,
-              generateQueryText(openQueryTabAction, collection.partitionKeyProperties),
-              openQueryTabAction.splitterDirection,
-              openQueryTabAction.queryViewSizePercent,
-            );
-            break;
-          }
-
-          if (
-            action.tabKind === ActionContracts.TabKind.MongoQuery ||
-            action.tabKind === ActionContracts.TabKind[ActionContracts.TabKind.MongoQuery]
-          ) {
-            const openQueryTabAction = action as ActionContracts.OpenQueryTab;
-            collection.onNewMongoQueryClick(
-              collection,
-              undefined,
-              generateQueryText(openQueryTabAction, collection.partitionKeyProperties),
-              openQueryTabAction.splitterDirection,
-              openQueryTabAction.queryViewSizePercent,
+              generateQueryText(action as ActionContracts.OpenQueryTab, collection.partitionKeyProperties),
             );
             break;
           }

src/Explorer/Tabs/DocumentsTabV2/DocumentsTabStateUtil.ts

@@ -3,11 +3,17 @@
 import { ColumnDefinition } from "Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent";
 import {
   AppStateComponentNames,
-  deleteSubComponentState,
-  readSubComponentState,
-  saveSubComponentState,
+  deleteState,
+  loadState,
+  saveState,
+  saveStateDebounced,
 } from "Shared/AppStatePersistenceUtility";
+import { userContext } from "UserContext";
 import * as ViewModels from "../../../Contracts/ViewModels";
+import { Action } from "../../../Shared/Telemetry/TelemetryConstants";
+import * as TelemetryProcessor from "../../../Shared/Telemetry/TelemetryProcessor";
+
+const componentName = AppStateComponentNames.DocumentsTab;
 
 export enum SubComponentName {
   ColumnSizes = "ColumnSizes",
@@ -15,7 +21,6 @@ export enum SubComponentName {
   MainTabDivider = "MainTabDivider",
   ColumnsSelection = "ColumnsSelection",
   ColumnSort = "ColumnSort",
-  CurrentFilter = "CurrentFilter",
 }
 
 export type ColumnSizesMap = { [columnId: string]: WidthDefinition };
@@ -25,22 +30,84 @@ export type TabDivider = { leftPaneWidthPercent: number };
 export type ColumnsSelection = { selectedColumnIds: string[]; columnDefinitions: ColumnDefinition[] };
 export type ColumnSort = { columnId: string; direction: "ascending" | "descending" };
 
-// Wrap the ...SubComponentState functions for type safety
-
-export const readDocumentsTabSubComponentState = <T>(
+/**
+ *
+ * @param subComponentName
+ * @param collection
+ * @param defaultValue Will be returned if persisted state is not found
+ * @returns
+ */
+export const readSubComponentState = <T>(
   subComponentName: SubComponentName,
   collection: ViewModels.CollectionBase,
   defaultValue: T,
-): T => readSubComponentState<T>(AppStateComponentNames.DocumentsTab, subComponentName, collection, defaultValue);
+): T => {
+  const globalAccountName = userContext.databaseAccount?.name;
+  if (!globalAccountName) {
+    const message = "Database account name not found in userContext";
+    console.error(message);
+    TelemetryProcessor.traceFailure(Action.ReadPersistedTabState, { message, componentName });
+    return defaultValue;
+  }
 
-export const saveDocumentsTabSubComponentState = <T>(
+  const state = loadState({
+    componentName: componentName,
+    subComponentName,
+    globalAccountName,
+    databaseName: collection.databaseId,
+    containerName: collection.id(),
+  }) as T;
+
+  return state || defaultValue;
+};
+
+/**
+ *
+ * @param subComponentName
+ * @param collection
+ * @param state State to save
+ * @param debounce true for high-frequency calls (e.g mouse drag events)
+ */
+export const saveSubComponentState = <T>(
   subComponentName: SubComponentName,
   collection: ViewModels.CollectionBase,
   state: T,
   debounce?: boolean,
-): void => saveSubComponentState<T>(AppStateComponentNames.DocumentsTab, subComponentName, collection, state, debounce);
+): void => {
+  const globalAccountName = userContext.databaseAccount?.name;
+  if (!globalAccountName) {
+    const message = "Database account name not found in userContext";
+    console.error(message);
+    TelemetryProcessor.traceFailure(Action.SavePersistedTabState, { message, componentName });
+    return;
+  }
 
-export const deleteDocumentsTabSubComponentState = (
-  subComponentName: SubComponentName,
-  collection: ViewModels.CollectionBase,
-) => deleteSubComponentState(AppStateComponentNames.DocumentsTab, subComponentName, collection);
+  (debounce ? saveStateDebounced : saveState)(
+    {
+      componentName: componentName,
+      subComponentName,
+      globalAccountName,
+      databaseName: collection.databaseId,
+      containerName: collection.id(),
+    },
+    state,
+  );
+};
+
+export const deleteSubComponentState = (subComponentName: SubComponentName, collection: ViewModels.CollectionBase) => {
+  const globalAccountName = userContext.databaseAccount?.name;
+  if (!globalAccountName) {
+    const message = "Database account name not found in userContext";
+    console.error(message);
+    TelemetryProcessor.traceFailure(Action.DeletePersistedTabState, { message, componentName });
+    return;
+  }
+
+  deleteState({
+    componentName: componentName,
+    subComponentName,
+    globalAccountName,
+    databaseName: collection.databaseId,
+    containerName: collection.id(),
+  });
+};

src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx

@@ -21,7 +21,6 @@ import { queryDocuments } from "Common/dataAccess/queryDocuments";
 import { readDocument } from "Common/dataAccess/readDocument";
 import { updateDocument } from "Common/dataAccess/updateDocument";
 import { Platform, configContext } from "ConfigContext";
-import { ActionType, OpenCollectionTab, TabKind } from "Contracts/ActionContracts";
 import { CommandButtonComponentProps } from "Explorer/Controls/CommandButton/CommandButtonComponent";
 import { useDialog } from "Explorer/Controls/Dialog";
 import { EditorReact } from "Explorer/Controls/Editor/EditorReact";
@@ -35,9 +34,8 @@ import {
   FilterHistory,
   SubComponentName,
   TabDivider,
-  deleteDocumentsTabSubComponentState,
-  readDocumentsTabSubComponentState,
-  saveDocumentsTabSubComponentState,
+  readSubComponentState,
+  saveSubComponentState,
 } from "Explorer/Tabs/DocumentsTabV2/DocumentsTabStateUtil";
 import { usePrevious } from "Explorer/Tabs/DocumentsTabV2/SelectionHelper";
 import { CosmosFluentProvider, LayoutConstants, cosmosShorthands, tokens } from "Explorer/Theme/ThemeUtil";
@@ -142,8 +140,6 @@ export class DocumentsTabV2 extends TabsBase {
   private title: string;
   private resourceTokenPartitionKey: string;
 
-  protected persistedState: OpenCollectionTab;
-
   constructor(options: ViewModels.DocumentsTabOptions) {
     super(options);
 
@@ -151,13 +147,6 @@ export class DocumentsTabV2 extends TabsBase {
     this.title = options.title;
     this.partitionKey = options.partitionKey;
     this.resourceTokenPartitionKey = options.resourceTokenPartitionKey;
-
-    this.persistedState = {
-      actionType: ActionType.OpenCollectionTab,
-      tabKind: options.isPreferredApiMongoDB ? TabKind.MongoDocuments : TabKind.SQLDocuments,
-      databaseResourceId: options.collection.databaseId,
-      collectionResourceId: options.collection.id(),
-    };
   }
 
   public render(): JSX.Element {
@@ -586,10 +575,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
   onIsExecutingChange,
   isTabActive,
 }): JSX.Element => {
-  const [filterContent, setFilterContent] = useState<string>(() =>
-    readDocumentsTabSubComponentState<string>(SubComponentName.CurrentFilter, _collection, ""),
-  );
-
+  const [filterContent, setFilterContent] = useState<string>("");
   const [documentIds, setDocumentIds] = useState<ExtendedDocumentId[]>([]);
   const [isExecuting, setIsExecuting] = useState<boolean>(false);
   const styles = useDocumentsTabStyles();
@@ -620,7 +606,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
 
   // State
   const [tabStateData, setTabStateData] = useState<TabDivider>(() =>
-    readDocumentsTabSubComponentState<TabDivider>(SubComponentName.MainTabDivider, _collection, {
+    readSubComponentState<TabDivider>(SubComponentName.MainTabDivider, _collection, {
       leftPaneWidthPercent: 35,
     }),
   );
@@ -635,7 +621,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
 
   // User's filter history
   const [lastFilterContents, setLastFilterContents] = useState<FilterHistory>(() =>
-    readDocumentsTabSubComponentState<FilterHistory>(SubComponentName.FilterHistory, _collection, [] as FilterHistory),
+    readSubComponentState<FilterHistory>(SubComponentName.FilterHistory, _collection, [] as FilterHistory),
   );
 
   // For progress bar for bulk delete (noSql)
@@ -777,7 +763,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
   };
 
   const [selectedColumnIds, setSelectedColumnIds] = useState<string[]>(() => {
-    const persistedColumnsSelection = readDocumentsTabSubComponentState<ColumnsSelection>(
+    const persistedColumnsSelection = readSubComponentState<ColumnsSelection>(
       SubComponentName.ColumnsSelection,
       _collection,
       undefined,
@@ -822,7 +808,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
   useEffect(() => {
     setKeyboardActions({
       [KeyboardAction.CLEAR_SEARCH]: () => {
-        updateFilterContent("");
+        setFilterContent("");
         refreshDocumentsGrid(true);
         return true;
       },
@@ -1659,7 +1645,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
 
   // Column definition is a map<id, ColumnDefinition> to garantee uniqueness
   const [columnDefinitions, setColumnDefinitions] = useState<ColumnDefinition[]>(() => {
-    const persistedColumnsSelection = readDocumentsTabSubComponentState<ColumnsSelection>(
+    const persistedColumnsSelection = readSubComponentState<ColumnsSelection>(
       SubComponentName.ColumnsSelection,
       _collection,
       undefined,
@@ -1970,7 +1956,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     const limitedLastFilterContents = lastFilterContents.slice(0, MAX_FILTER_HISTORY_COUNT);
 
     setLastFilterContents(limitedLastFilterContents);
-    saveDocumentsTabSubComponentState<FilterHistory>(SubComponentName.FilterHistory, _collection, lastFilterContents);
+    saveSubComponentState<FilterHistory>(SubComponentName.FilterHistory, _collection, lastFilterContents);
   };
 
   const refreshDocumentsGrid = useCallback(
@@ -2027,7 +2013,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
 
     setSelectedColumnIds(newSelectedColumnIds);
 
-    saveDocumentsTabSubComponentState<ColumnsSelection>(SubComponentName.ColumnsSelection, _collection, {
+    saveSubComponentState<ColumnsSelection>(SubComponentName.ColumnsSelection, _collection, {
       selectedColumnIds: newSelectedColumnIds,
       columnDefinitions,
     });
@@ -2077,15 +2063,6 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     return options;
   };
 
-  const updateFilterContent = (filter: string): void => {
-    if (filter === "" || filter === undefined) {
-      deleteDocumentsTabSubComponentState(SubComponentName.CurrentFilter, _collection);
-    } else {
-      saveDocumentsTabSubComponentState<string>(SubComponentName.CurrentFilter, _collection, filter, true);
-    }
-    setFilterContent(filter);
-  };
-
   return (
     <CosmosFluentProvider className={styles.container}>
       <div className="tab-pane active" role="tabpanel" style={{ display: "flex" }}>
@@ -2100,7 +2077,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
             }
             title="Type a query predicate or choose one from the list."
             value={filterContent}
-            onChange={updateFilterContent}
+            onChange={(value) => setFilterContent(value)}
             onKeyDown={onFilterKeyDown}
             bottomLink={{ text: "Learn more", url: DATA_EXPLORER_DOC_URL }}
           />
@@ -2126,7 +2103,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
         <Allotment
           onDragEnd={(sizes: number[]) => {
             tabStateData.leftPaneWidthPercent = (100 * sizes[0]) / (sizes[0] + sizes[1]);
-            saveDocumentsTabSubComponentState<TabDivider>(SubComponentName.MainTabDivider, _collection, tabStateData);
+            saveSubComponentState<TabDivider>(SubComponentName.MainTabDivider, _collection, tabStateData);
             setTabStateData(tabStateData);
           }}
         >

src/Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent.tsx

@@ -42,9 +42,9 @@ import { TableColumnSelectionPane } from "Explorer/Panes/TableColumnSelectionPan
 import {
   ColumnSizesMap,
   ColumnSort,
-  deleteDocumentsTabSubComponentState,
-  readDocumentsTabSubComponentState,
-  saveDocumentsTabSubComponentState,
+  deleteSubComponentState,
+  readSubComponentState,
+  saveSubComponentState,
   SubComponentName,
 } from "Explorer/Tabs/DocumentsTabV2/DocumentsTabStateUtil";
 import { INITIAL_SELECTED_ROW_INDEX, useDocumentsTabStyles } from "Explorer/Tabs/DocumentsTabV2/DocumentsTabV2";
@@ -118,11 +118,7 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
   const sortedRowsRef = React.useRef(null);
 
   const [columnSizingOptions, setColumnSizingOptions] = React.useState<TableColumnSizingOptions>(() => {
-    const columnSizesMap: ColumnSizesMap = readDocumentsTabSubComponentState(
-      SubComponentName.ColumnSizes,
-      collection,
-      {},
-    );
+    const columnSizesMap: ColumnSizesMap = readSubComponentState(SubComponentName.ColumnSizes, collection, {});
     const columnSizesPx: TableColumnSizingOptions = {};
     selectedColumnIds.forEach((columnId) => {
       if (
@@ -146,7 +142,7 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
     sortDirection: "ascending" | "descending";
     sortColumn: TableColumnId | undefined;
   }>(() => {
-    const sort = readDocumentsTabSubComponentState<ColumnSort>(SubComponentName.ColumnSort, collection, undefined);
+    const sort = readSubComponentState<ColumnSort>(SubComponentName.ColumnSort, collection, undefined);
 
     if (!sort) {
       return {
@@ -178,12 +174,7 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
         return acc;
       }, {} as ColumnSizesMap);
 
-      saveDocumentsTabSubComponentState<ColumnSizesMap>(
-        SubComponentName.ColumnSizes,
-        collection,
-        persistentSizes,
-        true,
-      );
+      saveSubComponentState<ColumnSizesMap>(SubComponentName.ColumnSizes, collection, persistentSizes, true);
 
       return newSizingOptions;
     });
@@ -195,14 +186,11 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
     setColumnSort(event, columnId, direction);
 
     if (columnId === undefined || direction === undefined) {
-      deleteDocumentsTabSubComponentState(SubComponentName.ColumnSort, collection);
+      deleteSubComponentState(SubComponentName.ColumnSort, collection);
       return;
     }
 
-    saveDocumentsTabSubComponentState<ColumnSort>(SubComponentName.ColumnSort, collection, {
-      columnId,
-      direction,
-    });
+    saveSubComponentState<ColumnSort>(SubComponentName.ColumnSort, collection, { columnId, direction });
   };
 
   // Columns must be a static object and cannot change on re-renders otherwise React will complain about too many refreshes

src/Explorer/Tabs/MongoQueryTab/MongoQueryTab.tsx

@@ -1,4 +1,3 @@
-import { ActionType, TabKind } from "Contracts/ActionContracts";
 import React from "react";
 import MongoUtility from "../../../Common/MongoUtility";
 import * as ViewModels from "../../../Contracts/ViewModels";
@@ -21,7 +20,7 @@ export class NewMongoQueryTab extends NewQueryTab {
     private mongoQueryTabProps: IMongoQueryTabProps,
   ) {
     super(options, mongoQueryTabProps);
-    this.queryText = options.queryText ?? "";
+    this.queryText = "";
     this.iMongoQueryTabComponentProps = {
       collection: options.collection,
       isExecutionError: this.isExecutionError(),
@@ -29,8 +28,6 @@ export class NewMongoQueryTab extends NewQueryTab {
       tabsBaseInstance: this,
       queryText: this.queryText,
       partitionKey: this.partitionKey,
-      splitterDirection: options.splitterDirection,
-      queryViewSizePercent: options.queryViewSizePercent,
       container: this.mongoQueryTabProps.container,
       onTabAccessor: (instance: ITabAccessor): void => {
         this.iTabAccessor = instance;
@@ -38,26 +35,6 @@ export class NewMongoQueryTab extends NewQueryTab {
       isPreferredApiMongoDB: true,
       monacoEditorSetting: "plaintext",
       viewModelcollection: this.mongoQueryTabProps.viewModelcollection,
-      onUpdatePersistedState: (state: {
-        queryText: string;
-        splitterDirection: string;
-        queryViewSizePercent: number;
-      }): void => {
-        this.persistedState = {
-          actionType: ActionType.OpenCollectionTab,
-          tabKind: TabKind.SQLQuery,
-          databaseResourceId: options.collection.databaseId,
-          collectionResourceId: options.collection.id(),
-          query: {
-            text: state.queryText,
-          },
-          splitterDirection: state.splitterDirection as "vertical" | "horizontal",
-          queryViewSizePercent: state.queryViewSizePercent,
-        };
-        if (this.triggerPersistState) {
-          this.triggerPersistState();
-        }
-      },
     };
   }
 

src/Explorer/Tabs/QueryTab/QueryTab.tsx

@@ -1,5 +1,4 @@
 import { sendMessage } from "Common/MessageHandler";
-import { ActionType, OpenQueryTab, TabKind } from "Contracts/ActionContracts";
 import { MessageTypes } from "Contracts/MessageTypes";
 import { CopilotProvider } from "Explorer/QueryCopilot/QueryCopilotContext";
 import { userContext } from "UserContext";
@@ -27,8 +26,6 @@ export class NewQueryTab extends TabsBase {
   public iQueryTabComponentProps: IQueryTabComponentProps;
   public iTabAccessor: ITabAccessor;
 
-  protected persistedState: OpenQueryTab;
-
   constructor(
     options: QueryTabOptions,
     private props: IQueryTabProps,
@@ -42,41 +39,12 @@ export class NewQueryTab extends TabsBase {
       tabsBaseInstance: this,
       queryText: options.queryText,
       partitionKey: this.partitionKey,
-      splitterDirection: options.splitterDirection,
-      queryViewSizePercent: options.queryViewSizePercent,
       container: this.props.container,
       onTabAccessor: (instance: ITabAccessor): void => {
         this.iTabAccessor = instance;
       },
       isPreferredApiMongoDB: false,
-      onUpdatePersistedState: (state: {
-        queryText: string;
-        splitterDirection: string;
-        queryViewSizePercent: number;
-      }): void => {
-        this.persistedState = {
-          actionType: ActionType.OpenCollectionTab,
-          tabKind: TabKind.SQLQuery,
-          databaseResourceId: options.collection.databaseId,
-          collectionResourceId: options.collection.id(),
-          query: {
-            text: state.queryText,
-          },
-          splitterDirection: state.splitterDirection as "vertical" | "horizontal",
-          queryViewSizePercent: state.queryViewSizePercent,
-        };
-        if (this.triggerPersistState) {
-          this.triggerPersistState();
-        }
-      },
     };
-
-    // set initial state
-    this.iQueryTabComponentProps.onUpdatePersistedState({
-      queryText: options.queryText,
-      splitterDirection: options.splitterDirection,
-      queryViewSizePercent: options.queryViewSizePercent,
-    });
   }
 
   public render(): JSX.Element {

src/Explorer/Tabs/QueryTab/QueryTabComponent.test.tsx

@@ -34,7 +34,6 @@ jest.mock("Shared/AppStatePersistenceUtility", () => ({
   AppStateComponentNames: {
     QueryCopilot: "QueryCopilot",
   },
-  readSubComponentState: jest.fn(),
 }));
 
 describe("QueryTabComponent", () => {

src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx

@@ -18,7 +18,13 @@ import { CosmosFluentProvider } from "Explorer/Theme/ThemeUtil";
 import { useSelectedNode } from "Explorer/useSelectedNode";
 import { KeyboardAction } from "KeyboardShortcuts";
 import { QueryConstants } from "Shared/Constants";
-import { LocalStorageUtility, StorageKey, getRUThreshold, ruThresholdEnabled } from "Shared/StorageUtility";
+import {
+  LocalStorageUtility,
+  StorageKey,
+  getDefaultQueryResultsView,
+  getRUThreshold,
+  ruThresholdEnabled,
+} from "Shared/StorageUtility";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { Allotment } from "allotment";
 import { QueryCopilotState, useQueryCopilot } from "hooks/useQueryCopilot";
@@ -93,13 +99,6 @@ export interface IQueryTabComponentProps {
   copilotEnabled?: boolean;
   isSampleCopilotActive?: boolean;
   copilotStore?: Partial<QueryCopilotState>;
-  splitterDirection?: "horizontal" | "vertical";
-  queryViewSizePercent?: number;
-  onUpdatePersistedState: (state: {
-    queryText: string;
-    splitterDirection: "vertical" | "horizontal";
-    queryViewSizePercent: number;
-  }) => void;
 }
 
 interface IQueryTabStates {
@@ -119,13 +118,11 @@ interface IQueryTabStates {
   queryResultsView: SplitterDirection;
   errors?: QueryError[];
   modelMarkers?: monaco.editor.IMarkerData[];
-  queryViewSizePercent: number;
 }
 
 export const QueryTabCopilotComponent = (props: IQueryTabComponentProps): any => {
   const styles = useQueryTabStyles();
   const copilotStore = useCopilotStore();
-
   const isSampleCopilotActive = useSelectedNode.getState().isQueryCopilotCollectionSelected();
   const queryTabProps = {
     ...props,
@@ -135,12 +132,12 @@ export const QueryTabCopilotComponent = (props: IQueryTabComponentProps): any =>
     isSampleCopilotActive: isSampleCopilotActive,
     copilotStore: copilotStore,
   };
-  return <QueryTabComponentImpl styles={styles} {...queryTabProps} />;
+  return <QueryTabComponentImpl styles={styles} {...queryTabProps}></QueryTabComponentImpl>;
 };
 
 export const QueryTabComponent = (props: IQueryTabComponentProps): any => {
   const styles = useQueryTabStyles();
-  return <QueryTabComponentImpl styles={styles} {...{ ...props }} />;
+  return <QueryTabComponentImpl styles={styles} {...props}></QueryTabComponentImpl>;
 };
 
 type QueryTabComponentImplProps = IQueryTabComponentProps & {
@@ -149,8 +146,6 @@ type QueryTabComponentImplProps = IQueryTabComponentProps & {
 
 // Inner (legacy) class component. We only use this component via one of the two functional components above (since we need to use the `useQueryTabStyles` hook).
 class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps, IQueryTabStates> {
-  private static readonly DEBOUNCE_DELAY_MS = 1000;
-
   public queryEditorId: string;
   public executeQueryButton: Button;
   public saveQueryButton: Button;
@@ -162,10 +157,10 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
   private _iterator: MinimalQueryIterator;
   private queryAbortController: AbortController;
   queryEditor: React.RefObject<EditorReact>;
-  private timeoutId: NodeJS.Timeout | undefined;
 
   constructor(props: QueryTabComponentImplProps) {
     super(props);
+
     this.queryEditor = createRef<EditorReact>();
 
     this.state = {
@@ -181,9 +176,7 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
       cancelQueryTimeoutID: undefined,
       copilotActive: this._queryCopilotActive(),
       currentTabActive: true,
-      queryResultsView:
-        props.splitterDirection === "vertical" ? SplitterDirection.Vertical : SplitterDirection.Horizontal,
-      queryViewSizePercent: props.queryViewSizePercent,
+      queryResultsView: getDefaultQueryResultsView(),
     };
     this.isCloseClicked = false;
     this.splitterId = this.props.tabId + "_splitter";
@@ -214,23 +207,6 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
     });
   }
 
-  /**
-   * Helper function to save the query text in the query tab state
-   * Since it reads and writes to the same state, it is debounced
-   */
-  private saveQueryTabStateDebounced = () => {
-    if (this.timeoutId) {
-      clearTimeout(this.timeoutId);
-    }
-    this.timeoutId = setTimeout(async () => {
-      this.props.onUpdatePersistedState({
-        queryText: this.state.sqlQueryEditorContent,
-        splitterDirection: this.state.queryResultsView,
-        queryViewSizePercent: this.state.queryViewSizePercent,
-      });
-    }, QueryTabComponentImpl.DEBOUNCE_DELAY_MS);
-  };
-
   private _queryCopilotActive(): boolean {
     if (this.props.copilotEnabled) {
       return readCopilotToggleStatus(userContext.databaseAccount);
@@ -591,7 +567,7 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
     };
   }
   private _setViewLayout(direction: SplitterDirection): void {
-    this.setState({ queryResultsView: direction }, () => this.saveQueryTabStateDebounced());
+    this.setState({ queryResultsView: direction });
 
     // We'll need to refresh the context buttons to update the selected state of the view buttons
     setTimeout(() => {
@@ -623,16 +599,13 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
     if (this.state.copilotActive) {
       this.props.copilotStore?.setQuery(newContent);
     }
-    this.setState(
-      {
-        sqlQueryEditorContent: newContent,
-        queryCopilotGeneratedQuery: "",
+    this.setState({
+      sqlQueryEditorContent: newContent,
+      queryCopilotGeneratedQuery: "",
 
-        // Clear the markers when the user edits the document.
-        modelMarkers: [],
-      },
-      () => this.saveQueryTabStateDebounced(),
-    );
+      // Clear the markers when the user edits the document.
+      modelMarkers: [],
+    });
     if (this.isPreferredApiMongoDB) {
       if (newContent.length > 0) {
         this.executeQueryButton = {
@@ -731,20 +704,8 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
             ></QueryCopilotPromptbar>
           )}
           {/* Set 'key' to the value of vertical to force re-rendering when vertical changes, to work around https://github.com/johnwalley/allotment/issues/457 */}
-          <Allotment
-            key={vertical.toString()}
-            vertical={vertical}
-            onDragEnd={(sizes: number[]) => {
-              const queryViewSizePercent = (100 * sizes[0]) / (sizes[0] + sizes[1]);
-              this.setState({ queryViewSizePercent }, () => this.saveQueryTabStateDebounced());
-            }}
-          >
-            <Allotment.Pane
-              data-test="QueryTab/EditorPane"
-              preferredSize={
-                this.state.queryViewSizePercent !== undefined ? `${this.state.queryViewSizePercent}%` : undefined
-              }
-            >
+          <Allotment key={vertical.toString()} vertical={vertical}>
+            <Allotment.Pane data-test="QueryTab/EditorPane">
               <EditorReact
                 ref={this.queryEditor}
                 className={this.props.styles.queryEditor}

src/Explorer/Tabs/SettingsTabV2.tsx

@@ -1,4 +1,3 @@
-import { ActionType, OpenCollectionTab, TabKind } from "Contracts/ActionContracts";
 import React from "react";
 import * as ViewModels from "../../Contracts/ViewModels";
 import { SettingsComponent } from "../Controls/Settings/SettingsComponent";
@@ -11,18 +10,6 @@ export class SettingsTabV2 extends TabsBase {
 }
 
 export class CollectionSettingsTabV2 extends SettingsTabV2 {
-  protected persistedState: OpenCollectionTab;
-
-  constructor(options: ViewModels.TabOptions) {
-    super(options);
-    this.persistedState = {
-      actionType: ActionType.OpenCollectionTab,
-      tabKind: TabKind.ScaleSettings,
-      databaseResourceId: options.collection.databaseId,
-      collectionResourceId: options.collection.id(),
-    };
-  }
-
   public onActivate(): void {
     super.onActivate();
     this.collection.selectedSubnodeKind(ViewModels.CollectionTabKind.CollectionSettingsV2);

src/Explorer/Tabs/TabsBase.ts

@@ -1,4 +1,3 @@
-import { OpenTab } from "Contracts/ActionContracts";
 import { KeyboardActionGroup, clearKeyboardActionGroup } from "KeyboardShortcuts";
 import * as ko from "knockout";
 import * as Constants from "../../Common/Constants";
@@ -31,8 +30,6 @@ export default class TabsBase extends WaitsForTemplateViewModel {
   protected _theme: string;
   public onLoadStartKey: number;
 
-  protected persistedState: OpenTab | undefined = undefined; // Used to store state of tab for persistence
-
   constructor(options: ViewModels.TabOptions) {
     super();
     this.index = options.index;
@@ -58,10 +55,6 @@ export default class TabsBase extends WaitsForTemplateViewModel {
     };
   }
 
-  // Called by useTabs to persist
-  public getPersistedState = (): OpenTab | null => this.persistedState;
-  public triggerPersistState: () => void = undefined;
-
   public onCloseTabButtonClick(): void {
     useTabs.getState().closeTab(this);
     TelemetryProcessor.trace(Action.Tab, ActionModifiers.Close, {

src/Explorer/Tree/Collection.ts

@@ -630,13 +630,7 @@ export default class Collection implements ViewModels.Collection {
     }
   };
 
-  public onNewQueryClick(
-    source: any,
-    event: MouseEvent,
-    queryText?: string,
-    splitterDirection?: "horizontal" | "vertical",
-    queryViewSizePercent?: number,
-  ) {
+  public onNewQueryClick(source: any, event: MouseEvent, queryText?: string) {
     const collection: ViewModels.Collection = source.collection || source;
     const id = useTabs.getState().getTabs(ViewModels.CollectionTabKind.Query).length + 1;
     const title = "Query " + id;
@@ -659,21 +653,13 @@ export default class Collection implements ViewModels.Collection {
           queryText: queryText,
           partitionKey: collection.partitionKey,
           onLoadStartKey: startKey,
-          splitterDirection,
-          queryViewSizePercent,
         },
         { container: this.container },
       ),
     );
   }
 
-  public onNewMongoQueryClick(
-    source: any,
-    event: MouseEvent,
-    queryText?: string,
-    splitterDirection?: "horizontal" | "vertical",
-    queryViewSizePercent?: number,
-  ) {
+  public onNewMongoQueryClick(source: any, event: MouseEvent, queryText?: string) {
     const collection: ViewModels.Collection = source.collection || source;
     const id = useTabs.getState().getTabs(ViewModels.CollectionTabKind.Query).length + 1;
 
@@ -695,9 +681,6 @@ export default class Collection implements ViewModels.Collection {
         node: this,
         partitionKey: collection.partitionKey,
         onLoadStartKey: startKey,
-        queryText,
-        splitterDirection,
-        queryViewSizePercent,
       },
       {
         container: this.container,

src/Platform/Hosted/extractFeatures.ts

@@ -38,7 +38,6 @@ export type Features = {
   readonly copilotChatFixedMonacoEditorHeight: boolean;
   readonly enablePriorityBasedExecution: boolean;
   readonly disableConnectionStringLogin: boolean;
-  readonly restoreTabs: boolean;
 
   // can be set via both flight and feature flag
   autoscaleDefault: boolean;
@@ -109,7 +108,6 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear
     copilotChatFixedMonacoEditorHeight: "true" === get("copilotchatfixedmonacoeditorheight"),
     enablePriorityBasedExecution: "true" === get("enableprioritybasedexecution"),
     disableConnectionStringLogin: "true" === get("disableconnectionstringlogin"),
-    restoreTabs: "true" === get("restoretabs"),
   };
 }
 

src/Shared/AppStatePersistenceUtility.ts

@@ -1,20 +1,12 @@
 import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
-import { Action } from "Shared/Telemetry/TelemetryConstants";
-import { userContext } from "UserContext";
-import * as ViewModels from "../Contracts/ViewModels";
-import * as TelemetryProcessor from "../Shared/Telemetry/TelemetryProcessor";
 
 // The component name whose state is being saved. Component name must not include special characters.
 export enum AppStateComponentNames {
   DocumentsTab = "DocumentsTab",
   MostRecentActivity = "MostRecentActivity",
   QueryCopilot = "QueryCopilot",
-  DataExplorerAction = "DataExplorerAction",
 }
 
-// Subcomponent for DataExplorerAction
-export const OPEN_TABS_SUBCOMPONENT_NAME = "OpenTabs";
-
 export const PATH_SEPARATOR = "/"; // export for testing purposes
 const SCHEMA_VERSION = 1;
 
@@ -80,18 +72,12 @@ export const hasState = (path: StorePath): boolean => {
 };
 
 // This is for high-frequency state changes
-// Keep track of timeouts per path
-const pathToTimeoutIdMap = new Map<string, NodeJS.Timeout>();
+let timeoutId: NodeJS.Timeout | undefined;
 export const saveStateDebounced = (path: StorePath, state: unknown, debounceDelayMs = 1000): void => {
-  const key = createKeyFromPath(path);
-  const timeoutId = pathToTimeoutIdMap.get(key);
   if (timeoutId) {
     clearTimeout(timeoutId);
   }
-  pathToTimeoutIdMap.set(
-    key,
-    setTimeout(() => saveState(path, state), debounceDelayMs),
-  );
+  timeoutId = setTimeout(() => saveState(path, state), debounceDelayMs);
 };
 
 interface ApplicationState {
@@ -126,93 +112,3 @@ export const createKeyFromPath = (path: StorePath): string => {
 export const deleteAllStates = (): void => {
   LocalStorageUtility.removeEntry(StorageKey.AppState);
 };
-
-// Convenience functions
-
-/**
- *
- * @param subComponentName
- * @param collection
- * @param defaultValue Will be returned if persisted state is not found
- * @returns
- */
-export const readSubComponentState = <T>(
-  componentName: AppStateComponentNames,
-  subComponentName: string,
-  collection: ViewModels.CollectionBase | undefined,
-  defaultValue: T,
-): T => {
-  const globalAccountName = userContext.databaseAccount?.name;
-  if (!globalAccountName) {
-    const message = "Database account name not found in userContext";
-    console.error(message);
-    TelemetryProcessor.traceFailure(Action.ReadPersistedTabState, { message, componentName });
-    return defaultValue;
-  }
-
-  const state = loadState({
-    componentName: componentName,
-    subComponentName,
-    globalAccountName,
-    databaseName: collection ? collection.databaseId : "",
-    containerName: collection ? collection.id() : "",
-  }) as T;
-
-  return state || defaultValue;
-};
-
-/**
- *
- * @param subComponentName
- * @param collection
- * @param state State to save
- * @param debounce true for high-frequency calls (e.g mouse drag events)
- */
-export const saveSubComponentState = <T>(
-  componentName: AppStateComponentNames,
-  subComponentName: string,
-  collection: ViewModels.CollectionBase | undefined,
-  state: T,
-  debounce?: boolean,
-): void => {
-  const globalAccountName = userContext.databaseAccount?.name;
-  if (!globalAccountName) {
-    const message = "Database account name not found in userContext";
-    console.error(message);
-    TelemetryProcessor.traceFailure(Action.SavePersistedTabState, { message, componentName });
-    return;
-  }
-
-  (debounce ? saveStateDebounced : saveState)(
-    {
-      componentName: componentName,
-      subComponentName,
-      globalAccountName,
-      databaseName: collection ? collection.databaseId : "",
-      containerName: collection ? collection.id() : "",
-    },
-    state,
-  );
-};
-
-export const deleteSubComponentState = (
-  componentName: AppStateComponentNames,
-  subComponentName: string,
-  collection: ViewModels.CollectionBase,
-) => {
-  const globalAccountName = userContext.databaseAccount?.name;
-  if (!globalAccountName) {
-    const message = "Database account name not found in userContext";
-    console.error(message);
-    TelemetryProcessor.traceFailure(Action.DeletePersistedTabState, { message, componentName });
-    return;
-  }
-
-  deleteState({
-    componentName: componentName,
-    subComponentName,
-    globalAccountName,
-    databaseName: collection.databaseId,
-    containerName: collection.id(),
-  });
-};

src/UserContext.ts

@@ -81,6 +81,7 @@ export interface UserContext {
   readonly endpoint?: string;
   readonly aadToken?: string;
   readonly accessToken?: string;
+  readonly armToken?: string;
   readonly authorizationToken?: string;
   readonly resourceToken?: string;
   readonly subscriptionType?: SubscriptionType;

src/hooks/useAADAuth.ts

@@ -3,6 +3,7 @@ import { useBoolean } from "@fluentui/react-hooks";
 import * as React from "react";
 import { configContext } from "../ConfigContext";
 import { acquireTokenWithMsal, getMsalInstance } from "../Utils/AuthorizationUtils";
+import { updateUserContext } from "UserContext";
 
 const msalInstance = await getMsalInstance();
 
@@ -79,7 +80,7 @@ export function useAADAuth(): ReturnType {
         authority: `${configContext.AAD_ENDPOINT}${tenantId}`,
         scopes: [`${configContext.ARM_ENDPOINT}/.default`],
       });
-
+      updateUserContext({ armToken: armToken });
       setArmToken(armToken);
       setAuthFailure(null);
     } catch (error) {

src/hooks/useDatabaseAccounts.tsx

@@ -1,6 +1,8 @@
 import { HttpHeaders } from "Common/Constants";
 import { QueryRequestOptions, QueryResponse } from "Contracts/AzureResourceGraph";
 import useSWR from "swr";
+import { updateUserContext, userContext } from "UserContext";
+import { acquireTokenWithMsal, getMsalInstance } from "Utils/AuthorizationUtils";
 import { configContext } from "../ConfigContext";
 import { DatabaseAccount } from "../Contracts/DataModels";
 /* eslint-disable  @typescript-eslint/no-explicit-any */
@@ -33,12 +35,9 @@ export async function fetchDatabaseAccounts(subscriptionId: string, accessToken:
   return accounts.sort((a, b) => a.name.localeCompare(b.name));
 }
 
-export async function fetchDatabaseAccountsFromGraph(
-  subscriptionId: string,
-  accessToken: string,
-): Promise<DatabaseAccount[]> {
+export async function fetchDatabaseAccountsFromGraph(subscriptionId: string): Promise<DatabaseAccount[]> {
   const headers = new Headers();
-  const bearer = `Bearer ${accessToken}`;
+  const bearer = `Bearer ${userContext.armToken}`;
 
   headers.append("Authorization", bearer);
   headers.append(HttpHeaders.contentType, "application/json");
@@ -46,8 +45,9 @@ export async function fetchDatabaseAccountsFromGraph(
   const apiVersion = "2021-03-01";
   const managementResourceGraphAPIURL = `${configContext.ARM_ENDPOINT}providers/Microsoft.ResourceGraph/resources?api-version=${apiVersion}`;
 
-  const databaseAccounts: DatabaseAccount[] = [];
+  let databaseAccounts: DatabaseAccount[] = [];
   let skipToken: string;
+  console.log("Old ARM Token - fetchDatabaseAccountsFromGraph function", userContext.armToken);
   do {
     const body = {
       query: databaseAccountsQuery,
@@ -85,10 +85,81 @@ export async function fetchDatabaseAccountsFromGraph(
   return databaseAccounts.sort((a, b) => a.name.localeCompare(b.name));
 }
 
-export function useDatabaseAccounts(subscriptionId: string, armToken: string): DatabaseAccount[] | undefined {
+export function useDatabaseAccounts(subscriptionId: string): DatabaseAccount[] | undefined {
   const { data } = useSWR(
-    () => (armToken && subscriptionId ? ["databaseAccounts", subscriptionId, armToken] : undefined),
-    (_, subscriptionId, armToken) => fetchDatabaseAccountsFromGraph(subscriptionId, armToken),
+    () => (subscriptionId ? ["databaseAccounts", subscriptionId] : undefined),
+    (_, subscriptionId) => runCommand(fetchDatabaseAccountsFromGraph, subscriptionId),
   );
   return data;
 }
+
+// Define the types for your responses
+interface DatabaseAccount {
+  name: string;
+  id: string;
+  // Add other relevant fields as per your use case
+}
+
+interface QueryRequestOptions {
+  $top?: number;
+  $skipToken?: string;
+  $allowPartialScopes?: boolean;
+}
+
+// Define the configuration context and headers if not already defined
+const configContext = {
+  ARM_ENDPOINT: "https://management.azure.com/",
+  AAD_ENDPOINT: "https://login.microsoftonline.com/",
+};
+
+interface QueryResponse {
+  data?: any[];
+  $skipToken?: string;
+}
+
+export async function runCommand<T>(fn: (...args: any[]) => Promise<T>, ...args: any[]): Promise<T> {
+  try {
+    // Attempt to execute the function passed as an argument
+    const result = await fn(...args);
+    console.log("Successfully executed function:", fn.name, result);
+    return result;
+  } catch (error) {
+    // Handle any error that is thrown during the execution of the function
+    if (error) {
+      console.log("Creating new token");
+      const msalInstance = await getMsalInstance();
+
+      const cachedAccount = msalInstance.getAllAccounts()?.[0];
+      const cachedTenantId = localStorage.getItem("cachedTenantId");
+
+      msalInstance.setActiveAccount(cachedAccount);
+
+      // TODO: Add condition to check if the ARM token needs to be renewed, then we need to run the code below for creating the ARM token
+
+      console.log("Creating new ARM token");
+      const newAccessToken = await acquireTokenWithMsal(msalInstance, {
+        authority: `${configContext.AAD_ENDPOINT}${cachedTenantId}`,
+        scopes: [`${configContext.ARM_ENDPOINT}/.default`],
+      });
+      updateUserContext({ armToken: newAccessToken });
+
+      // TODO: add condition to check if AAD token needs to be renewed (i.e) Token provider has failed with expired AAD token and create a new AAD Token using the below code
+
+      // const hrefEndpoint = new URL(userContext.databaseAccount.properties.documentEndpoint).href.replace(/\/$/, "/.default");
+      // console.log('Creating new AAD token');
+      // let aadToken = await acquireTokenWithMsal(msalInstance, {
+      //   forceRefresh: true,
+      //   authority: `${configContext.AAD_ENDPOINT}${cachedTenantId}`,
+      //   scopes: [hrefEndpoint],
+      // });
+      // updateUserContext({aadToken: aadToken});
+
+      //console.log('Latest AAD Token', fn.name, userContext.aadToken);
+      const result = await fn(...args);
+      return result;
+    } else {
+      console.error("An error occurred:", error.message);
+      throw new error();
+    }
+  }
+}

src/hooks/useKnockoutExplorer.ts

@@ -7,11 +7,6 @@ import Explorer from "Explorer/Explorer";
 import { useDataPlaneRbac } from "Explorer/Panes/SettingsPane/SettingsPane";
 import { useSelectedNode } from "Explorer/useSelectedNode";
 import { scheduleRefreshDatabaseResourceToken } from "Platform/Fabric/FabricUtil";
-import {
-  AppStateComponentNames,
-  OPEN_TABS_SUBCOMPONENT_NAME,
-  readSubComponentState,
-} from "Shared/AppStatePersistenceUtility";
 import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
 import { useNewPortalBackendEndpoint } from "Utils/EndpointUtils";
 import { getNetworkSettingsWarningMessage } from "Utils/NetworkUtility";
@@ -85,11 +80,6 @@ export function useKnockoutExplorer(platform: Platform): Explorer {
           await updateContextForCopilot(explorer);
           await updateContextForSampleData(explorer);
         }
-
-        if (userContext.features.restoreTabs) {
-          restoreOpenTabs();
-        }
-
         setExplorer(explorer);
       }
     };
@@ -826,17 +816,3 @@ async function updateContextForSampleData(explorer: Explorer): Promise<void> {
 interface SampledataconnectionResponse {
   connectionString: string;
 }
-
-const restoreOpenTabs = () => {
-  const openTabsState = readSubComponentState<(DataExplorerAction | undefined)[]>(
-    AppStateComponentNames.DataExplorerAction,
-    OPEN_TABS_SUBCOMPONENT_NAME,
-    undefined,
-    [],
-  );
-  openTabsState.forEach((openTabState) => {
-    if (openTabState) {
-      handleOpenAction(openTabState, useDatabases.getState().databases, this);
-    }
-  });
-};

src/hooks/useSubscriptions.tsx

@@ -3,6 +3,8 @@ import { QueryRequestOptions, QueryResponse } from "Contracts/AzureResourceGraph
 import useSWR from "swr";
 import { configContext } from "../ConfigContext";
 import { Subscription } from "../Contracts/DataModels";
+import { runCommand } from "hooks/useDatabaseAccounts";
+import { userContext } from "UserContext";
 /* eslint-disable  @typescript-eslint/no-explicit-any */
 
 interface SubscriptionListResult {
@@ -35,9 +37,9 @@ export async function fetchSubscriptions(accessToken: string): Promise<Subscript
   return subscriptions.sort((a, b) => a.displayName.localeCompare(b.displayName));
 }
 
-export async function fetchSubscriptionsFromGraph(accessToken: string): Promise<Subscription[]> {
+export async function fetchSubscriptionsFromGraph(): Promise<Subscription[]> {
   const headers = new Headers();
-  const bearer = `Bearer ${accessToken}`;
+  const bearer = `Bearer ${userContext.armToken}`;
 
   headers.append("Authorization", bearer);
   headers.append(HttpHeaders.contentType, "application/json");
@@ -48,6 +50,7 @@ export async function fetchSubscriptionsFromGraph(accessToken: string): Promise<
 
   const subscriptions: Subscription[] = [];
   let skipToken: string;
+  console.log("Old ARM Token fetchSubscriptionsFromGraph fn", userContext.armToken);
   do {
     const body = {
       query: subscriptionsQuery,
@@ -86,9 +89,14 @@ export async function fetchSubscriptionsFromGraph(accessToken: string): Promise<
 }
 
 export function useSubscriptions(armToken: string): Subscription[] | undefined {
-  const { data } = useSWR(
+  const { data, error } = useSWR(
     () => (armToken ? ["subscriptions", armToken] : undefined),
-    (_, armToken) => fetchSubscriptionsFromGraph(armToken),
+    (_) => runCommand(fetchSubscriptionsFromGraph),
   );
+
+  if (error) {
+    console.error("Error fetching subscriptions:", error);
+  }
+
   return data;
 }

src/hooks/useTabs.ts

@@ -1,11 +1,5 @@
 import { clamp } from "@fluentui/react";
-import { OpenTab } from "Contracts/ActionContracts";
 import { useSelectedNode } from "Explorer/useSelectedNode";
-import {
-  AppStateComponentNames,
-  OPEN_TABS_SUBCOMPONENT_NAME,
-  saveSubComponentState,
-} from "Shared/AppStatePersistenceUtility";
 import create, { UseStore } from "zustand";
 import * as ViewModels from "../Contracts/ViewModels";
 import { CollectionTabKind } from "../Contracts/ViewModels";
@@ -42,7 +36,6 @@ export interface TabsState {
   selectLeftTab: () => void;
   selectRightTab: () => void;
   closeActiveTab: () => void;
-  persistTabsState: () => void;
 }
 
 export enum ReactTabKind {
@@ -80,9 +73,7 @@ export const useTabs: UseStore<TabsState> = create((set, get) => ({
   },
   activateNewTab: (tab: TabsBase): void => {
     set((state) => ({ openedTabs: [...state.openedTabs, tab], activeTab: tab, activeReactTab: undefined }));
-    tab.triggerPersistState = get().persistTabsState;
     tab.onActivate();
-    get().persistTabsState();
   },
   activateReactTab: (tabKind: ReactTabKind): void => {
     // Clear the selected node when switching to a react tab.
@@ -139,8 +130,6 @@ export const useTabs: UseStore<TabsState> = create((set, get) => ({
     }
 
     set({ openedTabs: updatedTabs });
-
-    get().persistTabsState();
   },
   closeAllNotebookTabs: (hardClose): void => {
     const isNotebook = (tabKind: CollectionTabKind): boolean => {
@@ -237,15 +226,4 @@ export const useTabs: UseStore<TabsState> = create((set, get) => ({
       state.closeTab(state.activeTab);
     }
   },
-  persistTabsState: () => {
-    const state = get();
-    const openTabsStates = state.openedTabs.map((tab) => tab.getPersistedState());
-
-    saveSubComponentState<OpenTab[]>(
-      AppStateComponentNames.DataExplorerAction,
-      OPEN_TABS_SUBCOMPONENT_NAME,
-      undefined,
-      openTabsStates,
-    );
-  },
 }));