draft redux integration for requesting fakerjs mappings

Author: kpamaran

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,
@@ -32,6 +33,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/mock-data-generator-types.ts

@@ -0,0 +1,33 @@
+import type { MockDataSchemaResponse } from '@mongodb-js/compass-generative-ai';
+
+export const MOCK_DATA_GENERATOR_STATE_IDLE = 'idle';
+export const MOCK_DATA_GENERATOR_STATE_GENERATING = 'generating';
+export const MOCK_DATA_GENERATOR_STATE_COMPLETED = 'completed';
+export const MOCK_DATA_GENERATOR_STATE_ERROR = 'error';
+
+type MockDataGeneratorIdleState = {
+  status: typeof MOCK_DATA_GENERATOR_STATE_IDLE;
+};
+
+type MockDataGeneratorGeneratingState = {
+  status: typeof MOCK_DATA_GENERATOR_STATE_GENERATING;
+  requestId: string;
+};
+
+type MockDataGeneratorCompletedState = {
+  status: typeof MOCK_DATA_GENERATOR_STATE_COMPLETED;
+  fakerSchema: MockDataSchemaResponse;
+  requestId: string;
+};
+
+type MockDataGeneratorErrorState = {
+  status: typeof MOCK_DATA_GENERATOR_STATE_ERROR;
+  error: unknown;
+  requestId: string;
+};
+
+export type MockDataGeneratorState =
+  | MockDataGeneratorIdleState
+  | MockDataGeneratorGeneratingState
+  | MockDataGeneratorCompletedState
+  | MockDataGeneratorErrorState;

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

@@ -4,12 +4,18 @@ import { analyzeDocuments } 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 { 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 {
+  MockDataSchemaRequestShape,
+  type MockDataSchemaResponse,
+} from '@mongodb-js/compass-generative-ai';
 import { isInternalFieldPath } from 'hadron-document';
 import toNS from 'mongodb-ns';
 import {
@@ -24,11 +30,18 @@ import {
 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 {
+  MOCK_DATA_GENERATOR_STATE_IDLE,
+  MOCK_DATA_GENERATOR_STATE_GENERATING,
+  MOCK_DATA_GENERATOR_STATE_COMPLETED,
+  MOCK_DATA_GENERATOR_STATE_ERROR,
+} from '../mock-data-generator-types';
+import type { MockDataGeneratorState } from '../mock-data-generator-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,
@@ -61,6 +74,7 @@ type CollectionThunkAction<R, A extends AnyAction = AnyAction> = ThunkAction<
   {
     localAppRegistry: AppRegistry;
     dataService: DataService;
+    atlasAiService: AtlasAiService;
     workspaces: ReturnType<typeof workspacesServiceLocator>;
     experimentationServices: ReturnType<typeof experimentationServiceLocator>;
     logger: Logger;
@@ -79,6 +93,7 @@ export type CollectionState = {
     isModalOpen: boolean;
     currentStep: MockDataGeneratorStep;
   };
+  fakerSchemaGeneration: MockDataGeneratorState;
 };
 
 enum CollectionActions {
@@ -91,6 +106,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 {
@@ -137,6 +155,23 @@ interface MockDataGeneratorPreviousButtonClickedAction {
   type: CollectionActions.MockDataGeneratorPreviousButtonClicked;
 }
 
+interface FakerMappingGenerationStartedAction {
+  type: CollectionActions.FakerMappingGenerationStarted;
+  requestId: string;
+}
+
+interface FakerMappingGenerationCompletedAction {
+  type: CollectionActions.FakerMappingGenerationCompleted;
+  fakerSchema: MockDataSchemaResponse;
+  requestId: string;
+}
+
+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
@@ -150,6 +185,9 @@ const reducer: Reducer<CollectionState, Action> = (
       isModalOpen: false,
       currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
     },
+    fakerSchemaGeneration: {
+      status: MOCK_DATA_GENERATOR_STATE_IDLE,
+    },
   },
   action
 ) => {
@@ -332,6 +370,53 @@ const reducer: Reducer<CollectionState, Action> = (
     };
   }
 
+  if (
+    isAction<FakerMappingGenerationStartedAction>(
+      action,
+      CollectionActions.FakerMappingGenerationStarted
+    )
+  ) {
+    return {
+      ...state,
+      fakerSchemaGeneration: {
+        status: MOCK_DATA_GENERATOR_STATE_GENERATING,
+        requestId: action.requestId,
+      },
+    };
+  }
+
+  if (
+    isAction<FakerMappingGenerationCompletedAction>(
+      action,
+      CollectionActions.FakerMappingGenerationCompleted
+    )
+  ) {
+    return {
+      ...state,
+      fakerSchemaGeneration: {
+        status: MOCK_DATA_GENERATOR_STATE_COMPLETED,
+        fakerSchema: action.fakerSchema,
+        requestId: action.requestId,
+      },
+    };
+  }
+
+  if (
+    isAction<FakerMappingGenerationFailedAction>(
+      action,
+      CollectionActions.FakerMappingGenerationFailed
+    )
+  ) {
+    return {
+      ...state,
+      fakerSchemaGeneration: {
+        status: MOCK_DATA_GENERATOR_STATE_ERROR,
+        error: action.error,
+        requestId: action.requestId,
+      },
+    };
+  }
+
   return state;
 };
 
@@ -458,6 +543,65 @@ export const analyzeCollectionSchema = (): CollectionThunkAction<
   };
 };
 
