diff --git a/packages/atlas-service/src/atlas-service.ts b/packages/atlas-service/src/atlas-service.ts index 183a42012cb..87167baa42c 100644 --- a/packages/atlas-service/src/atlas-service.ts +++ b/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); diff --git a/packages/compass-preferences-model/src/feature-flags.ts b/packages/compass-preferences-model/src/feature-flags.ts index 5d97e8ef0ca..e089ff1a406 100644 --- a/packages/compass-preferences-model/src/feature-flags.ts +++ b/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', + }, + }, }; diff --git a/packages/compass-user-data/src/index.ts b/packages/compass-user-data/src/index.ts index 8d8ae1df65e..d7f7b5dd09b 100644 --- a/packages/compass-user-data/src/index.ts +++ b/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'; diff --git a/packages/compass-web/src/entrypoint.tsx b/packages/compass-web/src/entrypoint.tsx index 888e0574452..eecb7f006e5 100644 --- a/packages/compass-web/src/entrypoint.tsx +++ b/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( - {children} + + {children} + diff --git a/packages/compass-workspaces/src/index.ts b/packages/compass-workspaces/src/index.ts index 6a3dccbd716..ff6dd635d1e 100644 --- a/packages/compass-workspaces/src/index.ts +++ b/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,6 +41,11 @@ 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 '@mongodb-js/compass-user-data'; export type WorkspacesServices = { globalAppRegistry: AppRegistry; @@ -46,6 +53,7 @@ export type WorkspacesServices = { connections: ConnectionsService; logger: Logger; preferences: PreferencesAccess; + userData: IUserData; }; 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'; diff --git a/packages/compass-workspaces/src/services/index.ts b/packages/compass-workspaces/src/services/index.ts new file mode 100644 index 00000000000..bfcd01dd20e --- /dev/null +++ b/packages/compass-workspaces/src/services/index.ts @@ -0,0 +1,2 @@ +export { WorkspacesStorageServiceProviderDesktop } from './workspaces-storage-desktop'; +export { WorkspacesStorageServiceProviderWeb } from './workspaces-storage-web'; diff --git a/packages/compass-workspaces/src/services/workspaces-storage-desktop.tsx b/packages/compass-workspaces/src/services/workspaces-storage-desktop.tsx new file mode 100644 index 00000000000..aef7401d2dd --- /dev/null +++ b/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>( + new FileUserData(WorkspacesStateSchema, 'WorkspacesState', { + serialize: (content) => + EJSON.stringify(content, { + relaxed: false, + }), + deserialize: (content: string) => EJSON.parse(content), + }) as IUserData + ); + return ( + + {children} + + ); + }; diff --git a/packages/compass-workspaces/src/services/workspaces-storage-web.tsx b/packages/compass-workspaces/src/services/workspaces-storage-web.tsx new file mode 100644 index 00000000000..621b428e9e5 --- /dev/null +++ b/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; +}> = ({ orgId, projectId, getResourceUrl, authenticatedFetch, children }) => { + const storageRef = useRef>( + new AtlasUserData(WorkspacesStateSchema, 'WorkspacesState', { + orgId, + projectId, + getResourceUrl, + authenticatedFetch, + serialize: (content) => + EJSON.stringify(content, { + relaxed: false, + }), + deserialize: (content: string) => EJSON.parse(content), + }) + ); + return ( + + {children} + + ); +}; diff --git a/packages/compass-workspaces/src/services/workspaces-storage.tsx b/packages/compass-workspaces/src/services/workspaces-storage.tsx new file mode 100644 index 00000000000..21fca25891c --- /dev/null +++ b/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; +export type WorkspacesStateData = z.output; + +const throwIfNotTestEnv = () => { + if (process.env.NODE_ENV !== 'test') { + throw new Error("Can't find Workspaces storage service in React context"); + } +}; + +export class noopUserData extends IUserData { + write(): Promise { + throwIfNotTestEnv(); + return Promise.resolve(true); + } + delete(): Promise { + throwIfNotTestEnv(); + return Promise.resolve(true); + } + readAll(): Promise> { + throwIfNotTestEnv(); + return Promise.resolve({ data: [], errors: [] }); + } + readOne(): Promise> { + throwIfNotTestEnv(); + return Promise.resolve(undefined); + } + updateAttributes(): Promise { + throwIfNotTestEnv(); + return Promise.resolve(true); + } +} + +export const noopWorkspacesStorageService: IUserData< + typeof WorkspacesStateSchema +> = new noopUserData(WorkspacesStateSchema, 'WorkspacesState'); + +export const WorkspacesStorageServiceContext = React.createContext< + IUserData +>(noopWorkspacesStorageService); + +export const workspacesStorageServiceLocator = createServiceLocator(() => { + const service = useContext(WorkspacesStorageServiceContext); + return service; +}, 'workspacesStorageServiceLocator'); diff --git a/packages/compass-workspaces/src/stores/workspaces-middleware.ts b/packages/compass-workspaces/src/stores/workspaces-middleware.ts new file mode 100644 index 00000000000..21c2774fccc --- /dev/null +++ b/packages/compass-workspaces/src/stores/workspaces-middleware.ts @@ -0,0 +1,139 @@ +import type { Middleware, AnyAction } from 'redux'; +import type { WorkspacesState } from './workspaces'; +import type { + WorkspacesStateData, + WorkspaceTabData, +} from '../services/workspaces-storage'; +import type { WorkspacesServices } from '..'; +import { mongoLogId } from '@mongodb-js/compass-logging/provider'; + +/** + * Debounced handler to save the workspaces state. + */ +const handleWorkspacesStateChange = (() => { + let timeoutId: ReturnType; + return (state: WorkspacesState, services: WorkspacesServices) => { + clearTimeout(timeoutId); + timeoutId = setTimeout(() => { + // Fire and forget - don't await to avoid blocking the action + void saveWorkspaceStateToUserData(state, services); + }, 250); + }; +})(); + +/** + * Middleware that runs a callback whenever the workspaces state changes. + * This allows you to perform side effects when the state is updated. + */ +export function workspacesStateChangeMiddleware( + services: WorkspacesServices +): Middleware, WorkspacesState> { + return (store) => (next) => (action: AnyAction) => { + const prevState = store.getState(); + const result = next(action); + const nextState = store.getState(); + + // Only call the callback if the workspaces state actually changed + if (prevState !== nextState) { + if (services.preferences.getPreferences().enableRestoreWorkspaces) { + handleWorkspacesStateChange(nextState, services); + } + } + + return result; + }; +} + +/** + * Saves the workspace state to persistent storage using UserData + */ +async function saveWorkspaceStateToUserData( + state: WorkspacesState, + services: WorkspacesServices +) { + try { + // Transform the state to the format we want to save + const stateToSave: WorkspacesStateData = { + tabs: state.tabs.map((tab) => { + const { type, id } = tab; + + switch (type) { + case 'Welcome': + case 'My Queries': + case 'Data Modeling': + return { id, type }; + case 'Databases': + case 'Performance': + return { + id, + type, + connectionId: tab.connectionId, + }; + case 'Collections': + return { + id, + type, + connectionId: tab.connectionId, + namespace: tab.namespace, + }; + case 'Shell': { + const result: WorkspaceTabData = { + id, + type, + connectionId: tab.connectionId, + }; + if ('initialEvaluate' in tab) { + result.initialEvaluate = tab.initialEvaluate; + } + if ('initialInput' in tab) { + result.initialInput = tab.initialInput; + } + return result; + } + case 'Collection': { + const result: WorkspaceTabData = { + id, + type, + connectionId: tab.connectionId, + namespace: tab.namespace, + subTab: tab.subTab, + }; + if ('initialQuery' in tab) { + result.initialQuery = tab.initialQuery as Record; + } + if ('initialAggregation' in tab) { + result.initialAggregation = tab.initialAggregation as Record< + string, + unknown + >; + } + if ('editViewName' in tab) { + result.editViewName = tab.editViewName; + } + if ('initialPipeline' in tab) { + result.initialPipeline = tab.initialPipeline as Array< + Record + >; + } + if ('initialPipelineText' in tab) { + result.initialPipelineText = tab.initialPipelineText; + } + return result; + } + } + }), + activeTabId: state.activeTabId, + timestamp: Date.now(), + }; + + // Save to UserData with a fixed ID + await services.userData.write('saved-workspaces', stateToSave); + } catch (error) { + services.logger.log.error( + mongoLogId(1_001_000_229), + 'Workspaces middleware', + 'Failed to save workspace state to UserData', + { error } + ); + } +} diff --git a/packages/compass-workspaces/src/stores/workspaces.ts b/packages/compass-workspaces/src/stores/workspaces.ts index 6367f9f2c34..77b6da1dc16 100644 --- a/packages/compass-workspaces/src/stores/workspaces.ts +++ b/packages/compass-workspaces/src/stores/workspaces.ts @@ -14,6 +14,7 @@ import { } from '../components/workspace-close-handler'; import { type ConnectionInfo } from '@mongodb-js/compass-connections/provider'; import { showConfirmation } from '@mongodb-js/compass-components'; +import type { WorkspacesStateData } from '../services/workspaces-storage'; const LocalAppRegistryMap = new Map(); @@ -55,6 +56,7 @@ export enum WorkspacesActions { SelectNextTab = 'compass-workspaces/SelectNextTab', MoveTab = 'compass-workspaces/MoveTab', OpenTabFromCurrentActive = 'compass-workspaces/OpenTabFromCurrentActive', + RestoreWorkspaces = 'compass-workspaces/RestoreWorkspaces', DuplicateTab = 'compass-workspaces/DuplicateTab', CloseTabs = 'compass-workspaces/CloseTabs', CollectionRenamed = 'compass-workspaces/CollectionRenamed', @@ -347,6 +349,18 @@ const reducer: Reducer = ( }; } + if ( + isAction( + action, + WorkspacesActions.RestoreWorkspaces + ) + ) { + return { + ...state, + tabs: [...state.tabs, ...action.tabs.map(getInitialTabState)], + }; + } + if ( isAction( action, @@ -853,6 +867,11 @@ type OpenTabFromCurrentActiveAction = { defaultTab: OpenWorkspaceOptions; }; +type RestoreWorkspacesAction = { + type: WorkspacesActions.RestoreWorkspaces; + tabs: OpenWorkspaceOptions[]; +}; + export const openTabFromCurrent = ( defaultTab?: OpenWorkspaceOptions | null ): OpenTabFromCurrentActiveAction => { @@ -973,6 +992,127 @@ type DatabaseRemovedAction = { namespace: string; }; +/** + * Converts saved workspace state data back to OpenWorkspaceOptions format + * for initializing the store + */ +function convertSavedStateToOpenWorkspaceOptions( + savedState: WorkspacesStateData +): OpenWorkspaceOptions[] { + return savedState.tabs.map((tab) => { + const type = tab.type; + + switch (type) { + case 'Welcome': + case 'My Queries': + case 'Data Modeling': + return { type }; + case 'Databases': + case 'Performance': + return { + type, + connectionId: tab.connectionId, + }; + case 'Collections': + return { + type, + connectionId: tab.connectionId, + namespace: tab.namespace, + }; + case 'Shell': { + const result: OpenWorkspaceOptions = { + type, + connectionId: tab.connectionId, + }; + if ('initialEvaluate' in tab) { + result.initialEvaluate = tab.initialEvaluate; + } + if ('initialInput' in tab) { + result.initialInput = tab.initialInput; + } + return result; + } + case 'Collection': { + const result: OpenWorkspaceOptions = { + type, + connectionId: tab.connectionId, + namespace: tab.namespace, + }; + if ('subTab' in tab) { + result.initialSubtab = tab.subTab; + } + if ('initialQuery' in tab) { + result.initialQuery = tab.initialQuery; + } + + if ('initialAggregation' in tab) { + result.initialAggregation = tab.initialAggregation; + } + if ('editViewName' in tab) { + result.editViewName = tab.editViewName; + } + if ('initialPipeline' in tab) { + result.initialPipeline = tab.initialPipeline; + } + if ('initialPipelineText' in tab) { + result.initialPipelineText = tab.initialPipelineText; + } + return result; + } + } + }); +} + +export const loadSavedWorkspaces = (): WorkspacesThunkAction< + Promise, + RestoreWorkspacesAction +> => { + return async (dispatch, getState, { connections, userData, preferences }) => { + if (!preferences.getPreferences().enableRestoreWorkspaces) { + return; + } + const savedState = await userData.readOne('saved-workspaces', { + ignoreErrors: true, + }); + if (savedState && savedState.tabs.length > 0) { + const confirm = await showConfirmation({ + title: 'Reopen closed tabs?', + description: + 'Your connection and tabs were closed, this action will reopen your previous session', + buttonText: 'Reopen tabs', + }); + + const workspacesToRestore: OpenWorkspaceOptions[] = []; + if (confirm) { + for (const workspace of convertSavedStateToOpenWorkspaceOptions( + savedState + )) { + // 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 = connections.getConnectionById( + workspace.connectionId + )?.info; + + if (!connectionInfo) { + return; + } + + void connections.connect(connectionInfo); + } + + workspacesToRestore.push(workspace); + } + } + + dispatch({ + type: WorkspacesActions.RestoreWorkspaces, + tabs: workspacesToRestore, + }); + } + }; +}; + export const databaseRemoved = ( namespace: string ): WorkspacesThunkAction => { diff --git a/packages/compass-workspaces/src/types.ts b/packages/compass-workspaces/src/types.ts index 4c9c5ebaf62..59408589534 100644 --- a/packages/compass-workspaces/src/types.ts +++ b/packages/compass-workspaces/src/types.ts @@ -1,13 +1,16 @@ import type { CompassPluginComponent } from '@mongodb-js/compass-app-registry'; import type { WorkspaceTabCoreProps } from '@mongodb-js/compass-components'; -export type CollectionSubtab = - | 'Documents' - | 'Aggregations' - | 'Schema' - | 'Indexes' - | 'Validation' - | 'GlobalWrites'; +export const collectionSubtabValues = [ + 'Documents', + 'Aggregations', + 'Schema', + 'Indexes', + 'Validation', + 'GlobalWrites', +] as const; + +export type CollectionSubtab = (typeof collectionSubtabValues)[number]; export type WelcomeWorkspace = { type: 'Welcome'; diff --git a/packages/compass/src/app/components/entrypoint.tsx b/packages/compass/src/app/components/entrypoint.tsx index 0d20b5fc483..ec6911bdd65 100644 --- a/packages/compass/src/app/components/entrypoint.tsx +++ b/packages/compass/src/app/components/entrypoint.tsx @@ -31,6 +31,7 @@ import { createIpcSendTrack, } from '@mongodb-js/compass-telemetry'; import { DataModelStorageServiceProviderElectron } from '@mongodb-js/compass-data-modeling/renderer'; +import { WorkspacesStorageServiceProviderDesktop } from '@mongodb-js/compass-workspaces'; const WithPreferencesAndLoggerProviders: React.FC = ({ children }) => { const loggerProviderValue = useRef({ @@ -89,13 +90,16 @@ export const WithStorageProviders: React.FC = ({ children }) => { return createElectronRecentQueryStorage({ basepath: options?.basepath }); }, }); + return ( - - {children} - + + + {children} + +