feat(compass-data-modeling): add relationship via dragging COMPASS-9332

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

@@ -12,9 +12,11 @@ function renderDiagramEditorToolbar(
       step="EDITING"
       hasUndo={true}
       hasRedo={true}
+      isInRelationshipDrawingMode={false}
       onUndoClick={() => {}}
       onRedoClick={() => {}}
       onExportClick={() => {}}
+      onRelationshipDrawingToggle={() => {}}
       {...props}
     />
   );
@@ -63,6 +65,36 @@ describe('DiagramEditorToolbar', function () {
     });
   });
 
+  context('add relationship button', function () {
+    it('renders it active if isInRelationshipDrawingMode is true', function () {
+      renderDiagramEditorToolbar({ isInRelationshipDrawingMode: true });
+      const addButton = screen.getByRole('button', {
+        name: 'Exit Relationship Drawing Mode',
+      });
+      expect(addButton).to.have.attribute('aria-pressed', 'true');
+    });
+
+    it('does not render it active if isInRelationshipDrawingMode is false', function () {
+      renderDiagramEditorToolbar({ isInRelationshipDrawingMode: false });
+      const addButton = screen.getByRole('button', {
+        name: 'Add Relationship',
+      });
+      expect(addButton).to.have.attribute('aria-pressed', 'false');
+    });
+
+    it('clicking on it calls onRelationshipDrawingToggle', function () {
+      const relationshipDrawingToggleSpy = sinon.spy();
+      renderDiagramEditorToolbar({
+        onRelationshipDrawingToggle: relationshipDrawingToggleSpy,
+      });
+      const addRelationshipButton = screen.getByRole('button', {
+        name: 'Add Relationship',
+      });
+      userEvent.click(addRelationshipButton);
+      expect(relationshipDrawingToggleSpy).to.have.been.calledOnce;
+    });
+  });
+
   it('renders export button and calls onExportClick', function () {
     const exportSpy = sinon.spy();
     renderDiagramEditorToolbar({ onExportClick: exportSpy });

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

@@ -13,6 +13,7 @@ import {
   spacing,
   useDarkMode,
   transparentize,
+  Tooltip,
 } from '@mongodb-js/compass-components';
 
 const containerStyles = css({
@@ -44,10 +45,21 @@ export const DiagramEditorToolbar: React.FunctionComponent<{
   step: DataModelingState['step'];
   hasUndo: boolean;
   hasRedo: boolean;
+  isInRelationshipDrawingMode: boolean;
   onUndoClick: () => void;
   onRedoClick: () => void;
   onExportClick: () => void;
-}> = ({ step, hasUndo, onUndoClick, hasRedo, onRedoClick, onExportClick }) => {
+  onRelationshipDrawingToggle: () => void;
+}> = ({
+  step,
+  hasUndo,
+  onUndoClick,
+  hasRedo,
+  onRedoClick,
+  onExportClick,
+  onRelationshipDrawingToggle,
+  isInRelationshipDrawingMode,
+}) => {
   const darkmode = useDarkMode();
   if (step !== 'EDITING') {
     return null;
@@ -58,6 +70,24 @@ export const DiagramEditorToolbar: React.FunctionComponent<{
       data-testid="diagram-editor-toolbar"
     >
       <div className={toolbarGroupStyles}>
+        <Tooltip
+          trigger={
+            <IconButton
+              aria-label={
+                !isInRelationshipDrawingMode
+                  ? 'Add Relationship'
+                  : 'Exit Relationship Drawing Mode'
+              }
+              onClick={onRelationshipDrawingToggle}
+              active={isInRelationshipDrawingMode}
+              aria-pressed={isInRelationshipDrawingMode}
+            >
+              <Icon glyph="Relationship"></Icon>
+            </IconButton>
+          }
+        >
+          Drag from one collection to another to create a relationship.
+        </Tooltip>
         <IconButton aria-label="Undo" disabled={!hasUndo} onClick={onUndoClick}>
           <Icon glyph="Undo"></Icon>
         </IconButton>

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

@@ -1,4 +1,10 @@
-import React, { useCallback, useEffect, useMemo, useRef } from 'react';
+import React, {
+  useCallback,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from 'react';
 import { connect } from 'react-redux';
 import type { DataModelingState } from '../store/reducer';
 import {
@@ -8,6 +14,7 @@ import {
   selectBackground,
   type DiagramState,
   selectCurrentModelFromState,
+  createNewRelationship,
 } from '../store/diagram';
 import {
   Banner,
@@ -85,26 +92,38 @@ const modelPreviewContainerStyles = css({
 
 const modelPreviewStyles = css({
   minHeight: 0,
+
+  /** reactflow handles this normally, but there is a `* { userSelect: 'text' }` in this project,
+   *  which overrides inherited userSelect */
+  ['.connectablestart']: {
+    userSelect: 'none',
+  },
 });
 
 type SelectedItems = NonNullable<DiagramState>['selectedItems'];
 
 const DiagramContent: React.FunctionComponent<{
   diagramLabel: string;
   model: StaticModel | null;
+  isInRelationshipDrawingMode: boolean;
   editErrors?: string[];
   onMoveCollection: (ns: string, newPosition: [number, number]) => void;
   onCollectionSelect: (namespace: string) => void;
   onRelationshipSelect: (rId: string) => void;
   onDiagramBackgroundClicked: () => void;
   selectedItems: SelectedItems;
+  onCreateNewRelationship: (source: string, target: string) => void;
+  onRelationshipDrawn: () => void;
 }> = ({
   diagramLabel,
   model,
+  isInRelationshipDrawingMode,
   onMoveCollection,
   onCollectionSelect,
   onRelationshipSelect,
   onDiagramBackgroundClicked,
+  onCreateNewRelationship,
+  onRelationshipDrawn,
   selectedItems,
 }) => {
   const isDarkMode = useDarkMode();
@@ -138,9 +157,18 @@ const DiagramContent: React.FunctionComponent<{
         !!selectedItems &&
         selectedItems.type === 'collection' &&
         selectedItems.id === coll.ns;
-      return collectionToDiagramNode(coll, selectedFields, selected);
+      return collectionToDiagramNode(coll, {
+        selectedFields,
+        selected,
+        isInRelationshipDrawingMode,
+      });
     });
-  }, [model?.collections, model?.relationships, selectedItems]);
+  }, [
+    model?.collections,
+    model?.relationships,
+    selectedItems,
+    isInRelationshipDrawingMode,
+  ]);
 
   // Fit to view on initial mount
   useEffect(() => {
@@ -155,6 +183,14 @@ const DiagramContent: React.FunctionComponent<{
     });
   }, []);
 
+  const handleNodesConnect = useCallback(
+    (source: string, target: string) => {
+      onCreateNewRelationship(source, target);
+      onRelationshipDrawn();
+    },
+    [onRelationshipDrawn, onCreateNewRelationship]
+  );
+
   return (
     <div
       ref={setDiagramContainerRef}
@@ -191,6 +227,9 @@ const DiagramContent: React.FunctionComponent<{
           onNodeDragStop={(evt, node) => {
             onMoveCollection(node.id, [node.position.x, node.position.y]);
           }}
+          onConnect={({ source, target }) => {
+            handleNodesConnect(source, target);
+          }}
         />
       </div>
     </div>
@@ -211,6 +250,7 @@ const ConnectedDiagramContent = connect(
     onCollectionSelect: selectCollection,
     onRelationshipSelect: selectRelationship,
     onDiagramBackgroundClicked: selectBackground,
+    onCreateNewRelationship: createNewRelationship,
   }
 )(DiagramContent);
 
@@ -222,6 +262,17 @@ const DiagramEditor: React.FunctionComponent<{
 }> = ({ step, diagramId, onRetryClick, onCancelClick }) => {
   let content;
 
+  const [isInRelationshipDrawingMode, setIsInRelationshipDrawingMode] =
+    useState(false);
+
+  const handleRelationshipDrawingToggle = useCallback(() => {
+    setIsInRelationshipDrawingMode((prev) => !prev);
+  }, []);
+
+  const onRelationshipDrawn = useCallback(() => {
+    setIsInRelationshipDrawingMode(false);
+  }, []);
+
   if (step === 'NO_DIAGRAM_SELECTED') {
     return null;
   }
@@ -257,12 +308,23 @@ const DiagramEditor: React.FunctionComponent<{
 
   if (step === 'EDITING' && diagramId) {
     content = (
-      <ConnectedDiagramContent key={diagramId}></ConnectedDiagramContent>
+      <ConnectedDiagramContent
+        key={diagramId}
+        isInRelationshipDrawingMode={isInRelationshipDrawingMode}
+        onRelationshipDrawn={onRelationshipDrawn}
+      ></ConnectedDiagramContent>
     );
   }
 
   return (
-    <WorkspaceContainer toolbar={<DiagramEditorToolbar />}>
+    <WorkspaceContainer
+      toolbar={
+        <DiagramEditorToolbar
+          onRelationshipDrawingToggle={handleRelationshipDrawingToggle}
+          isInRelationshipDrawingMode={isInRelationshipDrawingMode}
+        />
+      }
+    >
       {content}
       <ExportDiagramModal />
     </WorkspaceContainer>

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

@@ -209,9 +209,11 @@ export function startAnalysis(
       const positioned = await applyLayout(
         collections.map((coll) => {
           return collectionToDiagramNode({
-            ns: coll.ns,
-            jsonSchema: coll.schema,
-            displayPosition: [0, 0],
+            coll: {
+              ns: coll.ns,
+              jsonSchema: coll.schema,
+              displayPosition: [0, 0],
+            },
           });
         }),
         [],

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

@@ -284,7 +284,8 @@ export function selectBackground(): DiagramBackgroundSelectedAction {
 }
 
 export function createNewRelationship(
-  namespace: string
+  localNamespace: string,
+  foreignNamespace: string | null = null
 ): DataModelingThunkAction<void, RelationSelectedAction> {
   return (dispatch, getState, { track }) => {
     const relationshipId = new UUID().toString();
@@ -297,8 +298,8 @@ export function createNewRelationship(
         relationship: {
           id: relationshipId,
           relationship: [
-            { ns: namespace, cardinality: 1, fields: null },
-            { ns: null, cardinality: 1, fields: null },
+            { ns: localNamespace, cardinality: 1, fields: null },
+            { ns: foreignNamespace, cardinality: 1, fields: null },
           ],
           isInferred: false,
         },

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

@@ -144,9 +144,18 @@ export const getFieldsFromSchema = (
 
 export function collectionToDiagramNode(
   coll: Pick<DataModelCollection, 'ns' | 'jsonSchema' | 'displayPosition'>,
-  selectedFields: Record<string, string[][] | undefined> = {},
-  selected = false
+  options: {
+    selectedFields?: Record<string, string[][] | undefined>;
+    selected?: boolean;
+    isInRelationshipDrawingMode?: boolean;
+  } = {}
 ): NodeProps {
+  const {
+    selectedFields = {},
+    selected = false,
+    isInRelationshipDrawingMode = false,
+  } = options;
+
   return {
     id: coll.ns,
     type: 'collection',
@@ -161,6 +170,8 @@ export function collectionToDiagramNode(
       0
     ),
     selected,
+    connectable: isInRelationshipDrawingMode,
+    draggable: !isInRelationshipDrawingMode,
   };
 }