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 9d64a77e369..b6c1aa726b0 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 @@ -10,6 +10,7 @@ import { } from '@mongodb-js/compass-components'; import React from 'react'; import { UNRECOGNIZED_FAKER_METHOD } from '../../modules/collection-tab'; +import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai'; const fieldMappingSelectorsStyles = css({ width: '50%', @@ -26,7 +27,7 @@ const labelStyles = css({ interface Props { activeJsonType: string; activeFakerFunction: string; - onJsonTypeSelect: (jsonType: string) => void; + onJsonTypeSelect: (jsonType: MongoDBFieldType) => void; onFakerFunctionSelect: (fakerFunction: string) => void; } @@ -43,7 +44,7 @@ const FakerMappingSelector = ({ label="JSON Type" allowDeselect={false} value={activeJsonType} - onChange={onJsonTypeSelect} + 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) => ( 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 9629d0a251e..39b6a5d7a27 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 @@ -12,7 +12,8 @@ import { import React from 'react'; import FieldSelector from './schema-field-selector'; import FakerMappingSelector from './faker-mapping-selector'; -import type { FakerSchemaMapping, MockDataGeneratorState } from './types'; +import type { FakerSchema, MockDataGeneratorState } from './types'; +import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai'; const containerStyles = css({ display: 'flex', @@ -49,55 +50,49 @@ const schemaEditorLoaderStyles = css({ }); const FakerSchemaEditorContent = ({ - fakerSchemaMappings, + fakerSchema, onSchemaConfirmed, }: { - fakerSchemaMappings: FakerSchemaMapping[]; + fakerSchema: FakerSchema; onSchemaConfirmed: (isConfirmed: boolean) => void; }) => { const [fakerSchemaFormValues, setFakerSchemaFormValues] = - React.useState>(fakerSchemaMappings); - const [activeField, setActiveField] = React.useState( - fakerSchemaFormValues[0].fieldPath - ); + React.useState(fakerSchema); + + const fieldPaths = Object.keys(fakerSchemaFormValues); + const [activeField, setActiveField] = React.useState(fieldPaths[0]); - const activeJsonType = fakerSchemaFormValues.find( - (mapping) => mapping.fieldPath === activeField - )?.mongoType; - const activeFakerFunction = fakerSchemaFormValues.find( - (mapping) => mapping.fieldPath === activeField - )?.fakerMethod; + const activeJsonType = fakerSchemaFormValues[activeField]?.mongoType; + const activeFakerFunction = fakerSchemaFormValues[activeField]?.fakerMethod; const resetIsSchemaConfirmed = () => { onSchemaConfirmed(false); }; - const onJsonTypeSelect = (newJsonType: string) => { - const updatedFakerFieldMapping = fakerSchemaFormValues.find( - (mapping) => mapping.fieldPath === activeField - ); - if (updatedFakerFieldMapping) { - updatedFakerFieldMapping.mongoType = newJsonType; - setFakerSchemaFormValues( - fakerSchemaFormValues.map((mapping) => - mapping.fieldPath === activeField ? updatedFakerFieldMapping : mapping - ) - ); + const onJsonTypeSelect = (newJsonType: MongoDBFieldType) => { + const currentMapping = fakerSchemaFormValues[activeField]; + if (currentMapping) { + setFakerSchemaFormValues({ + ...fakerSchemaFormValues, + [activeField]: { + ...currentMapping, + mongoType: newJsonType, + }, + }); resetIsSchemaConfirmed(); } }; const onFakerFunctionSelect = (newFakerFunction: string) => { - const updatedFakerFieldMapping = fakerSchemaFormValues.find( - (mapping) => mapping.fieldPath === activeField - ); - if (updatedFakerFieldMapping) { - updatedFakerFieldMapping.fakerMethod = newFakerFunction; - setFakerSchemaFormValues( - fakerSchemaFormValues.map((mapping) => - mapping.fieldPath === activeField ? updatedFakerFieldMapping : mapping - ) - ); + const currentMapping = fakerSchemaFormValues[activeField]; + if (currentMapping) { + setFakerSchemaFormValues({ + ...fakerSchemaFormValues, + [activeField]: { + ...currentMapping, + fakerMethod: newFakerFunction, + }, + }); resetIsSchemaConfirmed(); } }; @@ -107,7 +102,7 @@ const FakerSchemaEditorContent = ({
mapping.fieldPath)} + fields={fieldPaths} onFieldSelect={setActiveField} /> {activeJsonType && activeFakerFunction && ( @@ -163,7 +158,7 @@ const FakerSchemaEditorScreen = ({ )} {fakerSchemaGenerationState.status === 'completed' && ( )} 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 1b651c2f2b6..e758e4dbac6 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 @@ -5,6 +5,7 @@ import { renderWithActiveConnection, waitFor, userEvent, + waitForElementToBeRemoved, } from '@mongodb-js/testing-library-compass'; import { Provider } from 'react-redux'; import { createStore, applyMiddleware } from 'redux'; @@ -16,6 +17,20 @@ import type { CollectionState } from '../../modules/collection-tab'; import { default as collectionTabReducer } from '../../modules/collection-tab'; import type { ConnectionInfo } from '@mongodb-js/connection-info'; import type { MockDataSchemaResponse } from '@mongodb-js/compass-generative-ai'; +import type { SchemaAnalysisState } from '../../schema-analysis-types'; + +const defaultSchemaAnalysisState: SchemaAnalysisState = { + status: 'complete', + processedSchema: { + name: { + type: 'String', + probability: 1.0, + sample_values: ['John', 'Jane'], + }, + }, + sampleDocument: { name: 'John' }, + schemaMetadata: { maxNestingDepth: 1, validationRules: null }, +}; describe('MockDataGeneratorModal', () => { async function renderModal({ @@ -23,6 +38,7 @@ describe('MockDataGeneratorModal', () => { currentStep = MockDataGeneratorStep.SCHEMA_CONFIRMATION, enableGenAISampleDocumentPassing = false, mockServices = createMockServices(), + schemaAnalysis = defaultSchemaAnalysisState, connectionInfo, }: { isOpen?: boolean; @@ -30,23 +46,13 @@ describe('MockDataGeneratorModal', () => { currentStep?: MockDataGeneratorStep; mockServices?: any; connectionInfo?: ConnectionInfo; + schemaAnalysis?: SchemaAnalysisState; } = {}) { const initialState: CollectionState = { workspaceTabId: 'test-workspace-tab-id', namespace: 'test.collection', metadata: null, - schemaAnalysis: { - status: 'complete', - processedSchema: { - name: { - type: 'String', - probability: 1.0, - sample_values: ['John', 'Jane'], - }, - }, - sampleDocument: { name: 'John' }, - schemaMetadata: { maxNestingDepth: 1, validationRules: null }, - }, + schemaAnalysis, fakerSchemaGeneration: { status: 'idle', }, @@ -84,7 +90,7 @@ describe('MockDataGeneratorModal', () => { fields: [ { fieldPath: 'name', - mongoType: 'string', + mongoType: 'String', fakerMethod: 'person.firstName', fakerArgs: [], }, @@ -280,31 +286,58 @@ describe('MockDataGeneratorModal', () => { }); describe('on the schema editor step', () => { + const mockSchemaAnalysis: SchemaAnalysisState = { + ...defaultSchemaAnalysisState, + processedSchema: { + name: { + type: 'String', + probability: 1.0, + }, + age: { + type: 'Int32', + probability: 1.0, + }, + email: { + type: 'String', + probability: 1.0, + }, + username: { + type: 'String', + probability: 1.0, + }, + }, + sampleDocument: { + name: 'Jane', + age: 99, + email: 'Jane@email.com', + username: 'JaneDoe123', + }, + }; const mockServicesWithMockDataResponse = createMockServices(); mockServicesWithMockDataResponse.atlasAiService.getMockDataSchema = () => Promise.resolve({ fields: [ { fieldPath: 'name', - mongoType: 'string', + mongoType: 'String', fakerMethod: 'person.firstName', fakerArgs: [], }, { fieldPath: 'age', - mongoType: 'int', + mongoType: 'Int32', fakerMethod: 'number.int', fakerArgs: [], }, { fieldPath: 'email', - mongoType: 'string', + mongoType: 'String', fakerMethod: 'internet', fakerArgs: [], }, { fieldPath: 'username', - mongoType: 'string', + mongoType: 'String', fakerMethod: 'noSuchMethod', fakerArgs: [], }, @@ -320,11 +353,11 @@ describe('MockDataGeneratorModal', () => { resolve({ fields: [], }), - 1000 + 1 ) ); - await renderModal(); + await renderModal({ mockServices }); // advance to the schema editor step userEvent.click(screen.getByText('Confirm')); @@ -332,7 +365,10 @@ describe('MockDataGeneratorModal', () => { }); it('shows the faker schema editor when the faker schema generation is completed', async () => { - await renderModal({ mockServices: mockServicesWithMockDataResponse }); + await renderModal({ + mockServices: mockServicesWithMockDataResponse, + schemaAnalysis: mockSchemaAnalysis, + }); // advance to the schema editor step userEvent.click(screen.getByText('Confirm')); @@ -343,7 +379,10 @@ describe('MockDataGeneratorModal', () => { }); it('shows correct values for the faker schema editor', async () => { - await renderModal({ mockServices: mockServicesWithMockDataResponse }); + await renderModal({ + mockServices: mockServicesWithMockDataResponse, + schemaAnalysis: mockSchemaAnalysis, + }); // advance to the schema editor step userEvent.click(screen.getByText('Confirm')); @@ -352,21 +391,21 @@ describe('MockDataGeneratorModal', () => { }); // the "name" field should be selected by default expect(screen.getByText('name')).to.exist; - expect(screen.getByLabelText('JSON Type')).to.have.value('string'); + 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.getByText('age')).to.exist; - expect(screen.getByLabelText('JSON Type')).to.have.value('int'); + expect(screen.getByLabelText('JSON Type')).to.have.value('Int32'); expect(screen.getByLabelText('Faker Function')).to.have.value( 'number.int' ); // select the "email" field userEvent.click(screen.getByText('email')); expect(screen.getByText('email')).to.exist; - expect(screen.getByLabelText('JSON Type')).to.have.value('string'); + 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' @@ -380,7 +419,120 @@ describe('MockDataGeneratorModal', () => { // 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('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 () => { + const mockServices = createMockServices(); + mockServices.atlasAiService.getMockDataSchema = () => + Promise.resolve({ + fields: [ + { + fieldPath: 'name', + mongoType: 'String', + fakerMethod: 'person.firstName', + fakerArgs: [], + isArray: false, + probability: 1.0, + }, + { + fieldPath: 'email', + mongoType: 'String', + fakerMethod: 'internet.email', + fakerArgs: [], + isArray: false, + probability: 1.0, + }, + ], + }); + await renderModal({ + mockServices, + }); + + // advance to the schema editor step + userEvent.click(screen.getByText('Confirm')); + + await waitFor(() => { + expect(screen.getByTestId('faker-schema-editor')).to.exist; + }); + + expect(screen.getByText('name')).to.exist; + 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, + sample_values: ['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' ); diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.spec.ts b/packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.spec.ts index b7c24288dcf..82dcf6cb1da 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.spec.ts +++ b/packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.spec.ts @@ -1,9 +1,7 @@ import { expect } from 'chai'; import { faker } from '@faker-js/faker/locale/en'; -import { - generateScript, - type FakerFieldMapping, -} from './script-generation-utils'; +import { generateScript } from './script-generation-utils'; +import type { FakerFieldMapping } from './types'; /** * Helper function to test that generated document code is executable diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.ts b/packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.ts index 1fcc7b09f8f..6d465e6628a 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.ts +++ b/packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.ts @@ -1,17 +1,11 @@ -import type { MongoDBFieldType } from '../../schema-analysis-types'; +import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai'; +import type { FakerFieldMapping } from './types'; export type FakerArg = string | number | boolean | { json: string }; const DEFAULT_ARRAY_LENGTH = 3; const INDENT_SIZE = 2; -export interface FakerFieldMapping { - mongoType: MongoDBFieldType; - fakerMethod: string; - fakerArgs: FakerArg[]; - probability?: number; // 0.0 - 1.0 frequency of field (defaults to 1.0) -} - // Array length configuration for different array types export type ArrayLengthMap = { [fieldName: string]: 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 947911eabc1..af5150fbd55 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 @@ -1,4 +1,6 @@ import type { MockDataSchemaResponse } from '@mongodb-js/compass-generative-ai'; +import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai'; +import type { FakerArg } from './script-generation-utils'; export enum MockDataGeneratorStep { SCHEMA_CONFIRMATION = 'SCHEMA_CONFIRMATION', @@ -19,7 +21,7 @@ type MockDataGeneratorInProgressState = { type MockDataGeneratorCompletedState = { status: 'completed'; - fakerSchema: FakerSchemaMapping[]; + fakerSchema: FakerSchema; requestId: string; }; @@ -35,4 +37,13 @@ export type MockDataGeneratorState = | MockDataGeneratorCompletedState | MockDataGeneratorErrorState; -export type FakerSchemaMapping = MockDataSchemaResponse['fields'][number]; +export type LlmFakerMapping = MockDataSchemaResponse['fields'][number]; + +export interface FakerFieldMapping { + mongoType: MongoDBFieldType; + fakerMethod: string; + fakerArgs: FakerArg[]; + probability?: number; // 0.0 - 1.0 frequency of field (defaults to 1.0) +} + +export type FakerSchema = Record; diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts index 74ba34dd5a1..77720bd4bf3 100644 --- a/packages/compass-collection/src/modules/collection-tab.ts +++ b/packages/compass-collection/src/modules/collection-tab.ts @@ -16,10 +16,7 @@ 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, - MockDataSchemaResponse, -} from '@mongodb-js/compass-generative-ai'; +import type { MockDataSchemaRequest } from '@mongodb-js/compass-generative-ai'; import { isInternalFieldPath } from 'hadron-document'; import toNS from 'mongodb-ns'; import { @@ -39,7 +36,8 @@ import { import type { Document, MongoError } from 'mongodb'; import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types'; import type { - FakerSchemaMapping, + LlmFakerMapping, + FakerSchema, MockDataGeneratorState, } from '../components/mock-data-generator-modal/types'; @@ -184,7 +182,7 @@ export interface FakerMappingGenerationStartedAction { export interface FakerMappingGenerationCompletedAction { type: CollectionActions.FakerMappingGenerationCompleted; - fakerSchema: FakerSchemaMapping[]; + fakerSchema: FakerSchema; requestId: string; } @@ -698,32 +696,113 @@ export const cancelSchemaAnalysis = (): CollectionThunkAction => { }; }; +/** + * Transforms LLM array format to keyed object structure. + * Moves fieldPath from object property to object key. + */ +function transformFakerSchemaToObject( + fakerSchema: LlmFakerMapping[] +): FakerSchema { + const result: FakerSchema = {}; + + for (const field of fakerSchema) { + const { fieldPath, ...fieldMapping } = field; + result[fieldPath] = { + ...fieldMapping, + mongoType: fieldMapping.mongoType, + }; + } + + return result; +} + +/** + * Checks if the method exists and is callable on the faker object. + * + * Note: Only supports the format `module.method` (e.g., `internet.email`). + * Nested modules or other formats are not supported. + * @see {@link https://fakerjs.dev/api/} + */ +function isValidFakerMethod(fakerMethod: string): boolean { + const parts = fakerMethod.split('.'); + + // Validate format: exactly module.method + if (parts.length !== 2) { + return false; + } + + const [moduleName, methodName] = parts; + + try { + const fakerModule = (faker as unknown as Record)[ + moduleName + ]; + return ( + fakerModule !== null && + fakerModule !== undefined && + typeof fakerModule === 'object' && + typeof (fakerModule as Record)[methodName] === 'function' + ); + } catch { + return false; + } +} + +/** + * Validates a given faker schema against an input schema. + * + * - Validates the `fakerMethod` for each field, marking it as unrecognized if invalid + * - Adds any unmapped input schema fields to the result with an unrecognized faker method + * + * @param inputSchema - The schema definition for the input, mapping field names to their metadata. + * @param fakerSchemaArray - The array of faker schema mappings from LLM to validate and map. + * @param logger - Logger instance used to log warnings for invalid faker methods. + * @returns A keyed object of validated faker schema mappings, with one-to-one fields with input schema. + */ const validateFakerSchema = ( - fakerSchema: MockDataSchemaResponse, + inputSchema: Record, + fakerSchemaRaw: FakerSchema, logger: Logger -) => { - return fakerSchema.fields.map((field) => { - const { fakerMethod } = field; - - const [moduleName, methodName, ...rest] = fakerMethod.split('.'); - if ( - rest.length > 0 || - typeof (faker as any)[moduleName]?.[methodName] !== 'function' - ) { - logger.log.warn( - mongoLogId(1_001_000_372), - 'Collection', - 'Invalid faker method', - { fakerMethod } - ); - return { - ...field, +): FakerSchema => { + const result: FakerSchema = {}; + + // Process all input schema fields in a single O(n) pass + for (const fieldPath of Object.keys(inputSchema)) { + if (fakerSchemaRaw[fieldPath]) { + // input schema field exists in faker schema + const fakerMapping = { + ...fakerSchemaRaw[fieldPath], + probability: inputSchema[fieldPath].probability, + }; + // Validate the faker method + if (isValidFakerMethod(fakerMapping.fakerMethod)) { + result[fieldPath] = fakerMapping; + } else { + logger.log.warn( + mongoLogId(1_001_000_372), + 'Collection', + 'Invalid faker method', + { fakerMethod: fakerMapping.fakerMethod } + ); + result[fieldPath] = { + mongoType: fakerMapping.mongoType, + fakerMethod: UNRECOGNIZED_FAKER_METHOD, + fakerArgs: [], + probability: fakerMapping.probability, + }; + } + } else { + // Field not mapped by LLM - add default + result[fieldPath] = { + mongoType: inputSchema[fieldPath].type, fakerMethod: UNRECOGNIZED_FAKER_METHOD, + fakerArgs: [], + probability: inputSchema[fieldPath].probability, }; } + } - return field; - }); + return result; }; export const generateFakerMappings = (): CollectionThunkAction< @@ -792,7 +871,16 @@ export const generateFakerMappings = (): CollectionThunkAction< connectionInfoRef.current ); - const validatedFakerSchema = validateFakerSchema(response, logger); + // Transform to keyed object structure + const transformedFakerSchema = transformFakerSchemaToObject( + response.fields + ); + + const validatedFakerSchema = validateFakerSchema( + schemaAnalysis.processedSchema, + transformedFakerSchema, + logger + ); fakerSchemaGenerationAbortControllerRef.current = undefined; dispatch({ diff --git a/packages/compass-collection/src/schema-analysis-types.ts b/packages/compass-collection/src/schema-analysis-types.ts index 286fc166158..954080599af 100644 --- a/packages/compass-collection/src/schema-analysis-types.ts +++ b/packages/compass-collection/src/schema-analysis-types.ts @@ -1,5 +1,5 @@ +import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai'; import type { Document } from 'mongodb'; -import type { PrimitiveSchemaType } from 'mongodb-schema'; export const SCHEMA_ANALYSIS_STATE_INITIAL = 'initial'; export const SCHEMA_ANALYSIS_STATE_ANALYZING = 'analyzing'; @@ -30,11 +30,6 @@ export type SchemaAnalysisErrorState = { error: SchemaAnalysisError; }; -/** - * MongoDB schema type - */ -export type MongoDBFieldType = PrimitiveSchemaType['name']; - /** * Primitive values that can appear in sample_values after BSON-to-primitive conversion. * These are the JavaScript primitive equivalents of BSON values. diff --git a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts index a093b730e47..008a6ad772a 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts @@ -393,12 +393,10 @@ describe('AtlasAiService', function () { name: { type: 'string', sampleValues: ['John', 'Jane', 'Bob'], - probability: 0.9, }, age: { type: 'number', sampleValues: [25, 30, 35], - probability: 0.8, }, }, includeSampleValues: false, @@ -431,13 +429,13 @@ describe('AtlasAiService', function () { fields: [ { fieldPath: 'name', - mongoType: 'string', + mongoType: 'String', fakerMethod: 'person.fullName', fakerArgs: [], }, { fieldPath: 'age', - mongoType: 'int', + mongoType: 'Int32', fakerMethod: 'number.int', fakerArgs: [{ json: '{"min": 18, "max": 65}' }], }, @@ -466,13 +464,13 @@ describe('AtlasAiService', function () { fields: [ { fieldPath: 'name', - mongoType: 'string', + mongoType: 'String', fakerMethod: 'person.fullName', fakerArgs: [], }, { fieldPath: 'age', - mongoType: 'int', + mongoType: 'Int32', fakerMethod: 'number.int', fakerArgs: [{ json: '{"min": 18, "max": 122}' }], }, @@ -506,13 +504,13 @@ describe('AtlasAiService', function () { fields: [ { fieldPath: 'name', - mongoType: 'string', + mongoType: 'String', fakerMethod: 'person.fullName', fakerArgs: [], }, { fieldPath: 'age', - mongoType: 'int', + mongoType: 'Int32', fakerMethod: 'number.int', fakerArgs: [{ json: '{"min": 18, "max": 65}' }], }, @@ -536,7 +534,6 @@ describe('AtlasAiService', function () { ); expect(requestBody.schema.age).to.not.have.property('sampleValues'); expect(requestBody.schema.name.type).to.equal('string'); - expect(requestBody.schema.age.probability).to.equal(0.8); }); it('makes POST request with correct headers and body structure', async function () { @@ -544,13 +541,13 @@ describe('AtlasAiService', function () { fields: [ { fieldPath: 'name', - mongoType: 'string', + mongoType: 'String', fakerMethod: 'person.fullName', fakerArgs: [], }, { fieldPath: 'age', - mongoType: 'int', + mongoType: 'Int32', fakerMethod: 'number.int', fakerArgs: [{ json: '{"min": 18, "max": 65}' }], }, diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index 480beb00cd5..a490d5119a1 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.ts @@ -1,4 +1,4 @@ -import type { SimplifiedSchema } from 'mongodb-schema'; +import type { PrimitiveSchemaType, SimplifiedSchema } from 'mongodb-schema'; import { type PreferencesAccess, isAIFeatureEnabled, @@ -218,7 +218,6 @@ const aiURLConfig = { export interface MockDataSchemaRawField { type: string; sampleValues?: unknown[]; - probability?: number; } export interface MockDataSchemaRequest { @@ -231,11 +230,38 @@ export interface MockDataSchemaRequest { signal: AbortSignal; } +/** + * MongoDB schema type + */ +export type MongoDBFieldType = PrimitiveSchemaType['name']; + +// TODO(CLOUDP-346699): Export this from mongodb-schema +enum MongoDBFieldTypeValues { + String = 'String', + Number = 'Number', + Boolean = 'Boolean', + Date = 'Date', + Int32 = 'Int32', + Decimal128 = 'Decimal128', + Long = 'Long', + ObjectId = 'ObjectId', + RegExp = 'RegExp', + Symbol = 'Symbol', + MaxKey = 'MaxKey', + MinKey = 'MinKey', + Binary = 'Binary', + Code = 'Code', + Timestamp = 'Timestamp', + DBRef = 'DBRef', +} + export const MockDataSchemaResponseShape = z.object({ fields: z.array( z.object({ fieldPath: z.string(), - mongoType: z.string(), + mongoType: z.custom((val) => + Object.values(MongoDBFieldTypeValues).includes(val) + ), fakerMethod: z.string(), fakerArgs: z.array( z.union([ @@ -470,7 +496,7 @@ export class AtlasAiService { Omit > = {}; for (const [k, v] of Object.entries(schema)) { - newSchema[k] = { type: v.type, probability: v.probability }; + newSchema[k] = { type: v.type }; } schema = newSchema; } diff --git a/packages/compass-generative-ai/src/index.ts b/packages/compass-generative-ai/src/index.ts index 525d5b12008..26218310a07 100644 --- a/packages/compass-generative-ai/src/index.ts +++ b/packages/compass-generative-ai/src/index.ts @@ -35,4 +35,5 @@ export type { MockDataSchemaRequest, MockDataSchemaRawField, MockDataSchemaResponse, + MongoDBFieldType, } from './atlas-ai-service';