fix: store diagram with database COMPASS-9718

packages/compass-data-modeling/src/components/diagram-card.spec.tsx

@@ -9,6 +9,7 @@ describe('DiagramCard', () => {
     diagram: {
       id: 'test-diagram',
       connectionId: 'test-connection',
+      database: 'someDatabase',
       name: 'Test Diagram',
       createdAt: '2021-10-01T00:00:00.000Z',
       updatedAt: '2023-10-03T00:00:00.000',
@@ -31,7 +32,6 @@ describe('DiagramCard', () => {
           },
         },
       ] as [Edit],
-      databases: 'someDatabase',
     },
     onOpen: () => {},
     onDelete: () => {},

packages/compass-data-modeling/src/components/diagram-card.tsx

@@ -78,9 +78,7 @@ export function DiagramCard({
   onRename,
   onDelete,
 }: {
-  diagram: MongoDBDataModelDescription & {
-    databases: string;
-  };
+  diagram: MongoDBDataModelDescription;
   onOpen: (diagram: MongoDBDataModelDescription) => void;
   onRename: (id: string) => void;
   onDelete: (id: string) => void;
@@ -132,7 +130,7 @@ export function DiagramCard({
             color={palette.gray.dark1}
             className={namespaceIconStyles}
           ></Icon>
-          <span className={namespaceNameStyles}>{diagram.databases}</span>
+          <span className={namespaceNameStyles}>{diagram.database}</span>
         </div>
         <div className={lastModifiedLabel}>
           Last&nbsp;modified: {formattedDate}

packages/compass-data-modeling/src/components/diagram-editor.spec.tsx

@@ -58,6 +58,7 @@ const storageItems: MongoDBDataModelDescription[] = [
       },
     ],
     connectionId: null,
+    database: 'db1',
   },
   {
     id: 'new-diagram-id',
@@ -91,6 +92,7 @@ const storageItems: MongoDBDataModelDescription[] = [
       },
     ],
     connectionId: null,
+    database: 'db1',
   },
 ];
 

packages/compass-data-modeling/src/components/saved-diagrams-list.spec.tsx

@@ -37,6 +37,7 @@ const storageItems: MongoDBDataModelDescription[] = [
       },
     ],
     connectionId: null,
+    database: 'db1',
   },
   {
     id: '2',
@@ -63,6 +64,7 @@ const storageItems: MongoDBDataModelDescription[] = [
       },
     ],
     connectionId: null,
+    database: 'db2',
   },
   {
     id: '3',
@@ -89,6 +91,7 @@ const storageItems: MongoDBDataModelDescription[] = [
       },
     ],
     connectionId: null,
+    database: 'db3',
   },
 ];
 

packages/compass-data-modeling/src/components/saved-diagrams-list.tsx

@@ -16,7 +16,6 @@ import {
 import { useDataModelSavedItems } from '../provider';
 import {
   deleteDiagram,
-  selectCurrentModel,
   openDiagram,
   openDiagramFromFile,
   renameDiagram,
@@ -27,7 +26,6 @@ import SchemaVisualizationIcon from './icons/schema-visualization';
 import FlexibilityIcon from './icons/flexibility';
 import { CARD_HEIGHT, CARD_WIDTH, DiagramCard } from './diagram-card';
 import { DiagramListToolbar } from './diagram-list-toolbar';
-import toNS from 'mongodb-ns';
 import { ImportDiagramButton } from './import-diagram-button';
 
 const sortBy = [
@@ -181,18 +179,6 @@ const DiagramListEmptyContent: React.FunctionComponent<{
   );
 };
 
-const getDatabase = (
-  modelDescription: MongoDBDataModelDescription
-): string[] => {
-  try {
-    return selectCurrentModel(modelDescription.edits).collections.map(
-      ({ ns }) => toNS(ns).database
-    );
-  } catch {
-    return [];
-  }
-};
-
 export const SavedDiagramsList: React.FunctionComponent<{
   onCreateDiagramClick: () => void;
   onOpenDiagramClick: (diagram: MongoDBDataModelDescription) => void;
@@ -207,30 +193,15 @@ export const SavedDiagramsList: React.FunctionComponent<{
   onImportDiagramClick,
 }) => {
   const { items, status } = useDataModelSavedItems();
-  const decoratedItems = useMemo<
-    (MongoDBDataModelDescription & {
-      databases: string;
-    })[]
-  >(() => {
-    return items.map((item) => {
-      const databases = new Set(getDatabase(item));
-      return {
-        ...item,
-        databases: Array.from(databases).join(', '),
-      };
-    });
-  }, [items]);
   const [search, setSearch] = useState('');
   const filteredItems = useMemo(() => {
     try {
       const regex = new RegExp(search, 'i');
-      return decoratedItems.filter(
-        (x) => regex.test(x.name) || (x.databases && regex.test(x.databases))
-      );
+      return items.filter((x) => regex.test(x.name) || regex.test(x.database));
     } catch {
-      return decoratedItems;
+      return items;
     }
-  }, [decoratedItems, search]);
+  }, [items, search]);
   const [sortControls, sortState] = useSortControls(sortBy);
   const sortedItems = useSortedItems(filteredItems, sortState);
 

packages/compass-data-modeling/src/services/data-model-storage.ts

@@ -1,4 +1,5 @@
 import { z } from '@mongodb-js/compass-user-data';
+import toNS from 'mongodb-ns';
 import type { MongoDBJSONSchema } from 'mongodb-schema';
 
 export const FieldPathSchema = z.array(z.string());
@@ -170,20 +171,35 @@ export const validateEdit = (
   }
 };
 
