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

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

@@ -81,7 +81,11 @@ export class AtlasService {
   userDataEndpoint(
     orgId: string,
     groupId: string,
-    type: 'favoriteQueries' | 'recentQueries' | 'favoriteAggregations',
+    type:
+      | 'favoriteQueries'
+      | 'recentQueries'
+      | 'favoriteAggregations'
+      | 'savedWorkspaces',
     id?: string
   ): string {
     const encodedOrgId = encodeURIComponent(orgId);

packages/compass-preferences-model/src/feature-flags.ts

@@ -32,6 +32,7 @@ export type FeatureFlags = {
   enableAIAssistant: boolean;
   enablePerformanceInsightsEntrypoints: boolean;
   enableAutomaticRelationshipInference: boolean;
+  enableRestoreWorkspaces: boolean;
 };
 
 export const featureFlags: Required<{
@@ -198,4 +199,11 @@ export const featureFlags: Required<{
         'Enable automatic relationship inference during data model generation',
     },
   },
+
+  enableRestoreWorkspaces: {
+    stage: 'development',
+    description: {
+      short: 'Enable restoring previous workspace tabs on startup',
+    },
+  },
 };

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

@@ -1,3 +1,3 @@
 export type { ReadAllResult } from './user-data';
-export { type IUserData, FileUserData, AtlasUserData } from './user-data';
+export { IUserData, FileUserData, AtlasUserData } from './user-data';
 export { z } from 'zod';

packages/compass-web/src/entrypoint.tsx

@@ -13,6 +13,7 @@ import type {
 } from '@mongodb-js/compass-workspaces';
 import WorkspacesPlugin, {
   WorkspacesProvider,
+  WorkspacesStorageServiceProviderWeb,
 } from '@mongodb-js/compass-workspaces';
 import {
   CollectionsWorkspaceTab,
@@ -122,7 +123,8 @@ const WithStorageProviders = createServiceProvider(
       const type = pathParts[0] as
         | 'favoriteQueries'
         | 'recentQueries'
-        | 'favoriteAggregations';
+        | 'favoriteAggregations'
+        | 'savedWorkspaces';
       const pathOrgId = pathParts[1];
       const pathProjectId = pathParts[2];
       const id = pathParts[3];
@@ -173,7 +175,14 @@ const WithStorageProviders = createServiceProvider(
       <PipelineStorageProvider value={pipelineStorage.current}>
         <FavoriteQueryStorageProvider value={favoriteQueryStorage.current}>
           <RecentQueryStorageProvider value={recentQueryStorage.current}>
-            {children}
+            <WorkspacesStorageServiceProviderWeb
+              orgId={orgId}
+              projectId={projectId}
+              getResourceUrl={getResourceUrl}
+              authenticatedFetch={authenticatedFetch}
+            >
+              {children}
+            </WorkspacesStorageServiceProviderWeb>
           </RecentQueryStorageProvider>
         </FavoriteQueryStorageProvider>
       </PipelineStorageProvider>

packages/compass-workspaces/src/index.ts

@@ -16,11 +16,13 @@ import workspacesReducer, {
   connectionDisconnected,
   updateDatabaseInfo,
   updateCollectionInfo,
+  loadSavedWorkspaces,
   beforeUnloading,
 } from './stores/workspaces';
 import Workspaces from './components';
 import { applyMiddleware, createStore } from 'redux';
 import thunk from 'redux-thunk';
+import { workspacesStateChangeMiddleware } from './stores/workspaces-middleware';
 import type { MongoDBInstance } from '@mongodb-js/compass-app-stores/provider';
 import { mongoDBInstancesManagerLocator } from '@mongodb-js/compass-app-stores/provider';
 import type Collection from 'mongodb-collection-model';
@@ -39,13 +41,19 @@ import {
 } from '@mongodb-js/compass-app-stores/provider';
 import type { PreferencesAccess } from 'compass-preferences-model/provider';
 import { preferencesLocator } from 'compass-preferences-model/provider';
+import {
+  type WorkspacesStateSchema,
+  workspacesStorageServiceLocator,
+} from './services/workspaces-storage';
+import { type IUserData } from '../../compass-user-data/dist/user-data';
 
 export type WorkspacesServices = {
   globalAppRegistry: AppRegistry;
   instancesManager: MongoDBInstancesManager;
   connections: ConnectionsService;
   logger: Logger;
   preferences: PreferencesAccess;
+  userData: IUserData<typeof WorkspacesStateSchema>;
 };
 
 export function configureStore(
@@ -67,7 +75,10 @@ export function configureStore(
       collectionInfo: {},
       databaseInfo: {},
     },
-    applyMiddleware(thunk.withExtraArgument(services))
+    applyMiddleware(
+      thunk.withExtraArgument(services),
+      workspacesStateChangeMiddleware(services)
+    )
   );
 
   return store;
@@ -87,6 +98,7 @@ export function activateWorkspacePlugin(
     connections,
     logger,
     preferences,
+    userData,
   }: WorkspacesServices,
   { on, cleanup, addCleanup }: ActivateHelpers
 ) {
@@ -96,8 +108,11 @@ export function activateWorkspacePlugin(
     connections,
     logger,
     preferences,
+    userData,
   });
 
+  void store.dispatch(loadSavedWorkspaces());
+
   addCleanup(cleanupLocalAppRegistries);
 
   const setupInstanceListeners = (
@@ -230,9 +245,11 @@ const WorkspacesPlugin = registerCompassPlugin(
     connections: connectionsLocator,
     logger: createLoggerLocator('COMPASS-WORKSPACES-UI'),
     preferences: preferencesLocator,
+    userData: workspacesStorageServiceLocator,
   }
 );
 
+export { WorkspacesStateSchema } from './services/workspaces-storage';
 export default WorkspacesPlugin;
 export { WorkspacesProvider } from './components/workspaces-provider';
 export type { OpenWorkspaceOptions, CollectionTabInfo };
@@ -252,3 +269,6 @@ export type {
   CollectionSubtab,
   WorkspacePluginProps,
 } from './types';
+
+export { WorkspacesStorageServiceProviderDesktop } from './services/workspaces-storage-desktop';
+export { WorkspacesStorageServiceProviderWeb } from './services/workspaces-storage-web';

packages/compass-workspaces/src/services/index.ts

@@ -0,0 +1,2 @@
+export { WorkspacesStorageServiceProviderDesktop } from './workspaces-storage-desktop';
+export { WorkspacesStorageServiceProviderWeb } from './workspaces-storage-web';

packages/compass-workspaces/src/services/workspaces-storage-desktop.tsx

@@ -0,0 +1,25 @@
+import React, { useRef } from 'react';
+import { FileUserData, type IUserData } from '@mongodb-js/compass-user-data';
+import {
+  WorkspacesStateSchema,
+  WorkspacesStorageServiceContext,
+} from './workspaces-storage';
+import { EJSON } from 'bson';
+
+export const WorkspacesStorageServiceProviderDesktop: React.FunctionComponent =
+  ({ children }) => {
+    const storageRef = useRef<IUserData<typeof WorkspacesStateSchema>>(
+      new FileUserData(WorkspacesStateSchema, 'WorkspacesState', {
+        serialize: (content) =>
+          EJSON.stringify(content, {
+            relaxed: false,
+          }),
+        deserialize: (content: string) => EJSON.parse(content),
+      }) as IUserData<typeof WorkspacesStateSchema>
+    );
+    return (
+      <WorkspacesStorageServiceContext.Provider value={storageRef.current}>
+        {children}
+      </WorkspacesStorageServiceContext.Provider>
+    );
+  };

packages/compass-workspaces/src/services/workspaces-storage-web.tsx

@@ -0,0 +1,36 @@
+import React, { useRef } from 'react';
+import { AtlasUserData, type IUserData } from '@mongodb-js/compass-user-data';
+import {
+  WorkspacesStateSchema,
+  WorkspacesStorageServiceContext,
+} from './workspaces-storage';
+import { EJSON } from 'bson';
+
+export const WorkspacesStorageServiceProviderWeb: React.FunctionComponent<{
+  orgId: string;
+  projectId: string;
+  getResourceUrl: (path?: string) => string;
+  authenticatedFetch: (
+    url: RequestInfo | URL,
+    options?: RequestInit
+  ) => Promise<Response>;
+}> = ({ orgId, projectId, getResourceUrl, authenticatedFetch, children }) => {
+  const storageRef = useRef<IUserData<typeof WorkspacesStateSchema>>(
+    new AtlasUserData(WorkspacesStateSchema, 'WorkspacesState', {
+      orgId,
+      projectId,
+      getResourceUrl,
+      authenticatedFetch,
+      serialize: (content) =>
+        EJSON.stringify(content, {
+          relaxed: false,
+        }),
+      deserialize: (content: string) => EJSON.parse(content),
+    })
+  );
+  return (
+    <WorkspacesStorageServiceContext.Provider value={storageRef.current}>
+      {children}
+    </WorkspacesStorageServiceContext.Provider>
+  );
+};

packages/compass-workspaces/src/services/workspaces-storage.tsx

@@ -0,0 +1,112 @@
+import { createServiceLocator } from '@mongodb-js/compass-app-registry';
+import {
+  IUserData,
+  type ReadAllResult,
+  z,
+} from '@mongodb-js/compass-user-data';
+import React, { useContext } from 'react';
+import { collectionSubtabValues } from '../types';
+
+const CollectionSubtabSchema = z.enum(collectionSubtabValues);
+
+export const WorkspaceTabSchema = z
+  .discriminatedUnion('type', [
+    z.object({
+      type: z.literal('Welcome'),
+    }),
+    z.object({
+      type: z.literal('My Queries'),
+    }),
+    z.object({
+      type: z.literal('Data Modeling'),
+    }),
+    z.object({
+      type: z.literal('Databases'),
+      connectionId: z.string(),
+    }),
+    z.object({
+      type: z.literal('Performance'),
+      connectionId: z.string(),
+    }),
+    z.object({
+      type: z.literal('Shell'),
+      connectionId: z.string(),
+      initialEvaluate: z.union([z.string(), z.array(z.string())]).optional(),
+      initialInput: z.string().optional(),
+    }),
+    z.object({
+      type: z.literal('Collections'),
+      connectionId: z.string(),
+      namespace: z.string(),
+      inferredFromPrivileges: z.boolean().optional(),
+    }),
+    z.object({
+      type: z.literal('Collection'),
+      subTab: CollectionSubtabSchema,
+      initialQuery: z.record(z.any()).optional(),
+      initialPipeline: z.array(z.record(z.any())).optional(),
+      initialPipelineText: z.string().optional(),
+      initialAggregation: z.record(z.any()).optional(),
+      editViewName: z.string().optional(),
+      connectionId: z.string(),
+      namespace: z.string(),
+      inferredFromPrivileges: z.boolean().optional(),
+    }),
+  ])
+  .and(
+    z.object({
+      id: z.string(),
+    })
+  );
+
+export const WorkspacesStateSchema = z.object({
+  tabs: z.array(WorkspaceTabSchema),
+  activeTabId: z.string().nullable(),
+  timestamp: z.number(),
+});
+
+// TypeScript types derived from the schemas
+export type WorkspaceTabData = z.output<typeof WorkspaceTabSchema>;
+export type WorkspacesStateData = z.output<typeof WorkspacesStateSchema>;
+
+const throwIfNotTestEnv = () => {
+  if (process.env.NODE_ENV !== 'test') {
+    throw new Error("Can't find Workspaces storage service in React context");
+  }
+};
+
+export class noopUserData<T extends z.Schema> extends IUserData<T> {
+  write(): Promise<boolean> {
+    throwIfNotTestEnv();
+    return Promise.resolve(true);
+  }
+  delete(): Promise<boolean> {
+    throwIfNotTestEnv();
+    return Promise.resolve(true);
+  }
+  readAll(): Promise<ReadAllResult<T>> {
+    throwIfNotTestEnv();
+    return Promise.resolve({ data: [], errors: [] });
+  }
+  readOne(): Promise<z.output<T>> {
+    throwIfNotTestEnv();
+    return Promise.resolve(undefined);
+  }
+  updateAttributes(): Promise<boolean> {
+    throwIfNotTestEnv();
+    return Promise.resolve(true);
+  }
+}
+
+export const noopWorkspacesStorageService: IUserData<
+  typeof WorkspacesStateSchema
+> = new noopUserData(WorkspacesStateSchema, 'WorkspacesState');
+
+export const WorkspacesStorageServiceContext = React.createContext<
+  IUserData<typeof WorkspacesStateSchema>
+>(noopWorkspacesStorageService);
+
+export const workspacesStorageServiceLocator = createServiceLocator(() => {
+  const service = useContext(WorkspacesStorageServiceContext);
+  return service;
+}, 'workspacesStorageServiceLocator');