feat(data-modeling): add add field button to collection node COMPASS-9697 (#7221)

Author: Anemy

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

@@ -93,7 +93,7 @@ const mockDiagramming = {
     <div data-testid="mock-diagram">
       {Object.entries(props).map(([key, value]) => (
         <div key={key} data-testid={`diagram-prop-${key}`}>
-          {JSON.stringify(value)}
+          {typeof value === 'object' ? 'object' : JSON.stringify(value)}
         </div>
       ))}
     </div>

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

@@ -8,6 +8,7 @@ import React, {
 import { connect } from 'react-redux';
 import type { DataModelingState } from '../store/reducer';
 import {
+  addNewFieldToCollection,
   moveCollection,
   selectCollection,
   selectRelationship,
@@ -111,6 +112,7 @@ const DiagramContent: React.FunctionComponent<{
   isInRelationshipDrawingMode: boolean;
   editErrors?: string[];
   newCollection?: string;
+  onAddNewFieldToCollection: (ns: string) => void;
   onMoveCollection: (ns: string, newPosition: [number, number]) => void;
   onCollectionSelect: (namespace: string) => void;
   onRelationshipSelect: (rId: string) => void;
@@ -130,6 +132,7 @@ const DiagramContent: React.FunctionComponent<{
   model,
   isInRelationshipDrawingMode,
   newCollection,
+  onAddNewFieldToCollection,
   onMoveCollection,
   onCollectionSelect,
   onRelationshipSelect,
@@ -171,17 +174,21 @@ const DiagramContent: React.FunctionComponent<{
         !!selectedItems &&
         selectedItems.type === 'collection' &&
         selectedItems.id === coll.ns;
-      return collectionToDiagramNode(coll, {
+      return collectionToDiagramNode({
+        ...coll,
         highlightedFields,
         selectedField:
           selectedItems?.type === 'field' && selectedItems.namespace === coll.ns
             ? selectedItems.fieldPath
             : undefined,
+        onClickAddNewFieldToCollection: () =>
+          onAddNewFieldToCollection(coll.ns),
         selected,
         isInRelationshipDrawingMode,
       });
     });
   }, [
+    onAddNewFieldToCollection,
     model?.collections,
     model?.relationships,
     selectedItems,
@@ -301,6 +308,7 @@ const ConnectedDiagramContent = connect(
     };
   },
   {
+    onAddNewFieldToCollection: addNewFieldToCollection,
     onMoveCollection: moveCollection,
     onCollectionSelect: selectCollection,
     onRelationshipSelect: selectRelationship,

packages/compass-data-modeling/src/components/icons/plus-with-square.tsx

@@ -0,0 +1,27 @@
+import React, { useMemo } from 'react';
+import { palette, useDarkMode } from '@mongodb-js/compass-components';
+
+const PlusWithSquare: React.FunctionComponent = () => {
+  const darkMode = useDarkMode();
+  const strokeColor = useMemo(
+    () => (darkMode ? palette.gray.light1 : palette.gray.dark1),
+    [darkMode]
+  );
+
+  return (
+    <svg
+      width="14"
+      height="14"
+      viewBox="0 0 14 14"
+      fill="none"
+      xmlns="http://www.w3.org/2000/svg"
+    >
+      <path
+        d="M12 0.75H2C1.66848 0.75 1.35054 0.881696 1.11612 1.11612C0.881696 1.35054 0.75 1.66848 0.75 2V12C0.75 12.3315 0.881696 12.6495 1.11612 12.8839C1.35054 13.1183 1.66848 13.25 2 13.25H12C12.3315 13.25 12.6495 13.1183 12.8839 12.8839C13.1183 12.6495 13.25 12.3315 13.25 12V2C13.25 1.66848 13.1183 1.35054 12.8839 1.11612C12.6495 0.881696 12.3315 0.75 12 0.75ZM11.75 11.75H2.25V2.25H11.75V11.75ZM3.75 7C3.75 6.80109 3.82902 6.61032 3.96967 6.46967C4.11032 6.32902 4.30109 6.25 4.5 6.25H6.25V4.5C6.25 4.30109 6.32902 4.11032 6.46967 3.96967C6.61032 3.82902 6.80109 3.75 7 3.75C7.19891 3.75 7.38968 3.82902 7.53033 3.96967C7.67098 4.11032 7.75 4.30109 7.75 4.5V6.25H9.5C9.69891 6.25 9.88968 6.32902 10.0303 6.46967C10.171 6.61032 10.25 6.80109 10.25 7C10.25 7.19891 10.171 7.38968 10.0303 7.53033C9.88968 7.67098 9.69891 7.75 9.5 7.75H7.75V9.5C7.75 9.69891 7.67098 9.88968 7.53033 10.0303C7.38968 10.171 7.19891 10.25 7 10.25C6.80109 10.25 6.61032 10.171 6.46967 10.0303C6.32902 9.88968 6.25 9.69891 6.25 9.5V7.75H4.5C4.30109 7.75 4.11032 7.67098 3.96967 7.53033C3.82902 7.38968 3.75 7.19891 3.75 7Z"
+        fill={strokeColor}
+      />
+    </svg>
+  );
+};
+
+export default PlusWithSquare;

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

@@ -7,7 +7,7 @@ import type { Document } from 'bson';
 import type { AggregationCursor } from 'mongodb';
 import type { Relationship } from '../services/data-model-storage';
 import { applyLayout } from '@mongodb-js/diagramming';
-import { collectionToDiagramNode } from '../utils/nodes-and-edges';
+import { collectionToBaseNodeForLayout } from '../utils/nodes-and-edges';
 
 export type AnalysisProcessState = {
   currentAnalysisOptions:
@@ -207,13 +207,13 @@ export function startAnalysis(
       }
 
       const positioned = await applyLayout(
-        collections.map((coll) => {
-          return collectionToDiagramNode({
+        collections.map((coll) =>
+          collectionToBaseNodeForLayout({
             ns: coll.ns,
             jsonSchema: coll.schema,
             displayPosition: [0, 0],
-          });
-        }),
+          })
+        ),
         [],
         'LEFT_RIGHT'
       );

packages/compass-data-modeling/src/store/apply-edit.ts

@@ -4,6 +4,7 @@ import type {
   Relationship,
   StaticModel,
 } from '../services/data-model-storage';
+import { addFieldToJSONSchema } from '../utils/schema';
 import { updateSchema } from '../utils/schema-traversal';
 import {
   isRelationshipInvolvingField,
@@ -149,6 +150,24 @@ export function applyEdit(edit: Edit, model?: StaticModel): StaticModel {
         }),
       };
     }
+    case 'AddField': {
+      return {
+        ...model,
+        collections: model.collections.map((collection) => {
+          if (collection.ns === edit.ns) {
+            return {
+              ...collection,
+              jsonSchema: addFieldToJSONSchema(
+                collection.jsonSchema,
+                edit.field,
+                edit.jsonSchema
+              ),
+            };
+          }
+          return collection;
+        }),
+      };
+    }
     case 'RemoveField': {
       return {
         ...model,

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

@@ -27,10 +27,11 @@ import {
 } from '../services/open-and-download-diagram';
 import type { MongoDBJSONSchema } from 'mongodb-schema';
 import { getCoordinatesForNewNode } from '@mongodb-js/diagramming';
-import { collectionToDiagramNode } from '../utils/nodes-and-edges';
+import { collectionToBaseNodeForLayout } from '../utils/nodes-and-edges';
 import toNS from 'mongodb-ns';
 import { traverseSchema } from '../utils/schema-traversal';
 import { applyEdit as _applyEdit } from './apply-edit';
+import { getNewUnusedFieldName } from '../utils/schema';
 
 function isNonEmptyArray<T>(arr: T[]): arr is [T, ...T[]] {
   return Array.isArray(arr) && arr.length > 0;
@@ -394,6 +395,13 @@ const updateSelectedItemsFromAppliedEdit = (
         ],
       };
     }
+    case 'AddField': {
+      return {
+        type: 'field',
+        namespace: edit.ns,
+        fieldPath: edit.field,
+      };
+    }
   }
 
   return currentSelection;
@@ -485,6 +493,34 @@ export function redoEdit(): DataModelingThunkAction<void, RedoEditAction> {
   };
 }
 
+export function addNewFieldToCollection(
+  ns: string
+): DataModelingThunkAction<void, ApplyEditAction | ApplyEditFailedAction> {
+  return (dispatch, getState) => {
+    const modelState = selectCurrentModelFromState(getState());
+
+    const collection = modelState.collections.find((c) => c.ns === ns);
+    if (!collection) {
+      throw new Error('Collection to add field to not found');
+    }
+
+    const edit: Omit<
+      Extract<Edit, { type: 'AddField' }>,
+      'id' | 'timestamp'
+    > = {
+      type: 'AddField',
+      ns,
+      // Use the first unique field name we can use.
+      field: [getNewUnusedFieldName(collection.jsonSchema)],
+      jsonSchema: {
+        bsonType: 'string',
+      },
+    };
+
+    return dispatch(applyEdit(edit));
+  };
+}
+
 export function moveCollection(
   ns: string,
   newPosition: [number, number]
@@ -507,18 +543,16 @@ export function renameCollection(
   void,
   ApplyEditAction | ApplyEditFailedAction | CollectionSelectedAction
 > {
-  return (dispatch) => {
-    const edit: Omit<
-      Extract<Edit, { type: 'RenameCollection' }>,
-      'id' | 'timestamp'
-    > = {
-      type: 'RenameCollection',
-      fromNS,
-      toNS,
-    };
-
-    dispatch(applyEdit(edit));
+  const edit: Omit<
+    Extract<Edit, { type: 'RenameCollection' }>,
+    'id' | 'timestamp'
+  > = {
+    type: 'RenameCollection',
+    fromNS,
+    toNS,
   };
+
+  return applyEdit(edit);
 }
 
 export function applyEdit(
@@ -689,9 +723,9 @@ function getPositionForNewCollection(
   newCollection: Omit<DataModelCollection, 'displayPosition'>
 ): [number, number] {
   const existingNodes = existingCollections.map((collection) =>
-    collectionToDiagramNode(collection)
+    collectionToBaseNodeForLayout(collection)
   );
-  const newNode = collectionToDiagramNode({
+  const newNode = collectionToBaseNodeForLayout({
     ns: newCollection.ns,
     jsonSchema: newCollection.jsonSchema,
     displayPosition: [0, 0],

packages/compass-data-modeling/src/utils/nodes-and-edges.tsx

@@ -1,7 +1,12 @@
 import React from 'react';
 import toNS from 'mongodb-ns';
-import { InlineDefinition, Body, css } from '@mongodb-js/compass-components';
-import type { NodeProps, EdgeProps } from '@mongodb-js/diagramming';
+import {
+  Body,
+  IconButton,
+  InlineDefinition,
+  css,
+} from '@mongodb-js/compass-components';
+import type { NodeProps, EdgeProps, BaseNode } from '@mongodb-js/diagramming';
 import type { MongoDBJSONSchema } from 'mongodb-schema';
 import type { SelectedItems } from '../store/diagram';
 import type {
@@ -11,6 +16,7 @@ import type {
 } from '../services/data-model-storage';
 import { traverseSchema } from './schema-traversal';
 import { areFieldPathsEqual } from './utils';
+import PlusWithSquare from '../components/icons/plus-with-square';
 
 function getBsonTypeName(bsonType: string) {
   switch (bsonType) {
@@ -21,6 +27,10 @@ function getBsonTypeName(bsonType: string) {
   }
 }
 
+const addNewFieldStyles = css({
+  marginLeft: 'auto',
+});
+
 const mixedTypeTooltipContentStyles = css({
   overflowWrap: 'anywhere',
   textWrap: 'wrap',
@@ -116,38 +126,77 @@ export const getFieldsFromSchema = ({
   return fields;
 };
 
-export function collectionToDiagramNode(
-  coll: Pick<DataModelCollection, 'ns' | 'jsonSchema' | 'displayPosition'>,
-  options: {
-    highlightedFields?: Record<string, FieldPath[] | undefined>;
-    selectedField?: FieldPath;
-    selected?: boolean;
-    isInRelationshipDrawingMode?: boolean;
-  } = {}
-): NodeProps {
-  const {
-    highlightedFields = {},
-    selectedField,
-    selected = false,
-    isInRelationshipDrawingMode = false,
-  } = options;
+/**
+ * Create a base node to be used for positioning and measuring in node layouts.
+ */
+export function collectionToBaseNodeForLayout({
+  ns,
+  jsonSchema,
+  displayPosition,
+}: Pick<
+  DataModelCollection,
+  'ns' | 'jsonSchema' | 'displayPosition'
+>): BaseNode & Pick<NodeProps, 'fields'> {
+  return {
+    id: ns,
+    position: {
+      x: displayPosition[0],
+      y: displayPosition[1],
+    },
+    fields: getFieldsFromSchema({ jsonSchema }),
+  };
+}
 
+type CollectionWithRenderOptions = Pick<
+  DataModelCollection,
+  'ns' | 'jsonSchema' | 'displayPosition'
+> & {
+  highlightedFields: Record<string, FieldPath[] | undefined>;
+  selectedField?: FieldPath;
+  selected: boolean;
+  isInRelationshipDrawingMode: boolean;
+  onClickAddNewFieldToCollection: () => void;
+};
+
+export function collectionToDiagramNode({
+  ns,
+  jsonSchema,
+  displayPosition,
+  selectedField,
+  highlightedFields,
+  selected,
+  isInRelationshipDrawingMode,
+  onClickAddNewFieldToCollection,
+}: CollectionWithRenderOptions): NodeProps {
   return {
-    id: coll.ns,
+    id: ns,
     type: 'collection',
     position: {
-      x: coll.displayPosition[0],
-      y: coll.displayPosition[1],
+      x: displayPosition[0],
+      y: displayPosition[1],
     },
-    title: toNS(coll.ns).collection,
+    title: toNS(ns).collection,
     fields: getFieldsFromSchema({
-      jsonSchema: coll.jsonSchema,
-      highlightedFields: highlightedFields[coll.ns] ?? undefined,
+      jsonSchema: jsonSchema,
+      highlightedFields: highlightedFields[ns] ?? undefined,
       selectedField,
     }),
     selected,
     connectable: isInRelationshipDrawingMode,
     draggable: !isInRelationshipDrawingMode,
+    actions: onClickAddNewFieldToCollection ? (
+      <IconButton
+        aria-label="Add Field"
+        className={addNewFieldStyles}
+        onClick={(event: React.MouseEvent<HTMLButtonElement>) => {
+          event.stopPropagation();
+          onClickAddNewFieldToCollection();
+        }}
+        title="Add Field"
+      >
+        <PlusWithSquare />
+      </IconButton>
+    ) : undefined,
   };
 }