feat(data-modeling): implement data model storage that uses atlas user data COMPASS-9949

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'
+      | 'dataModelDescriptions',
     id?: string
   ): string {
     const encodedOrgId = encodeURIComponent(orgId);

packages/compass-data-modeling/package.json

@@ -26,13 +26,13 @@
     ".": "./dist/index.js",
     "./provider": "./dist/provider/index.js",
     "./renderer": "./dist/services/data-model-storage-electron.js",
-    "./web": "./dist/services/data-model-storage-in-memory.js"
+    "./web": "./dist/services/data-model-storage-atlas.js"
   },
   "compass:exports": {
     ".": "./src/index.ts",
     "./provider": "./src/provider/index.tsx",
     "./renderer": "./src/services/data-model-storage-electron.tsx",
-    "./web": "./src/services/data-model-storage-in-memory.tsx"
+    "./web": "./src/services/data-model-storage-atlas.tsx"
   },
   "types": "./dist/index.d.ts",
   "scripts": {

packages/compass-data-modeling/src/services/data-model-storage-atlas.tsx

@@ -0,0 +1,119 @@
+import { AtlasUserData } from '@mongodb-js/compass-user-data';
+import type {
+  DataModelStorage,
+  MongoDBDataModelDescription,
+} from './data-model-storage';
+import { MongoDBDataModelDescriptionSchema } from './data-model-storage';
+import dataModelStorageInMemory from './data-model-storage-in-memory';
+import {
+  atlasServiceLocator,
+  type AtlasService,
+} from '@mongodb-js/atlas-service/provider';
+import { createServiceProvider } from '@mongodb-js/compass-app-registry';
+import { DataModelStorageServiceProvider } from '../provider';
+import React, { useRef } from 'react';
+import { mongoLogId, useLogger } from '@mongodb-js/compass-logging/provider';
+
+class DataModelStorageAtlas implements DataModelStorage {
+  private readonly userData: AtlasUserData<
+    typeof MongoDBDataModelDescriptionSchema
+  >;
+  constructor(orgId: string, projectId: string, atlasService: AtlasService) {
+    this.userData = new AtlasUserData(
+      MongoDBDataModelDescriptionSchema,
+      'dataModelDescriptions',
+      {
+        orgId,
+        projectId,
+        getResourceUrl(path) {
+          // TODO(COMPASS-9960): this is copied from compass-web entrypoint for
+          // brevity, but shouldn't be defined outside of AtlasUserData, this
+          // logic is literally the same between all user data instances and can
+          // be encapsulated inside of the AtlasUserData class implementation
+          const [type, pathOrgId, pathProjectId, id] =
+            path?.split('/').filter(Boolean) || [];
+
+          if (
+            !type ||
+            !pathOrgId ||
+            !pathProjectId ||
+            type !== 'dataModelDescriptions'
+          ) {
+            throw new Error(
+              'DataModelStorageAtlas is used outside of Atlas Cloud context'
+            );
+          }
+
+          return atlasService.userDataEndpoint(
+            pathOrgId,
+            pathProjectId,
+            type,
+            id
+          );
+        },
+        authenticatedFetch: atlasService.authenticatedFetch.bind(atlasService),
+      }
+    );
+  }
+  save(description: MongoDBDataModelDescription) {
+    return this.userData.write(description.id, description);
+  }
+  delete(id: MongoDBDataModelDescription['id']) {
+    return this.userData.delete(id);
+  }
+  async loadAll(): Promise<MongoDBDataModelDescription[]> {
+    try {
+      const res = await this.userData.readAll();
+      return res.data;
+    } catch {
+      return [];
+    }
+  }
+  async load(id: string): Promise<MongoDBDataModelDescription | null> {
+    return (
+      (await this.loadAll()).find((item) => {
+        return item.id === id;
+      }) ?? null
+    );
+  }
+}
+
+export const DataModelStorageServiceProviderAtlas = createServiceProvider(
+  function DataModelStorageServiceProviderAtlas({
+    children,
+    orgId,
+    projectId,
+  }: {
+    children?: React.ReactNode;
+    orgId?: string;
+    projectId?: string;
+  }) {
+    const storageRef = useRef<DataModelStorage>();
+    const atlasService = atlasServiceLocator();
+    const logger = useLogger('DATA-MODEL-STORAGE');
+
+    if (!storageRef.current) {
+      if (orgId && projectId) {
+        storageRef.current = new DataModelStorageAtlas(
+          orgId,
+          projectId,
+          atlasService
+        );
+      } else {
+        logger.log.warn(
+          mongoLogId(1_001_000_378),
+          'DataModelStorageServiceProviderAtlas',
+          'Falling back to in memory storage because orgId or projectId is missing'
+        );
+        // Fallback to in-memory if we're outside of Atlas Cloud
+        storageRef.current = dataModelStorageInMemory;
+      }
+    }
+
+    return (
+      <DataModelStorageServiceProvider storage={storageRef.current}>
+        {children}
+      </DataModelStorageServiceProvider>
+    );
+  }
+);

packages/compass-data-modeling/web.d.ts

@@ -1 +1 @@
-export * from './dist/services/data-model-storage-in-memory.d';
+export * from './dist/services/data-model-storage-atlas.d';

packages/compass-data-modeling/web.js

@@ -1,2 +1,2 @@
 'use strict';
-module.exports = require('./dist/services/data-model-storage-in-memory');
+module.exports = require('./dist/services/data-model-storage-atlas');

packages/compass-web/sandbox/index.tsx

@@ -38,6 +38,7 @@ const App = () => {
   const [currentTab, updateCurrentTab] = useWorkspaceTabRouter();
   const { status, projectParams } = useAtlasProxySignIn();
   const {
+    orgId,
     projectId,
     csrfToken,
     csrfTime,
@@ -99,7 +100,7 @@ const App = () => {
       <SandboxPreferencesUpdateProvider>
         <Body as="div" className={sandboxContainerStyles}>
           <CompassWeb
-            orgId={''}
+            orgId={orgId ?? ''}
             projectId={projectId ?? ''}
             onActiveWorkspaceTabChange={updateCurrentTab}
             initialWorkspace={currentTab ?? undefined}

packages/compass-web/sandbox/sandbox-atlas-sign-in.tsx

@@ -13,6 +13,7 @@ console.info(
 type SignInStatus = 'checking' | 'signed-in' | 'signed-out';
 
 type ProjectParams = {
+  orgId: string;
   projectId: string;
   csrfToken: string;
   csrfTime: string;
@@ -127,10 +128,12 @@ export function useAtlasProxySignIn(): AtlasLoginReturnValue {
             currentOrganization: { genAIFeaturesEnabled },
             featureFlags: { groupEnabledFeatureFlags },
             userRoles,
+            currentOrganization,
           } = params;
           const overrideGenAIFeatures =
             process.env.COMPASS_OVERRIDE_ENABLE_AI_FEATURES === 'true';
           setProjectParams({
+            orgId: currentOrganization.id,
             projectId,
             csrfToken,
             csrfTime,

packages/compass-web/src/entrypoint.tsx

@@ -63,7 +63,7 @@ import { WebWorkspaceTab as WelcomeWorkspaceTab } from '@mongodb-js/compass-welc
 import { WorkspaceTab as MyQueriesWorkspace } from '@mongodb-js/compass-saved-aggregations-queries';
 import { useCompassWebPreferences } from './preferences';
 import { DataModelingWorkspaceTab as DataModelingWorkspace } from '@mongodb-js/compass-data-modeling';
-import { DataModelStorageServiceProviderInMemory } from '@mongodb-js/compass-data-modeling/web';
+import { DataModelStorageServiceProviderAtlas } from '@mongodb-js/compass-data-modeling/web';
 import {
   createWebRecentQueryStorage,
   createWebFavoriteQueryStorage,
@@ -507,7 +507,10 @@ const CompassWeb = ({
               <TelemetryProvider options={telemetryOptions.current}>
                 <WithAtlasProviders>
                   <WithStorageProviders orgId={orgId} projectId={projectId}>
-                    <DataModelStorageServiceProviderInMemory>
+                    <DataModelStorageServiceProviderAtlas
+                      orgId={orgId}
+                      projectId={projectId}
+                    >
                       <AtlasCloudConnectionStorageProvider
                         orgId={orgId}
                         projectId={projectId}
@@ -576,7 +579,7 @@ const CompassWeb = ({
                           </CompassConnections>
                         </CompassAssistantProvider>
                       </AtlasCloudConnectionStorageProvider>
-                    </DataModelStorageServiceProviderInMemory>
+                    </DataModelStorageServiceProviderAtlas>
                   </WithStorageProviders>
                 </WithAtlasProviders>
               </TelemetryProvider>