feat(schema, data-modeling): use iterable cursor in schema analysis COMPASS-9150 COMPASS-9315 (#6894)

Author: Anemy

package-lock.json

@@ -66,6 +66,7 @@
     "compass-preferences-model": "^2.36.0",
     "hadron-app-registry": "^9.4.8",
     "lodash": "^4.17.21",
+    "mongodb": "^6.14.1",
     "mongodb-ns": "^2.4.2",
     "mongodb-schema": "^12.5.2",
     "react": "^17.0.2",

packages/compass-data-modeling/package.json

@@ -3,6 +3,8 @@ import { isAction } from './util';
 import type { DataModelingThunkAction } from './reducer';
 import { analyzeDocuments } from 'mongodb-schema';
 import { getCurrentDiagramFromState } from './diagram';
+import type { Document } from 'bson';
+import type { AggregationCursor } from 'mongodb';
 
 export type AnalysisProcessState = {
   currentAnalysisOptions:
@@ -155,32 +157,32 @@ export function startAnalysis(
     try {
       const dataService =
         services.connections.getDataServiceForConnection(connectionId);
-      const samples = await Promise.all(
+      const schema = await Promise.all(
         namespaces.map(async (ns) => {
-          // TODO
-          const sample = await dataService.sample(
+          const sample: AggregationCursor<Document> = dataService.sampleCursor(
             ns,
             { size: 100 },
-            undefined,
             {
-              abortSignal: cancelController.signal,
+              signal: cancelController.signal,
+              promoteValues: false,
+            },
+            {
+              fallbackReadPreference: 'secondaryPreferred',
             }
           );
+
+          const accessor = await analyzeDocuments(sample, {
+            signal: cancelController.signal,
+          });
+
+          // TODO(COMPASS-9314): Update how we show analysis progress.
           dispatch({
             type: AnalysisProcessActionTypes.NAMESPACE_SAMPLE_FETCHED,
             namespace: ns,
           });
-          return { ns, sample };
-        })
-      );
-      const schema = await Promise.all(
-        samples.map(async ({ ns, sample }) => {
-          const schema = await analyzeDocuments(sample, {
+
+          const schema = await accessor.getMongoDBJsonSchema({
             signal: cancelController.signal,
-          }).then((accessor) => {
-            return accessor.getMongoDBJsonSchema({
-              signal: cancelController.signal,
-            });
           });
           dispatch({
             type: AnalysisProcessActionTypes.NAMESPACE_SCHEMA_ANALYZED,

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

@@ -13,6 +13,7 @@ import { createLoggerLocator } from '@mongodb-js/compass-logging/provider';
 import { telemetryLocator } from '@mongodb-js/compass-telemetry/provider';
 import { SchemaValidationTabTitle } from './plugin-title';
 import { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provider';
+import type { RequiredDataServiceProps } from './modules';
 
 const CompassSchemaValidationHadronPlugin = registerHadronPlugin(
   {
@@ -23,13 +24,8 @@ const CompassSchemaValidationHadronPlugin = registerHadronPlugin(
     activate: onActivated,
   },
   {
-    dataService: dataServiceLocator as DataServiceLocator<
-      | 'aggregate'
-      | 'collectionInfo'
-      | 'updateCollection'
-      | 'sample'
-      | 'isCancelError'
-    >,
+    dataService:
+      dataServiceLocator as DataServiceLocator<RequiredDataServiceProps>,
     connectionInfoRef: connectionInfoRefLocator,
     instance: mongoDBInstanceLocator,
     preferences: preferencesLocator,

packages/compass-schema-validation/src/index.ts

@@ -67,14 +67,13 @@ export type RootAction =
   | EditModeAction
   | ResetAction;
 
-export type DataService = Pick<
-  OriginalDataService,
+export type RequiredDataServiceProps =
   | 'aggregate'
   | 'collectionInfo'
   | 'updateCollection'
-  | 'sample'
-  | 'isCancelError'
->;
+  | 'sampleCursor'
+  | 'isCancelError';
+export type DataService = Pick<OriginalDataService, RequiredDataServiceProps>;
 
 export type SchemaValidationExtraArgs = {
   dataService: DataService;

packages/compass-schema-validation/src/modules/index.ts

@@ -42,7 +42,13 @@ const fakeDataService = {
       /* never resolves */
     }),
   isCancelError: () => false,
-  sample: () => [{ prop1: 'abc' }],
+  sampleCursor: () =>
+    ({
+      async *[Symbol.asyncIterator]() {
+        await new Promise((resolve) => setTimeout(resolve, 0));
+        yield* [{ prop1: 'abc' }];
+      },
+    } as any),
 } as any;
 
 const fakeWorkspaces = {

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

@@ -8,6 +8,7 @@ import {
 import CompassSchema from './components/compass-schema';
 import { registerHadronPlugin } from 'hadron-app-registry';
 import { activateSchemaPlugin } from './stores/store';
+import type { RequiredDataServiceProps } from './stores/store';
 import { createLoggerLocator } from '@mongodb-js/compass-logging/provider';
 import { telemetryLocator } from '@mongodb-js/compass-telemetry/provider';
 import { preferencesLocator } from 'compass-preferences-model/provider';
@@ -31,9 +32,8 @@ const CompassSchemaHadronPlugin = registerHadronPlugin(
     activate: activateSchemaPlugin,
   },
   {
-    dataService: dataServiceLocator as DataServiceLocator<
-      'sample' | 'isCancelError'
-    >,
+    dataService:
+      dataServiceLocator as DataServiceLocator<RequiredDataServiceProps>,
     logger: createLoggerLocator('COMPASS-SCHEMA-UI'),
     track: telemetryLocator,
     preferences: preferencesLocator,

packages/compass-schema/src/index.ts

@@ -5,13 +5,13 @@ import mongoDBSchemaAnalyzeSchema from 'mongodb-schema';
 import type { Schema } from 'mongodb-schema';
 import { createNoopLogger } from '@mongodb-js/compass-logging/provider';
 import { isInternalFieldPath } from 'hadron-document';
-
-import { analyzeSchema, calculateSchemaMetadata } from './schema-analysis';
 import {
   createSandboxFromDefaultPreferences,
   type PreferencesAccess,
 } from 'compass-preferences-model';
 
+import { analyzeSchema, calculateSchemaMetadata } from './schema-analysis';
+
 const testDocs = [
   {
     someFields: {
@@ -81,13 +81,18 @@ describe('schema-analysis', function () {
 
   describe('#getResult', function () {
     it('returns the schema', async function () {
+      const docs = [
+        { x: 1 },
+        { y: 2, __safeContent__: [bson.Binary.createFromBase64('aaaa')] },
+      ];
       const dataService = {
-        sample: () =>
-          Promise.resolve([
-            { x: 1 },
-            { y: 2, __safeContent__: [bson.Binary.createFromBase64('aaaa')] },
-          ]),
-        isCancelError: () => false,
+        sampleCursor: () =>
+          ({
+            async *[Symbol.asyncIterator]() {
+              await new Promise((resolve) => setTimeout(resolve, 0));
+              yield* docs;
+            },
+          } as any),
       };
       const abortController = new AbortController();
       const abortSignal = abortController.signal;
@@ -179,10 +184,17 @@ describe('schema-analysis', function () {
 
     it('adds promoteValues: false so the analyzer can report more accurate types', async function () {
       const dataService = {
-        sample: () => Promise.resolve([]),
-        isCancelError: () => false,
+        sampleCursor: () =>
+          ({
+            async *[Symbol.asyncIterator]() {
+              await new Promise((resolve) => setTimeout(resolve, 0));
+              yield {
+                a: 123,
+              };
+            },
+          } as any),
       };
-      const sampleSpy = sinon.spy(dataService, 'sample');
+      const sampleSpy = sinon.spy(dataService, 'sampleCursor');
       const abortController = new AbortController();
       const abortSignal = abortController.signal;
 
@@ -199,19 +211,31 @@ describe('schema-analysis', function () {
       expect(sampleSpy).to.have.been.calledWith(
         'db.coll',
         {},
-        { promoteValues: false }
+        { signal: abortSignal, promoteValues: false },
+        { fallbackReadPreference: 'secondaryPreferred' }
       );
     });
 
     it('returns undefined if is cancelled', async function () {
-      const dataService = {
-        sample: () => Promise.reject(new Error('test error')),
-        isCancelError: () => true,
-      };
-
       const abortController = new AbortController();
       const abortSignal = abortController.signal;
 
+      const dataService = {
+        sampleCursor: () =>
+          ({
+            async *[Symbol.asyncIterator]() {
+              await new Promise((resolve) => setTimeout(resolve, 0));
+              yield {
+                a: 123,
+              };
+              abortController.abort();
+              yield {
+                a: 345,
+              };
+            },
+          } as any),
+      };
+
       const result = await analyzeSchema(
         dataService,
         abortSignal,
@@ -228,13 +252,19 @@ describe('schema-analysis', function () {
     it('throws if sample throws', async function () {
       const error: Error & {
         code?: any;
-      } = new Error('should have been thrown');
+      } = new Error('pineapple');
       error.name = 'MongoError';
       error.code = new bson.Int32(1000);
 
       const dataService = {
-        sample: () => Promise.reject(error),
-        isCancelError: () => false,
+        sampleCursor: () =>
+          ({
+            async *[Symbol.asyncIterator]() {
+              await new Promise((resolve) => setTimeout(resolve, 0));
+              yield {};
+              throw error;
+            },
+          } as any),
       };
 
       const abortController = new AbortController();
@@ -251,7 +281,7 @@ describe('schema-analysis', function () {
           preferences
         );
       } catch (err: any) {
-        expect(err.message).to.equal('should have been thrown');
+        expect(err.message).to.equal('pineapple');
         expect(err.code).to.equal(1000);
         return;
       }

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

@@ -50,15 +50,15 @@ export const analyzeSchema = async (
       ns,
     });
 
-    const docs = await dataService.sample(
+    const sampleCursor = dataService.sampleCursor(
       ns,
       query,
       {
         ...aggregateOptions,
         promoteValues: false,
+        signal: abortSignal,
       },
       {
-        abortSignal,
         fallbackReadPreference: 'secondaryPreferred',
       }
     );
@@ -72,7 +72,10 @@ export const analyzeSchema = async (
       : {
           signal: abortSignal,
         };
-    const schemaAccessor = await analyzeDocuments(docs, schemaParseOptions);
+    const schemaAccessor = await analyzeDocuments(
+      sampleCursor,
+      schemaParseOptions
+    );
     log.info(mongoLogId(1001000090), 'Schema', 'Schema analysis completed', {
       ns,
     });
@@ -81,8 +84,14 @@ export const analyzeSchema = async (
     log.error(mongoLogId(1001000091), 'Schema', 'Schema analysis failed', {
       ns,
       error: err.message,
+      aborted: abortSignal.aborted,
+      ...(abortSignal.aborted
+        ? { abortReason: abortSignal.reason?.message ?? abortSignal.reason }
+        : {}),
     });
-    if (dataService.isCancelError(err)) {
+
+    if (abortSignal.aborted) {
+      // The operation was aborted, so we don't throw an error.
       debug('caught background operation terminated error', err);
       return;
     }