fixup: distinct stop analysis and cleanup actions, improve traversal value setting

Author: Anemy

packages/compass-schema/src/components/compass-schema.tsx

@@ -446,7 +446,7 @@ export default connect(
   }),
   {
     onStartAnalysis: startAnalysis,
-    onStopAnalysis: () => stopAnalysis(true),
+    onStopAnalysis: () => stopAnalysis(),
     onExportSchemaClicked: openExportSchema,
   }
 )(Schema);

packages/compass-schema/src/modules/schema-analysis.spec.ts

@@ -573,7 +573,6 @@ describe('schema-analysis', function () {
             Number: 1,
             Boolean: 1,
             String: 4,
-            Undefined: 3,
           });
         });
       });
@@ -604,7 +603,6 @@ describe('schema-analysis', function () {
           const { field_types } = await calculateSchemaMetadata(schema);
           expect(field_types).to.deep.equal({
             Number: 3,
-            Undefined: 1,
           });
         });
       });
@@ -636,7 +634,6 @@ describe('schema-analysis', function () {
             Array: 4,
             String: 2,
             Number: 3,
-            Undefined: 1,
           });
         });
       });

packages/compass-schema/src/modules/schema-analysis.ts

@@ -139,22 +139,23 @@ export async function calculateSchemaMetadata(schema: Schema): Promise<{
   let variableTypeCount = 0;
   let optionalFieldCount = 0;
   let unblockThreadCounter = 0;
+  let deepestPath = 0;
 
   async function traverseSchemaTree(
     fieldsOrTypes: SchemaField[] | SchemaType[],
-    isRoot = false
-  ): Promise<number> {
+    depth: number
+  ): Promise<void> {
     unblockThreadCounter++;
     if (unblockThreadCounter === UNBLOCK_INTERVAL_COUNT) {
       unblockThreadCounter = 0;
       await unblockThread();
     }
 
     if (!fieldsOrTypes || fieldsOrTypes.length === 0) {
-      return 0;
+      return;
     }
 
-    let deepestPath = 1;
+    deepestPath = Math.max(depth, deepestPath);
 
     for (const fieldOrType of fieldsOrTypes) {
       if (
@@ -170,9 +171,11 @@ export async function calculateSchemaMetadata(schema: Schema): Promise<{
       }
 
       if (isSchemaType(fieldOrType)) {
-        fieldTypes[fieldOrType.bsonType] =
-          (fieldTypes[fieldOrType.bsonType] ?? 0) + 1;
-      } else if (isRoot) {
+        if (fieldOrType.bsonType !== 'Undefined') {
+          fieldTypes[fieldOrType.bsonType] =
+            (fieldTypes[fieldOrType.bsonType] ?? 0) + 1;
+        }
+      } else if (depth === 1 /* is root level */) {
         // Count variable types (more than one unique type excluding undefined).
         if (
           fieldOrType.types &&
@@ -188,41 +191,31 @@ export async function calculateSchemaMetadata(schema: Schema): Promise<{
       }
 
       if ((fieldOrType as DocumentSchemaType).bsonType === 'Document') {
-        const deepestFieldPath =
-          (await traverseSchemaTree(
-            (fieldOrType as DocumentSchemaType).fields
-          )) + 1; /* Increment by one when we go a level deeper. */
-
-        deepestPath = Math.max(deepestFieldPath, deepestPath);
+        await traverseSchemaTree(
+          (fieldOrType as DocumentSchemaType).fields,
+          depth + 1 // Increment by one when we go a level deeper.
+        );
       } else if (
         (fieldOrType as ArraySchemaType).bsonType === 'Array' ||
         (fieldOrType as SchemaField).types
       ) {
-        // Increment by one when we go a level deeper.
         const increment =
           (fieldOrType as ArraySchemaType).bsonType === 'Array' ? 1 : 0;
-        const deepestFieldPath =
-          (await traverseSchemaTree(
-            (fieldOrType as ArraySchemaType | SchemaField).types
-          )) + increment;
-
-        deepestPath = Math.max(deepestFieldPath, deepestPath);
+        await traverseSchemaTree(
+          (fieldOrType as ArraySchemaType | SchemaField).types,
+          depth + increment // Increment by one when we go a level deeper.
+        );
       }
     }
-
-    return deepestPath;
   }
 
-  const schemaDepth = await traverseSchemaTree(
-    schema.fields,
-    true /* isRoot */
-  );
+  await traverseSchemaTree(schema.fields, 1);
 
   return {
     field_types: fieldTypes,
     geo_data: hasGeoData,
     optional_field_count: optionalFieldCount,
-    schema_depth: schemaDepth,
+    schema_depth: deepestPath,
     variable_type_count: variableTypeCount,
   };
 }

packages/compass-schema/src/stores/schema-analysis-reducer.ts

@@ -29,6 +29,7 @@ const ERROR_CODE_MAX_TIME_MS_EXPIRED = 50;
 
 export type SchemaAnalysisState = {
   analysisState: AnalysisState;
+  analysisStartTime?: number;
   errorMessage: string;
   schema: Schema | null;
   resultId: string;
@@ -42,6 +43,7 @@ export const enum SchemaAnalysisActions {
 
 export type AnalysisStartedAction = {
   type: SchemaAnalysisActions.analysisStarted;
+  analysisStartTime: number;
 };
 
 export type AnalysisFinishedAction = {
@@ -66,6 +68,7 @@ export const schemaAnalysisReducer: Reducer<SchemaAnalysisState, Action> = (
   ) {
     return {
       ...state,
+      analysisStartTime: action.analysisStartTime,
       analysisState: ANALYSIS_STATE_ANALYZING,
       errorMessage: '',
       schema: null,
@@ -164,19 +167,34 @@ export const geoLayersDeleted = (
   };
 };
 
-// We track when the user cancels the analysis. We use the userCancelled
-// flag to indicate if we should track, as we also automatically stop
-// analysis when deactivating the plugin.
-const userCancelledAnalysisAbortReason = 'Cancelled analysis';
+export const stopAnalysis = (): SchemaThunkAction<void> => {
+  return (
+    dispatch,
+    getState,
+    { analysisAbortControllerRef, connectionInfoRef, queryBar, track }
+  ) => {
+    if (!analysisAbortControllerRef.current) return;
+    const analysisTime =
+      Date.now() - (getState().schemaAnalysis.analysisStartTime ?? 0);
+
+    const query = queryBar.getLastAppliedQuery('schema');
+    track(
+      'Schema Analysis Cancelled',
+      {
+        analysis_time_ms: analysisTime,
+        with_filter: Object.entries(query.filter ?? {}).length > 0,
+      },
+      connectionInfoRef.current
+    );
+
+    analysisAbortControllerRef.current?.abort('Analysis cancelled');
+  };
+};
 
-export const stopAnalysis = (
-  userCancelled = false
-): SchemaThunkAction<void> => {
+export const cleanupAnalysis = (): SchemaThunkAction<void> => {
   return (dispatch, getState, { analysisAbortControllerRef }) => {
     if (!analysisAbortControllerRef.current) return;
-    analysisAbortControllerRef.current?.abort(
-      userCancelled ? userCancelledAnalysisAbortReason : undefined
-    );
+    analysisAbortControllerRef.current?.abort();
   };
 };
 
@@ -271,7 +289,10 @@ export const startAnalysis = (): SchemaThunkAction<
     try {
       debug('analysis started');
 
-      dispatch({ type: SchemaAnalysisActions.analysisStarted });
+      dispatch({
+        type: SchemaAnalysisActions.analysisStarted,
+        analysisStartTime,
+      });
 
       const schemaAccessor = await analyzeSchema(
         dataService,
@@ -316,21 +337,6 @@ export const startAnalysis = (): SchemaThunkAction<
 
       geoLayersRef.current = {};
     } catch (err: any) {
-      if (
-        abortSignal.aborted &&
-        abortSignal.reason === userCancelledAnalysisAbortReason
-      ) {
-        const analysisTime = Date.now() - analysisStartTime;
-        track(
-          'Schema Analysis Cancelled',
-          {
-            analysis_time_ms: analysisTime,
-            with_filter: Object.entries(query.filter ?? {}).length > 0,
-          },
-          connectionInfoRef.current
-        );
-      }
-
       log.error(
         mongoLogId(1_001_000_188),
         'Schema analysis',

packages/compass-schema/src/stores/store.spec.ts

@@ -112,10 +112,17 @@ describe('Schema Store', function () {
       sampleStub.resolves([{ name: 'Hans' }, { name: 'Greta' }]);
       await store.dispatch(startAnalysis());
       expect(sampleStub).to.have.been.called;
-      const { analysisState, errorMessage, schema, resultId } =
-        store.getState().schemaAnalysis;
+      const {
+        analysisState,
+        errorMessage,
+        schema,
+        resultId,
+        analysisStartTime,
+      } = store.getState().schemaAnalysis;
       expect(analysisState).to.equal('complete');
       expect(!!errorMessage).to.be.false;
+      expect(analysisStartTime).to.not.be.undefined;
+      expect(analysisStartTime).to.be.greaterThan(1000);
       expect(schema).not.to.be.null;
       expect(resultId).not.to.equal(oldResultId);
     });

packages/compass-schema/src/stores/store.ts

@@ -18,7 +18,10 @@ import type { FieldStoreService } from '@mongodb-js/compass-field-store';
 import type { QueryBarService } from '@mongodb-js/compass-query-bar';
 import type { TrackFunction } from '@mongodb-js/compass-telemetry';
 import type { SchemaAccessor } from 'mongodb-schema';
-import { schemaAnalysisReducer, stopAnalysis } from './schema-analysis-reducer';
+import {
+  schemaAnalysisReducer,
+  cleanupAnalysis,
+} from './schema-analysis-reducer';
 import {
   cancelExportSchema,
   confirmedExportLegacySchemaToClipboard,
@@ -88,7 +91,7 @@ export function activateSchemaPlugin(
     store.dispatch(confirmedExportLegacySchemaToClipboard());
   });
 
-  addCleanup(() => store.dispatch(stopAnalysis()));
+  addCleanup(() => store.dispatch(cleanupAnalysis()));
   addCleanup(() => store.dispatch(cancelExportSchema()));
 
   return {