CLOUDP-333846 Schema Analysis Redux Integration for Collection Plugin

ncarbon · ncarbon · commit 27f28ee20537 · 2025-08-13T11:33:00.000-04:00
diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts
@@ -1,11 +1,25 @@
 import type { Reducer, AnyAction, Action } from 'redux';
+import {
+  analyzeDocuments,
+  SchemaParseOptions,
+  type Schema,
+} from 'mongodb-schema';
+
 import type { CollectionMetadata } from 'mongodb-collection-model';
 import type { ThunkAction } from 'redux-thunk';
 import type AppRegistry from '@mongodb-js/compass-app-registry';
 import type { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provider';
 import type { CollectionSubtab } from '@mongodb-js/compass-workspaces';
 import type { DataService } from '@mongodb-js/compass-connections/provider';
 import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider';
+import { calculateSchemaMetadata } from '@mongodb-js/compass-schema';
+import type { Logger } from '@mongodb-js/compass-logging/provider';
+import { type PreferencesAccess } from 'compass-preferences-model/provider';
+import { isInternalFieldPath } from 'hadron-document';
+import { mongoLogId } from '@mongodb-js/compass-logging';
+import toNS from 'mongodb-ns';
+
+const DEFAULT_SAMPLE_SIZE = 100;
 
 function isAction<A extends AnyAction>(
   action: AnyAction,
@@ -22,32 +36,79 @@ type CollectionThunkAction<R, A extends AnyAction = AnyAction> = ThunkAction<
     dataService: DataService;
     workspaces: ReturnType<typeof workspacesServiceLocator>;
     experimentationServices: ReturnType<typeof experimentationServiceLocator>;
+    logger: Logger;
+    preferences: PreferencesAccess;
+    analysisAbortControllerRef: { current?: AbortController };
   },
   A
 >;
 
+export enum SchemaAnalysisStatus {
+  INITIAL = 'initial',
+  ANALYZING = 'analyzing',
+  COMPLETED = 'completed',
+  ERROR = 'error',
+}
+
+type SchemaAnalysis = {
+  status: SchemaAnalysisStatus;
+  schema: Schema | null;
+  sampleDocument: Document | null;
+  schemaMetadata: {
+    maxNestingDepth: number;
+    validationRules: Document;
+  } | null;
+  error: string | null;
+};
+
 export type CollectionState = {
   workspaceTabId: string;
   namespace: string;
   metadata: CollectionMetadata | null;
   editViewName?: string;
+  schemaAnalysis: SchemaAnalysis;
 };
 
-enum CollectionActions {
+export enum CollectionActions {
   CollectionMetadataFetched = 'compass-collection/CollectionMetadataFetched',
+  SchemaAnalysisStarted = 'compass-collection/SchemaAnalysisStarted',
+  SchemaAnalysisFinished = 'compass-collection/SchemaAnalysisFinished',
+  SchemaAnalysisFailed = 'compass-collection/SchemaAnalysisFailed',
 }
 
 interface CollectionMetadataFetchedAction {
   type: CollectionActions.CollectionMetadataFetched;
   metadata: CollectionMetadata;
 }
 
+interface SchemaAnalysisStartedAction {
+  type: CollectionActions.SchemaAnalysisStarted;
+  analysisStartTime: number;
+}
+
+interface SchemaAnalysisFinishedAction {
+  type: CollectionActions.SchemaAnalysisFinished;
+  schemaAnalysis: SchemaAnalysis;
+}
+
+interface SchemaAnalysisFailedAction {
+  type: CollectionActions.SchemaAnalysisFailed;
+  error: Error;
+}
+
 const reducer: Reducer<CollectionState, Action> = (
   state = {
     // TODO(COMPASS-7782): use hook to get the workspace tab id instead
     workspaceTabId: '',
     namespace: '',
     metadata: null,
+    schemaAnalysis: {
+      status: SchemaAnalysisStatus.INITIAL,
+      schema: null,
+      sampleDocument: null,
+      schemaMetadata: null,
+      error: null,
+    },
   },
   action
 ) => {
@@ -62,6 +123,53 @@ const reducer: Reducer<CollectionState, Action> = (
       metadata: action.metadata,
     };
   }
+
+  if (
+    isAction<SchemaAnalysisStartedAction>(
+      action,
+      CollectionActions.SchemaAnalysisStarted
+    )
+  ) {
+    return {
+      ...state,
+      schemaAnalysis: {
+        status: SchemaAnalysisStatus.ANALYZING,
+        schema: null,
+        sampleDocument: null,
+        schemaMetadata: null,
+        error: null,
+      },
+    };
+  }
+
+  if (
+    isAction<SchemaAnalysisFinishedAction>(
+      action,
+      CollectionActions.SchemaAnalysisFinished
+    )
+  ) {
+    return {
+      ...state,
+      schemaAnalysis: action.schemaAnalysis,
+    };
+  }
+
+  if (
+    isAction<SchemaAnalysisFailedAction>(
+      action,
+      CollectionActions.SchemaAnalysisFailed
+    )
+  ) {
+    return {
+      ...state,
+      schemaAnalysis: {
+        ...state.schemaAnalysis,
+        status: SchemaAnalysisStatus.ERROR,
+        error: action.error.message,
+      },
+    };
+  }
+
   return state;
 };
 
@@ -82,6 +190,115 @@ export const selectTab = (
   };
 };
 