+export const generateFakerMappings = (
+  connectionInfo: ConnectionInfo
+): CollectionThunkAction<Promise<void>> => {
+  return async (dispatch, getState, { logger, atlasAiService }) => {
+    const { schemaAnalysis, fakerSchemaGeneration } = getState();
+    if (schemaAnalysis.status !== SCHEMA_ANALYSIS_STATE_COMPLETE) {
+      logger.log.error(
+        mongoLogId(1_001_000_305),
+        'Collection',
+        'Cannot call `generateFakeMappings` unless schema analysis is complete'
+      );
+      return;
+    }
+
+    if (fakerSchemaGeneration.status === MOCK_DATA_GENERATOR_STATE_GENERATING) {
+      logger.debug(
+        'Faker mapping generation is already in progress, skipping new generation.'
+      );
+      return;
+    }
+
+    // todo: dedup/abort around requestId
+    const requestId = 'some-request-id';
+    try {
+      logger.debug('Generating faker mappings');
+
+      dispatch({
+        type: CollectionActions.FakerMappingGenerationStarted,
+        requestId: requestId,
+      });
+
+      const mockDataSchemaRequest = MockDataSchemaRequestShape.parse(
+        schemaAnalysis.processedSchema
+      );
+
+      const response = await atlasAiService.getMockDataSchema(
+        mockDataSchemaRequest,
+        connectionInfo
+      );
+      dispatch({
+        type: CollectionActions.FakerMappingGenerationCompleted,
+        fakerSchema: response,
+        requestId: requestId,
+      });
+    } catch {
+      logger.log.error(
+        mongoLogId(1_001_000_312),
+        'Collection',
+        'Failed to generate faker mappings'
+      );
+      dispatch({
+        type: CollectionActions.FakerMappingGenerationFailed,
+        error: 'Failed to generate faker mappings',
+        requestId: requestId,
+      });
+    }
+  };
+};
+
 export type CollectionTabPluginMetadata = CollectionMetadata & {
   /**
    * Initial query for the query bar

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'];
 
@@ -99,6 +100,7 @@ describe('Collection Tab Content store', function () {
       },
       {
         dataService,
+        atlasAiService,
         localAppRegistry,
         collection: mockCollection as any,
         workspaces: workspaces as any,

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

@@ -9,13 +9,15 @@ import reducer, {
   analyzeCollectionSchema,
 } from '../modules/collection-tab';
 import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
+import { MOCK_DATA_GENERATOR_STATE_IDLE } from '../mock-data-generator-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 { 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,
@@ -43,6 +45,7 @@ export type CollectionTabServices = {
   dataService: DataService;
   collection: Collection;
   localAppRegistry: AppRegistry;
+  atlasAiService: AtlasAiService;
   workspaces: ReturnType<typeof workspacesServiceLocator>;
   experimentationServices: ReturnType<typeof experimentationServiceLocator>;
   connectionInfoRef: ReturnType<typeof connectionInfoRefLocator>;
@@ -62,6 +65,7 @@ export function activatePlugin(
     dataService,
     collection: collectionModel,
     localAppRegistry,
+    atlasAiService,
     workspaces,
     experimentationServices,
     connectionInfoRef,
@@ -89,10 +93,14 @@ export function activatePlugin(
         isModalOpen: false,
         currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
       },
+      fakerSchemaGeneration: {
+        status: MOCK_DATA_GENERATOR_STATE_IDLE,
+      },
     },
     applyMiddleware(
       thunk.withExtraArgument({
         dataService,
+        atlasAiService,
         workspaces,
         localAppRegistry,
         experimentationServices,