-export const MongoDBDataModelDescriptionSchema = z.object({
-  id: z.string(),
-  name: z.string(),
-  /**
-   * Connection id associated with the data model at the moment of configuring
-   * and analyzing. No connection id means diagram was imported and not attached
-   * to a connection. Practically speaking it just means that we can't do
-   * anything that would require re-fetching data associated with the diagram
-   */
-  connectionId: z.string().nullable(),
-  edits: EditListSchema,
-  createdAt: z.string().datetime(),
-  updatedAt: z.string().datetime(),
-});
+export const MongoDBDataModelDescriptionSchema = z.preprocess(
+  (val) => {
+    const model = val as Record<string, unknown>;
+    if (
+      !model.database &&
+      Array.isArray(model.edits) &&
+      model.edits.length > 0
+    ) {
+      // Infer database from the first collection's namespace in the 'SetModel' edit
+      model.database = toNS(model.edits?.[0].model.collections[0].ns).database;
+    }
+    return model;
+  },
+  z.object({
+    id: z.string(),
+    name: z.string(),
+    /**
+     * Connection id associated with the data model at the moment of configuring
+     * and analyzing. No connection id means diagram was imported and not attached
+     * to a connection. Practically speaking it just means that we can't do
+     * anything that would require re-fetching data associated with the diagram
+     */
+    connectionId: z.string().nullable(),
+    database: z.string(),
+    edits: EditListSchema,
+    createdAt: z.string().datetime(),
+    updatedAt: z.string().datetime(),
+  })
+);
 
 export type MongoDBDataModelDescription = z.output<
   typeof MongoDBDataModelDescriptionSchema

packages/compass-data-modeling/src/services/open-and-download-diagram.spec.ts

