feat(compass-collection): Connect schema redux state to modal and filter collections based on nesting depth - CLOUDP-342665

packages/compass-collection/src/stores/collection-tab.spec.ts

@@ -7,12 +7,38 @@ import Sinon from 'sinon';
 import AppRegistry from '@mongodb-js/compass-app-registry';
 import { expect } from 'chai';
 import type { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provider';
-import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider';
+import type { ExperimentationServices } from '@mongodb-js/compass-telemetry/provider';
 import type { connectionInfoRefLocator } from '@mongodb-js/compass-connections/provider';
 import { createNoopLogger } from '@mongodb-js/compass-logging/provider';
 import { ReadOnlyPreferenceAccess } from 'compass-preferences-model/provider';
-import { ExperimentTestName } from '@mongodb-js/compass-telemetry/provider';
+import {
+  ExperimentTestName,
+  ExperimentTestGroup,
+} from '@mongodb-js/compass-telemetry/provider';
 import { type CollectionMetadata } from 'mongodb-collection-model';
+import type { types } from '@mongodb-js/mdb-experiment-js';
+
+// Helper function to create proper mock assignment objects for testing
+const createMockAssignment = (
+  variant: ExperimentTestGroup
+): types.SDKAssignment<ExperimentTestName, string> => ({
+  assignmentData: {
+    variant,
+    isInSample: true,
+  },
+  experimentData: {
+    assignmentDate: '2024-01-01T00:00:00Z',
+    entityType: 'USER' as types.EntityType,
+    id: 'test-assignment-id',
+    tag: 'test-tag',
+    testGroupId: 'test-group-id',
+    entityId: 'test-user-id',
+    testId: 'test-id',
+    testName: ExperimentTestName.mockDataGenerator,
+    testGroupDatabaseId: 'test-group-db-id',
+    meta: { isLaunchedExperiment: true },
+  },
+});
 
 const defaultMetadata = {
   namespace: 'test.foo',
@@ -71,9 +97,7 @@ describe('Collection Tab Content store', function () {
   const configureStore = async (
     options: Partial<CollectionTabOptions> = {},
     workspaces: Partial<ReturnType<typeof workspacesServiceLocator>> = {},
-    experimentationServices: Partial<
-      ReturnType<typeof experimentationServiceLocator>
-    > = {},
+    experimentationServices: Partial<ExperimentationServices> = {},
     connectionInfoRef: Partial<
       ReturnType<typeof connectionInfoRefLocator>
     > = {},
@@ -242,9 +266,21 @@ 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();
+      const getAssignment = sandbox.spy(() =>
+        Promise.resolve(
+          createMockAssignment(ExperimentTestGroup.mockDataGeneratorVariant)
+        )
+      );
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      await configureStore(undefined, undefined, {
+        getAssignment,
+        assignExperiment,
+      });
 
-      expect(analyzeCollectionSchemaStub).to.have.been.calledOnce;
+      await waitFor(() => {
+        expect(analyzeCollectionSchemaStub).to.have.been.calledOnce;
+      });
     });
 
     it('should not start schema analysis if collection is read-only', async function () {
@@ -274,5 +310,111 @@ describe('Collection Tab Content store', function () {
 
       expect(analyzeCollectionSchemaStub).to.not.have.been.called;
     });
+
+    it('should not start schema analysis in non-Atlas environments', async function () {
+      const getAssignment = sandbox.spy(() => Promise.resolve(null));
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+      const mockConnectionInfoRef = {
+        current: {
+          id: 'test-connection',
+          title: 'Test Connection',
+          connectionOptions: {
+            connectionString: 'mongodb://localhost:27017',
+          },
+          // No atlasMetadata (non-Atlas environment)
+        },
+      };
+
+      await configureStore(
+        undefined,
+        undefined,
+        { getAssignment, assignExperiment },
+        mockConnectionInfoRef
+      );
+
+      await waitFor(() => {
+        expect(getAssignment).to.have.been.calledOnceWith(
+          ExperimentTestName.mockDataGenerator,
+          false
+        );
+      });
+
+      // Wait a bit to ensure schema analysis would not have been called
+      await new Promise((resolve) => setTimeout(resolve, 50));
+      expect(analyzeCollectionSchemaStub).to.not.have.been.called;
+    });
+
+    it('should start schema analysis in Atlas when user is in treatment variant', async function () {
+      const getAssignment = sandbox.spy(() =>
+        Promise.resolve(
+          createMockAssignment(ExperimentTestGroup.mockDataGeneratorVariant)
+        )
+      );
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      await configureStore(
+        undefined,
+        undefined,
+        { getAssignment, assignExperiment },
+        mockAtlasConnectionInfo
+      );
+
+      await waitFor(() => {
+        expect(getAssignment).to.have.been.calledOnceWith(
+          ExperimentTestName.mockDataGenerator,
+          false // Don't track "Experiment Viewed" event
+        );
+        expect(analyzeCollectionSchemaStub).to.have.been.calledOnce;
+      });
+    });
+
+    it('should not start schema analysis in Atlas when user is in control variant', async function () {
+      const getAssignment = sandbox.spy(() =>
+        Promise.resolve(
+          createMockAssignment(ExperimentTestGroup.mockDataGeneratorControl)
+        )
+      );
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      await configureStore(
+        undefined,
+        undefined,
+        { getAssignment, assignExperiment },
+        mockAtlasConnectionInfo
+      );
+
+      await waitFor(() => {
+        expect(getAssignment).to.have.been.calledOnceWith(
+          ExperimentTestName.mockDataGenerator,
+          false
+        );
+      });
+
+      // Wait a bit to ensure schema analysis would not have been called
+      await new Promise((resolve) => setTimeout(resolve, 50));
+      expect(analyzeCollectionSchemaStub).to.not.have.been.called;
+    });
+
+    it('should not start schema analysis when getAssignment fails', async function () {
+      const getAssignment = sandbox.spy(() =>
+        Promise.reject(new Error('Assignment failed'))
+      );
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      await configureStore(
+        undefined,
+        undefined,
+        { getAssignment, assignExperiment },
+        mockAtlasConnectionInfo
+      );
+
+      await waitFor(() => {
+        expect(getAssignment).to.have.been.calledOnce;
+      });
+
+      // Wait a bit to ensure schema analysis would not have been called
+      await new Promise((resolve) => setTimeout(resolve, 50));
+      expect(analyzeCollectionSchemaStub).to.not.have.been.called;
+    });
   });
 });

packages/compass-collection/src/stores/collection-tab.ts

@@ -13,15 +13,18 @@ import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/t
 import type { Collection } from '@mongodb-js/compass-app-stores/provider';
 import type { ActivateHelpers } from '@mongodb-js/compass-app-registry';
 import type { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provider';
-import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider';
+import type { ExperimentationServices } from '@mongodb-js/compass-telemetry/provider';
 import type { connectionInfoRefLocator } from '@mongodb-js/compass-connections/provider';
 import type { Logger } from '@mongodb-js/compass-logging/provider';
 import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider';
 import {
   isAIFeatureEnabled,
   type PreferencesAccess,
 } from 'compass-preferences-model/provider';
-import { ExperimentTestName } from '@mongodb-js/compass-telemetry/provider';
+import {
+  ExperimentTestName,
+  ExperimentTestGroup,
+} from '@mongodb-js/compass-telemetry/provider';
 import { SCHEMA_ANALYSIS_STATE_INITIAL } from '../schema-analysis-types';
 
 export type CollectionTabOptions = {
@@ -46,7 +49,7 @@ export type CollectionTabServices = {
   localAppRegistry: AppRegistry;
   atlasAiService: AtlasAiService;
   workspaces: ReturnType<typeof workspacesServiceLocator>;
-  experimentationServices: ReturnType<typeof experimentationServiceLocator>;
+  experimentationServices: ExperimentationServices;
   connectionInfoRef: ReturnType<typeof connectionInfoRefLocator>;
   logger: Logger;
   preferences: PreferencesAccess;
@@ -153,8 +156,37 @@ export function activatePlugin(
     }
 
     if (!metadata.isReadonly && !metadata.isTimeSeries) {
-      // TODO: Consider checking experiment variant
-      void store.dispatch(analyzeCollectionSchema());
+      // Check experiment variant before running schema analysis
+      // Only run schema analysis if user is in treatment variant
+      const shouldRunSchemaAnalysis = async () => {
+        try {
+          const assignment = await experimentationServices.getAssignment(
+            ExperimentTestName.mockDataGenerator,
+            false // Don't track "Experiment Viewed" event here
+          );
+          return (
+            assignment?.assignmentData?.variant ===
+            ExperimentTestGroup.mockDataGeneratorVariant
+          );
+        } catch (error) {
+          // On error, default to not running schema analysis
+          logger.debug(
+            'Failed to get Mock Data Generator experiment assignment',
+            {
+              experiment: ExperimentTestName.mockDataGenerator,
+              namespace: namespace,
+              error: error instanceof Error ? error.message : String(error),
+            }
+          );
+          return false;
+        }
+      };
+
+      void shouldRunSchemaAnalysis().then((shouldRun) => {
+        if (shouldRun) {
+          void store.dispatch(analyzeCollectionSchema());
+        }
+      });
     }
   });
 

packages/compass-telemetry/src/experimentation-provider.tsx

@@ -14,9 +14,16 @@ type AssignExperimentFn = (
   options?: types.AssignOptions<string>
 ) => Promise<types.AsyncStatus | null>;
 
+type GetAssignmentFn = (
+  experimentName: ExperimentTestName,
+  trackIsInSample: boolean,
+  options?: types.GetAssignmentOptions<types.TypeData>
+) => Promise<types.SDKAssignment<ExperimentTestName, string> | null>;
+
 interface CompassExperimentationProviderContextValue {
   useAssignment: UseAssignmentHook;
   assignExperiment: AssignExperimentFn;
+  getAssignment: GetAssignmentFn;
 }
 
 const initialContext: CompassExperimentationProviderContextValue = {
@@ -33,6 +40,9 @@ const initialContext: CompassExperimentationProviderContextValue = {
   assignExperiment() {
     return Promise.resolve(null);
   },
+  getAssignment() {
+    return Promise.resolve(null);
+  },
 };
 
 export const ExperimentationContext =
@@ -43,12 +53,18 @@ export const CompassExperimentationProvider: React.FC<{
   children: React.ReactNode;
   useAssignment: UseAssignmentHook;
   assignExperiment: AssignExperimentFn;
-}> = ({ children, useAssignment, assignExperiment }) => {
+  getAssignment: GetAssignmentFn;
+}> = ({ children, useAssignment, assignExperiment, getAssignment }) => {
   // Use useRef to keep the functions up-to-date; Use mutation pattern to maintain the
   // same object reference to prevent unnecessary re-renders of consuming components
-  const { current: contextValue } = useRef({ useAssignment, assignExperiment });
+  const { current: contextValue } = useRef({
+    useAssignment,
+    assignExperiment,
+    getAssignment,
+  });
   contextValue.useAssignment = useAssignment;
   contextValue.assignExperiment = assignExperiment;
+  contextValue.getAssignment = getAssignment;
 
   return (
     <ExperimentationContext.Provider value={contextValue}>

packages/compass-telemetry/src/provider.tsx

@@ -54,13 +54,20 @@ export interface ExperimentationServices {
     experimentName: ExperimentTestName,
     options?: types.AssignOptions<string>
   ) => Promise<types.AsyncStatus | null>;
+  getAssignment: (
+    experimentName: ExperimentTestName,
+    trackIsInSample: boolean,
+    options?: types.GetAssignmentOptions<types.TypeData>
+  ) => Promise<types.SDKAssignment<ExperimentTestName, string> | null>;
 }
 
 // Service locator for experimentation services (non-component access)
 export const experimentationServiceLocator = createServiceLocator(
   function useExperimentationServices(): ExperimentationServices {
-    const { assignExperiment } = useContext(ExperimentationContext);
-    return { assignExperiment };
+    const { assignExperiment, getAssignment } = useContext(
+      ExperimentationContext
+    );
+    return { assignExperiment, getAssignment };
   },
   'experimentationServiceLocator'
 );