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

packages/compass-data-modeling/package.json

@@ -63,6 +63,7 @@
     "@mongodb-js/compass-telemetry": "^1.10.0",
     "@mongodb-js/compass-user-data": "^0.7.2",
     "@mongodb-js/compass-workspaces": "^0.42.0",
+    "html-to-image": "1.11.11",
     "bson": "^6.10.3",
     "compass-preferences-model": "^2.41.0",
     "@mongodb-js/compass-app-registry": "^9.4.11",

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

@@ -1,4 +1,4 @@
-import React, { useCallback, useState } from 'react';
+import React, { useCallback, useMemo, useRef, useState } from 'react';
 import {
   Button,
   css,
@@ -12,6 +12,7 @@ import {
   Radio,
   RadioGroup,
   spacing,
+  SpinLoader,
 } from '@mongodb-js/compass-components';
 import {
   closeExportModal,
@@ -21,7 +22,8 @@ import {
 import { connect } from 'react-redux';
 import type { DataModelingState } from '../store/reducer';
 import type { StaticModel } from '../services/data-model-storage';
-import { exportToJson } from '../services/export-diagram';
+import { exportToJson, exportToPng } from '../services/export-diagram';
+import { getNodesBounds, useDiagram } from '@mongodb-js/diagramming';
 
 const nbsp = '\u00a0';
 
@@ -59,15 +61,25 @@ const ExportDiagramModal = ({
   model,
   onCloseClick,
 }: ExportDiagramModalProps) => {
-  const [exportFormat, setExportFormat] = useState<'json' | null>(null);
+  const [exportFormat, setExportFormat] = useState<'png' | 'json' | null>(null);
+  const diagram = useDiagram();
+  const [isExporting, setIsExporting] = useState(false);
+  const exportDiagramContainerRef = useRef<HTMLDivElement>(null);
+  const bounds = useMemo(() => getNodesBounds(diagram.getNodes()), [diagram]);
 
-  const onExport = useCallback(() => {
+  const onExport = useCallback(async () => {
     if (!exportFormat || !model) {
       return;
     }
-    exportToJson(diagramLabel, model);
+    setIsExporting(true);
+    if (exportFormat === 'json') {
+      exportToJson(diagramLabel, model);
+    } else if (exportFormat === 'png') {
+      await exportToPng(diagramLabel, exportDiagramContainerRef, diagram);
+    }
     onCloseClick();
-  }, [exportFormat, onCloseClick, model, diagramLabel]);
+    setIsExporting(false);
+  }, [exportFormat, onCloseClick, model, diagram, diagramLabel]);
 
   return (
     <Modal
@@ -95,6 +107,17 @@ const ExportDiagramModal = ({
         <div className={contentContainerStyles}>
           <Label htmlFor="">Select file format:</Label>
           <RadioGroup className={contentContainerStyles} value={exportFormat}>
+            <div className={radioItemStyles}>
+              <Icon glyph="Diagram2" />
+              <Radio
+                checked={exportFormat === 'png'}
+                value="png"
+                aria-label="PNG"
+                onClick={() => setExportFormat('png')}
+              >
+                PNG
+              </Radio>
+            </div>
             <div className={radioItemStyles}>
               <Icon glyph="CurlyBraces" />
               <Radio
@@ -108,12 +131,26 @@ const ExportDiagramModal = ({
             </div>
           </RadioGroup>
         </div>
+        {/* Container where we render export diagram when exporting png */}
+        <div
+          style={{
+            width: bounds.width,
+            height: bounds.height,
+            // Fixed at bottom right corner
+            position: 'fixed',
+            top: '100vh',
+            left: '100vw',
+          }}
+          ref={exportDiagramContainerRef}
+        />
       </ModalBody>
       <ModalFooter className={footerStyles}>
         <Button
           variant="primary"
           onClick={() => void onExport()}
           data-testid="export-button"
+          disabled={isExporting}
+          loadingIndicator={<SpinLoader />}
         >
           Export
         </Button>

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

@@ -0,0 +1,184 @@
+import React from 'react';
+import {
+  getNodesBounds,
+  getViewportForBounds,
+  DiagramProvider,
+  Diagram,
+} from '@mongodb-js/diagramming';
+import type {
+  useDiagram,
+  NodeProps,
+  NodeType,
+  EdgeProps,
+  Marker,
+} 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,
+  };
+}
+
+function moveSvgDefsToViewportElement(
+  container: Element,
+  targetElement: Element
+) {
+  const markerDef = container.querySelector('svg defs');
+  if (!markerDef) {
+    return;
+  }
+  const diagramSvgElements = targetElement.querySelectorAll('svg');
+  diagramSvgElements.forEach((svg) => {
+    const pathsWithMarkers = svg.querySelectorAll(
+      'path[marker-end], path[marker-start]'
+    );
+    if (pathsWithMarkers.length !== 0) {
+      const clonedDef = markerDef.cloneNode(true) as SVGMarkerElement;
+      svg.insertBefore(clonedDef, svg.firstChild);
+    }
+  });
+  markerDef.remove();
+}
+
+export async function exportToPng(
+  fileName: string,
+  containerRef: React.RefObject<HTMLDivElement>,
+  diagram: DiagramInstance
+) {
+  const container = containerRef.current;
+  if (!container) {
+    throw new Error('Container reference is not set');
+  }
+  const dataUri = await getExportPngDataUri(container, diagram);
+  downloadFile(dataUri, fileName, () => {
+    ReactDOM.unmountComponentAtNode(container);
+  });
+}
+
+export function getExportPngDataUri(
+  container: HTMLDivElement,
+  diagram: DiagramInstance
+): Promise<string> {
+  return new Promise<string>((resolve, reject) => {
+    const nodes = diagram.getNodes();
+    const edges = diagram.getEdges();
+    ReactDOM.render(
+      <DiagramProvider>
+        <Diagram
+          edges={edges.map(mapEdgeToDiagramEdge)}
+          nodes={nodes.map(mapNodeToDiagramNode)}
+          onlyRenderVisibleElements={false}
+        />
+      </DiagramProvider>,
+      container,
+      () => {
+        rafraf(() => {
+          // For export we are selecting react-flow__viewport element,
+          // which contains the export canvas. It excludes diagram
+          // title, minmap, and other UI elements. However, it also
+          // excludes the svg defs that are currently outside of this element.
+          // So, when exporting, we need to include those defs as well so that
+          // edge markers are exported correctly.
+          const viewportElement = container.querySelector(
+            '.react-flow__viewport'
+          );
+          if (!viewportElement) {
+            throw new Error('Diagram element not found');
+          }
+
+          const bounds = getNodesBounds(nodes);
+          const transform = getViewportForBounds(
+            bounds,
+            bounds.width,
+            bounds.height,
+            0.5, // Minimum zoom
+            2, // Maximum zoom
+            `${spacing[400]}px` // 16px padding
+          );
+          // Moving svg defs to the viewport element
+          moveSvgDefsToViewportElement(container, viewportElement);
+          rafraf(() => {
+            toPng(viewportElement as HTMLElement, {
+              backgroundColor: '#fff',
+              pixelRatio: 2,
+              width: bounds.width,
+              height: bounds.height,
+              style: {
+                width: `${bounds.width}px`,
+                height: `${bounds.height}px`,
+                transform: `translate(${transform.x}px, ${transform.y}px) scale(${transform.zoom})`,
+              },
+            })
+              .then(resolve)
+              .catch(reject);
+          });
+        });
+      }
+    );
+  });
+}
+
+export function exportToJson(fileName: string, model: StaticModel) {
+  const json = getExportJsonFromModel(model);
+  const blob = new Blob([JSON.stringify(json, null, 2)], {
+    type: 'application/json',
+  });
+  const url = window.URL.createObjectURL(blob);
+  downloadFile(url, fileName, () => {
+    window.URL.revokeObjectURL(url);
+  });
+}
+
+export function getExportJsonFromModel({
+  collections,
+  relationships,
+}: StaticModel) {
+  return {
+    collections: Object.fromEntries(
+      collections.map((collection) => {
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        const { ns, jsonSchema, ...ignoredProps } = collection;
+        return [ns, { ns, jsonSchema }];
+      })
+    ),
+    relationships,
+  };
+}
+
+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();
+  }, 0);
+}