feat(data-modeling): export diagram to png COMPASS-9449

packages/compass-data-modeling/package.json

@@ -64,6 +64,7 @@
     "@mongodb-js/compass-user-data": "^0.7.2",
     "@mongodb-js/compass-workspaces": "^0.42.0",
     "html-to-image": "1.11.11",
+    "@mongodb-js/diagramming": "^1.0.3",
     "bson": "^6.10.3",
     "compass-preferences-model": "^2.41.0",
     "lodash": "^4.17.21",

packages/compass-data-modeling/src/components/export-diagram-modal.tsx

@@ -1,4 +1,4 @@
-import React, { useCallback, useState } from 'react';
+import React, { useCallback, useEffect, useRef, useState } from 'react';
 import {
   Button,
   css,
@@ -13,6 +13,7 @@ import {
   RadioGroup,
   spacing,
   SpinLoader,
+  useToast,
 } from '@mongodb-js/compass-components';
 import {
   closeExportModal,
@@ -24,6 +25,7 @@ import type { DataModelingState } from '../store/reducer';
 import type { StaticModel } from '../services/data-model-storage';
 import { exportToJson, exportToPng } from '../services/export-diagram';
 import { useDiagram } from '@mongodb-js/diagramming';
+import { isCancelError } from '@mongodb-js/compass-utils';
 
 const nbsp = '\u00a0';
 
@@ -64,25 +66,68 @@ const ExportDiagramModal = ({
   const [exportFormat, setExportFormat] = useState<'png' | 'json' | null>(null);
   const diagram = useDiagram();
   const [isExporting, setIsExporting] = useState(false);
+  const abortControllerRef = useRef<AbortController | null>(null);
+  const toast = useToast();
+  useEffect(() => {
+    const cleanup = () => {
+      if (abortControllerRef.current) {
+        abortControllerRef.current.abort();
+        abortControllerRef.current = null;
+      }
+    };
+    const abortController = new AbortController();
+    if (isModalOpen) {
+      abortControllerRef.current = abortController;
+    } else {
+      cleanup();
+    }
+    return cleanup;
+  }, [isModalOpen]);
+
+  const onClose = useCallback(() => {
+    setExportFormat(null);
+    setIsExporting(false);
+    abortControllerRef.current?.abort();
+    abortControllerRef.current = null;
+    onCloseClick();
+  }, [onCloseClick]);
 
   const onExport = useCallback(async () => {
-    if (!exportFormat || !model) {
-      return;
+    try {
+      if (!exportFormat || !model) {
+        return;
+      }
+      setIsExporting(true);
+      if (exportFormat === 'json') {
+        exportToJson(diagramLabel, model);
+      } else if (exportFormat === 'png') {
+        await exportToPng(
+          diagramLabel,
+          diagram,
+          abortControllerRef.current?.signal
+        );
+      }
+    } catch (error) {
+      if (isCancelError(error)) {
+        return;
+      }
+      toast.pushToast({
+        id: 'export-diagram-error',
+        variant: 'warning',
+        title: 'Export failed',
+        description: `An error occurred while exporting the diagram: ${
+          (error as Error).message
+        }`,
+      });
+    } finally {
+      onClose();
     }
-    setIsExporting(true);
-    if (exportFormat === 'json') {
-      exportToJson(diagramLabel, model);
-    } else if (exportFormat === 'png') {
-      await exportToPng(diagramLabel, diagram);
-    }
-    onCloseClick();
-    setIsExporting(false);
-  }, [exportFormat, onCloseClick, model, diagram, diagramLabel]);
+  }, [exportFormat, onClose, model, diagram, diagramLabel, toast]);
 
   return (
     <Modal
       open={isModalOpen}
-      setOpen={onCloseClick}
+      setOpen={onClose}
       data-testid="export-diagram-modal"
     >
       <ModalHeader
@@ -140,11 +185,7 @@ const ExportDiagramModal = ({
         >
           Export
         </Button>
-        <Button
-          variant="default"
-          onClick={onCloseClick}
-          data-testid="cancel-button"
-        >
+        <Button variant="default" onClick={onClose} data-testid="cancel-button">
           Cancel
         </Button>
       </ModalFooter>

packages/compass-data-modeling/src/services/export-diagram.tsx

@@ -4,47 +4,15 @@ import {
   getViewportForBounds,
   DiagramProvider,
   Diagram,
+  mapEdgeToDiagramEdge,
+  mapNodeToDiagramNode,
 } from '@mongodb-js/diagramming';
-import type {
-  useDiagram,
-  NodeProps,
-  NodeType,
-  EdgeProps,
-  Marker,
-} from '@mongodb-js/diagramming';
+import type { DiagramInstance } from '@mongodb-js/diagramming';
 import type { StaticModel } from './data-model-storage';
 import ReactDOM from 'react-dom';
 import { toPng } from 'html-to-image';
 import { rafraf, spacing } from '@mongodb-js/compass-components';
-
-// TODO: Export these methods (and type) from the diagramming package
-type DiagramInstance = ReturnType<typeof useDiagram>;
-function mapNodeToDiagramNode(
-  node: ReturnType<DiagramInstance['getNodes']>[number]
-): NodeProps {
-  const { data, type, ...restOfNode } = node;
-  return {
-    ...restOfNode,
-    ...(data as any), // TODO: Type data (or expose these methods from the diagramming package)
-    type: type as NodeType,
-  };
-}
-function mapEdgeToDiagramEdge(
-  edge: ReturnType<DiagramInstance['getEdges']>[number]
-): EdgeProps {
-  const { markerStart, markerEnd, ...restOfEdge } = edge;
-
-  // The diagramming package only allows string based markers
-  if (typeof markerStart !== 'string' || typeof markerEnd !== 'string') {
-    throw new Error('Unexpected edge with non-string markers');
-  }
-
-  return {
-    ...restOfEdge,
-    markerEnd: markerEnd.replace('end-', '') as Marker,
-    markerStart: markerStart.replace('start-', '') as Marker,
-  };
-}
+import { raceWithAbort } from '@mongodb-js/compass-utils';
 
 function moveSvgDefsToViewportElement(
   container: Element,
@@ -67,33 +35,50 @@ function moveSvgDefsToViewportElement(
   markerDef.remove();
 }
 
-export async function exportToPng(fileName: string, diagram: DiagramInstance) {
-  const container = document.createElement('div');
-  container.setAttribute('data-testid', 'export-diagram-container');
-  // Push it out of the viewport
-  container.style.position = 'fixed';
-  container.style.top = '100vh';
-  container.style.left = '100vw';
-  document.body.appendChild(container);
-
-  const dataUri = await getExportPngDataUri(container, diagram);
-  downloadFile(dataUri, fileName, () => {
-    container.remove();
-  });
+export async function exportToPng(
+  fileName: string,
+  diagram: DiagramInstance,
+  signal?: AbortSignal
+) {
+  const dataUri = await raceWithAbort(
+    getExportPngDataUri(diagram),
+    signal ?? new AbortController().signal
+  );
+  downloadFile(dataUri, fileName);
 }
 
-export function getExportPngDataUri(
-  container: HTMLDivElement,
-  diagram: DiagramInstance
-): Promise<string> {
-  return new Promise<string>((resolve, reject) => {
+export function getExportPngDataUri(diagram: DiagramInstance): Promise<string> {
+  return new Promise<string>((resolve, _reject) => {
     const nodes = diagram.getNodes();
     const edges = diagram.getEdges();
+    const bounds = getNodesBounds(nodes);
+
+    const container = document.createElement('div');
+    container.setAttribute('data-testid', 'export-diagram-container');
+    // Push it out of the viewport
+    container.style.position = 'fixed';
+    container.style.top = '100vh';
+    container.style.left = '100vw';
+    container.style.width = `${bounds.width}px`;
+    container.style.height = `${bounds.height}px`;
+    document.body.appendChild(container);
+
+    const diagramEdges = edges.map(mapEdgeToDiagramEdge);
+    const diagramNodes = nodes.map(mapNodeToDiagramNode).map((node) => ({
+      ...node,
+      selected: false, // Dont show selected state (blue border)
+    }));
+
+    const reject = (error: Error) => {
+      document.body.removeChild(container);
+      _reject(error);
+    };
+
     ReactDOM.render(
       <DiagramProvider>
         <Diagram
-          edges={edges.map(mapEdgeToDiagramEdge)}
-          nodes={nodes.map(mapNodeToDiagramNode)}
+          edges={diagramEdges}
+          nodes={diagramNodes}
           onlyRenderVisibleElements={false}
         />
       </DiagramProvider>,
@@ -110,10 +95,9 @@ export function getExportPngDataUri(
             '.react-flow__viewport'
           );
           if (!viewportElement) {
-            throw new Error('Diagram element not found');
+            return reject(new Error('Diagram element not found'));
           }
 
-          const bounds = getNodesBounds(nodes);
           const transform = getViewportForBounds(
             bounds,
             bounds.width,
@@ -122,6 +106,7 @@ export function getExportPngDataUri(
             2, // Maximum zoom
             `${spacing[400]}px` // 16px padding
           );
+
           // Moving svg defs to the viewport element
           moveSvgDefsToViewportElement(container, viewportElement);
           rafraf(() => {
@@ -172,13 +157,13 @@ export function getExportJsonFromModel({
   };
 }
 
-function downloadFile(uri: string, fileName: string, cleanup: () => void) {
+function downloadFile(uri: string, fileName: string, cleanup?: () => void) {
   const link = document.createElement('a');
   link.download = fileName;
   link.href = uri;
   link.click();
   setTimeout(() => {
     link.remove();
-    cleanup();
+    cleanup?.();
   }, 0);
 }