diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx index 1d7661a12b1..cc929355522 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx @@ -10,6 +10,7 @@ import { SpinLoaderWithLabel, } from '@mongodb-js/compass-components'; import React from 'react'; +import { connect } from 'react-redux'; import FieldSelector from './schema-field-selector'; import FakerMappingSelector from './faker-mapping-selector'; import type { @@ -172,7 +173,6 @@ const FakerSchemaEditorScreen = ({ onSchemaConfirmed, fakerSchemaGenerationState, }: { - isSchemaConfirmed: boolean; onSchemaConfirmed: (isConfirmed: boolean) => void; fakerSchemaGenerationState: MockDataGeneratorState; }) => { @@ -200,7 +200,7 @@ const FakerSchemaEditorScreen = ({ )} {fakerSchemaGenerationState.status === 'completed' && ( )} @@ -208,4 +208,4 @@ const FakerSchemaEditorScreen = ({ ); }; -export default FakerSchemaEditorScreen; +export default connect()(FakerSchemaEditorScreen); diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx index 96492f25e3f..f2a436e78d6 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx @@ -824,7 +824,15 @@ describe('MockDataGeneratorModal', () => { currentStep: MockDataGeneratorStep.GENERATE_DATA, fakerSchemaGeneration: { status: 'completed', - fakerSchema: { + originalLlmResponse: { + name: { + fakerMethod: 'person.firstName', + fakerArgs: [], + probability: 1.0, + mongoType: 'String', + }, + }, + editedFakerSchema: { name: { fakerMethod: 'person.firstName', fakerArgs: [], @@ -848,7 +856,15 @@ describe('MockDataGeneratorModal', () => { currentStep: MockDataGeneratorStep.GENERATE_DATA, fakerSchemaGeneration: { status: 'completed', - fakerSchema: { + originalLlmResponse: { + name: { + fakerMethod: 'person.firstName', + fakerArgs: [], + probability: 1.0, + mongoType: 'String', + }, + }, + editedFakerSchema: { name: { fakerMethod: 'person.firstName', fakerArgs: [], @@ -872,7 +888,15 @@ describe('MockDataGeneratorModal', () => { currentStep: MockDataGeneratorStep.GENERATE_DATA, fakerSchemaGeneration: { status: 'completed', - fakerSchema: { + originalLlmResponse: { + name: { + fakerMethod: 'person.firstName', + fakerArgs: [], + probability: 1.0, + mongoType: 'String', + }, + }, + editedFakerSchema: { name: { fakerMethod: 'person.firstName', fakerArgs: [], @@ -919,7 +943,15 @@ describe('MockDataGeneratorModal', () => { connectionInfo: atlasConnectionInfo, fakerSchemaGeneration: { status: 'completed', - fakerSchema: { + originalLlmResponse: { + name: { + fakerMethod: 'person.firstName', + fakerArgs: [], + probability: 1.0, + mongoType: 'String', + }, + }, + editedFakerSchema: { name: { fakerMethod: 'person.firstName', fakerArgs: [], @@ -971,7 +1003,15 @@ describe('MockDataGeneratorModal', () => { currentStep: MockDataGeneratorStep.GENERATE_DATA, fakerSchemaGeneration: { status: 'completed', - fakerSchema: { + originalLlmResponse: { + name: { + fakerMethod: 'person.firstName', + fakerArgs: [], + probability: 1.0, + mongoType: 'String', + }, + }, + editedFakerSchema: { name: { fakerMethod: 'person.firstName', fakerArgs: [], @@ -1001,7 +1041,21 @@ describe('MockDataGeneratorModal', () => { currentStep: MockDataGeneratorStep.GENERATE_DATA, fakerSchemaGeneration: { status: 'completed', - fakerSchema: { + originalLlmResponse: { + name: { + fakerMethod: 'person.firstName', + fakerArgs: [], + probability: 1.0, + mongoType: 'String', + }, + email: { + fakerMethod: 'internet.email', + fakerArgs: [], + probability: 1.0, + mongoType: 'String', + }, + }, + editedFakerSchema: { name: { fakerMethod: 'person.firstName', fakerArgs: [], diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx index 4e7a7b2e991..74f48452bc9 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx @@ -83,7 +83,6 @@ const MockDataGeneratorModal = ({ case MockDataGeneratorStep.SCHEMA_EDITOR: return ( @@ -100,7 +99,7 @@ const MockDataGeneratorModal = ({ @@ -111,7 +110,6 @@ const MockDataGeneratorModal = ({ }, [ currentStep, fakerSchemaGenerationState, - isSchemaConfirmed, documentCount, setDocumentCount, ]); diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/script-screen.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/script-screen.tsx index 7ead7b87b94..5f6aa624c8f 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/script-screen.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/script-screen.tsx @@ -217,7 +217,7 @@ const mapStateToProps = (state: CollectionState) => { return { fakerSchema: fakerSchemaGeneration.status === 'completed' - ? fakerSchemaGeneration.fakerSchema + ? fakerSchemaGeneration.editedFakerSchema : null, namespace, arrayLengthMap: diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/types.ts b/packages/compass-collection/src/components/mock-data-generator-modal/types.ts index af5150fbd55..e206c97adeb 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/types.ts +++ b/packages/compass-collection/src/components/mock-data-generator-modal/types.ts @@ -21,7 +21,8 @@ type MockDataGeneratorInProgressState = { type MockDataGeneratorCompletedState = { status: 'completed'; - fakerSchema: FakerSchema; + originalLlmResponse: Readonly; // Immutable LLM response + editedFakerSchema: FakerSchema; // User-modified version requestId: string; }; diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts index a1bfe9866e5..76b3566ae87 100644 --- a/packages/compass-collection/src/modules/collection-tab.ts +++ b/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,6 +130,8 @@ export enum CollectionActions { FakerMappingGenerationStarted = 'compass-collection/FakerMappingGenerationStarted', FakerMappingGenerationCompleted = 'compass-collection/FakerMappingGenerationCompleted', FakerMappingGenerationFailed = 'compass-collection/FakerMappingGenerationFailed', + FakerFieldTypeChanged = 'compass-collection/FakerFieldTypeChanged', + FakerFieldMethodChanged = 'compass-collection/FakerFieldMethodChanged', } interface CollectionMetadataFetchedAction { @@ -196,6 +201,18 @@ export interface FakerMappingGenerationFailedAction { requestId: string; } +export interface FakerFieldTypeChangedAction { + type: CollectionActions.FakerFieldTypeChanged; + fieldPath: string; + mongoType: MongoDBFieldType; +} + +export interface FakerFieldMethodChangedAction { + type: CollectionActions.FakerFieldMethodChanged; + fieldPath: string; + fakerMethod: string; +} + const reducer: Reducer = ( state = { // TODO(COMPASS-7782): use hook to get the workspace tab id instead @@ -457,7 +474,8 @@ const reducer: Reducer = ( ...state, fakerSchemaGeneration: { status: 'completed', - fakerSchema: action.fakerSchema, + originalLlmResponse: action.fakerSchema, + editedFakerSchema: action.fakerSchema, // Initially same as LLM response requestId: action.requestId, }, }; @@ -487,6 +505,72 @@ const reducer: Reducer = ( }; } + if ( + isAction( + action, + 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: { + ...state.fakerSchemaGeneration.editedFakerSchema, + [fieldPath]: { + ...currentMapping, + mongoType, + }, + }, + }, + }; + } + + if ( + isAction( + 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, + }, + }, + }, + }; + } + return state; }; @@ -528,6 +612,28 @@ export const mockDataGeneratorPreviousButtonClicked = (): CollectionThunkAction< }; }; +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 = ( tabName: CollectionSubtab ): CollectionThunkAction => {