feat(data-modeling): add collection COMPASS-9655

packages/compass-components/src/components/drawer-portal.spec.tsx

@@ -9,6 +9,8 @@ import {
   DrawerContentProvider,
   DrawerSection,
   DrawerAnchor,
+  useDrawerState,
+  useDrawerActions,
 } from './drawer-portal';
 import { expect } from 'chai';
 
@@ -162,4 +164,68 @@ describe('DrawerSection', function () {
       screen.getByTestId('lg-drawer')
     ).to.have.attribute('aria-hidden', 'true');
   });
+
+  it('can control drawer state via the hooks', async function () {
+    const ControlElement = () => {
+      const { isDrawerOpen } = useDrawerState();
+      const { openDrawer, closeDrawer } = useDrawerActions();
+      return (
+        <div>
+          <span data-testid="drawer-state">
+            {isDrawerOpen ? 'open' : 'closed'}
+          </span>
+          <button
+            data-testid="toggle-drawer"
+            onClick={
+              isDrawerOpen
+                ? () => closeDrawer()
+                : () => openDrawer('controlled-section')
+            }
+          >
+            {isDrawerOpen ? 'Close drawer' : 'Open drawer'}
+          </button>
+        </div>
+      );
+    };
+    render(
+      <DrawerContentProvider>
+        <ControlElement />
+        <DrawerAnchor>
+          <DrawerSection
+            id="unrelated-section"
+            label="Test section 1"
+            title="Test section 1"
+            glyph="Trash"
+          >
+            This is an unrelated section
+          </DrawerSection>
+          <DrawerSection
+            id="controlled-section"
+            label="Test section 2"
+            title="Test section 2"
+            glyph="Bell"
+          >
+            This is the controlled section
+          </DrawerSection>
+        </DrawerAnchor>
+      </DrawerContentProvider>
+    );
+
+    // Drawer is closed by default
+    expect(screen.getByTestId('drawer-state')).to.have.text('closed');
+
+    // Open the drawer
+    userEvent.click(screen.getByTestId('toggle-drawer'));
+    await waitFor(() => {
+      expect(screen.getByTestId('drawer-state')).to.have.text('open');
+      expect(screen.getByText('This is the controlled section')).to.be.visible;
+    });
+
+    // Close the drawer
+    userEvent.click(screen.getByTestId('toggle-drawer'));
+    await waitFor(() => {
+      expect(screen.getByTestId('drawer-state')).to.have.text('closed');
+      expect(screen.queryByText('This is the controlled section')).not.to.exist;
+    });
+  });
 });

packages/compass-components/src/components/drawer-portal.tsx

@@ -32,6 +32,10 @@ type DrawerSectionProps = Omit<SectionData, 'content' | 'onClick'> & {
   order?: number;
 };
 
