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

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,
@@ -177,11 +180,14 @@ const DiagramContent: React.FunctionComponent<{
           selectedItems?.type === 'field' && selectedItems.namespace === coll.ns
             ? selectedItems.fieldPath
             : undefined,
+        onClickAddNewFieldToCollection: () =>
+          onAddNewFieldToCollection(coll.ns),
         selected,
         isInRelationshipDrawingMode,
       });
     });
   }, [
+    onAddNewFieldToCollection,
     model?.collections,
     model?.relationships,
     selectedItems,
@@ -301,6 +307,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/services/data-model-storage.ts

@@ -57,6 +57,12 @@ const EditSchemaVariants = z.discriminatedUnion('type', [
     type: z.literal('SetModel'),
     model: StaticModelSchema,
   }),
+  z.object({
+    type: z.literal('AddField'),
+    ns: z.string(),
+    field: z.array(z.string()),
+    jsonSchema: z.custom<MongoDBJSONSchema>(),
+  }),
   z.object({
     type: z.literal('AddRelationship'),
     relationship: RelationshipSchema,

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

@@ -30,6 +30,7 @@ import { getCoordinatesForNewNode } from '@mongodb-js/diagramming';
 import { collectionToDiagramNode } from '../utils/nodes-and-edges';
 import toNS from 'mongodb-ns';
 import { traverseSchema } from '../utils/schema-traversal';
+import { addFieldToJSONSchema, getNewUnusedFieldName } from '../utils/schema';
 
 function isNonEmptyArray<T>(arr: T[]): arr is [T, ...T[]] {
   return Array.isArray(arr) && arr.length > 0;
@@ -474,6 +475,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]
@@ -496,18 +525,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(
@@ -752,6 +779,24 @@ function _applyEdit(edit: Edit, model?: StaticModel): StaticModel {
         collections: [...model.collections, newCollection],
       };
     }
+    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 'AddRelationship': {
       return {
         ...model,

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

@@ -1,6 +1,11 @@
 import React from 'react';
 import toNS from 'mongodb-ns';
-import { InlineDefinition, Body, css } from '@mongodb-js/compass-components';
+import {
+  Body,
+  IconButton,
+  InlineDefinition,
+  css,
+} from '@mongodb-js/compass-components';
 import type { NodeProps, EdgeProps } from '@mongodb-js/diagramming';
 import type { MongoDBJSONSchema } from 'mongodb-schema';
 import type { SelectedItems } from '../store/diagram';
@@ -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',
@@ -121,13 +131,15 @@ export function collectionToDiagramNode(
   options: {
     highlightedFields?: Record<string, FieldPath[] | undefined>;
     selectedField?: FieldPath;
+    onClickAddNewFieldToCollection?: () => void;
     selected?: boolean;
     isInRelationshipDrawingMode?: boolean;
   } = {}
 ): NodeProps {
   const {
     highlightedFields = {},
     selectedField,
+    onClickAddNewFieldToCollection,
     selected = false,
     isInRelationshipDrawingMode = false,
   } = options;
@@ -145,6 +157,19 @@ export function collectionToDiagramNode(
       highlightedFields: highlightedFields[coll.ns] ?? undefined,
       selectedField,
     }),
+    actions: onClickAddNewFieldToCollection ? (
+      <IconButton
+        aria-label="Add Field"
+        className={addNewFieldStyles}
+        onClick={(event: React.MouseEvent<HTMLButtonElement>) => {
+          event.stopPropagation();
+          onClickAddNewFieldToCollection();
+        }}
+        title="Add Field"
+      >
+        <PlusWithSquare />
+      </IconButton>
+    ) : undefined,
     selected,
     connectable: isInRelationshipDrawingMode,
     draggable: !isInRelationshipDrawingMode,

packages/compass-data-modeling/src/utils/schema.spec.ts

@@ -0,0 +1,122 @@
+import { expect } from 'chai';
+import { addFieldToJSONSchema, getNewUnusedFieldName } from './schema';
+
+describe('schema diagram utils', function () {
+  describe('#getNewUnusedFieldName', function () {
+    it('should return a new unused field name', function () {
+      const jsonSchema = {
+        bsonType: 'object',
+        properties: {
+          a: {
+            bsonType: 'string',
+          },
+          b: {
+            bsonType: 'string',
+          },
+        },
+      };
+      const newFieldName = getNewUnusedFieldName(jsonSchema);
+      expect(newFieldName).to.equal('field-1');
+    });
+
+    it('should return a new unused field name when there are conflicts', function () {
+      const jsonSchema = {
+        bsonType: 'object',
+        properties: {
+          'field-1': {
+            bsonType: 'string',
+          },
+          'field-2': {
+            bsonType: 'string',
+          },
+        },
+      };
+      const newFieldName = getNewUnusedFieldName(jsonSchema);
+      expect(newFieldName).to.equal('field-3');
+    });
+  });
+
+  describe('#addFieldToJSONSchema', function () {
+    it('should add a field to the root of the schema', function () {
+      const jsonSchema = {
+        bsonType: 'object',
+        properties: {
+          a: {
+            bsonType: 'string',
+          },
+          b: {
+            bsonType: 'string',
+          },
+        },
+      };
+      const newFieldSchema = {
+        bsonType: 'string',
+      };
+      const newJsonSchema = addFieldToJSONSchema(
+        jsonSchema,
+        ['c'],
+        newFieldSchema
+      );
+      expect(newJsonSchema).to.deep.equal({
+        bsonType: 'object',
+        properties: {
+          a: {
+            bsonType: 'string',
+          },
+          b: {
+            bsonType: 'string',
+          },
+          c: {
+            bsonType: 'string',
+          },
+        },
+      });
+    });
+
+    it('should add a field to a nested object in the schema', function () {
+      const jsonSchema = {
+        bsonType: 'object',
+        properties: {
+          a: {
+            bsonType: 'string',
+          },
+          b: {
+            bsonType: 'object',
+            properties: {
+              c: {
+                bsonType: 'string',
+              },
+            },
+          },
+        },
+      };
+      const newFieldSchema = {
+        bsonType: 'string',
+      };
+      const newJsonSchema = addFieldToJSONSchema(
+        jsonSchema,
+        ['b', 'd'],
+        newFieldSchema
+      );
+      expect(newJsonSchema).to.deep.equal({
+        bsonType: 'object',
+        properties: {
+          a: {
+            bsonType: 'string',
+          },
+          b: {
+            bsonType: 'object',
+            properties: {
+              c: {
+                bsonType: 'string',
+              },
+              d: {
+                bsonType: 'string',
+              },
+            },
+          },
+        },
+      });
+    });
+  });
+});