feat(compass-collection): Add Disclaimer Screen - Mock Data Generator CLOUDP-333851

packages/compass-collection/package.json

@@ -52,6 +52,7 @@
     "@mongodb-js/compass-app-stores": "^7.56.0",
     "@mongodb-js/compass-components": "^1.48.0",
     "@mongodb-js/compass-connections": "^1.70.0",
+    "@mongodb-js/compass-generative-ai": "^0.50.0",
     "@mongodb-js/compass-logging": "^1.7.11",
     "@mongodb-js/compass-telemetry": "^1.13.0",
     "@mongodb-js/compass-workspaces": "^0.51.0",

packages/compass-collection/src/components/collection-header/collection-header.tsx

@@ -20,7 +20,7 @@ import { useConnectionInfo } from '@mongodb-js/compass-connections/provider';
 import { getConnectionTitle } from '@mongodb-js/connection-info';
 import MockDataGeneratorModal from '../mock-data-generator-modal/mock-data-generator-modal';
 import { connect } from 'react-redux';
-import { mockDataGeneratorModalOpened } from '../../modules/collection-tab';
+import { openMockDataGeneratorModal } from '../../modules/collection-tab';
 
 const collectionHeaderStyles = css({
   padding: spacing[400],
@@ -182,7 +182,7 @@ const CollectionHeader: React.FunctionComponent<CollectionHeaderProps> = ({
 };
 
 const ConnectedCollectionHeader = connect(undefined, {
-  onOpenMockDataModal: mockDataGeneratorModalOpened,
+  onOpenMockDataModal: openMockDataGeneratorModal,
 })(CollectionHeader);
 
 export default ConnectedCollectionHeader;

packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx

@@ -23,7 +23,7 @@ describe('MockDataGeneratorModal', () => {
 
   function renderModal({
     isOpen = true,
-    currentStep = MockDataGeneratorStep.AI_DISCLAIMER,
+    currentStep = MockDataGeneratorStep.SCHEMA_CONFIRMATION,
   } = {}) {
     return render(
       <MockDataGeneratorModal

packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx

@@ -70,7 +70,7 @@ const MockDataGeneratorModal = ({
       <ModalFooter className={footerStyles}>
         <Button
           onClick={onPreviousStep}
-          disabled={currentStep === MockDataGeneratorStep.AI_DISCLAIMER}
+          disabled={currentStep === MockDataGeneratorStep.SCHEMA_CONFIRMATION}
         >
           Back
         </Button>

packages/compass-collection/src/index.ts

@@ -18,6 +18,7 @@ import {
   CollectionWorkspaceTitle,
   CollectionPluginTitleComponent,
 } from './plugin-tab-title';
+import { atlasAiServiceLocator } from '@mongodb-js/compass-generative-ai/provider';
 
 export const WorkspaceTab: WorkspacePlugin<typeof CollectionWorkspaceTitle> = {
   name: CollectionWorkspaceTitle,
@@ -37,6 +38,7 @@ export const WorkspaceTab: WorkspacePlugin<typeof CollectionWorkspaceTitle> = {
       connectionInfoRef: connectionInfoRefLocator,
       logger: createLoggerLocator('COMPASS-COLLECTION'),
       preferences: preferencesLocator,
+      atlasAiService: atlasAiServiceLocator,
     }
   ),
   content: CollectionTab,

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

@@ -10,6 +10,8 @@ import type { DataService } from '@mongodb-js/compass-connections/provider';
 import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider';
 import { type Logger, mongoLogId } from '@mongodb-js/compass-logging/provider';
 import { type PreferencesAccess } from 'compass-preferences-model/provider';
+import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider';
+
 import { isInternalFieldPath } from 'hadron-document';
 import toNS from 'mongodb-ns';
 import {
@@ -22,11 +24,11 @@ import {
 } from '../schema-analysis-types';
 import { calculateSchemaDepth } from '../calculate-schema-depth';
 import type { Document, MongoError } from 'mongodb';
+import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
 
 const DEFAULT_SAMPLE_SIZE = 100;
 
 const NO_DOCUMENTS_ERROR = 'No documents found in the collection to analyze.';
-import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
 
 function isAction<A extends AnyAction>(
   action: AnyAction,
@@ -63,6 +65,7 @@ type CollectionThunkAction<R, A extends AnyAction = AnyAction> = ThunkAction<
     experimentationServices: ReturnType<typeof experimentationServiceLocator>;
     logger: Logger;
     preferences: PreferencesAccess;
+    atlasAiService: AtlasAiService;
   },
   A
 >;
@@ -146,7 +149,7 @@ const reducer: Reducer<CollectionState, Action> = (
     },
     mockDataGenerator: {
       isModalOpen: false,
-      currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
+      currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
     },
   },
   action
@@ -234,7 +237,7 @@ const reducer: Reducer<CollectionState, Action> = (
       mockDataGenerator: {
         ...state.mockDataGenerator,
         isModalOpen: true,
-        currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
+        currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
       },
     };
   }
@@ -370,6 +373,27 @@ export const selectTab = (
   };
 };
 
+export const openMockDataGeneratorModal = (): CollectionThunkAction<
+  Promise<void>
+> => {
+  return async (dispatch, _getState, { atlasAiService, logger }) => {
+    try {
+      if (process.env.COMPASS_E2E_SKIP_ATLAS_SIGNIN !== 'true') {
+        await atlasAiService.ensureAiFeatureAccess();
+      }
+      dispatch(mockDataGeneratorModalOpened());
+    } catch (error) {
+      // if failed or user canceled we just don't show the modal
+      logger.log.error(
+        mongoLogId(1_001_000_364),
+        'Collections',
+        'Failed to ensure AI feature access and open mock data generator modal',
+        error
+      );
+    }
+  };
+};
+
 export const analyzeCollectionSchema = (): CollectionThunkAction<
   Promise<void>
 > => {

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

@@ -63,6 +63,7 @@ describe('Collection Tab Content store', function () {
     .stub(collectionTabModule, 'analyzeCollectionSchema')
     .returns(async () => {});
   const dataService = {} as any;
+  const atlasAiService = {} as any;
   let store: ReturnType<typeof activatePlugin>['store'];
   let deactivate: ReturnType<typeof activatePlugin>['deactivate'];
 
@@ -106,6 +107,7 @@ describe('Collection Tab Content store', function () {
         connectionInfoRef: connectionInfoRef as any,
         logger,
         preferences,
+        atlasAiService,
       },
       { on() {}, cleanup() {} } as any
     ));

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

@@ -22,6 +22,7 @@ import {
 } from 'compass-preferences-model/provider';
 import { ExperimentTestName } from '@mongodb-js/compass-telemetry/provider';
 import { SCHEMA_ANALYSIS_STATE_INITIAL } from '../schema-analysis-types';
+import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider';
 
 export type CollectionTabOptions = {
   /**
@@ -48,6 +49,7 @@ export type CollectionTabServices = {
   connectionInfoRef: ReturnType<typeof connectionInfoRefLocator>;
   logger: Logger;
   preferences: PreferencesAccess;
+  atlasAiService: AtlasAiService;
 };
 
 export function activatePlugin(
@@ -67,6 +69,7 @@ export function activatePlugin(
     connectionInfoRef,
     logger,
     preferences,
+    atlasAiService,
   } = services;
 
   if (!collectionModel) {
@@ -87,7 +90,7 @@ export function activatePlugin(
       },
       mockDataGenerator: {
         isModalOpen: false,
-        currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
+        currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
       },
     },
     applyMiddleware(
@@ -98,6 +101,7 @@ export function activatePlugin(
         experimentationServices,
         logger,
         preferences,
+        atlasAiService,
       })
     )
   );