Persist and restore query text, tab position and splitter direction in QueryTabComponent (#1993)

* Save query text, tab splitter direction and position in QueryTabComponent

* Fix unit tests

src/Explorer/Tabs/DocumentsTabV2/DocumentsTabStateUtil.ts

@@ -3,17 +3,11 @@
 import { ColumnDefinition } from "Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent";
 import {
   AppStateComponentNames,
-  deleteState,
-  loadState,
-  saveState,
-  saveStateDebounced,
+  deleteSubComponentState,
+  readSubComponentState,
+  saveSubComponentState,
 } 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",
@@ -30,84 +24,22 @@ export type TabDivider = { leftPaneWidthPercent: number };
 export type ColumnsSelection = { selectedColumnIds: string[]; columnDefinitions: ColumnDefinition[] };
 export type ColumnSort = { columnId: string; direction: "ascending" | "descending" };
 
-/**
- *
- * @param subComponentName
- * @param collection
- * @param defaultValue Will be returned if persisted state is not found
- * @returns
- */
-export const readSubComponentState = <T>(
+// Wrap the ...SubComponentState functions for type safety
+
+export const readDocumentsTabSubComponentState = <T>(
   subComponentName: SubComponentName,
   collection: ViewModels.CollectionBase,
   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;
-  }
+): T => readSubComponentState<T>(AppStateComponentNames.DocumentsTab, subComponentName, collection, defaultValue);
 
-  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>(
+export const saveDocumentsTabSubComponentState = <T>(
   subComponentName: SubComponentName,
   collection: ViewModels.CollectionBase,
   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;
-  }
+): void => saveSubComponentState<T>(AppStateComponentNames.DocumentsTab, subComponentName, collection, state, debounce);
 
-  (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(),
-  });
-};
+export const deleteDocumentsTabSubComponentState = (
+  subComponentName: SubComponentName,
+  collection: ViewModels.CollectionBase,
+) => deleteSubComponentState(AppStateComponentNames.DocumentsTab, subComponentName, collection);

src/Explorer/Tabs/DocumentsTabV2/DocumentsTabV2.tsx

@@ -35,8 +35,8 @@ import {
   FilterHistory,
   SubComponentName,
   TabDivider,
-  readSubComponentState,
-  saveSubComponentState,
+  readDocumentsTabSubComponentState,
+  saveDocumentsTabSubComponentState,
 } from "Explorer/Tabs/DocumentsTabV2/DocumentsTabStateUtil";
 import { usePrevious } from "Explorer/Tabs/DocumentsTabV2/SelectionHelper";
 import { CosmosFluentProvider, LayoutConstants, cosmosShorthands, tokens } from "Explorer/Theme/ThemeUtil";
