Save and restore states for QueryTab and DocumentsTab for SQL and Mongo

src/Contracts/ActionContracts.ts

@@ -9,6 +9,7 @@ export enum TabKind {
   Graph,
   SQLQuery,
   ScaleSettings,
+  MongoQuery,
 }
 
 /**

src/Contracts/ViewModels.ts

@@ -155,7 +155,13 @@ export interface Collection extends CollectionBase {
   onSettingsClick: () => Promise<void>;
 
   onNewGraphClick(): void;
-  onNewMongoQueryClick(source: any, event?: MouseEvent, queryText?: string): void;
+  onNewMongoQueryClick(
+    source: any,
+    event?: MouseEvent,
+    queryText?: string,
+    stringsplitterDirection?: "horizontal" | "vertical",
+    queryViewSizePercent?: number,
+  ): void;
   onNewMongoShellClick(): void;
   onNewStoredProcedureClick(source: Collection, event?: MouseEvent): void;
   onNewUserDefinedFunctionClick(source: Collection, event?: MouseEvent): void;

src/Explorer/Explorer.tsx

@@ -9,10 +9,13 @@ import { MessageTypes } from "Contracts/ExplorerContracts";
 import { handleOpenAction } from "Explorer/OpenActions/OpenActions";
 import { useDataPlaneRbac } from "Explorer/Panes/SettingsPane/SettingsPane";
 import { getCopilotEnabled, isCopilotFeatureRegistered } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
-import { OPEN_TABS_SUBCOMPONENT_NAME } from "Explorer/Tabs/QueryTab/QueryTabStateUtil";
 import { IGalleryItem } from "Juno/JunoClient";
 import { scheduleRefreshDatabaseResourceToken } from "Platform/Fabric/FabricUtil";
-import { AppStateComponentNames, readSubComponentState } from "Shared/AppStatePersistenceUtility";
+import {
+  AppStateComponentNames,
+  OPEN_TABS_SUBCOMPONENT_NAME,
+  readSubComponentState,
+} from "Shared/AppStatePersistenceUtility";
 import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility";
 import { acquireMsalTokenForAccount } from "Utils/AuthorizationUtils";
 import { allowedNotebookServerUrls, validateEndpoint } from "Utils/EndpointUtils";
@@ -1212,7 +1215,7 @@ export default class Explorer {
   }
 
   private restoreOpenTabs() {
-    const openTabsState = readSubComponentState<DataExplorerAction[]>(
+    const openTabsState = readSubComponentState<(DataExplorerAction | undefined)[]>(
       AppStateComponentNames.DataExplorerAction,
       OPEN_TABS_SUBCOMPONENT_NAME,
       undefined,

src/Explorer/OpenActions/OpenActions.tsx

@@ -132,6 +132,21 @@ function openCollectionTab(
             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,
+            );
+            break;
+          }
+
           if (
             action.tabKind === ActionContracts.TabKind.ScaleSettings ||
             action.tabKind === ActionContracts.TabKind[ActionContracts.TabKind.ScaleSettings]

src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx

@@ -23,6 +23,7 @@ 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";
@@ -146,6 +147,8 @@ export class DocumentsTabV2 extends TabsBase {
   private title: string;
   private resourceTokenPartitionKey: string;
 
+  protected persistedState: OpenCollectionTab;
+
   constructor(options: ViewModels.DocumentsTabOptions) {
     super(options);
 
@@ -153,6 +156,13 @@ 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 {

src/Explorer/Tabs/MongoQueryTab/MongoQueryTab.tsx

@@ -1,3 +1,4 @@
+import { ActionType, TabKind } from "Contracts/ActionContracts";
 import React from "react";
 import MongoUtility from "../../../Common/MongoUtility";
 import * as ViewModels from "../../../Contracts/ViewModels";
@@ -20,7 +21,7 @@ export class NewMongoQueryTab extends NewQueryTab {
     private mongoQueryTabProps: IMongoQueryTabProps,
   ) {
     super(options, mongoQueryTabProps);
-    this.queryText = "";
+    this.queryText = options.queryText ?? "";
     this.iMongoQueryTabComponentProps = {
       collection: options.collection,
       isExecutionError: this.isExecutionError(),
@@ -28,6 +29,8 @@ export class NewMongoQueryTab extends NewQueryTab {
       tabsBaseInstance: this,
       queryText: this.queryText,
       partitionKey: this.partitionKey,
+      stringsplitterDirection: options.stringsplitterDirection,
+      queryViewSizePercent: options.queryViewSizePercent,
       container: this.mongoQueryTabProps.container,
       onTabAccessor: (instance: ITabAccessor): void => {
         this.iTabAccessor = instance;
@@ -35,6 +38,26 @@ 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,4 +1,5 @@
 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";
@@ -26,6 +27,8 @@ export class NewQueryTab extends TabsBase {
   public iQueryTabComponentProps: IQueryTabComponentProps;
   public iTabAccessor: ITabAccessor;
 
+  protected persistedState: OpenQueryTab;
+
   constructor(
     options: QueryTabOptions,
     private props: IQueryTabProps,
@@ -46,7 +49,34 @@ export class NewQueryTab extends TabsBase {
         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.stringsplitterDirection,
+      queryViewSizePercent: options.queryViewSizePercent,
+    });
   }
 
   public render(): JSX.Element {

src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx

@@ -13,7 +13,6 @@ import { readCopilotToggleStatus, saveCopilotToggleStatus } from "Explorer/Query
 import { OnExecuteQueryClick, QueryDocumentsPerPage } from "Explorer/QueryCopilot/Shared/QueryCopilotClient";
 import { QueryCopilotSidebar } from "Explorer/QueryCopilot/V2/Sidebar/QueryCopilotSidebar";
 import { QueryResultSection } from "Explorer/Tabs/QueryTab/QueryResultSection";
-import { deleteQueryTabState, saveQueryTabState } from "Explorer/Tabs/QueryTab/QueryTabStateUtil";
 import { QueryTabStyles, useQueryTabStyles } from "Explorer/Tabs/QueryTab/Styles";
 import { CosmosFluentProvider } from "Explorer/Theme/ThemeUtil";
 import { useSelectedNode } from "Explorer/useSelectedNode";
@@ -96,6 +95,11 @@ export interface IQueryTabComponentProps {
   copilotStore?: Partial<QueryCopilotState>;
   stringsplitterDirection?: "horizontal" | "vertical";
   queryViewSizePercent?: number;
+  onUpdatePersistedState: (state: {
+    queryText: string;
+    splitterDirection: "vertical" | "horizontal";
+    queryViewSizePercent: number;
+  }) => void;
 }
 
 interface IQueryTabStates {
@@ -183,7 +187,7 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
       copilotActive: this._queryCopilotActive(),
       currentTabActive: true,
       queryResultsView:
-        props.stringsplitterDirection === "horizontal" ? SplitterDirection.Horizontal : SplitterDirection.Vertical,
+        props.stringsplitterDirection === "vertical" ? SplitterDirection.Vertical : SplitterDirection.Horizontal,
       queryViewSizePercent: props.queryViewSizePercent,
     };
     this.isCloseClicked = false;
@@ -213,10 +217,6 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
       onSaveClickEvent: this.getCurrentEditorQuery.bind(this),
       onCloseClickEvent: this.onCloseClick.bind(this),
     });
-
-    // Update persistence
-    this.saveQueryTabStateDebounced();
-    // DO THIS IN useTabs.activateNewTab() INSTEAD
   }
 
   /**
@@ -228,15 +228,11 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
       clearTimeout(this.timeoutId);
     }
     this.timeoutId = setTimeout(async () => {
-      saveQueryTabState(
-        this.props.collection,
-        {
-          queryText: this.state.sqlQueryEditorContent,
-          splitterDirection: this.state.queryResultsView,
-          queryViewSizePercent: this.state.queryViewSizePercent,
-        },
-        this.props.tabIndex,
-      );
+      this.props.onUpdatePersistedState({
+        queryText: this.state.sqlQueryEditorContent,
+        splitterDirection: this.state.queryResultsView,
+        queryViewSizePercent: this.state.queryViewSizePercent,
+      });
     }, QueryTabComponentImpl.DEBOUNCE_DELAY_MS);
   };
 
@@ -363,9 +359,9 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
       this._iterator = this.props.isPreferredApiMongoDB
         ? queryIterator(this.props.collection.databaseId, this.props.viewModelcollection, query)
         : queryDocuments(this.props.collection.databaseId, this.props.collection.id(), query, {
-          enableCrossPartitionQuery: HeadersUtility.shouldEnableCrossPartitionKey(),
-          abortSignal: this.queryAbortController.signal,
-        } as unknown as FeedOptions);
+            enableCrossPartitionQuery: HeadersUtility.shouldEnableCrossPartitionKey(),
+            abortSignal: this.queryAbortController.signal,
+          } as unknown as FeedOptions);
     }
 
     await this._queryDocumentsPage(firstItemIndex);
@@ -724,9 +720,6 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
 
   componentWillUnmount(): void {
     document.removeEventListener("keydown", this.handleCopilotKeyDown);
-
-    // Remove persistence
-    deleteQueryTabState(this.props.tabIndex);
   }
 
   private getEditorAndQueryResult(): JSX.Element {

src/Explorer/Tabs/QueryTab/QueryTabStateUtil.ts

@@ -1,65 +0,0 @@
-// Definitions of State data
-
-import { ActionType, OpenQueryTab, TabKind } from "Contracts/ActionContracts";
-import {
-  AppStateComponentNames,
-  readSubComponentState,
-  saveSubComponentState,
-} from "Shared/AppStatePersistenceUtility";
-import * as ViewModels from "../../../Contracts/ViewModels";
-
-export const OPEN_TABS_SUBCOMPONENT_NAME = "OpenTabs";
-
-export const saveQueryTabState = (
-  collection: ViewModels.CollectionBase,
-  state: {
-    queryText: string;
-    splitterDirection: "vertical" | "horizontal";
-    queryViewSizePercent: number;
-  },
-  tabIndex: number,
-): void => {
-  const openTabsState = readSubComponentState<OpenQueryTab[]>(
-    AppStateComponentNames.DataExplorerAction,
-    OPEN_TABS_SUBCOMPONENT_NAME,
-    undefined,
-    [],
-  );
-
-  openTabsState[tabIndex] = {
-    actionType: ActionType.OpenCollectionTab,
-    tabKind: TabKind.SQLQuery,
-    databaseResourceId: collection.databaseId,
-    collectionResourceId: collection.id(),
-    query: {
-      text: state.queryText,
-    },
-    splitterDirection: state.splitterDirection,
-    queryViewSizePercent: state.queryViewSizePercent,
-  };
-
-  saveSubComponentState<OpenQueryTab[]>(
-    AppStateComponentNames.DataExplorerAction,
-    OPEN_TABS_SUBCOMPONENT_NAME,
-    undefined,
-    openTabsState,
-  );
-};
-
-export const deleteQueryTabState = (tabIndex: number): void => {
-  const openTabsState = readSubComponentState<OpenQueryTab[]>(
-    AppStateComponentNames.DataExplorerAction,
-    OPEN_TABS_SUBCOMPONENT_NAME,
-    undefined,
-    [],
-  );
-
-  openTabsState.splice(tabIndex, 1);
-
-  saveSubComponentState<OpenQueryTab[]>(
-    AppStateComponentNames.DataExplorerAction,
-    OPEN_TABS_SUBCOMPONENT_NAME,
-    undefined,
-    openTabsState,
-  );
-};

src/Explorer/Tabs/TabsBase.ts

@@ -1,3 +1,4 @@
+import { OpenTab } from "Contracts/ActionContracts";
 import { KeyboardActionGroup, clearKeyboardActionGroup } from "KeyboardShortcuts";
 import * as ko from "knockout";
 import * as Constants from "../../Common/Constants";
@@ -30,6 +31,8 @@ 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;
@@ -55,6 +58,10 @@ 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

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

src/Shared/AppStatePersistenceUtility.ts

@@ -12,6 +12,9 @@ export enum AppStateComponentNames {
   DataExplorerAction = "DataExplorerAction",
 }
 
+// Subcomponent for DataExplorerAction
+export const OPEN_TABS_SUBCOMPONENT_NAME = "OpenTabs";
+
 export const PATH_SEPARATOR = "/"; // export for testing purposes
 const SCHEMA_VERSION = 1;
 

src/hooks/useTabs.ts

@@ -1,5 +1,11 @@
 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";
@@ -36,6 +42,7 @@ export interface TabsState {
   selectLeftTab: () => void;
   selectRightTab: () => void;
   closeActiveTab: () => void;
+  persistTabsState: () => void;
 }
 
 export enum ReactTabKind {
@@ -73,7 +80,9 @@ 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.
@@ -130,6 +139,8 @@ export const useTabs: UseStore<TabsState> = create((set, get) => ({
     }
 
     set({ openedTabs: updatedTabs });
+
+    get().persistTabsState();
   },
   closeAllNotebookTabs: (hardClose): void => {
     const isNotebook = (tabKind: CollectionTabKind): boolean => {
@@ -226,4 +237,15 @@ 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,
+    );
+  },
 }));