feat(compass-collection): manage faker mapping request state in redux CLOUDP-333850

packages/compass-collection/package.json

@@ -58,6 +58,7 @@
     "@mongodb-js/compass-workspaces": "^0.52.0",
     "@mongodb-js/connection-info": "^0.17.2",
     "@mongodb-js/mongodb-constants": "^0.14.0",
+    "bson": "^6.10.1",
     "compass-preferences-model": "^2.51.0",
     "hadron-document": "^8.9.6",
     "mongodb": "^6.17.0",
@@ -67,8 +68,7 @@
     "react": "^17.0.2",
     "react-redux": "^8.1.3",
     "redux": "^4.2.1",
-    "redux-thunk": "^2.4.2",
-    "bson": "^6.10.1"
+    "redux-thunk": "^2.4.2"
   },
   "devDependencies": {
     "@mongodb-js/eslint-config-compass": "^1.4.7",

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

@@ -1,7 +1,36 @@
+import type { MockDataSchemaResponse } from '@mongodb-js/compass-generative-ai';
+
 export enum MockDataGeneratorStep {
   SCHEMA_CONFIRMATION = 'SCHEMA_CONFIRMATION',
   SCHEMA_EDITOR = 'SCHEMA_EDITOR',
   DOCUMENT_COUNT = 'DOCUMENT_COUNT',
   PREVIEW_DATA = 'PREVIEW_DATA',
   GENERATE_DATA = 'GENERATE_DATA',
 }
+
+type MockDataGeneratorIdleState = {
+  status: 'idle';
+};
+
+type MockDataGeneratorInProgressState = {
+  status: 'in-progress';
+  requestId: string;
+};
+
+type MockDataGeneratorCompletedState = {
+  status: 'completed';
+  fakerSchema: MockDataSchemaResponse;
+  requestId: string;
+};
+
+type MockDataGeneratorErrorState = {
+  status: 'error';
+  error: unknown;
+  requestId: string;
+};
+
+export type MockDataGeneratorState =
+  | MockDataGeneratorIdleState
+  | MockDataGeneratorInProgressState
+  | MockDataGeneratorCompletedState
+  | MockDataGeneratorErrorState;

packages/compass-collection/src/index.ts

@@ -14,6 +14,7 @@ import { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provide
 import { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider';
 import { createLoggerLocator } from '@mongodb-js/compass-logging/provider';
 import { preferencesLocator } from 'compass-preferences-model/provider';
+import { atlasAiServiceLocator } from '@mongodb-js/compass-generative-ai/provider';
 import {
   CollectionWorkspaceTitle,
   CollectionPluginTitleComponent,
@@ -33,6 +34,7 @@ export const WorkspaceTab: WorkspacePlugin<typeof CollectionWorkspaceTitle> = {
     {
       dataService: dataServiceLocator as DataServiceLocator<keyof DataService>,
       collection: collectionModelLocator,
+      atlasAiService: atlasAiServiceLocator,
       workspaces: workspacesServiceLocator,
       experimentationServices: experimentationServiceLocator,
       connectionInfoRef: connectionInfoRefLocator,

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

@@ -1,17 +1,22 @@
 import type { Reducer, AnyAction, Action } from 'redux';
 import { analyzeDocuments } from 'mongodb-schema';
+import { UUID } from 'bson';
 
 import type { CollectionMetadata } from 'mongodb-collection-model';
 import type { ThunkAction } from 'redux-thunk';
 import type AppRegistry from '@mongodb-js/compass-app-registry';
+import type { ConnectionInfo } from '@mongodb-js/connection-info';
 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 { CollectionSubtab } from '@mongodb-js/compass-workspaces';
+import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/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 type {
+  MockDataSchemaRequest,
+  MockDataSchemaResponse,
+} from '@mongodb-js/compass-generative-ai';
 import { isInternalFieldPath } from 'hadron-document';
 import toNS from 'mongodb-ns';
 import {
@@ -27,6 +32,7 @@ import { calculateSchemaDepth } from '../calculate-schema-depth';
 import { processSchema } from '../transform-schema-to-field-info';
 import type { Document, MongoError } from 'mongodb';
 import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
+import type { MockDataGeneratorState } from '../components/mock-data-generator-modal/types';
 
 const DEFAULT_SAMPLE_SIZE = 100;
 
@@ -63,11 +69,11 @@ type CollectionThunkAction<R, A extends AnyAction = AnyAction> = ThunkAction<
   {
     localAppRegistry: AppRegistry;
     dataService: DataService;
+    atlasAiService: AtlasAiService;
     workspaces: ReturnType<typeof workspacesServiceLocator>;
     experimentationServices: ReturnType<typeof experimentationServiceLocator>;
     logger: Logger;
     preferences: PreferencesAccess;
-    atlasAiService: AtlasAiService;
   },
   A
 >;
@@ -82,9 +88,10 @@ export type CollectionState = {
     isModalOpen: boolean;
     currentStep: MockDataGeneratorStep;
   };
+  fakerSchemaGeneration: MockDataGeneratorState;
 };
 
-enum CollectionActions {
+export enum CollectionActions {
   CollectionMetadataFetched = 'compass-collection/CollectionMetadataFetched',
   SchemaAnalysisStarted = 'compass-collection/SchemaAnalysisStarted',
   SchemaAnalysisFinished = 'compass-collection/SchemaAnalysisFinished',
@@ -94,6 +101,9 @@ enum CollectionActions {
   MockDataGeneratorModalClosed = 'compass-collection/MockDataGeneratorModalClosed',
   MockDataGeneratorNextButtonClicked = 'compass-collection/MockDataGeneratorNextButtonClicked',
   MockDataGeneratorPreviousButtonClicked = 'compass-collection/MockDataGeneratorPreviousButtonClicked',
+  FakerMappingGenerationStarted = 'compass-collection/FakerMappingGenerationStarted',
+  FakerMappingGenerationCompleted = 'compass-collection/FakerMappingGenerationCompleted',
+  FakerMappingGenerationFailed = 'compass-collection/FakerMappingGenerationFailed',
 }
 
 interface CollectionMetadataFetchedAction {
@@ -140,6 +150,23 @@ interface MockDataGeneratorPreviousButtonClickedAction {
   type: CollectionActions.MockDataGeneratorPreviousButtonClicked;
 }
 
+export interface FakerMappingGenerationStartedAction {
+  type: CollectionActions.FakerMappingGenerationStarted;
+  requestId: string;
+}
+
+export interface FakerMappingGenerationCompletedAction {
+  type: CollectionActions.FakerMappingGenerationCompleted;
+  fakerSchema: MockDataSchemaResponse;
+  requestId: string;
+}
+
+export interface FakerMappingGenerationFailedAction {
+  type: CollectionActions.FakerMappingGenerationFailed;
+  error: string;
+  requestId: string;
+}
+
 const reducer: Reducer<CollectionState, Action> = (
   state = {
     // TODO(COMPASS-7782): use hook to get the workspace tab id instead
@@ -153,6 +180,9 @@ const reducer: Reducer<CollectionState, Action> = (
       isModalOpen: false,
       currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
     },
+    fakerSchemaGeneration: {
+      status: 'idle',
+    },
   },
   action
 ) => {
@@ -256,6 +286,9 @@ const reducer: Reducer<CollectionState, Action> = (
         ...state.mockDataGenerator,
         isModalOpen: false,
       },
+      fakerSchemaGeneration: {
+        status: 'idle',
+      },
     };
   }
 
@@ -333,6 +366,75 @@ const reducer: Reducer<CollectionState, Action> = (
     };
   }
 
+  if (
+    isAction<FakerMappingGenerationStartedAction>(
+      action,
+      CollectionActions.FakerMappingGenerationStarted
+    )
+  ) {
+    if (
+      state.mockDataGenerator.currentStep !==
+      MockDataGeneratorStep.SCHEMA_CONFIRMATION
+    ) {
+      return state;
+    }
+
+    if (
+      state.fakerSchemaGeneration.status === 'in-progress' ||
+      state.fakerSchemaGeneration.status === 'completed'
+    ) {
+      return state;
+    }
+
+    return {
+      ...state,
+      fakerSchemaGeneration: {
+        status: 'in-progress',
+        requestId: action.requestId,
+      },
+    };
+  }
+
+  if (
+    isAction<FakerMappingGenerationCompletedAction>(
+      action,
+      CollectionActions.FakerMappingGenerationCompleted
+    )
+  ) {
+    if (state.fakerSchemaGeneration.status !== 'in-progress') {
+      return state;
+    }
+
+    return {
+      ...state,
+      fakerSchemaGeneration: {
+        status: 'completed',
+        fakerSchema: action.fakerSchema,
+        requestId: action.requestId,
+      },
+    };
+  }
+
+  if (
+    isAction<FakerMappingGenerationFailedAction>(
+      action,
+      CollectionActions.FakerMappingGenerationFailed
+    )
+  ) {
+    if (state.fakerSchemaGeneration.status !== 'in-progress') {
+      return state;
+    }
+
+    return {
+      ...state,
+      fakerSchemaGeneration: {
+        status: 'error',
+        error: action.error,
+        requestId: action.requestId,
+      },
+    };
+  }
+
   return state;
 };
 
@@ -480,6 +582,86 @@ export const analyzeCollectionSchema = (): CollectionThunkAction<
   };
 };
 
+export const generateFakerMappings = (
+  connectionInfo: ConnectionInfo
+): CollectionThunkAction<Promise<void>> => {
+  return async (
+    dispatch,
+    getState,
+    { logger, atlasAiService, preferences }
+  ) => {
+    const { schemaAnalysis, fakerSchemaGeneration, namespace } = getState();
+    if (schemaAnalysis.status !== SCHEMA_ANALYSIS_STATE_COMPLETE) {
+      logger.log.warn(
+        mongoLogId(1_001_000_305),
+        'Collection',
+        'Cannot call `generateFakeMappings` unless schema analysis is complete'
+      );
+      return;
+    }
+
+    if (fakerSchemaGeneration.status === 'in-progress') {
+      logger.debug(
+        'Faker mapping generation is already in progress, skipping new generation.'
+      );
+      return;
+    }
+
+    const includeSampleValues =
+      preferences.getPreferences().enableGenAISampleDocumentPassing;
+
+    // todo: dedup/abort requests using requestId (CLOUDP-333850)
+    const requestId = new UUID().toString();
+
+    try {
+      logger.debug('Generating faker mappings');
+
+      const { database, collection } = toNS(namespace);
+
+      dispatch({
+        type: CollectionActions.FakerMappingGenerationStarted,
+        requestId: requestId,
+      });
+
+      const mockDataSchemaRequest: MockDataSchemaRequest = {
+        databaseName: database,
+        collectionName: collection,
+        schema: schemaAnalysis.processedSchema,
+        validationRules: schemaAnalysis.schemaMetadata.validationRules,
+        includeSampleValues,
+        requestId,
+      };
+
+      const response = await atlasAiService.getMockDataSchema(
+        mockDataSchemaRequest,
+        connectionInfo
+      );
+      dispatch({
+        type: CollectionActions.FakerMappingGenerationCompleted,
+        fakerSchema: response,
+        requestId: requestId,
+      });
+    } catch (e) {
+      const errorMessage = e instanceof Error ? e.stack : String(e);
+
+      logger.log.error(
+        mongoLogId(1_001_000_312),
+        'Collection',
+        'Failed to generate faker.js mappings',
+        {
+          message: errorMessage,
+          namespace,
+        }
+      );
+      dispatch({
+        type: CollectionActions.FakerMappingGenerationFailed,
+        error: 'faker mapping request failed',
+        requestId,
+      });
+    }
+  };
+};
+
 export type CollectionTabPluginMetadata = CollectionMetadata & {
   /**
    * Initial query for the query bar