@@ -9,15 +9,18 @@ import FlightDiagram from '../../test/fixtures/data-model-with-relationships.jso
 describe('open-and-download-diagram', function () {
   it('should return correct content to download', function () {
     const fileName = 'test-diagram';
+    const database = 'test-database';
 
     const { edits, ...restOfContent } = getDownloadDiagramContent(
       fileName,
-      FlightDiagram.edits as any
+      FlightDiagram.edits as any,
+      database
     );
     expect(restOfContent).to.deep.equal({
       type: 'Compass Data Modeling Diagram',
       version: 1,
       name: fileName,
+      database,
     });
 
     const decodedEdits = JSON.parse(
@@ -132,6 +135,7 @@ describe('open-and-download-diagram', function () {
             version: 1,
             type: 'Compass Data Modeling Diagram',
             name: 'Test Diagram',
+            database: 'test',
             edits: Buffer.from(
               JSON.stringify([{ type: 'NonExistent' }])
             ).toString('base64'),
@@ -146,6 +150,7 @@ describe('open-and-download-diagram', function () {
             version: 1,
             type: 'Compass Data Modeling Diagram',
             name: 'Test Diagram',
+            database: 'test',
             edits: Buffer.from(
               JSON.stringify([
                 {
@@ -179,6 +184,7 @@ describe('open-and-download-diagram', function () {
           version: 1,
           type: 'Compass Data Modeling Diagram',
           name: 'Test Diagram',
+          database: 'test-database',
           edits: Buffer.from(JSON.stringify(FlightDiagram.edits)).toString(
             'base64'
           ),

packages/compass-data-modeling/src/services/open-and-download-diagram.ts

@@ -6,9 +6,19 @@ import { z } from '@mongodb-js/compass-user-data';
 const kCurrentVersion = 1;
 const kFileTypeDescription = 'Compass Data Modeling Diagram';
 
-export function downloadDiagram(fileName: string, edits: Edit[]) {
+export function downloadDiagram(
+  fileName: string,
+  edits: Edit[],
+  database: string
+) {
   const blob = new Blob(
-    [JSON.stringify(getDownloadDiagramContent(fileName, edits), null, 2)],
+    [
+      JSON.stringify(
+        getDownloadDiagramContent(fileName, edits, database),
+        null,
+        2
+      ),
+    ],
     {
       type: 'application/json',
     }
@@ -19,18 +29,25 @@ export function downloadDiagram(fileName: string, edits: Edit[]) {
   });
 }
 
-export function getDownloadDiagramContent(name: string, edits: Edit[]) {
+export function getDownloadDiagramContent(
+  name: string,
+  edits: Edit[],
+  database: string
+) {
   return {
     type: kFileTypeDescription,
     version: kCurrentVersion,
     name,
+    database,
     edits: Buffer.from(JSON.stringify(edits)).toString('base64'),
   };
 }
 
-export async function getDiagramContentsFromFile(
-  file: File
-): Promise<{ name: string; edits: [SetModelEdit, ...Edit[]] }> {
+export async function getDiagramContentsFromFile(file: File): Promise<{
+  name: string;
+  edits: [SetModelEdit, ...Edit[]];
+  database: string;
+}> {
   const reader = new FileReader();
   return new Promise((resolve, reject) => {
     reader.onload = (event) => {
@@ -48,9 +65,9 @@ export async function getDiagramContentsFromFile(
           throw new Error('Unsupported diagram file format');
         }
 
-        const { name, edits } = parsedContent;
+        const { name, edits, database } = parsedContent;
 
-        if (!name || !edits || typeof edits !== 'string') {
+        if (!name || !edits || typeof edits !== 'string' || !database) {
           throw new Error('Diagram file is missing required fields');
         }
 
@@ -63,6 +80,7 @@ export async function getDiagramContentsFromFile(
         return resolve({
           name: parsedContent.name,
           edits: [validEdits[0] as SetModelEdit, ...validEdits.slice(1)],
+          database: parsedContent.database,
         });
       } catch (error) {
         const message =

packages/compass-data-modeling/src/store/analysis-process.ts

@@ -71,6 +71,7 @@ export type AnalysisFinishedAction = {
   type: AnalysisProcessActionTypes.ANALYSIS_FINISHED;
   name: string;
   connectionId: string;
+  database: string;
   collections: {
     ns: string;
     schema: MongoDBJSONSchema;
@@ -315,6 +316,7 @@ export function startAnalysis(
         type: AnalysisProcessActionTypes.ANALYSIS_FINISHED,
         name,
         connectionId,
+        database,
         collections: collections.map((coll) => {
           const node = positioned.nodes.find((node) => {
             return node.id === coll.ns;

packages/compass-data-modeling/src/store/diagram.spec.ts

@@ -19,6 +19,10 @@ import type {
 } from '../services/data-model-storage';
 import { UUID } from 'bson';
 import Sinon from 'sinon';
+import {
+  type AnalysisFinishedAction,
+  AnalysisProcessActionTypes,
+} from './analysis-process';
 
 const model: StaticModel = {
   collections: [
@@ -61,6 +65,7 @@ const loadedDiagram: MongoDBDataModelDescription = {
   id: 'diagram-id',
   name: 'diagram-name',
   connectionId: 'connection-id',
+  database: 'db',
   createdAt: '2023-10-01T00:00:00.000Z',
   updatedAt: '2023-10-05T00:00:00.000Z',
   edits: [{ type: 'SetModel', model } as Edit],
@@ -81,6 +86,7 @@ describe('Data Modeling store', function () {
       const newDiagram = {
         name: 'New Diagram',
         connectionId: 'connection-id',
+        database: 'db',
         collections: [
           {
             ns: 'db.collection1',
@@ -95,14 +101,16 @@ describe('Data Modeling store', function () {
         ],
         relations: model.relationships,
       };
-      store.dispatch({
-        type: 'data-modeling/analysis-stats/ANALYSIS_FINISHED',
+      const analysisFinishedAction: AnalysisFinishedAction = {
+        type: AnalysisProcessActionTypes.ANALYSIS_FINISHED,
         ...newDiagram,
-      });
+      };
+      store.dispatch(analysisFinishedAction);
 
       const initialDiagram = getCurrentDiagramFromState(store.getState());
       expect(initialDiagram.name).to.equal(newDiagram.name);
       expect(initialDiagram.connectionId).to.equal(newDiagram.connectionId);
+      expect(initialDiagram.database).to.equal(newDiagram.database);
       expect(initialDiagram.edits).to.have.length(1);
       expect(initialDiagram.edits[0].type).to.equal('SetModel');
       const initialEdit = initialDiagram.edits[0] as Extract<
@@ -133,6 +141,7 @@ describe('Data Modeling store', function () {
       expect(diagram.id).to.equal(loadedDiagram.id);
       expect(diagram.name).to.equal(loadedDiagram.name);
       expect(diagram.connectionId).to.equal(loadedDiagram.connectionId);
+      expect(diagram.database).to.equal(loadedDiagram.database);
       expect(diagram.edits).to.deep.equal(loadedDiagram.edits);
     });
   });