WIP

Author: jcobis

packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx

@@ -15,7 +15,16 @@ import FieldSelector from './schema-field-selector';
 import FakerMappingSelector from './faker-mapping-selector';
 import type { FakerSchema, MockDataGeneratorState } from './types';
 import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai';
-import { updateEditedFakerSchema } from '../../modules/collection-tab';
+import {
+  fakerFieldTypeChanged,
+  fakerFieldMethodChanged,
+  type FakerFieldTypeChangedAction,
+  type FakerFieldMethodChangedAction,
+} from '../../modules/collection-tab';
+
+type FakerSchemaDispatch = (
+  action: FakerFieldTypeChangedAction | FakerFieldMethodChangedAction
+) => void;
 
 const containerStyles = css({
   display: 'flex',
@@ -54,11 +63,11 @@ const schemaEditorLoaderStyles = css({
 const FakerSchemaEditorContent = ({
   fakerSchema,
   onSchemaConfirmed,
-  onSchemaChange,
+  dispatch,
 }: {
   fakerSchema: FakerSchema;
   onSchemaConfirmed: (isConfirmed: boolean) => void;
-  onSchemaChange: (updatedSchema: FakerSchema) => void;
+  dispatch: FakerSchemaDispatch;
 }) => {
   const [fakerSchemaFormValues, setFakerSchemaFormValues] =
     React.useState<FakerSchema>(fakerSchema);
@@ -77,31 +86,35 @@ const FakerSchemaEditorContent = ({
   const onJsonTypeSelect = (newJsonType: MongoDBFieldType) => {
     const currentMapping = fakerSchemaFormValues[activeField];
     if (currentMapping) {
-      const updatedSchema = {
+      // Update local form state
+      setFakerSchemaFormValues({
         ...fakerSchemaFormValues,
         [activeField]: {
           ...currentMapping,
           mongoType: newJsonType,
         },
-      };
-      setFakerSchemaFormValues(updatedSchema);
-      onSchemaChange(updatedSchema);
+      });
+
+      // Dispatch event to Redux
+      dispatch(fakerFieldTypeChanged(activeField, newJsonType));
       resetIsSchemaConfirmed();
     }
   };
 
   const onFakerFunctionSelect = (newFakerFunction: string) => {
     const currentMapping = fakerSchemaFormValues[activeField];
     if (currentMapping) {
-      const updatedSchema = {
+      // Update local form state
+      setFakerSchemaFormValues({
         ...fakerSchemaFormValues,
         [activeField]: {
           ...currentMapping,
           fakerMethod: newFakerFunction,
         },
-      };
-      setFakerSchemaFormValues(updatedSchema);
-      onSchemaChange(updatedSchema);
+      });
+
+      // Dispatch event to Redux
+      dispatch(fakerFieldMethodChanged(activeField, newFakerFunction));
       resetIsSchemaConfirmed();
     }
   };
@@ -144,7 +157,7 @@ const FakerSchemaEditorScreen = ({
   isSchemaConfirmed: boolean;
   onSchemaConfirmed: (isConfirmed: boolean) => void;
   fakerSchemaGenerationState: MockDataGeneratorState;
-  dispatch: (action: any) => void;
+  dispatch: FakerSchemaDispatch;
 }) => {
   return (
     <div data-testid="faker-schema-editor" className={containerStyles}>
@@ -172,9 +185,7 @@ const FakerSchemaEditorScreen = ({
         <FakerSchemaEditorContent
           fakerSchema={fakerSchemaGenerationState.editedFakerSchema}
           onSchemaConfirmed={onSchemaConfirmed}
-          onSchemaChange={(updatedSchema) =>
-            dispatch(updateEditedFakerSchema(updatedSchema))
-          }
+          dispatch={dispatch}
         />
       )}
     </div>

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

@@ -16,7 +16,10 @@ 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 { MockDataSchemaRequest } from '@mongodb-js/compass-generative-ai';
+import type {
+  MockDataSchemaRequest,
+  MongoDBFieldType,
+} from '@mongodb-js/compass-generative-ai';
 import { isInternalFieldPath } from 'hadron-document';
 import toNS from 'mongodb-ns';
 import {
@@ -127,7 +130,8 @@ export enum CollectionActions {
   FakerMappingGenerationStarted = 'compass-collection/FakerMappingGenerationStarted',
   FakerMappingGenerationCompleted = 'compass-collection/FakerMappingGenerationCompleted',
   FakerMappingGenerationFailed = 'compass-collection/FakerMappingGenerationFailed',
-  FakerSchemaEdited = 'compass-collection/FakerSchemaEdited',
+  FakerFieldTypeChanged = 'compass-collection/FakerFieldTypeChanged',
+  FakerFieldMethodChanged = 'compass-collection/FakerFieldMethodChanged',
 }
 
 interface CollectionMetadataFetchedAction {
@@ -197,9 +201,16 @@ export interface FakerMappingGenerationFailedAction {
   requestId: string;
 }
 
-export interface FakerSchemaEditedAction {
-  type: CollectionActions.FakerSchemaEdited;
-  editedFakerSchema: FakerSchema;
+export interface FakerFieldTypeChangedAction {
+  type: CollectionActions.FakerFieldTypeChanged;
+  fieldPath: string;
+  mongoType: MongoDBFieldType;
+}
+
+export interface FakerFieldMethodChangedAction {
+  type: CollectionActions.FakerFieldMethodChanged;
+  fieldPath: string;
+  fakerMethod: string;
 }
 
 const reducer: Reducer<CollectionState, Action> = (
@@ -495,20 +506,67 @@ const reducer: Reducer<CollectionState, Action> = (
   }
 
   if (
-    isAction<FakerSchemaEditedAction>(
+    isAction<FakerFieldTypeChangedAction>(
       action,
-      CollectionActions.FakerSchemaEdited
+      CollectionActions.FakerFieldTypeChanged
     )
   ) {
     if (state.fakerSchemaGeneration.status !== 'completed') {
       return state;
     }
 
+    const { fieldPath, mongoType } = action;
+    const currentMapping =
+      state.fakerSchemaGeneration.editedFakerSchema[fieldPath];
+
+    if (!currentMapping) {
+      return state;
+    }
+
     return {
       ...state,
       fakerSchemaGeneration: {
         ...state.fakerSchemaGeneration,
-        editedFakerSchema: action.editedFakerSchema,
+        editedFakerSchema: {
+          ...state.fakerSchemaGeneration.editedFakerSchema,
+          [fieldPath]: {
+            ...currentMapping,
+            mongoType,
+          },
+        },
+      },
+    };
+  }
+
+  if (
+    isAction<FakerFieldMethodChangedAction>(
+      action,
+      CollectionActions.FakerFieldMethodChanged
+    )
+  ) {
+    if (state.fakerSchemaGeneration.status !== 'completed') {
+      return state;
+    }
+
+    const { fieldPath, fakerMethod } = action;
+    const currentMapping =
+      state.fakerSchemaGeneration.editedFakerSchema[fieldPath];
+
+    if (!currentMapping) {
+      return state;
+    }
+
+    return {
+      ...state,
+      fakerSchemaGeneration: {
+        ...state.fakerSchemaGeneration,
+        editedFakerSchema: {
+          ...state.fakerSchemaGeneration.editedFakerSchema,
+          [fieldPath]: {
+            ...currentMapping,
+            fakerMethod,
+          },
+        },
       },
     };
   }
@@ -554,10 +612,26 @@ export const mockDataGeneratorPreviousButtonClicked = (): CollectionThunkAction<
   };
 };
 
-export const updateEditedFakerSchema = (
-  editedFakerSchema: FakerSchema
-): FakerSchemaEditedAction => {
-  return { type: CollectionActions.FakerSchemaEdited, editedFakerSchema };
+export const fakerFieldTypeChanged = (
+  fieldPath: string,
+  mongoType: MongoDBFieldType
+): FakerFieldTypeChangedAction => {
+  return {
+    type: CollectionActions.FakerFieldTypeChanged,
+    fieldPath,
+    mongoType,
+  };
+};
+
+export const fakerFieldMethodChanged = (
+  fieldPath: string,
+  fakerMethod: string
+): FakerFieldMethodChangedAction => {
+  return {
+    type: CollectionActions.FakerFieldMethodChanged,
+    fieldPath,
+    fakerMethod,
+  };
 };
 
 export const selectTab = (