store the layout

Author: paula-stacho

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

@@ -9,6 +9,7 @@ import { connect } from 'react-redux';
 import type { DataModelingState } from '../store/reducer';
 import {
   applyEdit,
+  applyInitialLayout,
   getCurrentDiagramFromState,
   redoEdit,
   selectCurrentModel,
@@ -127,6 +128,7 @@ const DiagramEditor: React.FunctionComponent<{
   onRetryClick: () => void;
   onCancelClick: () => void;
   onApplyClick: (edit: Omit<Edit, 'id' | 'timestamp'>) => void;
+  onApplyInitialLayout: (positions: Record<string, [number, number]>) => void;
 }> = ({
   diagramLabel,
   step,
@@ -139,6 +141,7 @@ const DiagramEditor: React.FunctionComponent<{
   onRetryClick,
   onCancelClick,
   onApplyClick,
+  onApplyInitialLayout,
 }) => {
   const isDarkMode = useDarkMode();
   const diagramContainerRef = useRef<HTMLDivElement | null>(null);
@@ -226,8 +229,13 @@ const DiagramEditor: React.FunctionComponent<{
           edges,
           'STAR'
         );
-        // TODO: save the new positions to the model
         setNodes(positionedNodes);
+        onApplyInitialLayout(
+          positionedNodes.reduce((obj, node) => {
+            obj[node.id] = [node.position.x, node.position.y];
+            return obj;
+          }, {} as Record<string, [number, number]>)
+        );
       } catch (err) {
         console.error('Error applying layout:', err);
       }
@@ -427,5 +435,6 @@ export default connect(
     onRetryClick: retryAnalysis,
     onCancelClick: cancelAnalysis,
     onApplyClick: applyEdit,
+    onApplyInitialLayout: applyInitialLayout,
   }
 )(DiagramEditor);

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

@@ -88,12 +88,20 @@ export const MongoDBDataModelDescriptionSchema = z.object({
    */
   connectionId: z.string().nullable(),
 
-  edits: z.array(EditSchema).nonempty(),
+  // Ensure first item exists and is 'SetModel'
+  edits: z
+    .array(EditSchema)
+    .nonempty()
+    .refine((edits) => edits[0]?.type === 'SetModel', {
+      message: "First edit must be of type 'SetModel'",
+    }),
 
   createdAt: z.string().datetime(),
   updatedAt: z.string().datetime(),
 });
 
+type SetModelEdit = Extract<Edit, { type: 'SetModel' }>;
+
 export type MongoDBDataModelDescription = z.output<
   typeof MongoDBDataModelDescriptionSchema
 >;

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

@@ -31,6 +31,7 @@ export enum DiagramActionTypes {
   OPEN_DIAGRAM = 'data-modeling/diagram/OPEN_DIAGRAM',
   DELETE_DIAGRAM = 'data-modeling/diagram/DELETE_DIAGRAM',
   RENAME_DIAGRAM = 'data-modeling/diagram/RENAME_DIAGRAM',
+  APPLY_INITIAL_LAYOUT = 'data-modeling/diagram/APPLY_INITIAL_LAYOUT',
   APPLY_EDIT = 'data-modeling/diagram/APPLY_EDIT',
   APPLY_EDIT_FAILED = 'data-modeling/diagram/APPLY_EDIT_FAILED',
   UNDO_EDIT = 'data-modeling/diagram/UNDO_EDIT',
@@ -53,6 +54,11 @@ export type RenameDiagramAction = {
   name: string;
 };
 
+export type ApplyInitialLayoutAction = {
+  type: DiagramActionTypes.APPLY_INITIAL_LAYOUT;
+  positions: Record<string, [number, number]>;
+};
+
 export type ApplyEditAction = {
   type: DiagramActionTypes.APPLY_EDIT;
   edit: Edit;
@@ -75,6 +81,7 @@ export type DiagramActions =
   | OpenDiagramAction
   | DeleteDiagramAction
   | RenameDiagramAction
+  | ApplyInitialLayoutAction
   | ApplyEditAction
   | ApplyEditFailedAction
   | UndoEditAction
@@ -141,6 +148,30 @@ export const diagramReducer: Reducer<DiagramState> = (
       updatedAt: new Date().toISOString(),
     };
   }
+  if (isAction(action, DiagramActionTypes.APPLY_INITIAL_LAYOUT)) {
+    const initialEdit = state.edits.current[0];
+    if (!initialEdit || initialEdit.type !== 'SetModel') {
+      throw new Error('No initial model edit found to apply layout to');
+    }
+    return {
+      ...state,
+      edits: {
+        ...state.edits,
+        current: [
+          {
+            ...initialEdit,
+            model: {
+              ...initialEdit.model,
+              collections: initialEdit.model.collections.map((collection) => ({
+                ...collection,
+                displayPosition: action.positions[collection.ns] || [-1, -1],
+              })),
+            },
+          },
+        ],
+      },
+    };
+  }
   if (isAction(action, DiagramActionTypes.APPLY_EDIT)) {
     return {
       ...state,
@@ -232,6 +263,18 @@ export function applyEdit(
   };
 }
 
+export function applyInitialLayout(
+  positions: Record<string, [number, number]>
+): DataModelingThunkAction<void, ApplyInitialLayoutAction> {
+  return (dispatch, getState, { dataModelStorage }) => {
+    dispatch({
+      type: DiagramActionTypes.APPLY_INITIAL_LAYOUT,
+      positions,
+    });
+    void dataModelStorage.save(getCurrentDiagramFromState(getState()));
+  };
+}
+
 export function openDiagram(diagram: MongoDBDataModelDescription) {
   return { type: DiagramActionTypes.OPEN_DIAGRAM, diagram };
 }