@@ -619,7 +619,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
 
   // State
   const [tabStateData, setTabStateData] = useState<TabDivider>(() =>
-    readSubComponentState<TabDivider>(SubComponentName.MainTabDivider, _collection, {
+    readDocumentsTabSubComponentState<TabDivider>(SubComponentName.MainTabDivider, _collection, {
       leftPaneWidthPercent: 35,
     }),
   );
@@ -634,7 +634,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
 
   // User's filter history
   const [lastFilterContents, setLastFilterContents] = useState<FilterHistory>(() =>
-    readSubComponentState<FilterHistory>(SubComponentName.FilterHistory, _collection, [] as FilterHistory),
+    readDocumentsTabSubComponentState<FilterHistory>(SubComponentName.FilterHistory, _collection, [] as FilterHistory),
   );
 
   // For progress bar for bulk delete (noSql)
@@ -804,7 +804,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
   };
 
   const [selectedColumnIds, setSelectedColumnIds] = useState<string[]>(() => {
-    const persistedColumnsSelection = readSubComponentState<ColumnsSelection>(
+    const persistedColumnsSelection = readDocumentsTabSubComponentState<ColumnsSelection>(
       SubComponentName.ColumnsSelection,
       _collection,
       undefined,
@@ -1714,7 +1714,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
 
   // Column definition is a map<id, ColumnDefinition> to garantee uniqueness
   const [columnDefinitions, setColumnDefinitions] = useState<ColumnDefinition[]>(() => {
-    const persistedColumnsSelection = readSubComponentState<ColumnsSelection>(
+    const persistedColumnsSelection = readDocumentsTabSubComponentState<ColumnsSelection>(
       SubComponentName.ColumnsSelection,
       _collection,
       undefined,
@@ -2025,7 +2025,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
     const limitedLastFilterContents = lastFilterContents.slice(0, MAX_FILTER_HISTORY_COUNT);
 
     setLastFilterContents(limitedLastFilterContents);
-    saveSubComponentState<FilterHistory>(SubComponentName.FilterHistory, _collection, lastFilterContents);
+    saveDocumentsTabSubComponentState<FilterHistory>(SubComponentName.FilterHistory, _collection, lastFilterContents);
   };
 
   const refreshDocumentsGrid = useCallback(
@@ -2086,7 +2086,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
 
     setSelectedColumnIds(newSelectedColumnIds);
 
-    saveSubComponentState<ColumnsSelection>(SubComponentName.ColumnsSelection, _collection, {
+    saveDocumentsTabSubComponentState<ColumnsSelection>(SubComponentName.ColumnsSelection, _collection, {
       selectedColumnIds: newSelectedColumnIds,
       columnDefinitions,
     });
@@ -2214,7 +2214,7 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
           <Allotment
             onDragEnd={(sizes: number[]) => {
               tabStateData.leftPaneWidthPercent = (100 * sizes[0]) / (sizes[0] + sizes[1]);
-              saveSubComponentState<TabDivider>(SubComponentName.MainTabDivider, _collection, tabStateData);
+              saveDocumentsTabSubComponentState<TabDivider>(SubComponentName.MainTabDivider, _collection, tabStateData);
               setTabStateData(tabStateData);
             }}
           >
@@ -2316,17 +2316,15 @@ export const DocumentsTabComponent: React.FunctionComponent<IDocumentsTabCompone
                 </MessageBarBody>
               </MessageBar>
             )}
-            {bulkDeleteProcess.hasBeenThrottled && (
-              <MessageBar intent="warning">
-                <MessageBarBody>
-                  <MessageBarTitle>Warning</MessageBarTitle>
-                  {get429WarningMessageNoSql()}{" "}
-                  <Link href={NO_SQL_THROTTLING_DOC_URL} target="_blank">
-                    Learn More
-                  </Link>
-                </MessageBarBody>
-              </MessageBar>
-            )}
+            <MessageBar intent="warning">
+              <MessageBarBody>
+                <MessageBarTitle>Warning</MessageBarTitle>
+                {get429WarningMessageNoSql()}{" "}
+                <Link href={NO_SQL_THROTTLING_DOC_URL} target="_blank">
+                  Learn More
+                </Link>
+              </MessageBarBody>
+            </MessageBar>
           </div>
         </ProgressModalDialog>
       )}

src/Explorer/Tabs/DocumentsTabV2/DocumentsTableComponent.tsx

@@ -42,9 +42,9 @@ import { TableColumnSelectionPane } from "Explorer/Panes/TableColumnSelectionPan
 import {
   ColumnSizesMap,
   ColumnSort,
-  deleteSubComponentState,
-  readSubComponentState,
-  saveSubComponentState,
+  deleteDocumentsTabSubComponentState,
+  readDocumentsTabSubComponentState,
+  saveDocumentsTabSubComponentState,
   SubComponentName,
 } from "Explorer/Tabs/DocumentsTabV2/DocumentsTabStateUtil";
 import { INITIAL_SELECTED_ROW_INDEX, useDocumentsTabStyles } from "Explorer/Tabs/DocumentsTabV2/DocumentsTabV2";
@@ -116,7 +116,11 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
   const styles = useDocumentsTabStyles();
 
   const [columnSizingOptions, setColumnSizingOptions] = React.useState<TableColumnSizingOptions>(() => {
-    const columnSizesMap: ColumnSizesMap = readSubComponentState(SubComponentName.ColumnSizes, collection, {});
+    const columnSizesMap: ColumnSizesMap = readDocumentsTabSubComponentState(
+      SubComponentName.ColumnSizes,
+      collection,
+      {},
+    );
     const columnSizesPx: TableColumnSizingOptions = {};
     selectedColumnIds.forEach((columnId) => {
       if (
@@ -140,7 +144,7 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
     sortDirection: "ascending" | "descending";
     sortColumn: TableColumnId | undefined;
   }>(() => {
-    const sort = readSubComponentState<ColumnSort>(SubComponentName.ColumnSort, collection, undefined);
+    const sort = readDocumentsTabSubComponentState<ColumnSort>(SubComponentName.ColumnSort, collection, undefined);
 
     if (!sort) {
       return {
@@ -172,7 +176,12 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
         return acc;
       }, {} as ColumnSizesMap);
 
-      saveSubComponentState<ColumnSizesMap>(SubComponentName.ColumnSizes, collection, persistentSizes, true);
+      saveDocumentsTabSubComponentState<ColumnSizesMap>(
+        SubComponentName.ColumnSizes,
+        collection,
+        persistentSizes,
+        true,
+      );
 
       return newSizingOptions;
     });
@@ -184,11 +193,14 @@ export const DocumentsTableComponent: React.FC<IDocumentsTableComponentProps> =
     setColumnSort(event, columnId, direction);
 
     if (columnId === undefined || direction === undefined) {
-      deleteSubComponentState(SubComponentName.ColumnSort, collection);
+      deleteDocumentsTabSubComponentState(SubComponentName.ColumnSort, collection);
       return;
     }
 
-    saveSubComponentState<ColumnSort>(SubComponentName.ColumnSort, collection, { columnId, direction });
+    saveDocumentsTabSubComponentState<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/QueryTab/QueryTabComponent.test.tsx

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

src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx

@@ -13,6 +13,11 @@ 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 {
+  SubComponentName,
+  readQueryTabSubComponentState,
+  saveQueryTabSubComponentState,
+} from "Explorer/Tabs/QueryTab/QueryTabStateUtil";
 import { QueryTabStyles, useQueryTabStyles } from "Explorer/Tabs/QueryTab/Styles";
 import { CosmosFluentProvider } from "Explorer/Theme/ThemeUtil";
 import { useSelectedNode } from "Explorer/useSelectedNode";
@@ -118,6 +123,7 @@ interface IQueryTabStates {
   queryResultsView: SplitterDirection;
   errors?: QueryError[];
   modelMarkers?: monaco.editor.IMarkerData[];
+  queryViewSizePercent: number;
 }
 
 export const QueryTabCopilotComponent = (props: IQueryTabComponentProps): any => {
@@ -165,7 +171,7 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
 
     this.state = {
       toggleState: ToggleState.Result,
-      sqlQueryEditorContent: props.isPreferredApiMongoDB ? "{}" : props.queryText || "SELECT * FROM c",
+      sqlQueryEditorContent: this._getDefaultQueryEditorContent(props),
       selectedContent: "",
       queryResults: undefined,
       errors: [],
@@ -176,7 +182,8 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
       cancelQueryTimeoutID: undefined,
       copilotActive: this._queryCopilotActive(),
       currentTabActive: true,
-      queryResultsView: getDefaultQueryResultsView(),
+      queryResultsView: this._getDefaultQUeryResultsViewDirection(props),
+      queryViewSizePercent: this._getQueryViewSizePercent(props),
     };
     this.isCloseClicked = false;
     this.splitterId = this.props.tabId + "_splitter";
@@ -207,6 +214,25 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
     });
   }
 
+  private _getQueryViewSizePercent(props: QueryTabComponentImplProps): number {
+    return readQueryTabSubComponentState<number>(SubComponentName.QueryViewSizePercent, props.collection, 50);
+  }
+
+  private _getDefaultQUeryResultsViewDirection(props: QueryTabComponentImplProps): SplitterDirection {
+    const defaultQueryResultsView = getDefaultQueryResultsView();
+    return readQueryTabSubComponentState<SplitterDirection>(
+      SubComponentName.SplitterDirection,
+      props.collection,
+      defaultQueryResultsView,
+    );
+  }
+
+  private _getDefaultQueryEditorContent(props: QueryTabComponentImplProps): string {
+    const defaultText = props.isPreferredApiMongoDB ? "{}" : props.queryText || "SELECT * FROM c";
+    // Retrieve from app state if available
+    return readQueryTabSubComponentState<string>(SubComponentName.QueryText, props.collection, defaultText);
+  }
+
   private _queryCopilotActive(): boolean {
     if (this.props.copilotEnabled) {
       return readCopilotToggleStatus(userContext.databaseAccount);
@@ -569,6 +595,13 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
   private _setViewLayout(direction: SplitterDirection): void {
     this.setState({ queryResultsView: direction });
 
+    // Store to local storage
+    saveQueryTabSubComponentState<SplitterDirection>(
+      SubComponentName.SplitterDirection,
+      this.props.collection,
+      direction,
+    );
+
     // We'll need to refresh the context buttons to update the selected state of the view buttons
     setTimeout(() => {
       useCommandBar.getState().setContextButtons(this.getTabsButtons());
@@ -623,6 +656,8 @@ class QueryTabComponentImpl extends React.Component<QueryTabComponentImplProps,
     this.saveQueryButton.enabled = newContent.length > 0;
 
     useCommandBar.getState().setContextButtons(this.getTabsButtons());
+
+    saveQueryTabSubComponentState<string>(SubComponentName.QueryText, this.props.collection, newContent, true);
   }
 
   public onSelectedContent(selectedContent: string, selection: monaco.Selection): void {
@@ -704,8 +739,21 @@ 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}>
-            <Allotment.Pane data-test="QueryTab/EditorPane">
+          <Allotment
+            key={vertical.toString()}
+            vertical={vertical}
+            onDragEnd={(sizes: number[]) => {
+              const queryViewSizePercent = (100 * sizes[0]) / (sizes[0] + sizes[1]);
+              saveQueryTabSubComponentState<number>(
+                SubComponentName.QueryViewSizePercent,
+                this.props.collection,
+                queryViewSizePercent,
+                true,
+              );
+              this.setState({ queryViewSizePercent });
+            }}
+          >
+            <Allotment.Pane data-test="QueryTab/EditorPane" preferredSize={`${this.state.queryViewSizePercent}%`}>
               <EditorReact
                 ref={this.queryEditor}
                 className={this.props.styles.queryEditor}

src/Explorer/Tabs/QueryTab/QueryTabStateUtil.ts

@@ -0,0 +1,37 @@
+// Definitions of State data
+
+import {
+  AppStateComponentNames,
+  deleteSubComponentState,
+  readSubComponentState,
+  saveSubComponentState,
+} from "Shared/AppStatePersistenceUtility";
+import * as ViewModels from "../../../Contracts/ViewModels";
+
+export enum SubComponentName {
+  SplitterDirection = "SplitterDirection",
+  QueryViewSizePercent = "QueryViewSizePercent",
+  QueryText = "QueryText",
+}
+
+export type QueryViewSizePercent = number;
+export type QueryText = string;
+
+// Wrap the ...SubComponentState functions for type safety
+export const readQueryTabSubComponentState = <T>(
+  subComponentName: SubComponentName,
+  collection: ViewModels.CollectionBase,
+  defaultValue: T,
+): T => readSubComponentState<T>(AppStateComponentNames.QueryTab, subComponentName, collection, defaultValue);
+
+export const saveQueryTabSubComponentState = <T>(
+  subComponentName: SubComponentName,
+  collection: ViewModels.CollectionBase,
+  state: T,
+  debounce?: boolean,
+): void => saveSubComponentState<T>(AppStateComponentNames.QueryTab, subComponentName, collection, state, debounce);
+
+export const deleteQueryTabSubComponentState = (
+  subComponentName: SubComponentName,
+  collection: ViewModels.CollectionBase,
+) => deleteSubComponentState(AppStateComponentNames.QueryTab, subComponentName, collection);

src/Shared/AppStatePersistenceUtility.ts

@@ -1,10 +1,15 @@
 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",
+  QueryTab = "QueryTab",
 }
 
 export const PATH_SEPARATOR = "/"; // export for testing purposes
@@ -112,3 +117,93 @@ 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,
+  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.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>(
+  componentName: AppStateComponentNames,
+  subComponentName: string,
+  collection: ViewModels.CollectionBase,
+  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.databaseId,
+      containerName: 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(),
+  });
+};