diff --git a/package-lock.json b/package-lock.json index 1a10f197185..b940df8c5d5 100644 --- a/package-lock.json +++ b/package-lock.json @@ -47722,6 +47722,7 @@ }, "devDependencies": { "@mongodb-js/eslint-config-compass": "^1.4.8", + "@mongodb-js/mdb-experiment-js": "1.9.0", "@mongodb-js/mocha-config-compass": "^1.7.1", "@mongodb-js/prettier-config-compass": "^1.2.8", "@mongodb-js/testing-library-compass": "^1.3.11", @@ -61182,6 +61183,7 @@ "@mongodb-js/compass-workspaces": "^0.53.1", "@mongodb-js/connection-info": "^0.18.0", "@mongodb-js/eslint-config-compass": "^1.4.8", + "@mongodb-js/mdb-experiment-js": "1.9.0", "@mongodb-js/mocha-config-compass": "^1.7.1", "@mongodb-js/mongodb-constants": "^0.14.0", "@mongodb-js/prettier-config-compass": "^1.2.8", diff --git a/packages/compass-collection/package.json b/packages/compass-collection/package.json index 588c09aedbe..31bdfbcd877 100644 --- a/packages/compass-collection/package.json +++ b/packages/compass-collection/package.json @@ -73,6 +73,7 @@ }, "devDependencies": { "@mongodb-js/eslint-config-compass": "^1.4.8", + "@mongodb-js/mdb-experiment-js": "1.9.0", "@mongodb-js/mocha-config-compass": "^1.7.1", "@mongodb-js/prettier-config-compass": "^1.2.8", "@mongodb-js/testing-library-compass": "^1.3.11", diff --git a/packages/compass-collection/src/components/collection-header-actions/collection-header-actions.spec.tsx b/packages/compass-collection/src/components/collection-header-actions/collection-header-actions.spec.tsx index 54f18833a16..1f98a73199d 100644 --- a/packages/compass-collection/src/components/collection-header-actions/collection-header-actions.spec.tsx +++ b/packages/compass-collection/src/components/collection-header-actions/collection-header-actions.spec.tsx @@ -48,6 +48,7 @@ describe('CollectionHeaderActions [Component]', function () { @@ -221,97 +225,21 @@ describe('CollectionHeaderActions [Component]', function () { }, }; - it('should not show Mock Data Generator button when user is in control group', async function () { - mockUseAssignment.returns({ - assignment: { - assignmentData: { - variant: 'mockDataGeneratorControl', - }, - }, - }); - - await renderCollectionHeaderActions( - { - namespace: 'test.collection', - isReadonly: false, - }, - {}, - atlasConnectionInfo - ); - - expect( - screen.queryByTestId('collection-header-generate-mock-data-button') - ).to.not.exist; - }); - - it('should not show Mock Data Generator button when not in Atlas', async function () { - mockUseAssignment.returns({ - assignment: { - assignmentData: { - variant: 'treatment', - }, - }, - }); - + it('should call useAssignment with correct parameters', async function () { await renderCollectionHeaderActions({ namespace: 'test.collection', isReadonly: false, - // Don't pass atlasConnectionInfo, to simulate not being in Atlas }); - expect( - screen.queryByTestId('collection-header-generate-mock-data-button') - ).to.not.exist; - }); - - it('should not show Mock Data Generator button for readonly collections', async function () { - mockUseAssignment.returns({ - assignment: { - assignmentData: { - variant: 'treatment', - }, - }, - }); - - await renderCollectionHeaderActions( - { - namespace: 'test.collection', - isReadonly: true, - }, - {}, - atlasConnectionInfo + expect(mockUseAssignment).to.have.been.calledWith( + ExperimentTestName.mockDataGenerator, + true // trackIsInSample - Experiment viewed analytics event ); - - expect( - screen.queryByTestId('collection-header-generate-mock-data-button') - ).to.not.exist; }); - it('should not show Mock Data Generator button for views (sourceName present)', async function () { - mockUseAssignment.returns({ - assignment: { - assignmentData: { - variant: 'treatment', - }, - }, - }); - - await renderCollectionHeaderActions( - { - namespace: 'test.collection', - isReadonly: false, - sourceName: 'source-collection', - }, - {}, - atlasConnectionInfo - ); - - expect( - screen.queryByTestId('collection-header-generate-mock-data-button') - ).to.not.exist; - }); + it('should call onOpenMockDataModal when CTA button is clicked', async function () { + const onOpenMockDataModal = sinon.stub(); - it('should show Mock Data Generator button when user is in treatment group and in Atlas', async function () { mockUseAssignment.returns({ assignment: { assignmentData: { @@ -324,35 +252,25 @@ describe('CollectionHeaderActions [Component]', function () { { namespace: 'test.collection', isReadonly: false, + onOpenMockDataModal, }, {}, atlasConnectionInfo ); - expect( - screen.getByTestId('collection-header-generate-mock-data-button') - ).to.exist; - }); - - it('should call useAssignment with correct parameters', async function () { - await renderCollectionHeaderActions({ - namespace: 'test.collection', - isReadonly: false, - }); - - expect(mockUseAssignment).to.have.been.calledWith( - ExperimentTestName.mockDataGenerator, - true // trackIsInSample - Experiment viewed analytics event + const button = screen.getByTestId( + 'collection-header-generate-mock-data-button' ); - }); + button.click(); - it('should call onOpenMockDataModal when CTA button is clicked', async function () { - const onOpenMockDataModal = sinon.stub(); + expect(onOpenMockDataModal).to.have.been.calledOnce; + }); + it('should disable button for deeply nested collections', async function () { mockUseAssignment.returns({ assignment: { assignmentData: { - variant: 'mockDataGeneratorVariant', + variant: 'mockDataGeneratorVariant', // Treatment variant }, }, }); @@ -361,7 +279,10 @@ describe('CollectionHeaderActions [Component]', function () { { namespace: 'test.collection', isReadonly: false, - onOpenMockDataModal, + hasSchemaAnalysisData: true, + analyzedSchemaDepth: 5, // Exceeds MAX_COLLECTION_NESTING_DEPTH (3) + schemaAnalysisStatus: 'complete', + onOpenMockDataModal: sinon.stub(), }, {}, atlasConnectionInfo @@ -370,9 +291,8 @@ describe('CollectionHeaderActions [Component]', function () { const button = screen.getByTestId( 'collection-header-generate-mock-data-button' ); - button.click(); - - expect(onOpenMockDataModal).to.have.been.calledOnce; + expect(button).to.exist; + expect(button).to.have.attribute('aria-disabled', 'true'); }); }); }); diff --git a/packages/compass-collection/src/components/collection-header-actions/collection-header-actions.tsx b/packages/compass-collection/src/components/collection-header-actions/collection-header-actions.tsx index e6f06f6525e..e8204823d5c 100644 --- a/packages/compass-collection/src/components/collection-header-actions/collection-header-actions.tsx +++ b/packages/compass-collection/src/components/collection-header-actions/collection-header-actions.tsx @@ -18,6 +18,15 @@ import { ExperimentTestName, ExperimentTestGroup, } from '@mongodb-js/compass-telemetry/provider'; +import { + SCHEMA_ANALYSIS_STATE_ANALYZING, + type SchemaAnalysisStatus, +} from '../../schema-analysis-types'; + +/** + * Maximum allowed nesting depth for collections to show Mock Data Generator + */ +const MAX_COLLECTION_NESTING_DEPTH = 3; const collectionHeaderActionsStyles = css({ display: 'flex', @@ -47,6 +56,9 @@ type CollectionHeaderActionsProps = { sourceName?: string; sourcePipeline?: unknown[]; onOpenMockDataModal: () => void; + hasSchemaAnalysisData: boolean; + analyzedSchemaDepth: number; + schemaAnalysisStatus: SchemaAnalysisStatus | null; }; const CollectionHeaderActions: React.FunctionComponent< @@ -58,6 +70,9 @@ const CollectionHeaderActions: React.FunctionComponent< sourceName, sourcePipeline, onOpenMockDataModal, + hasSchemaAnalysisData, + analyzedSchemaDepth, + schemaAnalysisStatus, }: CollectionHeaderActionsProps) => { const connectionInfo = useConnectionInfo(); const { id: connectionId, atlasMetadata } = connectionInfo; @@ -85,9 +100,13 @@ const CollectionHeaderActions: React.FunctionComponent< atlasMetadata && // Only show in Atlas !isReadonly && // Don't show for readonly collections (views) !sourceName; // sourceName indicates it's a view - // TODO: CLOUDP-337090: also filter out overly nested collections - const hasData = true; // TODO: CLOUDP-337090 + const exceedsMaxNestingDepth = + analyzedSchemaDepth > MAX_COLLECTION_NESTING_DEPTH; + + const isCollectionEmpty = + !hasSchemaAnalysisData && + schemaAnalysisStatus !== SCHEMA_ANALYSIS_STATE_ANALYZING; return (
} > - Please add data to your collection to generate similar mock documents + {exceedsMaxNestingDepth && + 'At this time we are unable to generate mock data for collections that have deeply nested documents'} + {isCollectionEmpty && + 'Please add data to your collection to generate similar mock documents'} )} {atlasMetadata && ( diff --git a/packages/compass-collection/src/components/collection-header/collection-header.spec.tsx b/packages/compass-collection/src/components/collection-header/collection-header.spec.tsx index 0a7b3118ce8..82553cc22af 100644 --- a/packages/compass-collection/src/components/collection-header/collection-header.spec.tsx +++ b/packages/compass-collection/src/components/collection-header/collection-header.spec.tsx @@ -3,6 +3,7 @@ import type { ComponentProps } from 'react'; import React from 'react'; import { renderWithConnections, + renderWithActiveConnection, screen, cleanup, within, @@ -16,19 +17,26 @@ import { type WorkspacesService, } from '@mongodb-js/compass-workspaces/provider'; import { MockDataGeneratorStep } from '../mock-data-generator-modal/types'; +import { SCHEMA_ANALYSIS_STATE_COMPLETE } from '../../schema-analysis-types'; +import { CompassExperimentationProvider } from '@mongodb-js/compass-telemetry'; +import type { ConnectionInfo } from '@mongodb-js/compass-connections/provider'; import Sinon from 'sinon'; function renderCollectionHeader( props: Partial> = {}, - workspaceService: Partial = {} + workspaceService: Partial = {}, + stateOverrides: any = {} ) { - const mockStore = createStore(() => ({ + const defaultState = { mockDataGenerator: { isModalOpen: false, currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION, }, - })); + ...stateOverrides, + }; + + const mockStore = createStore(() => defaultState); return renderWithConnections( @@ -294,4 +302,293 @@ describe('CollectionHeader [Component]', function () { }); }); }); + + it('should handle undefined schemaAnalysis gracefully and render collection header successfully', function () { + // Create a store with undefined schemaAnalysis to simulate initial state + const mockStoreWithUndefinedSchema = createStore(() => ({ + mockDataGenerator: { + isModalOpen: false, + currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION, + }, + // schemaAnalysis not provided + })); + + expect(() => { + renderWithConnections( + + + + + + ); + }).to.not.throw(); + + expect(screen.getByTestId('collection-header')).to.exist; + expect(screen.getByTestId('collection-header-actions')).to.exist; + }); + + describe('Mock Data Generator integration', function () { + let mockUseAssignment: Sinon.SinonStub; + + beforeEach(function () { + // Mock the useAssignment hook from compass-experimentation + mockUseAssignment = Sinon.stub().returns({ + assignment: { + assignmentData: { + variant: 'mockDataGeneratorVariant', + }, + }, + }); + }); + + afterEach(function () { + Sinon.restore(); + }); + + const atlasConnectionInfo: ConnectionInfo = { + id: 'test-atlas-connection', + connectionOptions: { + connectionString: 'mongodb://localhost:27017', + }, + atlasMetadata: { + orgId: 'test-org', + projectId: 'test-project', + clusterName: 'test-cluster', + clusterUniqueId: 'test-cluster-unique-id', + clusterType: 'REPLICASET', + clusterState: 'IDLE', + metricsId: 'test-metrics-id', + metricsType: 'replicaSet', + regionalBaseUrl: null, + instanceSize: 'M10', + supports: { + globalWrites: false, + rollingIndexes: true, + }, + }, + }; + + function renderCollectionHeaderWithExperimentation( + props: Partial> = {}, + workspaceService: Partial = {}, + stateOverrides: any = {}, + connectionInfo?: ConnectionInfo + ) { + const defaultState = { + mockDataGenerator: { + isModalOpen: false, + currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION, + }, + ...stateOverrides, + }; + + const mockStore = createStore(() => defaultState); + + return renderWithActiveConnection( + + + + + + + , + connectionInfo + ); + } + + it('should show Mock Data Generator button when all conditions are met', async function () { + await renderCollectionHeaderWithExperimentation( + { + isAtlas: true, // Atlas environment + isReadonly: false, // Not readonly + namespace: 'test.collection', + }, + {}, + { + schemaAnalysis: { + status: SCHEMA_ANALYSIS_STATE_COMPLETE, + processedSchema: { + field1: { type: 'String', sample_values: ['value1'] }, + }, + schemaMetadata: { + maxNestingDepth: 2, // Below the limit of 4 + }, + }, + }, + atlasConnectionInfo + ); + + expect(screen.getByTestId('collection-header-generate-mock-data-button')) + .to.exist; + expect( + screen.getByTestId('collection-header-generate-mock-data-button') + ).to.not.have.attribute('aria-disabled', 'true'); + }); + + it('should disable Mock Data Generator button when collection has no schema analysis data', async function () { + await renderCollectionHeaderWithExperimentation( + { + isAtlas: true, + isReadonly: false, + namespace: 'test.collection', + }, + {}, + { + schemaAnalysis: { + status: SCHEMA_ANALYSIS_STATE_COMPLETE, + processedSchema: {}, // Empty schema + schemaMetadata: { + maxNestingDepth: 2, + }, + }, + }, + atlasConnectionInfo + ); + + const button = screen.getByTestId( + 'collection-header-generate-mock-data-button' + ); + expect(button).to.exist; + expect(button).to.have.attribute('aria-disabled', 'true'); + }); + + it('should disable Mock Data Generator button for collections with excessive nesting depth', async function () { + await renderCollectionHeaderWithExperimentation( + { + isAtlas: true, + isReadonly: false, + namespace: 'test.collection', + }, + {}, + { + schemaAnalysis: { + status: SCHEMA_ANALYSIS_STATE_COMPLETE, + processedSchema: { + field1: { type: 'String', sample_values: ['value1'] }, + }, + schemaMetadata: { + maxNestingDepth: 4, // Exceeds the limit + }, + }, + }, + atlasConnectionInfo + ); + + const button = screen.getByTestId( + 'collection-header-generate-mock-data-button' + ); + expect(button).to.exist; + expect(button).to.have.attribute('aria-disabled', 'true'); + }); + + it('should not show Mock Data Generator button for readonly collections (views)', async function () { + await renderCollectionHeaderWithExperimentation( + { + isAtlas: true, + isReadonly: true, // Readonly (view) + namespace: 'test.view', + }, + {}, + { + schemaAnalysis: { + status: SCHEMA_ANALYSIS_STATE_COMPLETE, + processedSchema: { + field1: { type: 'String', sample_values: ['value1'] }, + }, + schemaMetadata: { + maxNestingDepth: 2, + }, + }, + }, + atlasConnectionInfo + ); + + expect( + screen.queryByTestId('collection-header-generate-mock-data-button') + ).to.not.exist; + }); + + it('should not show Mock Data Generator button in non-Atlas environments', async function () { + await renderCollectionHeaderWithExperimentation( + { + isAtlas: false, // Not Atlas + isReadonly: false, + namespace: 'test.collection', + }, + {}, + { + schemaAnalysis: { + status: SCHEMA_ANALYSIS_STATE_COMPLETE, + processedSchema: { + field1: { type: 'String', sample_values: ['value1'] }, + }, + schemaMetadata: { + maxNestingDepth: 2, + }, + }, + } + // Don't pass atlasConnectionInfo to simulate non-Atlas environment + ); + + expect( + screen.queryByTestId('collection-header-generate-mock-data-button') + ).to.not.exist; + }); + + it('should not show Mock Data Generator button when not in treatment variant', async function () { + mockUseAssignment.returns({ + assignment: { + assignmentData: { + variant: 'control', + }, + }, + }); + + await renderCollectionHeaderWithExperimentation( + { + isAtlas: true, + isReadonly: false, + namespace: 'test.collection', + }, + {}, + { + schemaAnalysis: { + status: SCHEMA_ANALYSIS_STATE_COMPLETE, + processedSchema: { + field1: { type: 'String', sample_values: ['value1'] }, + }, + schemaMetadata: { + maxNestingDepth: 2, + }, + }, + }, + atlasConnectionInfo + ); + + expect( + screen.queryByTestId('collection-header-generate-mock-data-button') + ).to.not.exist; + }); + }); }); diff --git a/packages/compass-collection/src/components/collection-header/collection-header.tsx b/packages/compass-collection/src/components/collection-header/collection-header.tsx index 6af9fea8dd7..7aca29141b7 100644 --- a/packages/compass-collection/src/components/collection-header/collection-header.tsx +++ b/packages/compass-collection/src/components/collection-header/collection-header.tsx @@ -21,6 +21,11 @@ import { getConnectionTitle } from '@mongodb-js/connection-info'; import MockDataGeneratorModal from '../mock-data-generator-modal/mock-data-generator-modal'; import { connect } from 'react-redux'; import { openMockDataGeneratorModal } from '../../modules/collection-tab'; +import type { CollectionState } from '../../modules/collection-tab'; +import { + SCHEMA_ANALYSIS_STATE_COMPLETE, + type SchemaAnalysisStatus, +} from '../../schema-analysis-types'; const collectionHeaderStyles = css({ padding: spacing[400], @@ -62,6 +67,9 @@ type CollectionHeaderProps = { editViewName?: string; sourcePipeline?: unknown[]; onOpenMockDataModal: () => void; + hasSchemaAnalysisData: boolean; + analyzedSchemaDepth: number; + schemaAnalysisStatus: SchemaAnalysisStatus | null; }; const getInsightsForPipeline = (pipeline: any[], isAtlas: boolean) => { @@ -97,6 +105,9 @@ const CollectionHeader: React.FunctionComponent = ({ editViewName, sourcePipeline, onOpenMockDataModal, + hasSchemaAnalysisData, + analyzedSchemaDepth, + schemaAnalysisStatus, }) => { const darkMode = useDarkMode(); const showInsights = usePreference('showInsights'); @@ -174,6 +185,9 @@ const CollectionHeader: React.FunctionComponent = ({ sourceName={sourceName} sourcePipeline={sourcePipeline} onOpenMockDataModal={onOpenMockDataModal} + hasSchemaAnalysisData={hasSchemaAnalysisData} + analyzedSchemaDepth={analyzedSchemaDepth} + schemaAnalysisStatus={schemaAnalysisStatus} /> @@ -181,7 +195,23 @@ const CollectionHeader: React.FunctionComponent = ({ ); }; -const ConnectedCollectionHeader = connect(undefined, { +const mapStateToProps = (state: CollectionState) => { + const { schemaAnalysis } = state; + + return { + hasSchemaAnalysisData: + schemaAnalysis && + schemaAnalysis.status === SCHEMA_ANALYSIS_STATE_COMPLETE && + Object.keys(schemaAnalysis.processedSchema).length > 0, + analyzedSchemaDepth: + schemaAnalysis && schemaAnalysis.status === SCHEMA_ANALYSIS_STATE_COMPLETE + ? schemaAnalysis.schemaMetadata.maxNestingDepth + : 0, + schemaAnalysisStatus: schemaAnalysis?.status || null, + }; +}; + +const ConnectedCollectionHeader = connect(mapStateToProps, { onOpenMockDataModal: openMockDataGeneratorModal, })(CollectionHeader); diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts index 7a001624dd4..e1f3afa2eac 100644 --- a/packages/compass-collection/src/modules/collection-tab.ts +++ b/packages/compass-collection/src/modules/collection-tab.ts @@ -79,6 +79,7 @@ type CollectionThunkAction = ThunkAction< preferences: PreferencesAccess; connectionInfoRef: ConnectionInfoRef; fakerSchemaGenerationAbortControllerRef: { current?: AbortController }; + schemaAnalysisAbortControllerRef: { current?: AbortController }; }, A >; @@ -101,6 +102,7 @@ export enum CollectionActions { SchemaAnalysisStarted = 'compass-collection/SchemaAnalysisStarted', SchemaAnalysisFinished = 'compass-collection/SchemaAnalysisFinished', SchemaAnalysisFailed = 'compass-collection/SchemaAnalysisFailed', + SchemaAnalysisCanceled = 'compass-collection/SchemaAnalysisCanceled', SchemaAnalysisReset = 'compass-collection/SchemaAnalysisReset', MockDataGeneratorModalOpened = 'compass-collection/MockDataGeneratorModalOpened', MockDataGeneratorModalClosed = 'compass-collection/MockDataGeneratorModalClosed', @@ -139,6 +141,10 @@ interface SchemaAnalysisFailedAction { error: Error; } +interface SchemaAnalysisCanceledAction { + type: CollectionActions.SchemaAnalysisCanceled; +} + interface MockDataGeneratorModalOpenedAction { type: CollectionActions.MockDataGeneratorModalOpened; } @@ -263,6 +269,20 @@ const reducer: Reducer = ( }; } + if ( + isAction( + action, + CollectionActions.SchemaAnalysisCanceled + ) + ) { + return { + ...state, + schemaAnalysis: { + status: SCHEMA_ANALYSIS_STATE_INITIAL, + }, + }; + } + if ( isAction( action, @@ -524,7 +544,11 @@ export const openMockDataGeneratorModal = (): CollectionThunkAction< export const analyzeCollectionSchema = (): CollectionThunkAction< Promise > => { - return async (dispatch, getState, { dataService, preferences, logger }) => { + return async ( + dispatch, + getState, + { dataService, preferences, logger, schemaAnalysisAbortControllerRef } + ) => { const { schemaAnalysis, namespace } = getState(); const analysisStatus = schemaAnalysis.status; if (analysisStatus === SCHEMA_ANALYSIS_STATE_ANALYZING) { @@ -534,6 +558,10 @@ export const analyzeCollectionSchema = (): CollectionThunkAction< return; } + // Create abort controller for this analysis + const abortController = new AbortController(); + schemaAnalysisAbortControllerRef.current = abortController; + try { logger.debug('Schema analysis started.'); @@ -552,8 +580,15 @@ export const analyzeCollectionSchema = (): CollectionThunkAction< driverOptions, { fallbackReadPreference: 'secondaryPreferred', + abortSignal: abortController.signal, } ); + + // Check if analysis was aborted after sampling + if (abortController.signal.aborted) { + logger.debug('Schema analysis was aborted during sampling'); + return; + } if (sampleDocuments.length === 0) { logger.debug(NO_DOCUMENTS_ERROR); dispatch({ @@ -565,6 +600,13 @@ export const analyzeCollectionSchema = (): CollectionThunkAction< // Analyze sampled documents const schemaAccessor = await analyzeDocuments(sampleDocuments); + + // Check if analysis was aborted after document analysis + if (abortController.signal.aborted) { + logger.debug('Schema analysis was aborted during document analysis'); + return; + } + const schema = await schemaAccessor.getInternalSchema(); // Filter out internal fields from the schema @@ -583,6 +625,13 @@ export const analyzeCollectionSchema = (): CollectionThunkAction< maxNestingDepth, validationRules, }; + + // Final check before dispatching results + if (abortController.signal.aborted) { + logger.debug('Schema analysis was aborted before completion'); + return; + } + dispatch({ type: CollectionActions.SchemaAnalysisFinished, processedSchema, @@ -590,6 +639,15 @@ export const analyzeCollectionSchema = (): CollectionThunkAction< schemaMetadata, }); } catch (err: any) { + // Check if the error is due to cancellation + if (isCancelError(err) || abortController.signal.aborted) { + logger.debug('Schema analysis was aborted'); + dispatch({ + type: CollectionActions.SchemaAnalysisCanceled, + }); + return; + } + logger.log.error( mongoLogId(1_001_000_363), 'Collection', @@ -603,6 +661,23 @@ export const analyzeCollectionSchema = (): CollectionThunkAction< type: CollectionActions.SchemaAnalysisFailed, error: err as Error, }); + } finally { + // Clean up abort controller + schemaAnalysisAbortControllerRef.current = undefined; + } + }; +}; + +export const cancelSchemaAnalysis = (): CollectionThunkAction => { + return ( + _dispatch, + _getState, + { schemaAnalysisAbortControllerRef, logger } + ) => { + if (schemaAnalysisAbortControllerRef.current) { + logger.debug('Canceling schema analysis'); + schemaAnalysisAbortControllerRef.current.abort(); + schemaAnalysisAbortControllerRef.current = undefined; } }; }; diff --git a/packages/compass-collection/src/stores/collection-tab.spec.ts b/packages/compass-collection/src/stores/collection-tab.spec.ts index 6de48f4426f..bcbcb52f763 100644 --- a/packages/compass-collection/src/stores/collection-tab.spec.ts +++ b/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 => ({ + 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 = {}, workspaces: Partial> = {}, - experimentationServices: Partial< - ReturnType - > = {}, + experimentationServices: Partial = {}, connectionInfoRef: Partial< ReturnType > = {}, @@ -110,7 +134,7 @@ describe('Collection Tab Content store', function () { logger, preferences, }, - { on() {}, cleanup() {} } as any + { on() {}, cleanup() {}, addCleanup() {} } as any )); await waitFor(() => { expect(store.getState()) @@ -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)); - expect(analyzeCollectionSchemaStub).to.have.been.calledOnce; + await configureStore(undefined, undefined, { + getAssignment, + assignExperiment, + }); + + await waitFor(() => { + expect(analyzeCollectionSchemaStub).to.have.been.calledOnce; + }); }); it('should not start schema analysis if collection is read-only', async function () { @@ -274,5 +310,135 @@ 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; + }); + }); + + describe('schema analysis cancellation', function () { + it('should cancel schema analysis when cancelSchemaAnalysis is dispatched', async function () { + const getAssignment = sandbox.spy(() => + Promise.resolve( + createMockAssignment(ExperimentTestGroup.mockDataGeneratorVariant) + ) + ); + const assignExperiment = sandbox.spy(() => Promise.resolve(null)); + + const store = await configureStore(undefined, undefined, { + getAssignment, + assignExperiment, + }); + + // Dispatch cancel action + store.dispatch(collectionTabModule.cancelSchemaAnalysis() as any); + + // Verify the state is reset to initial + expect((store.getState() as any).schemaAnalysis.status).to.equal( + 'initial' + ); + }); }); }); diff --git a/packages/compass-collection/src/stores/collection-tab.ts b/packages/compass-collection/src/stores/collection-tab.ts index 48b4240ff63..b49f4e90845 100644 --- a/packages/compass-collection/src/stores/collection-tab.ts +++ b/packages/compass-collection/src/stores/collection-tab.ts @@ -7,13 +7,14 @@ import reducer, { selectTab, collectionMetadataFetched, analyzeCollectionSchema, + cancelSchemaAnalysis, } from '../modules/collection-tab'; import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types'; 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'; @@ -21,7 +22,10 @@ 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 +50,7 @@ export type CollectionTabServices = { localAppRegistry: AppRegistry; atlasAiService: AtlasAiService; workspaces: ReturnType; - experimentationServices: ReturnType; + experimentationServices: ExperimentationServices; connectionInfoRef: ReturnType; logger: Logger; preferences: PreferencesAccess; @@ -55,7 +59,7 @@ export type CollectionTabServices = { export function activatePlugin( { namespace, editViewName, tabId }: CollectionTabOptions, services: CollectionTabServices, - { on, cleanup }: ActivateHelpers + { on, cleanup, addCleanup }: ActivateHelpers ): { store: ReturnType; deactivate: () => void; @@ -81,6 +85,9 @@ export function activatePlugin( const fakerSchemaGenerationAbortControllerRef = { current: undefined, }; + const schemaAnalysisAbortControllerRef = { + current: undefined, + }; const store = createStore( reducer, { @@ -110,6 +117,7 @@ export function activatePlugin( logger, preferences, fakerSchemaGenerationAbortControllerRef, + schemaAnalysisAbortControllerRef, }) ) ); @@ -153,11 +161,43 @@ 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()); + } + }); } }); + // Cancel schema analysis when plugin is deactivated + addCleanup(() => store.dispatch(cancelSchemaAnalysis())); + return { store, deactivate: cleanup, diff --git a/packages/compass-telemetry/src/experimentation-provider.tsx b/packages/compass-telemetry/src/experimentation-provider.tsx index 3055fd255a3..6bdbda73647 100644 --- a/packages/compass-telemetry/src/experimentation-provider.tsx +++ b/packages/compass-telemetry/src/experimentation-provider.tsx @@ -14,9 +14,16 @@ type AssignExperimentFn = ( options?: types.AssignOptions ) => Promise; +type GetAssignmentFn = ( + experimentName: ExperimentTestName, + trackIsInSample: boolean, + options?: types.GetAssignmentOptions +) => Promise | 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 ( diff --git a/packages/compass-telemetry/src/provider.tsx b/packages/compass-telemetry/src/provider.tsx index 301222d3838..4446dbbfb03 100644 --- a/packages/compass-telemetry/src/provider.tsx +++ b/packages/compass-telemetry/src/provider.tsx @@ -54,13 +54,20 @@ export interface ExperimentationServices { experimentName: ExperimentTestName, options?: types.AssignOptions ) => Promise; + getAssignment: ( + experimentName: ExperimentTestName, + trackIsInSample: boolean, + options?: types.GetAssignmentOptions + ) => Promise | 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' );