diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx
index edae8dedc29..962b5fd8212 100644
--- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx
@@ -9,7 +9,6 @@ import {
} from '@mongodb-js/testing-library-compass';
import sinon from 'sinon';
import FakerMappingSelector from './faker-mapping-selector';
-import { UNRECOGNIZED_FAKER_METHOD } from '../../modules/collection-tab';
import type { MongoDBFieldType } from '../../schema-analysis-types';
import {
MONGO_TYPE_TO_FAKER_METHODS,
@@ -122,39 +121,52 @@ describe('FakerMappingSelector', () => {
);
});
- it('should show warning banner when faker method is unrecognized', () => {
+ it('should always include the original LLM faker method in the dropdown', () => {
+ const originalLlmMethod = 'custom.llmMethod';
+
render(
);
- expect(
- screen.getByText(
- /Please select a function or we will default fill this field/
- )
- ).to.exist;
+ const fakerFunctionSelect = screen.getByLabelText('Faker Function');
+ userEvent.click(fakerFunctionSelect);
+
+ // Should include the original LLM method even though it's not in MONGO_TYPE_TO_FAKER_METHODS
+ expect(screen.getByRole('option', { name: originalLlmMethod })).to.exist;
+
+ // Should also include standard methods for String type
+ expect(screen.getByRole('option', { name: 'lorem.word' })).to.exist;
+ expect(screen.getByRole('option', { name: 'lorem.sentence' })).to.exist;
});
- it('should not show warning banner when faker method is recognized', () => {
+ it('should not duplicate the original LLM method if it is already in the standard methods', () => {
+ const originalLlmMethod = 'lorem.word';
+
render(
);
- expect(
- screen.queryByText(
- /Please select a function or we will default fill this field/
- )
- ).to.not.exist;
+ const fakerFunctionSelect = screen.getByLabelText('Faker Function');
+ userEvent.click(fakerFunctionSelect);
+
+ // Should only have one instance of 'lorem.word'
+ const loremWordOptions = screen.getAllByRole('option', {
+ name: 'lorem.word',
+ });
+ expect(loremWordOptions).to.have.length(1);
});
});
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx
index 5b6c289b5fa..4b5f2628aa1 100644
--- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx
@@ -1,6 +1,4 @@
import {
- Banner,
- BannerVariant,
Body,
Code,
css,
@@ -10,7 +8,6 @@ import {
spacing,
} from '@mongodb-js/compass-components';
import React, { useMemo } from 'react';
-import { UNRECOGNIZED_FAKER_METHOD } from '../../modules/collection-tab';
import {
MONGO_TYPE_TO_FAKER_METHODS,
MongoDBFieldTypeValues,
@@ -65,6 +62,7 @@ interface Props {
onJsonTypeSelect: (jsonType: MongoDBFieldType) => void;
activeFakerArgs: FakerArg[];
onFakerFunctionSelect: (fakerFunction: string) => void;
+ originalLlmFakerMethod?: string;
}
const FakerMappingSelector = ({
@@ -73,16 +71,18 @@ const FakerMappingSelector = ({
activeFakerArgs,
onJsonTypeSelect,
onFakerFunctionSelect,
+ originalLlmFakerMethod,
}: Props) => {
const fakerMethodOptions = useMemo(() => {
const methods = MONGO_TYPE_TO_FAKER_METHODS[activeJsonType] || [];
- if (methods.includes(activeFakerFunction)) {
- return methods;
+ // Include original LLM method if it's not already in the list of methods
+ if (originalLlmFakerMethod && !methods.includes(originalLlmFakerMethod)) {
+ return [originalLlmFakerMethod, ...methods];
}
- return [activeFakerFunction, ...methods];
- }, [activeJsonType, activeFakerFunction]);
+ return methods;
+ }, [activeJsonType, originalLlmFakerMethod]);
return (
@@ -111,29 +111,17 @@ const FakerMappingSelector = ({
))}
- {activeFakerFunction === UNRECOGNIZED_FAKER_METHOD ? (
-
- Please select a function or we will default fill this field with the
- string "Unrecognized"
-
- ) : (
- <>
-
-
- {formatFakerFunctionCallWithArgs(
- activeFakerFunction,
- activeFakerArgs
- )}
-
-
- )}
+
+
+ {formatFakerFunctionCallWithArgs(activeFakerFunction, activeFakerArgs)}
+
);
};
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 646262059a1..8ac01d98fb4 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
@@ -8,16 +8,17 @@ import {
} from '@mongodb-js/compass-components';
import React from 'react';
import { connect } from 'react-redux';
+import type { Dispatch } from 'redux';
import FieldSelector from './schema-field-selector';
import FakerMappingSelector from './faker-mapping-selector';
import { getDefaultFakerMethod } from './script-generation-utils';
-import type {
- FakerSchema,
- FakerFieldMapping,
- MockDataGeneratorState,
-} from './types';
+import type { FakerSchema, MockDataGeneratorState } from './types';
import type { MongoDBFieldType } from '../../schema-analysis-types';
import { useTelemetry } from '@mongodb-js/compass-telemetry/provider';
+import {
+ fakerFieldTypeChanged,
+ fakerFieldMethodChanged,
+} from '../../modules/collection-tab';
const containerStyles = css({
display: 'flex',
@@ -46,53 +47,39 @@ const schemaEditorLoaderStyles = css({
const FakerSchemaEditorContent = ({
fakerSchema,
+ originalLlmResponse,
+ onFieldTypeChanged,
+ onFieldMethodChanged,
}: {
fakerSchema: FakerSchema;
+ originalLlmResponse: FakerSchema;
+ onFieldTypeChanged: (fieldPath: string, mongoType: MongoDBFieldType) => void;
+ onFieldMethodChanged: (fieldPath: string, fakerMethod: string) => void;
}) => {
const track = useTelemetry();
- const [fakerSchemaFormValues, setFakerSchemaFormValues] =
- React.useState(fakerSchema);
-
- // Store original LLM mappings to restore when reselecting original methods
- const originalLlmMappings = React.useRef>(
- Object.fromEntries(
- Object.entries(fakerSchema).map(([field, mapping]) => [
- field,
- {
- ...mapping,
- },
- ])
- )
- );
- const fieldPaths = Object.keys(fakerSchemaFormValues);
+ const fieldPaths = Object.keys(fakerSchema);
const [activeField, setActiveField] = React.useState(fieldPaths[0]);
- const activeJsonType = fakerSchemaFormValues[activeField]?.mongoType;
- const activeFakerFunction = fakerSchemaFormValues[activeField]?.fakerMethod;
- const activeFakerArgs = fakerSchemaFormValues[activeField]?.fakerArgs;
+ const activeJsonType = fakerSchema[activeField]?.mongoType;
+ const activeFakerFunction = fakerSchema[activeField]?.fakerMethod;
+ const activeFakerArgs = fakerSchema[activeField]?.fakerArgs;
const onJsonTypeSelect = (newJsonType: MongoDBFieldType) => {
- const currentMapping = fakerSchemaFormValues[activeField];
- const originalLlmMapping = originalLlmMappings.current[activeField];
+ const currentMapping = fakerSchema[activeField];
+ const originalLlmMapping = originalLlmResponse[activeField];
if (currentMapping) {
const previousJsonType = currentMapping.mongoType;
const previousFakerMethod = currentMapping.fakerMethod;
- const isSwitchingToOriginalType =
- originalLlmMapping && newJsonType === originalLlmMapping.mongoType;
-
- const newMapping = isSwitchingToOriginalType
- ? { ...originalLlmMapping }
- : {
- ...currentMapping,
- mongoType: newJsonType,
- fakerMethod: getDefaultFakerMethod(newJsonType),
- fakerArgs: [],
- };
+ const newFakerMethod =
+ originalLlmMapping && newJsonType === originalLlmMapping.mongoType
+ ? originalLlmMapping.fakerMethod
+ : getDefaultFakerMethod(newJsonType);
- const newFakerMethod = newMapping.fakerMethod;
+ onFieldTypeChanged(activeField, newJsonType);
+ onFieldMethodChanged(activeField, newFakerMethod);
track('Mock Data JSON Type Changed', {
field_name: activeField,
@@ -101,33 +88,16 @@ const FakerSchemaEditorContent = ({
previous_faker_method: previousFakerMethod,
new_faker_method: newFakerMethod,
});
-
- setFakerSchemaFormValues({
- ...fakerSchemaFormValues,
- [activeField]: newMapping,
- });
}
};
const onFakerFunctionSelect = (newFakerFunction: string) => {
- const currentMapping = fakerSchemaFormValues[activeField];
- const originalLlmMapping = originalLlmMappings.current[activeField];
+ const currentMapping = fakerSchema[activeField];
if (currentMapping) {
const previousFakerMethod = currentMapping.fakerMethod;
- const isSwitchingToLlmSuggestion =
- originalLlmMapping &&
- currentMapping.mongoType === originalLlmMapping.mongoType &&
- newFakerFunction === originalLlmMapping.fakerMethod;
-
- const newMapping = isSwitchingToLlmSuggestion
- ? { ...originalLlmMapping }
- : {
- ...currentMapping,
- fakerMethod: newFakerFunction,
- fakerArgs: [],
- };
+ onFieldMethodChanged(activeField, newFakerFunction);
track('Mock Data Faker Method Changed', {
field_name: activeField,
@@ -135,11 +105,6 @@ const FakerSchemaEditorContent = ({
previous_faker_method: previousFakerMethod,
new_faker_method: newFakerFunction,
});
-
- setFakerSchemaFormValues({
- ...fakerSchemaFormValues,
- [activeField]: newMapping,
- });
}
};
@@ -150,7 +115,6 @@ const FakerSchemaEditorContent = ({
activeField={activeField}
fields={fieldPaths}
onFieldSelect={setActiveField}
- fakerSchema={fakerSchemaFormValues}
/>
{activeJsonType && activeFakerFunction && (
)}
@@ -168,8 +137,12 @@ const FakerSchemaEditorContent = ({
const FakerSchemaEditorScreen = ({
fakerSchemaGenerationState,
+ onFieldTypeChanged,
+ onFieldMethodChanged,
}: {
fakerSchemaGenerationState: MockDataGeneratorState;
+ onFieldTypeChanged: (fieldPath: string, mongoType: MongoDBFieldType) => void;
+ onFieldMethodChanged: (fieldPath: string, fakerMethod: string) => void;
}) => {
return (
@@ -197,10 +170,20 @@ const FakerSchemaEditorScreen = ({
{fakerSchemaGenerationState.status === 'completed' && (
)}
);
};
-export default connect()(FakerSchemaEditorScreen);
+const mapDispatchToProps = (dispatch: Dispatch) => ({
+ onFieldTypeChanged: (fieldPath: string, mongoType: MongoDBFieldType) =>
+ dispatch(fakerFieldTypeChanged(fieldPath, mongoType)),
+ onFieldMethodChanged: (fieldPath: string, fakerMethod: string) =>
+ dispatch(fakerFieldMethodChanged(fieldPath, fakerMethod)),
+});
+
+export default connect(null, mapDispatchToProps)(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 54fb81e7a7d..55f13900643 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
@@ -6,7 +6,6 @@ import {
renderWithActiveConnection,
waitFor,
userEvent,
- waitForElementToBeRemoved,
} from '@mongodb-js/testing-library-compass';
import { Provider } from 'react-redux';
import { createStore, applyMiddleware } from 'redux';
@@ -475,23 +474,14 @@ describe('MockDataGeneratorModal', () => {
userEvent.click(screen.getByText('email'));
expect(screen.getByText('email')).to.exist;
expect(screen.getByLabelText('JSON Type')).to.have.value('String');
- // the "email" field should have a warning banner since the faker method is invalid
expect(screen.getByLabelText('Faker Function')).to.have.value(
- 'Unrecognized'
+ 'lorem.word'
);
- expect(
- screen.getByText(
- 'Please select a function or we will default fill this field with the string "Unrecognized"'
- )
- ).to.exist;
// select the "username" field
userEvent.click(screen.getByText('username'));
expect(screen.getByText('username')).to.exist;
expect(screen.getByLabelText('JSON Type')).to.have.value('String');
- expect(screen.getByLabelText('Faker Function')).to.have.value(
- 'Unrecognized'
- );
});
it('does not show any fields that are not in the input schema', async () => {
@@ -532,81 +522,6 @@ describe('MockDataGeneratorModal', () => {
expect(screen.queryByText('email')).to.not.exist;
});
- it('shows unmapped fields as "Unrecognized"', async () => {
- const mockServices = createMockServices();
- mockServices.atlasAiService.getMockDataSchema = () =>
- Promise.resolve({
- fields: [
- {
- fieldPath: 'name',
- mongoType: 'String',
- fakerMethod: 'person.firstName',
- fakerArgs: [],
- isArray: false,
- probability: 1.0,
- },
- {
- fieldPath: 'age',
- mongoType: 'Int32',
- fakerMethod: 'number.int',
- fakerArgs: [],
- isArray: false,
- probability: 1.0,
- },
- ],
- });
-
- await renderModal({
- mockServices,
- schemaAnalysis: {
- ...defaultSchemaAnalysisState,
- processedSchema: {
- name: {
- type: 'String',
- probability: 1.0,
- },
- age: {
- type: 'Int32',
- probability: 1.0,
- },
- type: {
- type: 'String',
- probability: 1.0,
- sampleValues: ['cat', 'dog'],
- },
- },
- sampleDocument: { name: 'Peaches', age: 10, type: 'cat' },
- },
- });
-
- // advance to the schema editor step
- userEvent.click(screen.getByText('Confirm'));
- await waitForElementToBeRemoved(() =>
- screen.queryByTestId('faker-schema-editor-loader')
- );
-
- // select the "name" field
- userEvent.click(screen.getByText('name'));
- expect(screen.getByLabelText('JSON Type')).to.have.value('String');
- expect(screen.getByLabelText('Faker Function')).to.have.value(
- 'person.firstName'
- );
-
- // select the "age" field
- userEvent.click(screen.getByText('age'));
- expect(screen.getByLabelText('JSON Type')).to.have.value('Int32');
- expect(screen.getByLabelText('Faker Function')).to.have.value(
- 'number.int'
- );
-
- // select the "type" field
- userEvent.click(screen.getByText('type'));
- expect(screen.getByLabelText('JSON Type')).to.have.value('String');
- expect(screen.getByLabelText('Faker Function')).to.have.value(
- 'Unrecognized'
- );
- });
-
it('displays preview of the faker call without args when the args are invalid', async () => {
const largeLengthArgs = Array.from({ length: 11 }, () => 'testArg');
const mockServices = createMockServices();
@@ -818,6 +733,126 @@ describe('MockDataGeneratorModal', () => {
);
});
});
+
+ it('persists user modifications to faker mappings when navigating back to the screen', async () => {
+ await renderModal({
+ mockServices: mockServicesWithMockDataResponse,
+ schemaAnalysis: mockSchemaAnalysis,
+ });
+
+ userEvent.click(screen.getByText('Confirm'));
+ await waitFor(() => {
+ expect(screen.getByTestId('faker-schema-editor')).to.exist;
+ });
+
+ // Change the JSON type from String to Number
+ const jsonTypeSelect = screen.getByLabelText('JSON Type');
+ userEvent.click(jsonTypeSelect);
+ const numberOption = await screen.findByRole('option', {
+ name: 'Number',
+ });
+ userEvent.click(numberOption);
+
+ await waitFor(() => {
+ expect(screen.getByLabelText('JSON Type')).to.have.value('Number');
+ });
+
+ const fakerMethodSelect = screen.getByLabelText('Faker Function');
+ userEvent.click(fakerMethodSelect);
+ const floatOption = await screen.findByRole('option', {
+ name: 'number.float',
+ });
+ userEvent.click(floatOption);
+
+ await waitFor(() => {
+ expect(screen.getByLabelText('Faker Function')).to.have.value(
+ 'number.float'
+ );
+ });
+
+ // Advance to document count step
+ userEvent.click(screen.getByTestId('next-step-button'));
+ await waitFor(() => {
+ expect(screen.getByText('Specify Number of Documents to Generate')).to
+ .exist;
+ });
+
+ // Go back to schema editor
+ userEvent.click(screen.getByText('Back'));
+ await waitFor(() => {
+ expect(screen.getByTestId('faker-schema-editor')).to.exist;
+ });
+
+ // Verify both selections persisted
+ await waitFor(() => {
+ expect(screen.getByLabelText('JSON Type')).to.have.value('Number');
+ expect(screen.getByLabelText('Faker Function')).to.have.value(
+ 'number.float'
+ );
+ });
+ });
+
+ it('preserves original fakerArgs when selecting original LLM method', async () => {
+ // Mock response with fakerArgs
+ const mockServicesWithArgs = {
+ ...mockServicesWithMockDataResponse,
+ atlasAiService: {
+ getMockDataSchema: sinon.stub().resolves({
+ fields: [
+ {
+ fieldPath: 'name',
+ mongoType: 'String',
+ fakerMethod: 'person.firstName',
+ fakerArgs: [{ json: '{"locale":"en"}' }],
+ probability: 0.8,
+ },
+ ],
+ }),
+ },
+ };
+
+ await renderModal({
+ mockServices: mockServicesWithArgs,
+ schemaAnalysis: mockSchemaAnalysis,
+ });
+
+ userEvent.click(screen.getByText('Confirm'));
+ await waitFor(() => {
+ expect(screen.getByTestId('faker-schema-editor')).to.exist;
+ });
+
+ // Change to a different faker method (this resets fakerArgs)
+ const fakerMethodSelect = screen.getByLabelText('Faker Function');
+ userEvent.click(fakerMethodSelect);
+ const wordOption = await screen.findByRole('option', {
+ name: 'lorem.word',
+ });
+ userEvent.click(wordOption);
+
+ await waitFor(() => {
+ expect(screen.getByLabelText('Faker Function')).to.have.value(
+ 'lorem.word'
+ );
+ });
+
+ // Select the original LLM method again
+ userEvent.click(fakerMethodSelect);
+ const firstNameOption = await screen.findByRole('option', {
+ name: 'person.firstName',
+ });
+ userEvent.click(firstNameOption);
+
+ await waitFor(() => {
+ expect(screen.getByLabelText('Faker Function')).to.have.value(
+ 'person.firstName'
+ );
+ });
+
+ const preview = screen.getByTestId('faker-function-call-preview');
+ expect(preview.textContent).to.include(
+ 'faker.person.firstName({"locale":"en"})'
+ );
+ });
});
describe('on the document count step', () => {
diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/raw-schema-confirmation-screen.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/raw-schema-confirmation-screen.tsx
index 506e74178b2..f44f775fe38 100644
--- a/packages/compass-collection/src/components/mock-data-generator-modal/raw-schema-confirmation-screen.tsx
+++ b/packages/compass-collection/src/components/mock-data-generator-modal/raw-schema-confirmation-screen.tsx
@@ -11,9 +11,11 @@ import {
DocumentList,
useDarkMode,
cx,
+ Link,
} from '@mongodb-js/compass-components';
import { usePreference } from 'compass-preferences-model/provider';
+import { useConnectionInfo } from '@mongodb-js/compass-connections/provider';
import toSimplifiedFieldInfo from './to-simplified-field-info';
import type { CollectionState } from '../../modules/collection-tab';
import type { SchemaAnalysisState } from '../../schema-analysis-types';
@@ -44,6 +46,11 @@ const descriptionStyles = css({
marginBottom: spacing[200],
});
+const projectSettingsInfoStyles = css({
+ marginBottom: spacing[400],
+ color: palette.gray.dark1,
+});
+
const errorBannerStyles = css({
marginTop: spacing[400],
marginBottom: spacing[400],
@@ -61,6 +68,7 @@ const RawSchemaConfirmationScreen = ({
'enableGenAISampleDocumentPassing'
);
const isDarkMode = useDarkMode();
+ const connectionInfo = useConnectionInfo();
const subtitleText = enableSampleDocumentPassing
? 'Sample Documents Collected'
@@ -70,6 +78,19 @@ const RawSchemaConfirmationScreen = ({
? 'A sample of documents from your collection will be sent to an LLM for processing.'
: 'We have identified the following schema from your documents. This schema will be sent to an LLM for processing.';
+ const projectId = connectionInfo.atlasMetadata?.projectId;
+ const projectSettingsLink = projectId ? (
+
+ Project Settings
+
+ ) : (
+ 'Project Settings'
+ );
+
return (
{schemaAnalysis.status === 'complete' ? (
@@ -78,6 +99,11 @@ const RawSchemaConfirmationScreen = ({
{subtitleText}