controller in args and add comment for rafraf

Author: mabaasit

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

@@ -75,6 +75,9 @@ export function getExportPngDataUri(diagram: DiagramInstance): Promise<string> {
       </DiagramProvider>,
       container,
       () => {
+        // We skip some frames here to ensure that the DOM has fully rendered and React has
+        // committed all updates before we try to query for viewport element. Without this,
+        // the element may not exist yet or may not have the correct styles etc.
         rafraf(() => {
           // For export we are selecting react-flow__viewport element,
           // which contains the export canvas. It excludes diagram

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

@@ -144,7 +144,7 @@ export function startAnalysis(
     const namespaces = collections.map((collName) => {
       return `${database}.${collName}`;
     });
-    const cancelController = (services.cancelControllerRef.current =
+    const cancelController = (services.cancelAnalysisControllerRef.current =
       new AbortController());
     dispatch({
       type: AnalysisProcessActionTypes.ANALYZING_COLLECTIONS_START,
@@ -223,7 +223,7 @@ export function startAnalysis(
         });
       }
     } finally {
-      services.cancelControllerRef.current = null;
+      services.cancelAnalysisControllerRef.current = null;
     }
   };
 }
@@ -250,8 +250,8 @@ export function retryAnalysis(): DataModelingThunkAction<void, never> {
 }
 
 export function cancelAnalysis(): DataModelingThunkAction<void, never> {
-  return (_dispatch, _getState, { cancelControllerRef }) => {
-    cancelControllerRef.current?.abort();
-    cancelControllerRef.current = null;
+  return (_dispatch, _getState, { cancelAnalysisControllerRef }) => {
+    cancelAnalysisControllerRef.current?.abort();
+    cancelAnalysisControllerRef.current = null;
   };
 }

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

@@ -56,8 +56,6 @@ const INITIAL_STATE = {
   isExporting: false,
 };
 
-let cancelExportAbortController: AbortController | null = null;
-
 export const exportDiagramReducer: Reducer<ExportDiagramState> = (
   state = INITIAL_STATE,
   action
@@ -103,7 +101,7 @@ export function exportDiagram(
   Promise<void>,
   ExportStartedAction | ExportCompletedAction
 > {
-  return async (dispatch, getState, { track }) => {
+  return async (dispatch, getState, { track, cancelExportControllerRef }) => {
     const {
       exportDiagram: { exportFormat, isExporting },
       diagram,
@@ -117,7 +115,8 @@ export function exportDiagram(
         type: ExportDiagramActionTypes.EXPORT_STARTED,
       });
 
-      cancelExportAbortController = new AbortController();
+      const cancelController = (cancelExportControllerRef.current =
+        new AbortController());
 
       const model = selectCurrentModel(getCurrentDiagramFromState(getState()));
       if (exportFormat === 'json') {
@@ -126,7 +125,7 @@ export function exportDiagram(
         await exportToPng(
           diagram.name,
           diagramInstance,
-          cancelExportAbortController.signal
+          cancelController.signal
         );
       }
       track('Data Modeling Diagram Exported', {
@@ -143,7 +142,7 @@ export function exportDiagram(
         });
       }
     } finally {
-      cancelExportAbortController = null;
+      cancelExportControllerRef.current = null;
       dispatch({
         type: ExportDiagramActionTypes.EXPORT_COMPLETED,
       });
@@ -155,7 +154,7 @@ export function showExportModal(): DataModelingThunkAction<
   void,
   ModalOpenedAction
 > {
-  return (dispatch, getState, { track }) => {
+  return (dispatch, _getState, { track }) => {
     track('Screen', { name: 'export_diagram_modal' }, undefined);
     return dispatch({ type: ExportDiagramActionTypes.MODAL_OPENED });
   };
@@ -165,8 +164,8 @@ export function closeExportModal(): DataModelingThunkAction<
   void,
   ModalClosedAction
 > {
-  return (dispatch) => {
-    cancelExportAbortController?.abort();
+  return (dispatch, _getState, { cancelExportControllerRef }) => {
+    cancelExportControllerRef.current?.abort();
     dispatch({ type: ExportDiagramActionTypes.MODAL_CLOSED });
   };
 }

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

@@ -6,6 +6,7 @@ import type { MongoDBInstancesManager } from '@mongodb-js/compass-app-stores/pro
 import type { DataModelStorageService } from '../provider';
 import { applyMiddleware, createStore } from 'redux';
 import reducer from './reducer';
+import type { DataModelingExtraArgs } from './reducer';
 import thunk from 'redux-thunk';
 import type { ActivateHelpers } from '@mongodb-js/compass-app-registry';
 
@@ -25,11 +26,16 @@ export function activateDataModelingStore(
   services: DataModelingStoreServices,
   { cleanup }: ActivateHelpers
 ) {
-  const cancelControllerRef = { current: null };
+  const cancelAnalysisControllerRef = { current: null };
+  const cancelExportControllerRef = { current: null };
   const store = createStore(
     reducer,
     applyMiddleware(
-      thunk.withExtraArgument({ ...services, cancelControllerRef })
+      thunk.withExtraArgument<DataModelingExtraArgs>({
+        ...services,
+        cancelAnalysisControllerRef,
+        cancelExportControllerRef,
+      })
     )
   );
   return { store, deactivate: cleanup };

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

@@ -44,7 +44,8 @@ export type DataModelingActionTypes =
 export type DataModelingState = ReturnType<typeof reducer>;
 
 export type DataModelingExtraArgs = DataModelingStoreServices & {
-  cancelControllerRef: { current: AbortController | null };
+  cancelAnalysisControllerRef: { current: AbortController | null };
+  cancelExportControllerRef: { current: AbortController | null };
 };
 
 export type DataModelingThunkAction<R, A extends AnyAction> = ThunkAction<