+type DrawerOpenStateContextValue = boolean;
+
+type DrawerSetOpenStateContextValue = (isOpen: boolean) => void;
+
 type DrawerActionsContextValue = {
   current: {
     openDrawer: (id: string) => void;
@@ -43,6 +47,12 @@ type DrawerActionsContextValue = {
 
 const DrawerStateContext = React.createContext<DrawerSectionProps[]>([]);
 
+const DrawerOpenStateContext =
+  React.createContext<DrawerOpenStateContextValue>(false);
+
+const DrawerSetOpenStateContext =
+  React.createContext<DrawerSetOpenStateContextValue>(() => {});
+
 const DrawerActionsContext = React.createContext<DrawerActionsContextValue>({
   current: {
     openDrawer: () => undefined,
@@ -89,6 +99,8 @@ export const DrawerContentProvider: React.FunctionComponent = ({
   children,
 }) => {
   const [drawerState, setDrawerState] = useState<DrawerSectionProps[]>([]);
+  const [drawerOpenState, setDrawerOpenState] =
+    useState<DrawerOpenStateContextValue>(false);
   const drawerActions = useRef({
     openDrawer: () => undefined,
     closeDrawer: () => undefined,
@@ -116,18 +128,26 @@ export const DrawerContentProvider: React.FunctionComponent = ({
 
   return (
     <DrawerStateContext.Provider value={drawerState}>
-      <DrawerActionsContext.Provider value={drawerActions}>
-        {children}
-      </DrawerActionsContext.Provider>
+      <DrawerOpenStateContext.Provider value={drawerOpenState}>
+        <DrawerSetOpenStateContext.Provider value={setDrawerOpenState}>
+          <DrawerActionsContext.Provider value={drawerActions}>
+            {children}
+          </DrawerActionsContext.Provider>
+        </DrawerSetOpenStateContext.Provider>
+      </DrawerOpenStateContext.Provider>
     </DrawerStateContext.Provider>
   );
 };
 
 const DrawerContextGrabber: React.FunctionComponent = ({ children }) => {
   const drawerToolbarContext = useDrawerToolbarContext();
   const actions = useContext(DrawerActionsContext);
+  const openStateSetter = useContext(DrawerSetOpenStateContext);
   actions.current.openDrawer = drawerToolbarContext.openDrawer;
   actions.current.closeDrawer = drawerToolbarContext.closeDrawer;
+  useEffect(() => {
+    openStateSetter(drawerToolbarContext.isDrawerOpen);
+  }, [drawerToolbarContext.isDrawerOpen, openStateSetter]);
   return <>{children}</>;
 };
 
@@ -321,3 +341,14 @@ export function useDrawerActions() {
   });
   return stableActions.current;
 }
+
+export const useDrawerState = () => {
+  const drawerOpenStateContext = useContext(DrawerOpenStateContext);
+  const drawerState = useContext(DrawerStateContext);
+  return {
+    isDrawerOpen:
+      drawerOpenStateContext &&
+      // the second check is a workaround, because LG doesn't set isDrawerOpen to false when it's empty
+      drawerState.length > 0,
+  };
+};

packages/compass-data-modeling/package.json

@@ -63,7 +63,7 @@
     "@mongodb-js/compass-user-data": "^0.9.0",
     "@mongodb-js/compass-utils": "^0.9.10",
     "@mongodb-js/compass-workspaces": "^0.51.0",
-    "@mongodb-js/diagramming": "^1.3.3",
+    "@mongodb-js/diagramming": "^1.3.5",
     "bson": "^6.10.4",
     "compass-preferences-model": "^2.50.0",
     "html-to-image": "1.11.11",

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

@@ -17,6 +17,7 @@ function renderDiagramEditorToolbar(
       onRedoClick={() => {}}
       onExportClick={() => {}}
       onRelationshipDrawingToggle={() => {}}
+      onAddCollectionClick={() => {}}
       {...props}
     />
   );
@@ -65,6 +66,16 @@ describe('DiagramEditorToolbar', function () {
     });
   });
 
+  context('add collection button', function () {
+    it('starts adding collection', function () {
+      const addCollectionSpy = sinon.spy();
+      renderDiagramEditorToolbar({ onAddCollectionClick: addCollectionSpy });
+      const addButton = screen.getByRole('button', { name: 'Add Collection' });
+      userEvent.click(addButton);
+      expect(addCollectionSpy).to.have.been.calledOnce;
+    });
+  });
+
   context('add relationship button', function () {
     it('renders it active if isInRelationshipDrawingMode is true', function () {
       renderDiagramEditorToolbar({ isInRelationshipDrawingMode: true });

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

@@ -15,7 +15,7 @@ import {
   transparentize,
   Tooltip,
 } from '@mongodb-js/compass-components';
-
+import AddCollection from './icons/add-collection';
 const containerStyles = css({
   display: 'flex',
   justifyContent: 'space-between',
@@ -50,6 +50,7 @@ export const DiagramEditorToolbar: React.FunctionComponent<{
   onRedoClick: () => void;
   onExportClick: () => void;
   onRelationshipDrawingToggle: () => void;
+  onAddCollectionClick: () => void;
 }> = ({
   step,
   hasUndo,
@@ -58,6 +59,7 @@ export const DiagramEditorToolbar: React.FunctionComponent<{
   onRedoClick,
   onExportClick,
   onRelationshipDrawingToggle,
+  onAddCollectionClick,
   isInRelationshipDrawingMode,
 }) => {
   const darkmode = useDarkMode();
@@ -70,6 +72,15 @@ export const DiagramEditorToolbar: React.FunctionComponent<{
       data-testid="diagram-editor-toolbar"
     >
       <div className={toolbarGroupStyles}>
+        <IconButton aria-label="Undo" disabled={!hasUndo} onClick={onUndoClick}>
+          <Icon glyph="Undo"></Icon>
+        </IconButton>
+        <IconButton aria-label="Redo" disabled={!hasRedo} onClick={onRedoClick}>
+          <Icon glyph="Redo"></Icon>
+        </IconButton>
+        <IconButton aria-label="Add Collection" onClick={onAddCollectionClick}>
+          <AddCollection />
+        </IconButton>
         <Tooltip
           trigger={
             <IconButton
@@ -88,12 +99,6 @@ export const DiagramEditorToolbar: React.FunctionComponent<{
         >
           Drag from one collection to another to create a relationship.
         </Tooltip>
-        <IconButton aria-label="Undo" disabled={!hasUndo} onClick={onUndoClick}>
-          <Icon glyph="Undo"></Icon>
-        </IconButton>
-        <IconButton aria-label="Redo" disabled={!hasRedo} onClick={onRedoClick}>
-          <Icon glyph="Redo"></Icon>
-        </IconButton>
       </div>
       <div className={toolbarGroupStyles}>
         <Button size="xsmall" aria-label="Export" onClick={onExportClick}>

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

@@ -16,6 +16,7 @@ import sinon from 'sinon';
 import { DiagramProvider } from '@mongodb-js/diagramming';
 import { DataModelingWorkspaceTab } from '..';
 import { openDiagram } from '../store/diagram';
+import { DrawerAnchor } from '@mongodb-js/compass-components';
 
 const storageItems: MongoDBDataModelDescription[] = [
   {
@@ -143,9 +144,11 @@ const renderDiagramEditor = ({
   const {
     plugin: { store },
   } = renderWithConnections(
-    <DiagramProvider fitView>
-      <DiagramEditor />
-    </DiagramProvider>
+    <DrawerAnchor>
+      <DiagramProvider fitView>
+        <DiagramEditor />
+      </DiagramProvider>
+    </DrawerAnchor>
   );
   store.dispatch(openDiagram(renderedItem));
 

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

@@ -15,6 +15,7 @@ import {
   type DiagramState,
   selectCurrentModelFromState,
   createNewRelationship,
+  addCollection,
 } from '../store/diagram';
 import {
   Banner,
@@ -25,6 +26,7 @@ import {
   Button,
   useDarkMode,
   useDrawerActions,
+  useDrawerState,
   rafraf,
 } from '@mongodb-js/compass-components';
 import { cancelAnalysis, retryAnalysis } from '../store/analysis-process';
@@ -107,6 +109,7 @@ const DiagramContent: React.FunctionComponent<{
   model: StaticModel | null;
   isInRelationshipDrawingMode: boolean;
   editErrors?: string[];
+  newCollection?: string;
   onMoveCollection: (ns: string, newPosition: [number, number]) => void;
   onCollectionSelect: (namespace: string) => void;
   onRelationshipSelect: (rId: string) => void;
@@ -118,6 +121,7 @@ const DiagramContent: React.FunctionComponent<{
   diagramLabel,
   model,
   isInRelationshipDrawingMode,
+  newCollection,
   onMoveCollection,
   onCollectionSelect,
   onRelationshipSelect,
@@ -129,6 +133,7 @@ const DiagramContent: React.FunctionComponent<{
   const isDarkMode = useDarkMode();
   const diagram = useRef(useDiagram());
   const { openDrawer } = useDrawerActions();
+  const { isDrawerOpen } = useDrawerState();
 
   const setDiagramContainerRef = useCallback((ref: HTMLDivElement | null) => {
     if (ref) {
@@ -183,6 +188,35 @@ const DiagramContent: React.FunctionComponent<{
     });
   }, []);
 
+  // Center on a new collection when it is added
+  const previouslyOpenedDrawer = useRef<boolean>(false);
+  useEffect(() => {
+    const wasDrawerPreviouslyOpened = previouslyOpenedDrawer.current;
+    previouslyOpenedDrawer.current = !!isDrawerOpen;
+
+    if (!newCollection) return;
+    const node = nodes.find((n) => n.id === newCollection);
+    if (!node) return;
+
+    // For calculating the center, we're taking into account the drawer,
+    // so that the new node is centered in the visible part.
+    const drawerOffset = wasDrawerPreviouslyOpened ? 0 : 240;
+    const zoom = diagram.current.getViewport().zoom;
+    const drawerOffsetInDiagramCoords = drawerOffset / zoom;
+    const newNodeWidth = 244;
+    const newNodeHeight = 64;
+    return rafraf(() => {
+      void diagram.current.setCenter(
+        node.position.x + newNodeWidth / 2 + drawerOffsetInDiagramCoords,
+        node.position.y + newNodeHeight / 2,
+        {
+          duration: 500,
+          zoom,
+        }
+      );
+    });
+  }, [newCollection, nodes, isDrawerOpen]);
+
   const handleNodesConnect = useCallback(
     (source: string, target: string) => {
       onCreateNewRelationship(source, target);
@@ -241,6 +275,7 @@ const ConnectedDiagramContent = connect(
       model: diagram ? selectCurrentModelFromState(state) : null,
       diagramLabel: diagram?.name || 'Schema Preview',
       selectedItems: state.diagram?.selectedItems ?? null,
+      newCollection: diagram?.draftCollection,
     };
   },
   {
@@ -257,7 +292,15 @@ const DiagramEditor: React.FunctionComponent<{
   diagramId?: string;
   onRetryClick: () => void;
   onCancelClick: () => void;
-}> = ({ step, diagramId, onRetryClick, onCancelClick }) => {
+  onAddCollectionClick: () => void;
+}> = ({
+  step,
+  diagramId,
+  onRetryClick,
+  onCancelClick,
+  onAddCollectionClick,
+}) => {
+  const { openDrawer } = useDrawerActions();
   let content;
 
   const [isInRelationshipDrawingMode, setIsInRelationshipDrawingMode] =
@@ -271,6 +314,11 @@ const DiagramEditor: React.FunctionComponent<{
     setIsInRelationshipDrawingMode(false);
   }, []);
 
+  const handleAddCollectionClick = useCallback(() => {
+    onAddCollectionClick();
+    openDrawer(DATA_MODELING_DRAWER_ID);
+  }, [openDrawer, onAddCollectionClick]);
+
   if (step === 'NO_DIAGRAM_SELECTED') {
     return null;
   }
@@ -320,6 +368,7 @@ const DiagramEditor: React.FunctionComponent<{
         <DiagramEditorToolbar
           onRelationshipDrawingToggle={handleRelationshipDrawingToggle}
           isInRelationshipDrawingMode={isInRelationshipDrawingMode}
+          onAddCollectionClick={handleAddCollectionClick}
         />
       }
     >
@@ -341,5 +390,6 @@ export default connect(
   {
     onRetryClick: retryAnalysis,
     onCancelClick: cancelAnalysis,
+    onAddCollectionClick: addCollection,
   }
 )(DiagramEditor);