feat(workspaces): Save and restore tabs COMPASS-9499

packages/atlas-service/src/atlas-service.ts

@@ -78,6 +78,11 @@ export class AtlasService {
     // https://github.com/10gen/mms/blob/9f858bb987aac6aa80acfb86492dd74c89cbb862/client/packages/project/common/ajaxPrefilter.ts#L34-L49
     return this.cloudEndpoint(path);
   }
+  userDataEndpoint(path?: string): string {
+    return `https://cluster-connection.cloud-dev.mongodb.com/userData${normalizePath(
+      path
+    )}`;
+  }
   driverProxyEndpoint(path?: string): string {
     return `${this.config.ccsBaseUrl}${normalizePath(path)}`;
   }

packages/compass-user-data/src/user-data.ts

@@ -10,7 +10,7 @@ const { log, mongoLogId } = createLogger('COMPASS-USER-STORAGE');
 
 type SerializeContent<I> = (content: I) => string;
 type DeserializeContent = (content: string) => unknown;
-type GetResourceUrl = (path?: string) => Promise<string>;
+type GetResourceUrl = (path?: string) => string;
 type AuthenticatedFetch = (
   url: RequestInfo | URL,
   options?: RequestInit
@@ -65,6 +65,7 @@ export abstract class IUserData<T extends z.Schema> {
   abstract write(id: string, content: z.input<T>): Promise<boolean>;
   abstract delete(id: string): Promise<boolean>;
   abstract readAll(options?: ReadOptions): Promise<ReadAllResult<T>>;
+  abstract readOne(id: string, options: ReadOptions): Promise<z.output<T>>;
   abstract updateAttributes(
     id: string,
     data: Partial<z.input<T>>
@@ -295,7 +296,7 @@ export class AtlasUserData<T extends z.Schema> extends IUserData<T> {
   }
 
   async write(id: string, content: z.input<T>): Promise<boolean> {
-    const url = await this.getResourceUrl(
+    const url = this.getResourceUrl(
       `${this.dataType}/${this.orgId}/${this.projectId}`
     );
 
@@ -331,7 +332,7 @@ export class AtlasUserData<T extends z.Schema> extends IUserData<T> {
   }
 
   async delete(id: string): Promise<boolean> {
-    const url = await this.getResourceUrl(
+    const url = this.getResourceUrl(
       `${this.dataType}/${this.orgId}/${this.projectId}/${id}`
     );
 
@@ -361,9 +362,7 @@ export class AtlasUserData<T extends z.Schema> extends IUserData<T> {
     };
     try {
       const response = await this.authenticatedFetch(
-        await this.getResourceUrl(
-          `${this.dataType}/${this.orgId}/${this.projectId}`
-        ),
+        this.getResourceUrl(`${this.dataType}/${this.orgId}/${this.projectId}`),
         {
           method: 'GET',
         }
@@ -396,7 +395,7 @@ export class AtlasUserData<T extends z.Schema> extends IUserData<T> {
       };
 
       await this.authenticatedFetch(
-        await this.getResourceUrl(
+        this.getResourceUrl(
           `${this.dataType}/${this.orgId}/${this.projectId}/${id}`
         ),
         {
@@ -414,7 +413,7 @@ export class AtlasUserData<T extends z.Schema> extends IUserData<T> {
         'Atlas Backend',
         'Error updating data',
         {
-          url: await this.getResourceUrl(
+          url: this.getResourceUrl(
             `${this.dataType}/${this.orgId}/${this.projectId}/${id}`
           ),
           error: (error as Error).message,
@@ -426,7 +425,7 @@ export class AtlasUserData<T extends z.Schema> extends IUserData<T> {
 
   // TODO: change this depending on whether or not updateAttributes can provide all current data
   async readOne(id: string): Promise<z.output<T>> {
-    const url = await this.getResourceUrl(
+    const url = this.getResourceUrl(
       `${this.dataType}/${this.orgId}/${this.projectId}/${id}`
     );
 

packages/compass-web/src/entrypoint.tsx

@@ -14,6 +14,8 @@ import {
   DatabasesWorkspaceTab,
   CollectionsWorkspaceTab,
 } from '@mongodb-js/compass-databases-collections';
+import { EJSON } from 'bson';
+import { atlasServiceLocator } from '@mongodb-js/atlas-service/provider';
 import { CompassComponentsProvider, css } from '@mongodb-js/compass-components';
 import {
   WorkspaceTab as CollectionWorkspace,
@@ -62,6 +64,8 @@ import {
   CompassAssistantDrawer,
   CompassAssistantProvider,
 } from '@mongodb-js/compass-assistant';
+import { AtlasUserData, type IUserData } from '@mongodb-js/compass-user-data';
+import { WorkspacesStateSchema } from '@mongodb-js/compass-workspaces';
 
 export type TrackFunction = (
   event: string,
@@ -89,7 +93,10 @@ type CompassWorkspaceProps = Pick<
   Pick<
     React.ComponentProps<typeof CompassSidebarPlugin>,
     'onOpenConnectViaModal'
-  >;
+  > & {
+    orgId: string;
+    projectId: string;
+  };
 
 type CompassWebProps = {
   /**
@@ -140,6 +147,11 @@ type CompassWebProps = {
    */
   initialPreferences?: Partial<AllPreferences>;
 
+  /**
+   * UserData instance to use for persisting workspace state
+   */
+  userData: IUserData<typeof WorkspacesStateSchema>;
+
   /**
    * Callback prop called every time any code inside Compass logs something
    */
@@ -173,7 +185,23 @@ function CompassWorkspace({
   initialWorkspaceTabs,
   onActiveWorkspaceTabChange,
   onOpenConnectViaModal,
+  orgId,
+  projectId,
 }: CompassWorkspaceProps) {
+  const atlasService = atlasServiceLocator();
+  const workspacesUserData = useRef(
+    new AtlasUserData(WorkspacesStateSchema, 'savedWorkspaces', {
+      orgId,
+      projectId,
+      getResourceUrl: (path?: string) => {
+        const url = atlasService.userDataEndpoint(`/${path || ''}`);
+        return url;
+      },
+      authenticatedFetch: atlasService.authenticatedFetch.bind(atlasService),
+      serialize: (content) => EJSON.stringify(content),
+      deserialize: (content: string) => EJSON.parse(content),
+    })
+  );
   return (
     <WorkspacesProvider
       value={[
@@ -204,6 +232,7 @@ function CompassWorkspace({
           className={connectedContainerStyles}
         >
           <WorkspacesPlugin
+            userData={workspacesUserData.current}
             initialWorkspaceTabs={initialWorkspaceTabs}
             openOnEmptyWorkspace={{ type: 'Welcome' }}
             onActiveWorkspaceTabChange={onActiveWorkspaceTabChange}
@@ -419,6 +448,8 @@ const CompassWeb = ({
                             <FieldStorePlugin>
                               <WithConnectionsStore>
                                 <CompassWorkspace
+                                  orgId={orgId}
+                                  projectId={projectId}
                                   initialWorkspaceTabs={
                                     initialWorkspaceTabsRef.current
                                   }

packages/compass-workspaces/src/components/index.tsx

@@ -10,6 +10,8 @@ import type { WorkspaceTab } from '../types';
 import Workspaces from './workspaces';
 import { connect } from '../stores/context';
 import { WorkspacesServiceProvider } from '../provider';
+import type { IUserData } from '@mongodb-js/compass-user-data';
+import type { WorkspacesStateSchema } from '../stores/workspaces-storage';
 
 type WorkspacesWithSidebarProps = {
   /**
@@ -37,6 +39,11 @@ type WorkspacesWithSidebarProps = {
    * Initial workspace tab to show (by default no tabs will be shown initially)
    */
   initialWorkspaceTabs?: OpenWorkspaceOptions[] | null;
+
+  /**
+   * UserData instance to use for persisting workspace state
+   */
+  userData: IUserData<typeof WorkspacesStateSchema>;
   /**
    * Workspace configuration to be opened when all tabs are closed (defaults to
    * "My Queries")
@@ -93,13 +100,15 @@ const WorkspacesWithSidebar: React.FunctionComponent<
   onActiveWorkspaceTabChange,
   renderSidebar,
   renderModals,
+  userData,
 }) => {
   const darkMode = useDarkMode();
   const onChange = useRef(onActiveWorkspaceTabChange);
   onChange.current = onActiveWorkspaceTabChange;
   useEffect(() => {
     onChange.current(activeTab, activeTabCollectionInfo);
   }, [activeTab, activeTabCollectionInfo]);
+
   return (
     <WorkspacesServiceProvider>
       <div
@@ -110,7 +119,10 @@ const WorkspacesWithSidebar: React.FunctionComponent<
       >
         <div className={sidebarStyles}>{renderSidebar?.()}</div>
         <div className={workspacesStyles}>
-          <Workspaces openOnEmptyWorkspace={openOnEmptyWorkspace}></Workspaces>
+          <Workspaces
+            openOnEmptyWorkspace={openOnEmptyWorkspace}
+            userData={userData}
+          ></Workspaces>
         </div>
       </div>
       {renderModals?.()}

packages/compass-workspaces/src/components/workspaces.tsx

@@ -1,10 +1,11 @@
-import React, { useCallback, useMemo } from 'react';
+import React, { useCallback, useEffect, useMemo, useRef } from 'react';
 import {
   DrawerAnchor,
   ErrorBoundary,
   MongoDBLogoMark,
   WorkspaceTabs,
   css,
+  showConfirmation,
   spacing,
   useDarkMode,
   type WorkspaceTabCoreProps,
@@ -26,13 +27,22 @@ import {
   selectNextTab,
   selectPrevTab,
   selectTab,
+  restoreWorkspaces,
 } from '../stores/workspaces';
 import { useWorkspacePlugins } from './workspaces-provider';
+import type { ConnectionInfo } from '@mongodb-js/compass-connections/provider';
+import {
+  useConnectionActions,
+  useConnectionsListRef,
+} from '@mongodb-js/compass-connections/provider';
 import toNS from 'mongodb-ns';
 import { useLogger } from '@mongodb-js/compass-logging/provider';
 import { connect } from '../stores/context';
 import { WorkspaceTabContextProvider } from './workspace-tab-context-provider';
 import type { WorkspaceTab } from '../types';
+import { loadWorkspaceStateFromUserData } from '../stores/workspaces-middleware';
+import type { IUserData } from '@mongodb-js/compass-user-data';
+import type { WorkspacesStateSchema } from '../stores/workspaces-storage';
 
 const emptyWorkspaceStyles = css({
   margin: '0 auto',
@@ -71,6 +81,7 @@ type CompassWorkspacesProps = {
   collectionInfo: Record<string, CollectionTabInfo>;
   databaseInfo: Record<string, DatabaseTabInfo>;
   openOnEmptyWorkspace?: OpenWorkspaceOptions | null;
+  userData: IUserData<typeof WorkspacesStateSchema>;
 
   onSelectTab(at: number): void;
   onSelectNextTab(): void;
@@ -84,6 +95,7 @@ type CompassWorkspacesProps = {
     tab: Extract<WorkspaceTab, { namespace: string }>,
     fallbackNamespace: string | null
   ): void;
+  onRestoreTabs(tabs: OpenWorkspaceOptions[]): void;
 };
 
 const CompassWorkspaces: React.FunctionComponent<CompassWorkspacesProps> = ({
@@ -92,6 +104,7 @@ const CompassWorkspaces: React.FunctionComponent<CompassWorkspacesProps> = ({
   collectionInfo,
   databaseInfo,
   openOnEmptyWorkspace,
+  userData,
   onSelectTab,
   onSelectNextTab,
   onSelectPrevTab,
@@ -101,6 +114,7 @@ const CompassWorkspaces: React.FunctionComponent<CompassWorkspacesProps> = ({
   onCloseTab,
   onCloseAllOtherTabs,
   onNamespaceNotFound,
+  onRestoreTabs,
 }) => {
   const { log, mongoLogId } = useLogger('COMPASS-WORKSPACES');
   const { getWorkspacePluginByName } = useWorkspacePlugins();
@@ -111,6 +125,79 @@ const CompassWorkspaces: React.FunctionComponent<CompassWorkspacesProps> = ({
     onCreateTab(openOnEmptyWorkspace);
   }, [onCreateTab, openOnEmptyWorkspace]);
 
+  const connectionActions = useConnectionActions();
+  const { getConnectionById } = useConnectionsListRef();
+  const savedWorkspacesPromiseRef = useRef(
+    loadWorkspaceStateFromUserData(userData)
+  );
+
+  useEffect(() => {
+    savedWorkspacesPromiseRef.current.then(
+      (res) => {
+        if (res !== null) {
+          showConfirmation({
+            title: 'Reopen closed tabs?',
+            description:
+              'Your connection and tabs were closed, this action will reopen your previous session',
+            buttonText: 'Reopen tabs',
+          }).then(
+            (confirm) => {
+              if (confirm) {
+                const workspacesToRestore: OpenWorkspaceOptions[] = [];
+                const connectionsToRestore: Map<string, ConnectionInfo> =
+                  new Map();
+                res.forEach((workspace) => {
+                  // If the workspace is tied to a connection, check if the connection exists
+                  // and add it to the list of connections to restore if so.
+                  if ('connectionId' in workspace) {
+                    const connectionInfo = getConnectionById(
+                      workspace.connectionId
+                    )?.info;
+
+                    if (!connectionInfo) {
+                      return;
+                    }
+
+                    connectionsToRestore.set(
+                      workspace.connectionId,
+                      connectionInfo
+                    );
+                  }
+
+                  workspacesToRestore.push(workspace);
+                });
+
+                connectionsToRestore.forEach((connectionInfo) => {
+                  void connectionActions.connect(connectionInfo);
+                });
+
+                onRestoreTabs(workspacesToRestore);
+              }
+            },
+            (err) => {
+              throw err;
+            }
+          );
+        }
+      },
+      (err) => {
+        log.error(
+          mongoLogId(1_001_000_361),
+          'Workspaces',
+          'Failed to load saved workspaces from previous session',
+          { error: err }
+        );
+      }
+    );
+  }, [
+    savedWorkspacesPromiseRef,
+    onRestoreTabs,
+    connectionActions,
+    getConnectionById,
+    log,
+    mongoLogId,
+  ]);
+
   const workspaceTabs = useMemo(() => {
     return tabs.map((tab) => {
       const plugin = getWorkspacePluginByName(tab.type);
@@ -248,5 +335,6 @@ export default connect(
     onCloseTab: closeTab,
     onCloseAllOtherTabs: closeAllOtherTabs,
     onNamespaceNotFound: openFallbackWorkspace,
+    onRestoreTabs: restoreWorkspaces,
   }
 )(CompassWorkspaces);