refactor: update MongoDBFieldType import paths and improve type usage across mock data generator components

Author: ncarbon

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) => (

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

@@ -13,7 +13,7 @@ import React from 'react';
 import FieldSelector from './schema-field-selector';
 import FakerMappingSelector from './faker-mapping-selector';
 import type { FakerSchema, MockDataGeneratorState } from './types';
-import type { MongoDBFieldType } from '../../schema-analysis-types';
+import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai';
 
 const containerStyles = css({
   display: 'flex',
@@ -69,14 +69,14 @@ const FakerSchemaEditorContent = ({
     onSchemaConfirmed(false);
   };
 
-  const onJsonTypeSelect = (newJsonType: string) => {
+  const onJsonTypeSelect = (newJsonType: MongoDBFieldType) => {
     const currentMapping = fakerSchemaFormValues[activeField];
     if (currentMapping) {
       setFakerSchemaFormValues({
         ...fakerSchemaFormValues,
         [activeField]: {
           ...currentMapping,
-          mongoType: newJsonType as MongoDBFieldType,
+          mongoType: newJsonType,
         },
       });
       resetIsSchemaConfirmed();

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

@@ -90,7 +90,7 @@ describe('MockDataGeneratorModal', () => {
               fields: [
                 {
                   fieldPath: 'name',
-                  mongoType: 'string',
+                  mongoType: 'String',
                   fakerMethod: 'person.firstName',
                   fakerArgs: [],
                   isArray: false,
@@ -323,31 +323,31 @@ describe('MockDataGeneratorModal', () => {
           fields: [
             {
               fieldPath: 'name',
-              mongoType: 'string',
+              mongoType: 'String',
               fakerMethod: 'person.firstName',
               fakerArgs: [],
               isArray: false,
               probability: 1.0,
             },
             {
               fieldPath: 'age',
-              mongoType: 'int',
+              mongoType: 'Int32',
               fakerMethod: 'number.int',
               fakerArgs: [],
               isArray: false,
               probability: 1.0,
             },
             {
               fieldPath: 'email',
-              mongoType: 'string',
+              mongoType: 'String',
               fakerMethod: 'internet',
               fakerArgs: [],
               isArray: false,
               probability: 1.0,
             },
             {
               fieldPath: 'username',
-              mongoType: 'string',
+              mongoType: 'String',
               fakerMethod: 'noSuchMethod',
               fakerArgs: [],
               isArray: false,
@@ -406,21 +406,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'
@@ -434,7 +434,7 @@ 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'
       );
@@ -448,15 +448,15 @@ describe('MockDataGeneratorModal', () => {
             fields: [
               {
                 fieldPath: 'name',
-                mongoType: 'string',
+                mongoType: 'String',
                 fakerMethod: 'person.firstName',
                 fakerArgs: [],
                 isArray: false,
                 probability: 1.0,
               },
               {
                 fieldPath: 'email',
-                mongoType: 'string',
+                mongoType: 'String',
                 fakerMethod: 'internet.email',
                 fakerArgs: [],
                 isArray: false,

packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.ts

@@ -1,4 +1,4 @@
-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 };

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

@@ -1,5 +1,5 @@
 import type { MockDataSchemaResponse } from '@mongodb-js/compass-generative-ai';
-import type { MongoDBFieldType } from '../../schema-analysis-types';
+import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai';
 import type { FakerArg } from './script-generation-utils';
 
 export enum MockDataGeneratorStep {

packages/compass-collection/src/modules/collection-tab.ts

@@ -27,7 +27,6 @@ import {
   type SchemaAnalysisError,
   type SchemaAnalysisState,
   type FieldInfo,
-  type MongoDBFieldType,
 } from '../schema-analysis-types';
 import { calculateSchemaDepth } from '../calculate-schema-depth';
 import {
@@ -710,7 +709,7 @@ function transformFakerSchemaToObject(
     const { fieldPath, ...fieldMapping } = field;
     result[fieldPath] = {
       ...fieldMapping,
-      mongoType: fieldMapping.mongoType as MongoDBFieldType,
+      mongoType: fieldMapping.mongoType,
     };
   }
 

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.

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,
@@ -231,12 +231,17 @@ export interface MockDataSchemaRequest {
   signal: AbortSignal;
 }
 
+/**
+ * MongoDB schema type
+ */
+export type MongoDBFieldType = PrimitiveSchemaType['name'];
+
 export const MockDataSchemaResponseShape = z.object({
   content: z.object({
     fields: z.array(
       z.object({
         fieldPath: z.string(),
-        mongoType: z.string(),
+        mongoType: z.string() as z.ZodType<MongoDBFieldType>,
         fakerMethod: z.string(),
         fakerArgs: z.array(
           z.union([

packages/compass-generative-ai/src/index.ts

@@ -35,4 +35,5 @@ export type {
   MockDataSchemaRequest,
   MockDataSchemaRawField,
   MockDataSchemaResponse,
+  MongoDBFieldType,
 } from './atlas-ai-service';