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-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,10 +16,12 @@ import workspacesReducer, {
   connectionDisconnected,
   updateDatabaseInfo,
   updateCollectionInfo,
+  loadSavedWorkspaces,
 } 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';
@@ -38,13 +40,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(
@@ -66,7 +74,10 @@ export function configureStore(
       collectionInfo: {},
       databaseInfo: {},
     },
-    applyMiddleware(thunk.withExtraArgument(services))
+    applyMiddleware(
+      thunk.withExtraArgument(services),
+      workspacesStateChangeMiddleware(services)
+    )
   );
 
   return store;
@@ -75,13 +86,16 @@ export function configureStore(
 export function activateWorkspacePlugin(
   {
     initialWorkspaceTabs,
-  }: { initialWorkspaceTabs?: OpenWorkspaceOptions[] | null },
+  }: {
+    initialWorkspaceTabs?: OpenWorkspaceOptions[] | null;
+  },
   {
     globalAppRegistry,
     instancesManager,
     connections,
     logger,
     preferences,
+    userData,
   }: WorkspacesServices,
   { on, cleanup, addCleanup }: ActivateHelpers
 ) {
@@ -91,8 +105,11 @@ export function activateWorkspacePlugin(
     connections,
     logger,
     preferences,
+    userData,
   });
 
+  void store.dispatch(loadSavedWorkspaces());
+
   addCleanup(cleanupLocalAppRegistries);
 
   const setupInstanceListeners = (
@@ -236,9 +253,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 };
@@ -258,3 +277,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,22 @@
+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, undefined, 2),
+        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,33 @@
+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),
+      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,102 @@
+import { createServiceLocator } from '@mongodb-js/compass-app-registry';
+import {
+  IUserData,
+  type ReadAllResult,
+  z,
+} from '@mongodb-js/compass-user-data';
+import React, { useContext, useRef } from 'react';
+
+/**
+ * Schema for saving workspace tab state to user data
+ */
+
+// Define schema for collection subtab
+const CollectionSubtabSchema = z.enum([
+  'Documents',
+  'Aggregations',
+  'Schema',
+  'Indexes',
+  'Validation',
+  'GlobalWrites',
+]);
+
+// Define schema for workspace tab type
+const WorkspaceTabTypeSchema = z.enum([
+  'Welcome',
+  'My Queries',
+  'Data Modeling',
+  'Shell',
+  'Databases',
+  'Performance',
+  'Collections',
+  'Collection',
+]);
+
+// Define schema for a workspace tab
+export const WorkspaceTabSchema = z.object({
+  id: z.string(),
+  type: WorkspaceTabTypeSchema,
+  connectionId: z.string().optional(),
+  namespace: z.string().optional(),
+  initialQuery: z.record(z.any()).optional(),
+  initialAggregation: z.record(z.any()).optional(),
+  initialPipeline: z.array(z.record(z.any())).optional(),
+  initialPipelineText: z.string().optional(),
+  editViewName: z.string().optional(),
+  initialEvaluate: z.union([z.string(), z.array(z.string())]).optional(),
+  initialInput: z.string().optional(),
+  subTab: CollectionSubtabSchema.optional(),
+});
+
+// Define schema for the complete workspaces state
+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');