diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts b/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts index c5935479ee6..6f7f9d7bbb4 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts +++ b/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts @@ -1,5 +1,5 @@ -import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai'; import { MockDataGeneratorStep } from './types'; +import type { MongoDBFieldType } from '../../schema-analysis-types'; export const StepButtonLabelMap = { [MockDataGeneratorStep.SCHEMA_CONFIRMATION]: 'Confirm', @@ -12,6 +12,25 @@ export const StepButtonLabelMap = { export const DEFAULT_DOCUMENT_COUNT = 1000; export const MAX_DOCUMENT_COUNT = 100000; +export const MongoDBFieldTypeValues: MongoDBFieldType[] = [ + 'String', + 'Number', + 'Boolean', + 'Date', + 'Int32', + 'Decimal128', + 'Long', + 'ObjectId', + 'RegExp', + 'Symbol', + 'MaxKey', + 'MinKey', + 'Binary', + 'Code', + 'Timestamp', + 'DBRef', +]; + /** * Map of MongoDB types to available Faker v9 methods. * Not all Faker methods are included here. 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 84d7c129603..edae8dedc29 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 @@ -10,11 +10,14 @@ import { 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'; +import type { MongoDBFieldType } from '../../schema-analysis-types'; +import { + MONGO_TYPE_TO_FAKER_METHODS, + MongoDBFieldTypeValues, +} from './constants'; import type { FakerArg } from './script-generation-utils'; -const mockActiveJsonType = MongoDBFieldTypeValues.String; +const mockActiveJsonType: MongoDBFieldType = 'String'; const mockActiveFakerFunction = 'lorem.word'; const mockActiveFakerArgs: Array = []; const onJsonTypeSelectStub = sinon.stub(); @@ -26,9 +29,6 @@ describe('FakerMappingSelector', () => { }); 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( { const jsonTypeSelect = screen.getByLabelText('JSON Type'); userEvent.click(jsonTypeSelect); - for (const type of mongoTypes) { + for (const type of MongoDBFieldTypeValues) { await waitFor(() => { expect(screen.getByRole('option', { name: type })).to.exist; }); 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 ecff308c314..aa08248262b 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 @@ -12,9 +12,11 @@ import { } from '@mongodb-js/compass-components'; 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'; +import { + MONGO_TYPE_TO_FAKER_METHODS, + MongoDBFieldTypeValues, +} from './constants'; +import type { MongoDBFieldType } from '../../schema-analysis-types'; import type { FakerArg } from './script-generation-utils'; const fieldMappingSelectorsStyles = css({ @@ -93,7 +95,7 @@ const FakerMappingSelector = ({ value={activeJsonType} onChange={(value) => onJsonTypeSelect(value as MongoDBFieldType)} > - {Object.values(MongoDBFieldTypeValues).map((type) => ( + {MongoDBFieldTypeValues.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 cc929355522..c4b5a6cfd20 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 @@ -13,13 +13,13 @@ import React from 'react'; import { connect } from 'react-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 { MongoDBFieldType } from '@mongodb-js/compass-generative-ai'; -import { getDefaultFakerMethod } from './script-generation-utils'; +import type { MongoDBFieldType } from '../../schema-analysis-types'; const containerStyles = css({ display: 'flex', 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 f2a436e78d6..3c2ecf68f71 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 @@ -97,7 +97,6 @@ describe('MockDataGeneratorModal', () => { fields: [ { fieldPath: 'name', - mongoType: 'String', fakerMethod: 'person.firstName', fakerArgs: [], }, @@ -326,25 +325,21 @@ describe('MockDataGeneratorModal', () => { fields: [ { fieldPath: 'name', - mongoType: 'String', fakerMethod: 'person.firstName', fakerArgs: [], }, { fieldPath: 'age', - mongoType: 'Int32', fakerMethod: 'number.int', fakerArgs: [], }, { fieldPath: 'email', - mongoType: 'String', fakerMethod: 'internet', fakerArgs: [], }, { fieldPath: 'username', - mongoType: 'String', fakerMethod: 'noSuchMethod', fakerArgs: [], }, @@ -439,7 +434,6 @@ describe('MockDataGeneratorModal', () => { fields: [ { fieldPath: 'name', - mongoType: 'String', fakerMethod: 'person.firstName', fakerArgs: [], isArray: false, @@ -447,7 +441,6 @@ describe('MockDataGeneratorModal', () => { }, { fieldPath: 'email', - mongoType: 'String', fakerMethod: 'internet.email', fakerArgs: [], isArray: false, @@ -477,7 +470,6 @@ describe('MockDataGeneratorModal', () => { fields: [ { fieldPath: 'name', - mongoType: 'String', fakerMethod: 'person.firstName', fakerArgs: [], isArray: false, @@ -485,7 +477,6 @@ describe('MockDataGeneratorModal', () => { }, { fieldPath: 'age', - mongoType: 'Int32', fakerMethod: 'number.int', fakerArgs: [], isArray: false, @@ -553,7 +544,6 @@ describe('MockDataGeneratorModal', () => { fields: [ { fieldPath: 'name', - mongoType: 'String', fakerMethod: 'person.firstName', fakerArgs: largeLengthArgs, isArray: false, @@ -561,7 +551,6 @@ describe('MockDataGeneratorModal', () => { }, { fieldPath: 'age', - mongoType: 'Int32', fakerMethod: 'number.int', fakerArgs: [ { @@ -575,7 +564,6 @@ describe('MockDataGeneratorModal', () => { }, { fieldPath: 'username', - mongoType: 'String', fakerMethod: 'string.alpha', // large string fakerArgs: ['a'.repeat(1001)], @@ -584,7 +572,6 @@ describe('MockDataGeneratorModal', () => { }, { fieldPath: 'avatar', - mongoType: 'String', fakerMethod: 'image.url', fakerArgs: [ { @@ -653,6 +640,7 @@ describe('MockDataGeneratorModal', () => { it('disables the Next button when the faker schema mapping is not confirmed', async () => { await renderModal({ mockServices: mockServicesWithMockDataResponse, + schemaAnalysis: mockSchemaAnalysis, }); // advance to the schema editor step @@ -669,6 +657,7 @@ describe('MockDataGeneratorModal', () => { it('resets the confirm schema mapping state when the user clicks the back button then goes back to the schema editor step', async () => { await renderModal({ mockServices: mockServicesWithMockDataResponse, + schemaAnalysis: mockSchemaAnalysis, }); // advance to the schema editor step @@ -705,6 +694,7 @@ describe('MockDataGeneratorModal', () => { it('preserves the confirm schema mapping state when the user clicks the next button then goes back to the schema editor step', async () => { await renderModal({ mockServices: mockServicesWithMockDataResponse, + schemaAnalysis: mockSchemaAnalysis, }); // advance to the schema editor step 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 7179a08ad38..7d26ef9825b 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,5 +1,5 @@ -import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai'; import type { FakerFieldMapping } from './types'; +import type { MongoDBFieldType } from '../../schema-analysis-types'; import { prettify } from '@mongodb-js/compass-editor'; import { UNRECOGNIZED_FAKER_METHOD } from '../../modules/collection-tab'; import { faker } from '@faker-js/faker/locale/en'; 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 e206c97adeb..180b51e709b 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,6 +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'; +import type { MongoDBFieldType } from '../../schema-analysis-types'; export enum MockDataGeneratorStep { SCHEMA_CONFIRMATION = 'SCHEMA_CONFIRMATION', diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts index 76b3566ae87..569360e9841 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, - MongoDBFieldType, -} 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 { @@ -30,6 +27,7 @@ import { type SchemaAnalysisError, type SchemaAnalysisState, type FieldInfo, + type MongoDBFieldType, } from '../schema-analysis-types'; import { calculateSchemaDepth } from '../calculate-schema-depth'; import { @@ -818,9 +816,13 @@ export const cancelSchemaAnalysis = (): CollectionThunkAction => { /** * Transforms LLM array format to keyed object structure. * Moves fieldPath from object property to object key. + * + * @param fakerSchema - The faker schema array from LLM response + * @param inputSchema - The schema definition for the LLM input used to carry over `mongoType` data */ function transformFakerSchemaToObject( - fakerSchema: LlmFakerMapping[] + fakerSchema: LlmFakerMapping[], + inputSchema: Record ): FakerSchema { const result: FakerSchema = {}; @@ -828,7 +830,8 @@ function transformFakerSchemaToObject( const { fieldPath, ...fieldMapping } = field; result[fieldPath] = { ...fieldMapping, - mongoType: fieldMapping.mongoType, + // Note: `validateFakerSchema` already handles fields that are not present in `inputSchema` + mongoType: inputSchema[fieldPath]?.type, }; } @@ -968,7 +971,8 @@ export const generateFakerMappings = (): CollectionThunkAction< // Transform to keyed object structure const transformedFakerSchema = transformFakerSchemaToObject( - response.fields + response.fields, + schemaAnalysis.processedSchema ); const validatedFakerSchema = validateFakerSchema( diff --git a/packages/compass-collection/src/schema-analysis-types.ts b/packages/compass-collection/src/schema-analysis-types.ts index c975f136130..b1705b41cc6 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'; @@ -42,6 +42,8 @@ export type SampleValue = | null | undefined; +export type MongoDBFieldType = PrimitiveSchemaType['name']; + /** * Schema field information (for LLM processing) */ 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 008a6ad772a..dbcb793f307 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts @@ -429,13 +429,11 @@ describe('AtlasAiService', function () { fields: [ { fieldPath: 'name', - mongoType: 'String', fakerMethod: 'person.fullName', fakerArgs: [], }, { fieldPath: 'age', - mongoType: 'Int32', fakerMethod: 'number.int', fakerArgs: [{ json: '{"min": 18, "max": 65}' }], }, @@ -464,13 +462,11 @@ describe('AtlasAiService', function () { fields: [ { fieldPath: 'name', - mongoType: 'String', fakerMethod: 'person.fullName', fakerArgs: [], }, { fieldPath: 'age', - mongoType: 'Int32', fakerMethod: 'number.int', fakerArgs: [{ json: '{"min": 18, "max": 122}' }], }, @@ -504,13 +500,11 @@ describe('AtlasAiService', function () { fields: [ { fieldPath: 'name', - mongoType: 'String', fakerMethod: 'person.fullName', fakerArgs: [], }, { fieldPath: 'age', - mongoType: 'Int32', fakerMethod: 'number.int', fakerArgs: [{ json: '{"min": 18, "max": 65}' }], }, @@ -541,13 +535,11 @@ describe('AtlasAiService', function () { fields: [ { fieldPath: 'name', - mongoType: 'String', fakerMethod: 'person.fullName', fakerArgs: [], }, { fieldPath: 'age', - 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 cb19dae3d1a..262be30797a 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 { PrimitiveSchemaType, SimplifiedSchema } from 'mongodb-schema'; +import type { SimplifiedSchema } from 'mongodb-schema'; import { type PreferencesAccess, isAIFeatureEnabled, @@ -230,37 +230,10 @@ export interface MockDataSchemaRequest { signal: AbortSignal; } -/** - * MongoDB schema type - */ -export type MongoDBFieldType = PrimitiveSchemaType['name']; - -export 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.custom((val) => - Object.values(MongoDBFieldTypeValues).includes(val) - ), fakerMethod: z.string(), fakerArgs: z.array( z.union([ diff --git a/packages/compass-generative-ai/src/index.ts b/packages/compass-generative-ai/src/index.ts index 0089fe592d8..525d5b12008 100644 --- a/packages/compass-generative-ai/src/index.ts +++ b/packages/compass-generative-ai/src/index.ts @@ -35,7 +35,4 @@ export type { MockDataSchemaRequest, MockDataSchemaRawField, MockDataSchemaResponse, - MongoDBFieldType, } from './atlas-ai-service'; - -export { MongoDBFieldTypeValues } from './atlas-ai-service';