feat(compass-collection): CLOUDP-348131: Schema Editor Dropdowns: Add mongo type to faker method mapping CLOUDP-348131

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

@@ -1,3 +1,4 @@
+import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai';
 import { MockDataGeneratorStep } from './types';
 
 export const StepButtonLabelMap = {
@@ -10,3 +11,72 @@ export const StepButtonLabelMap = {
 
 export const DEFAULT_DOCUMENT_COUNT = 1000;
 export const MAX_DOCUMENT_COUNT = 100000;
+
+/**
+ * Map of MongoDB types to available Faker methods.
+ * Not all Faker methods are included here.
+ * More can be found in the Faker.js API: https://fakerjs.dev/api/
+ */
+export const MONGO_TYPE_TO_FAKER_METHODS: Record<
+  MongoDBFieldType,
+  Array<string>
+> = {
+  String: [
+    'lorem.word',
+    'lorem.words',
+    'lorem.sentence',
+    'lorem.paragraph',
+    'person.firstName',
+    'person.lastName',
+    'person.fullName',
+    'person.jobTitle',
+    'internet.email',
+    'internet.url',
+    'internet.domainName',
+    'internet.userName',
+    'phone.number',
+    'location.city',
+    'location.country',
+    'location.streetAddress',
+    'location.zipCode',
+    'location.state',
+    'company.name',
+    'company.catchPhrase',
+    'color.human',
+    'commerce.productName',
+    'commerce.department',
+    'finance.accountName',
+    'finance.currencyCode',
+    'git.commitSha',
+    'string.uuid',
+    'string.alpha',
+    'string.alphanumeric',
+  ],
+  Number: [
+    'number.int',
+    'number.float',
+    'finance.amount',
+    'location.latitude',
+    'location.longitude',
+  ],
+  Int32: ['number.int', 'finance.amount'],
+  Long: ['number.int', 'number.bigInt'],
+  Decimal128: ['number.float', 'finance.amount'],
+  Boolean: ['datatype.boolean'],
+  Date: [
+    'date.recent',
+    'date.past',
+    'date.future',
+    'date.anytime',
+    'date.birthdate',
+  ],
+  Timestamp: ['date.recent', 'date.past', 'date.future', 'date.anytime'],
+  ObjectId: ['database.mongodbObjectId'],
+  Binary: ['string.hexadecimal', 'string.binary'],
+  RegExp: ['lorem.word', 'string.alpha'],
+  Code: ['lorem.sentence', 'lorem.paragraph', 'git.commitMessage'],
+  MinKey: ['number.int'],
+  MaxKey: ['number.int'],
+  Symbol: ['lorem.word', 'string.symbol'],
+  DBRef: ['database.mongodbObjectId'],
+};

packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx

@@ -0,0 +1,152 @@
+import { expect } from 'chai';
+import React from 'react';
+import {
+  screen,
+  render,
+  cleanup,
+  waitFor,
+  userEvent,
+} 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 { MONGO_TYPE_TO_FAKER_METHODS } from './constants';
+import { MongoDBFieldTypeValues } from '@mongodb-js/compass-generative-ai';
+
+const mockActiveJsonType = MongoDBFieldTypeValues.String;
+const mockActiveFakerFunction = 'lorem.word';
+const onJsonTypeSelectStub = sinon.stub();
+const onFakerFunctionSelectStub = sinon.stub();
+
+describe('FakerMappingSelector', () => {
+  afterEach(() => {
+    cleanup();
+  });
+
+  it('should display all MongoDB types in the dropdown', async () => {
+    // Check that all MongoDB types from the constant are present
+    const mongoTypes = Object.keys(MongoDBFieldTypeValues);
+
+    render(
+      <FakerMappingSelector
+        activeJsonType={mockActiveJsonType}
+        activeFakerFunction="lorem.word"
+        onJsonTypeSelect={onJsonTypeSelectStub}
+        onFakerFunctionSelect={onFakerFunctionSelectStub}
+      />
+    );
+
+    const jsonTypeSelect = screen.getByLabelText('JSON Type');
+    userEvent.click(jsonTypeSelect);
+
+    for (const type of mongoTypes) {
+      await waitFor(() => {
+        expect(screen.getByRole('option', { name: type })).to.exist;
+      });
+    }
+  });
+
+  describe('should display faker methods for each MongoDB type', () => {
+    for (const [mongoType, methods] of Object.entries(
+      MONGO_TYPE_TO_FAKER_METHODS
+    )) {
+      it(`should display faker methods for ${mongoType}`, () => {
+        const firstMethod = methods[0];
+
+        render(
+          <FakerMappingSelector
+            activeJsonType={
+              mongoType as keyof typeof MONGO_TYPE_TO_FAKER_METHODS
+            }
+            activeFakerFunction={firstMethod}
+            onJsonTypeSelect={onJsonTypeSelectStub}
+            onFakerFunctionSelect={onFakerFunctionSelectStub}
+          />
+        );
+
+        const fakerFunctionSelect = screen.getByLabelText('Faker Function');
+        userEvent.click(fakerFunctionSelect);
+
+        methods.forEach((method) => {
+          expect(screen.getByRole('option', { name: method })).to.exist;
+        });
+      });
+    }
+  });
+
+  it('should call onJsonTypeSelect when MongoDB type changes', async () => {
+    render(
+      <FakerMappingSelector
+        activeJsonType={mockActiveJsonType}
+        activeFakerFunction={mockActiveFakerFunction}
+        onJsonTypeSelect={onJsonTypeSelectStub}
+        onFakerFunctionSelect={onFakerFunctionSelectStub}
+      />
+    );
+
+    const jsonTypeSelect = screen.getByLabelText('JSON Type');
+    userEvent.click(jsonTypeSelect);
+
+    const numberOption = await screen.findByRole('option', { name: 'Number' });
+    userEvent.click(numberOption);
+
+    expect(onJsonTypeSelectStub).to.have.been.calledOnceWith('Number');
+  });
+
+  it('should call onFakerFunctionSelect when faker function changes', async () => {
+    render(
+      <FakerMappingSelector
+        activeJsonType={mockActiveJsonType}
+        activeFakerFunction={mockActiveFakerFunction}
+        onJsonTypeSelect={onJsonTypeSelectStub}
+        onFakerFunctionSelect={onFakerFunctionSelectStub}
+      />
+    );
+
+    const fakerFunctionSelect = screen.getByLabelText('Faker Function');
+    userEvent.click(fakerFunctionSelect);
+
+    const emailOption = await screen.findByRole('option', {
+      name: 'internet.email',
+    });
+    userEvent.click(emailOption);
+
+    expect(onFakerFunctionSelectStub).to.have.been.calledOnceWith(
+      'internet.email'
+    );
+  });
+
+  it('should show warning banner when faker method is unrecognized', () => {
+    render(
+      <FakerMappingSelector
+        activeJsonType={mockActiveJsonType}
+        activeFakerFunction={UNRECOGNIZED_FAKER_METHOD}
+        onJsonTypeSelect={onJsonTypeSelectStub}
+        onFakerFunctionSelect={onFakerFunctionSelectStub}
+      />
+    );
+
+    expect(
+      screen.getByText(
+        /Please select a function or we will default fill this field/
+      )
+    ).to.exist;
+  });
+
+  it('should not show warning banner when faker method is recognized', () => {
+    render(
+      <FakerMappingSelector
+        activeJsonType={mockActiveJsonType}
+        activeFakerFunction={mockActiveFakerFunction}
+        onJsonTypeSelect={onJsonTypeSelectStub}
+        onFakerFunctionSelect={onFakerFunctionSelectStub}
+      />
+    );
+
+    expect(
+      screen.queryByText(
+        /Please select a function or we will default fill this field/
+      )
+    ).to.not.exist;
+  });
+});

packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx

@@ -8,9 +8,11 @@ import {
   Select,
   spacing,
 } from '@mongodb-js/compass-components';
-import React from 'react';
+import React, { useMemo } from 'react';
 import { UNRECOGNIZED_FAKER_METHOD } from '../../modules/collection-tab';
 import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai';
+import { MongoDBFieldTypeValues } from '@mongodb-js/compass-generative-ai';
+import { MONGO_TYPE_TO_FAKER_METHODS } from './constants';
 
 const fieldMappingSelectorsStyles = css({
   width: '50%',
@@ -25,7 +27,7 @@ const labelStyles = css({
 });
 
 interface Props {
-  activeJsonType: string;
+  activeJsonType: MongoDBFieldType;
   activeFakerFunction: string;
   onJsonTypeSelect: (jsonType: MongoDBFieldType) => void;
   onFakerFunctionSelect: (fakerFunction: string) => void;
@@ -37,6 +39,16 @@ const FakerMappingSelector = ({
   onJsonTypeSelect,
   onFakerFunctionSelect,
 }: Props) => {
+  const fakerMethodOptions = useMemo(() => {
+    const methods = MONGO_TYPE_TO_FAKER_METHODS[activeJsonType] || [];
+
+    if (methods.includes(activeFakerFunction)) {
+      return methods;
+    }
+
+    return [activeFakerFunction, ...methods];
+  }, [activeJsonType, activeFakerFunction]);
+
   return (
     <div className={fieldMappingSelectorsStyles}>
       <Body className={labelStyles}>Mapping</Body>
@@ -46,8 +58,7 @@ const FakerMappingSelector = ({
         value={activeJsonType}
         onChange={(value) => onJsonTypeSelect(value as MongoDBFieldType)}
       >
-        {/* TODO(CLOUDP-344400) : Make the select input editable and render other options depending on the JSON type selected */}
-        {[activeJsonType].map((type) => (
+        {Object.values(MongoDBFieldTypeValues).map((type) => (
           <Option key={type} value={type}>
             {type}
           </Option>
@@ -59,10 +70,9 @@ const FakerMappingSelector = ({
         value={activeFakerFunction}
         onChange={onFakerFunctionSelect}
       >
-        {/* TODO(CLOUDP-344400): Make the select input editable and render other JSON types */}
-        {[activeFakerFunction].map((field) => (
-          <Option key={field} value={field}>
-            {field}
+        {fakerMethodOptions.map((method) => (
+          <Option key={method} value={method}>
+            {method}
           </Option>
         ))}
       </Select>

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

@@ -12,8 +12,13 @@ import {
 import React from 'react';
 import FieldSelector from './schema-field-selector';
 import FakerMappingSelector from './faker-mapping-selector';
-import type { FakerSchema, MockDataGeneratorState } from './types';
+import type {
+  FakerSchema,
+  FakerFieldMapping,
+  MockDataGeneratorState,
+} from './types';
 import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai';
+import { getDefaultFakerMethod } from './script-generation-utils';
 
 const containerStyles = css({
   display: 'flex',
@@ -59,6 +64,20 @@ const FakerSchemaEditorContent = ({
   const [fakerSchemaFormValues, setFakerSchemaFormValues] =
     React.useState<FakerSchema>(fakerSchema);
 
+  // Store original LLM mappings to restore when reselecting original methods
+  const originalLlmMappings = React.useRef<Record<string, FakerFieldMapping>>(
+    Object.fromEntries(
+      Object.entries(fakerSchema).map(([field, mapping]) => [
+        field,
+        {
+          mongoType: mapping.mongoType,
+          fakerMethod: mapping.fakerMethod,
+          fakerArgs: mapping.fakerArgs,
+        },
+      ])
+    )
+  );
+
   const fieldPaths = Object.keys(fakerSchemaFormValues);
   const [activeField, setActiveField] = React.useState<string>(fieldPaths[0]);
 
@@ -71,27 +90,52 @@ const FakerSchemaEditorContent = ({
 
   const onJsonTypeSelect = (newJsonType: MongoDBFieldType) => {
     const currentMapping = fakerSchemaFormValues[activeField];
+    const originalLlmMapping = originalLlmMappings.current[activeField];
+
     if (currentMapping) {
+      const isSwitchingToOriginalType =
+        originalLlmMapping && newJsonType === originalLlmMapping.mongoType;
+
+      const newMapping =
+        isSwitchingToOriginalType && originalLlmMapping
+          ? { ...originalLlmMapping }
+          : {
+              ...currentMapping,
+              mongoType: newJsonType,
+              fakerMethod: getDefaultFakerMethod(newJsonType),
+              fakerArgs: [],
+            };
+
       setFakerSchemaFormValues({
         ...fakerSchemaFormValues,
-        [activeField]: {
-          ...currentMapping,
-          mongoType: newJsonType,
-        },
+        [activeField]: newMapping,
       });
       resetIsSchemaConfirmed();
     }
   };
 
   const onFakerFunctionSelect = (newFakerFunction: string) => {
     const currentMapping = fakerSchemaFormValues[activeField];
+    const originalLlmMapping = originalLlmMappings.current[activeField];
+
     if (currentMapping) {
+      const isSwitchingToLlmSuggestion =
+        originalLlmMapping &&
+        currentMapping.mongoType === originalLlmMapping.mongoType &&
+        newFakerFunction === originalLlmMapping.fakerMethod;
+
+      const newMapping =
+        isSwitchingToLlmSuggestion && originalLlmMapping
+          ? { ...originalLlmMapping }
+          : {
+              ...currentMapping,
+              fakerMethod: newFakerFunction,
+              fakerArgs: [],
+            };
+
       setFakerSchemaFormValues({
         ...fakerSchemaFormValues,
-        [activeField]: {
-          ...currentMapping,
-          fakerMethod: newFakerFunction,
-        },
+        [activeField]: newMapping,
       });
       resetIsSchemaConfirmed();
     }