+export const analyzeCollectionSchema = (): CollectionThunkAction<void> => {
+  return async (
+    dispatch,
+    getState,
+    { analysisAbortControllerRef, dataService, preferences, logger }
+  ) => {
+    const { schemaAnalysis, namespace } = getState();
+    const analysisStatus = schemaAnalysis.status;
+    if (analysisStatus === SchemaAnalysisStatus.ANALYZING) {
+      logger.debug(
+        'Schema analysis is already in progress, skipping new analysis.'
+      );
+      return;
+    }
+
+    analysisAbortControllerRef.current = new AbortController();
+    const abortSignal = analysisAbortControllerRef.current.signal;
+
+    const analysisStartTime = Date.now();
+
+    try {
+      logger.debug('Schema analysis started.');
+
+      dispatch({
+        type: CollectionActions.SchemaAnalysisStarted,
+        analysisStartTime,
+      });
+
+      // Sample documents
+      const samplingOptions = { size: DEFAULT_SAMPLE_SIZE };
+      const driverOptions = {
+        maxTimeMS: preferences.getPreferences().maxTimeMS,
+        signal: abortSignal,
+      };
+      const sampleCursor = dataService.sampleCursor(
+        namespace,
+        samplingOptions,
+        driverOptions,
+        {
+          fallbackReadPreference: 'secondaryPreferred',
+        }
+      );
+      const sampleDocuments = await sampleCursor.toArray();
+
+      // Analyze sampled documents
+      const schemaParseOptions: SchemaParseOptions = {
+        signal: abortSignal,
+      };
+      const schemaAccessor = await analyzeDocuments(
+        sampleDocuments,
+        schemaParseOptions
+      );
+      if (abortSignal?.aborted) {
+        throw new Error(abortSignal?.reason || new Error('Operation aborted'));
+      }
+
+      let schema: Schema | null = null;
+      if (schemaAccessor) {
+        schema = await schemaAccessor.getInternalSchema();
+        // Filter out internal fields from the schema
+        schema.fields = schema.fields.filter(
+          ({ path }) => !isInternalFieldPath(path[0])
+        );
+        // TODO: Transform schema to structure that will be used by the LLM.
+      }
+
+      let schemaMetadata = null;
+      if (schema !== null) {
+        const { schema_depth } = await calculateSchemaMetadata(schema);
+        const { database, collection } = toNS(namespace);
+        const collInfo = await dataService.collectionInfo(database, collection);
+        schemaMetadata = {
+          maxNestingDepth: schema_depth,
+          validationRules: collInfo?.validation?.validator || null,
+        };
+      }
+      dispatch({
+        type: CollectionActions.SchemaAnalysisFinished,
+        schemaAnalysis: {
+          status: SchemaAnalysisStatus.COMPLETED,
+          schema,
+          sampleDocument: sampleDocuments[0] ?? null,
+          schemaMetadata,
+        },
+      });
+    } catch (err: any) {
+      logger.log.error(
+        mongoLogId(1_001_000_363),
+        'Collection',
+        'Schema analysis failed',
+        {
+          namespace,
+          error: err.message,
+          aborted: abortSignal.aborted,
+          ...(abortSignal.aborted
+            ? { abortReason: abortSignal.reason?.message ?? abortSignal.reason }
+            : {}),
+        }
+      );
+      dispatch({
+        type: CollectionActions.SchemaAnalysisFailed,
+        error: err as Error,
+      });
+    } finally {
+      analysisAbortControllerRef.current = undefined;
+    }
+  };
+};
+
 export type CollectionTabPluginMetadata = CollectionMetadata & {
   /**
    * Initial query for the query bar
diff --git a/packages/compass-collection/src/stores/collection-tab.spec.ts b/packages/compass-collection/src/stores/collection-tab.spec.ts
@@ -1,6 +1,7 @@
 import type { CollectionTabOptions } from './collection-tab';
 import { activatePlugin } from './collection-tab';
 import { selectTab } from '../modules/collection-tab';
+import * as collectionTabModule from '../modules/collection-tab';
 import { waitFor } from '@mongodb-js/testing-library-compass';
 import Sinon from 'sinon';
 import AppRegistry from '@mongodb-js/compass-app-registry';
@@ -11,6 +12,7 @@ import type { connectionInfoRefLocator } from '@mongodb-js/compass-connections/p
 import { createNoopLogger } from '@mongodb-js/compass-logging/provider';
 import { ReadOnlyPreferenceAccess } from 'compass-preferences-model/provider';
 import { ExperimentTestName } from '@mongodb-js/compass-telemetry/provider';
+import { CollectionMetadata } from 'mongodb-collection-model';
 
 const defaultMetadata = {
   namespace: 'test.foo',
@@ -27,16 +29,6 @@ const defaultTabOptions = {
   namespace: defaultMetadata.namespace,
 };
 
-const mockCollection = {
-  _id: defaultMetadata.namespace,
-  fetchMetadata() {
-    return Promise.resolve(defaultMetadata);
-  },
-  toJSON() {
-    return this;
-  },
-};
-
 const mockAtlasConnectionInfo = {
   current: {
     id: 'test-connection',
@@ -67,6 +59,9 @@ describe('Collection Tab Content store', function () {
   const sandbox = Sinon.createSandbox();
 
   const localAppRegistry = sandbox.spy(new AppRegistry());
+  const analyzeCollectionSchemaStub = sandbox
+    .stub(collectionTabModule, 'analyzeCollectionSchema')
+    .returns(() => async () => {});
   const dataService = {} as any;
   let store: ReturnType<typeof activatePlugin>['store'];
   let deactivate: ReturnType<typeof activatePlugin>['deactivate'];
@@ -85,8 +80,19 @@ describe('Collection Tab Content store', function () {
       enableGenAIFeatures: true,
       enableGenAIFeaturesAtlasOrg: true,
       cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
-    })
+    }),
+    collectionMetadata: Partial<CollectionMetadata> = defaultMetadata,
+    analysisAbortControllerRef: { current?: AbortController } = {}
   ) => {
+    const mockCollection = {
+      _id: collectionMetadata.namespace,
+      fetchMetadata() {
+        return Promise.resolve(collectionMetadata);
+      },
+      toJSON() {
+        return this;
+      },
+    };
     ({ store, deactivate } = activatePlugin(
       {
         ...defaultTabOptions,
@@ -107,7 +113,7 @@ describe('Collection Tab Content store', function () {
     await waitFor(() => {
       expect(store.getState())
         .to.have.property('metadata')
-        .deep.eq(defaultMetadata);
+        .deep.eq(collectionMetadata);
     });
     return store;
   };
@@ -231,4 +237,40 @@ describe('Collection Tab Content store', function () {
       });
     });
   });
+
+  describe('schema analysis on collection load', function () {
+    it('should start schema analysis if collection is not read-only and not time-series', async function () {
+      await configureStore();
+
+      expect(analyzeCollectionSchemaStub).to.have.been.calledOnce;
+    });
+
+    it('should not start schema analysis if collection is read-only', async function () {
+      await configureStore(
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+        { ...defaultMetadata, isReadonly: true }
+      );
+
+      expect(analyzeCollectionSchemaStub).to.not.have.been.called;
+    });
+
+    it('should not start schema analysis if collection is time-series', async function () {
+      await configureStore(
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+        { ...defaultMetadata, isTimeSeries: true }
+      );
+
+      expect(analyzeCollectionSchemaStub).to.not.have.been.called;
+    });
+  });
 });
diff --git a/packages/compass-collection/src/stores/collection-tab.ts b/packages/compass-collection/src/stores/collection-tab.ts
