Skip to content

Commit d630729

Browse files
kpamarankpamaran
committed
simplify typing
1 parent ca06674 commit d630729

File tree

4 files changed

+58
-77
lines changed

4 files changed

+58
-77
lines changed

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

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,28 +9,23 @@ export enum MockDataGeneratorStep {
99
GENERATE_DATA = 'GENERATE_DATA',
1010
}
1111

12-
export const MOCK_DATA_GENERATOR_REQUEST_IDLE = 'idle';
13-
export const MOCK_DATA_GENERATOR_REQUEST_GENERATING = 'generating';
14-
export const MOCK_DATA_GENERATOR_REQUEST_COMPLETED = 'completed';
15-
export const MOCK_DATA_GENERATOR_REQUEST_ERROR = 'error';
16-
1712
type MockDataGeneratorIdleState = {
18-
status: typeof MOCK_DATA_GENERATOR_REQUEST_IDLE;
13+
status: 'idle';
1914
};
2015

2116
type MockDataGeneratorGeneratingState = {
22-
status: typeof MOCK_DATA_GENERATOR_REQUEST_GENERATING;
17+
status: 'generating';
2318
requestId: string;
2419
};
2520

2621
type MockDataGeneratorCompletedState = {
27-
status: typeof MOCK_DATA_GENERATOR_REQUEST_COMPLETED;
22+
status: 'completed';
2823
fakerSchema: MockDataSchemaResponse;
2924
requestId: string;
3025
};
3126

3227
type MockDataGeneratorErrorState = {
33-
status: typeof MOCK_DATA_GENERATOR_REQUEST_ERROR;
28+
status: 'error';
3429
error: unknown;
3530
requestId: string;
3631
};

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

Lines changed: 31 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetr
1313
import { type Logger, mongoLogId } from '@mongodb-js/compass-logging/provider';
1414
import { type PreferencesAccess } from 'compass-preferences-model/provider';
1515
import {
16-
MockDataSchemaRequestShape,
16+
type MockDataSchemaRequest,
1717
type MockDataSchemaResponse,
1818
} from '@mongodb-js/compass-generative-ai';
1919
import { isInternalFieldPath } from 'hadron-document';
@@ -30,13 +30,7 @@ import {
3030
import { calculateSchemaDepth } from '../calculate-schema-depth';
3131
import { processSchema } from '../transform-schema-to-field-info';
3232
import type { Document, MongoError } from 'mongodb';
33-
import {
34-
MockDataGeneratorStep,
35-
MOCK_DATA_GENERATOR_REQUEST_IDLE,
36-
MOCK_DATA_GENERATOR_REQUEST_GENERATING,
37-
MOCK_DATA_GENERATOR_REQUEST_COMPLETED,
38-
MOCK_DATA_GENERATOR_REQUEST_ERROR,
39-
} from '../components/mock-data-generator-modal/types';
33+
import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
4034
import type { MockDataGeneratorState } from '../components/mock-data-generator-modal/types';
4135

4236
const DEFAULT_SAMPLE_SIZE = 100;
@@ -186,7 +180,7 @@ const reducer: Reducer<CollectionState, Action> = (
186180
currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
187181
},
188182
fakerSchemaGeneration: {
189-
status: MOCK_DATA_GENERATOR_REQUEST_IDLE,
183+
status: 'idle',
190184
},
191185
},
192186
action
@@ -386,18 +380,16 @@ const reducer: Reducer<CollectionState, Action> = (
386380
}
387381

388382
if (
389-
state.fakerSchemaGeneration.status ===
390-
MOCK_DATA_GENERATOR_REQUEST_GENERATING ||
391-
state.fakerSchemaGeneration.status ===
392-
MOCK_DATA_GENERATOR_REQUEST_COMPLETED
383+
state.fakerSchemaGeneration.status === 'generating' ||
384+
state.fakerSchemaGeneration.status === 'completed'
393385
) {
394386
return state;
395387
}
396388

397389
return {
398390
...state,
399391
fakerSchemaGeneration: {
400-
status: MOCK_DATA_GENERATOR_REQUEST_GENERATING,
392+
status: 'generating',
401393
requestId: action.requestId,
402394
},
403395
};
@@ -409,17 +401,14 @@ const reducer: Reducer<CollectionState, Action> = (
409401
CollectionActions.FakerMappingGenerationCompleted
410402
)
411403
) {
412-
if (
413-
state.fakerSchemaGeneration.status !==
414-
MOCK_DATA_GENERATOR_REQUEST_GENERATING
415-
) {
404+
if (state.fakerSchemaGeneration.status !== 'generating') {
416405
return state;
417406
}
418407

419408
return {
420409
...state,
421410
fakerSchemaGeneration: {
422-
status: MOCK_DATA_GENERATOR_REQUEST_COMPLETED,
411+
status: 'completed',
423412
fakerSchema: action.fakerSchema,
424413
requestId: action.requestId,
425414
},
@@ -432,17 +421,14 @@ const reducer: Reducer<CollectionState, Action> = (
432421
CollectionActions.FakerMappingGenerationFailed
433422
)
434423
) {
435-
if (
436-
state.fakerSchemaGeneration.status !==
437-
MOCK_DATA_GENERATOR_REQUEST_GENERATING
438-
) {
424+
if (state.fakerSchemaGeneration.status !== 'generating') {
439425
return state;
440426
}
441427

442428
return {
443429
...state,
444430
fakerSchemaGeneration: {
445-
status: MOCK_DATA_GENERATOR_REQUEST_ERROR,
431+
status: 'error',
446432
error: action.error,
447433
requestId: action.requestId,
448434
},
@@ -589,9 +575,7 @@ export const generateFakerMappings = (
589575
return;
590576
}
591577

592-
if (
593-
fakerSchemaGeneration.status === MOCK_DATA_GENERATOR_REQUEST_GENERATING
594-
) {
578+
if (fakerSchemaGeneration.status === 'generating') {
595579
logger.debug(
596580
'Faker mapping generation is already in progress, skipping new generation.'
597581
);
@@ -611,12 +595,29 @@ export const generateFakerMappings = (
611595
requestId: requestId,
612596
});
613597

614-
const mockDataSchemaRequest = MockDataSchemaRequestShape.parse({
598+
// Convert FieldInfo objects to MockDataSchemaRawField format
599+
const schema: Record<
600+
string,
601+
{ type: string; sampleValues?: unknown[]; probability?: number }
602+
> = {};
603+
for (const [fieldName, fieldInfo] of Object.entries(
604+
schemaAnalysis.processedSchema
605+
)) {
606+
schema[fieldName] = {
607+
type: fieldInfo.type,
608+
sampleValues: fieldInfo.sample_values,
609+
probability: fieldInfo.probability,
610+
};
611+
}
612+
613+
const mockDataSchemaRequest: MockDataSchemaRequest = {
615614
databaseName: database,
616615
collectionName: collection,
617-
schema: schemaAnalysis.processedSchema,
616+
schema: schema,
618617
validationRules: schemaAnalysis.schemaMetadata.validationRules,
619-
});
618+
// todo
619+
includeSampleValues: true,
620+
};
620621

621622
const response = await atlasAiService.getMockDataSchema(
622623
mockDataSchemaRequest,

‎packages/compass-collection/src/stores/collection-tab.spec.ts‎

Lines changed: 22 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -25,16 +25,10 @@ import { type CollectionMetadata } from 'mongodb-collection-model';
2525
import {
2626
SCHEMA_ANALYSIS_STATE_COMPLETE,
2727
SCHEMA_ANALYSIS_STATE_INITIAL,
28-
SchemaAnalysisState,
2928
} from '../schema-analysis-types';
30-
import {
31-
MOCK_DATA_GENERATOR_REQUEST_COMPLETED,
32-
MOCK_DATA_GENERATOR_REQUEST_GENERATING,
33-
MOCK_DATA_GENERATOR_REQUEST_IDLE,
34-
MOCK_DATA_GENERATOR_REQUEST_ERROR,
35-
MockDataGeneratorState,
36-
MockDataGeneratorStep,
37-
} from '../components/mock-data-generator-modal/types';
29+
import type { SchemaAnalysisState } from '../schema-analysis-types';
30+
import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
31+
import type { MockDataGeneratorState } from '../components/mock-data-generator-modal/types';
3832
import { CollectionActions } from '../modules/collection-tab';
3933
import { type MockDataSchemaResponse } from '@mongodb-js/compass-generative-ai';
4034

@@ -330,7 +324,7 @@ describe('Collection Tab Content store', function () {
330324
validationRules: null,
331325
},
332326
},
333-
fakerSchemaGeneration: { status: MOCK_DATA_GENERATOR_REQUEST_IDLE },
327+
fakerSchemaGeneration: { status: 'idle' },
334328
});
335329
const logger = {
336330
log: { error: sandbox.spy() },
@@ -348,14 +342,6 @@ describe('Collection Tab Content store', function () {
348342
fakerArgs: [],
349343
isArray: false,
350344
},
351-
{
352-
fieldPath: 'age',
353-
probability: 1.0,
354-
mongoType: 'number',
355-
fakerMethod: 'number.int',
356-
fakerArgs: [],
357-
isArray: false,
358-
},
359345
{
360346
fieldPath: 'isActive',
361347
probability: 1.0,
@@ -404,7 +390,7 @@ describe('Collection Tab Content store', function () {
404390
status: SCHEMA_ANALYSIS_STATE_COMPLETE,
405391
processedSchema: undefined,
406392
},
407-
fakerSchemaGeneration: { status: MOCK_DATA_GENERATOR_REQUEST_IDLE },
393+
fakerSchemaGeneration: { status: 'idle' },
408394
});
409395
const logger = {
410396
log: { error: sandbox.spy() },
@@ -445,7 +431,7 @@ describe('Collection Tab Content store', function () {
445431
const dispatch = sandbox.spy();
446432
const getState = sandbox.stub().returns({
447433
schemaAnalysis: { status: SCHEMA_ANALYSIS_STATE_INITIAL },
448-
fakerSchemaGeneration: { status: MOCK_DATA_GENERATOR_REQUEST_IDLE },
434+
fakerSchemaGeneration: { status: 'idle' },
449435
});
450436
const logger = {
451437
log: { error: sandbox.spy() },
@@ -470,7 +456,7 @@ describe('Collection Tab Content store', function () {
470456
const getState = sandbox.stub().returns({
471457
schemaAnalysis: { status: SCHEMA_ANALYSIS_STATE_COMPLETE },
472458
fakerSchemaGeneration: {
473-
status: MOCK_DATA_GENERATOR_REQUEST_GENERATING,
459+
status: 'generating',
474460
},
475461
});
476462
const logger = {
@@ -503,7 +489,7 @@ describe('Collection Tab Content store', function () {
503489
isModalOpen: false,
504490
currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
505491
},
506-
fakerSchemaGeneration: { status: MOCK_DATA_GENERATOR_REQUEST_IDLE },
492+
fakerSchemaGeneration: { status: 'idle' },
507493
};
508494

509495
const completeSchemaState: SchemaAnalysisState = {
@@ -537,7 +523,7 @@ describe('Collection Tab Content store', function () {
537523
currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
538524
},
539525
fakerSchemaGeneration: {
540-
status: MOCK_DATA_GENERATOR_REQUEST_GENERATING,
526+
status: 'generating',
541527
requestId: 'existing_id',
542528
},
543529
},
@@ -549,7 +535,7 @@ describe('Collection Tab Content store', function () {
549535
currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
550536
},
551537
fakerSchemaGeneration: {
552-
status: MOCK_DATA_GENERATOR_REQUEST_COMPLETED,
538+
status: 'completed',
553539
fakerSchema: {
554540
content: {
555541
fields: [
@@ -583,13 +569,13 @@ describe('Collection Tab Content store', function () {
583569
isModalOpen: false,
584570
currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
585571
},
586-
fakerSchemaGeneration: { status: MOCK_DATA_GENERATOR_REQUEST_IDLE },
572+
fakerSchemaGeneration: { status: 'idle' },
587573
};
588574
const action = fakerMappingGenerationStartedAction;
589575
const newState = collectionTabReducer(state, action);
590576

591577
expect(newState.fakerSchemaGeneration).to.deep.equal({
592-
status: MOCK_DATA_GENERATOR_REQUEST_GENERATING,
578+
status: 'generating',
593579
requestId: 'some_request_id',
594580
});
595581

@@ -624,9 +610,9 @@ describe('Collection Tab Content store', function () {
624610

625611
it('should not transition to completed if in idle, completed, or error state', function () {
626612
const noOpStates: MockDataGeneratorState[] = [
627-
{ status: MOCK_DATA_GENERATOR_REQUEST_IDLE },
613+
{ status: 'idle' },
628614
{
629-
status: MOCK_DATA_GENERATOR_REQUEST_COMPLETED,
615+
status: 'completed',
630616
fakerSchema: {
631617
content: {
632618
fields: [
@@ -644,7 +630,7 @@ describe('Collection Tab Content store', function () {
644630
requestId: 'existing_id',
645631
},
646632
{
647-
status: MOCK_DATA_GENERATOR_REQUEST_ERROR,
633+
status: 'error',
648634
error: 'Some error',
649635
requestId: 'error_request_id',
650636
},
@@ -665,15 +651,15 @@ describe('Collection Tab Content store', function () {
665651
const state: CollectionState = {
666652
...baseState,
667653
fakerSchemaGeneration: {
668-
status: MOCK_DATA_GENERATOR_REQUEST_GENERATING,
654+
status: 'generating',
669655
requestId: 'generating_request_id',
670656
},
671657
};
672658
const action = fakerMappingGenerationCompletedAction;
673659
const newState = collectionTabReducer(state, action);
674660

675661
expect(newState.fakerSchemaGeneration).to.deep.equal({
676-
status: MOCK_DATA_GENERATOR_REQUEST_COMPLETED,
662+
status: 'completed',
677663
fakerSchema: action.fakerSchema,
678664
requestId: action.requestId,
679665
});
@@ -698,9 +684,9 @@ describe('Collection Tab Content store', function () {
698684

699685
it('should not transition to error if in idle, completed, or error state', function () {
700686
const noOpStates: MockDataGeneratorState[] = [
701-
{ status: MOCK_DATA_GENERATOR_REQUEST_IDLE },
687+
{ status: 'idle' },
702688
{
703-
status: MOCK_DATA_GENERATOR_REQUEST_COMPLETED,
689+
status: 'completed',
704690
fakerSchema: {
705691
content: {
706692
fields: [
@@ -718,7 +704,7 @@ describe('Collection Tab Content store', function () {
718704
requestId: 'existing_id',
719705
},
720706
{
721-
status: MOCK_DATA_GENERATOR_REQUEST_ERROR,
707+
status: 'error',
722708
error: 'Previous error',
723709
requestId: 'error_request_id',
724710
},
@@ -739,15 +725,15 @@ describe('Collection Tab Content store', function () {
739725
const state: CollectionState = {
740726
...baseState,
741727
fakerSchemaGeneration: {
742-
status: MOCK_DATA_GENERATOR_REQUEST_GENERATING,
728+
status: 'generating',
743729
requestId: 'generating_request_id',
744730
},
745731
};
746732
const action = fakerMappingGenerationFailedAction;
747733
const newState = collectionTabReducer(state, action);
748734

749735
expect(newState.fakerSchemaGeneration).to.deep.equal({
750-
status: MOCK_DATA_GENERATOR_REQUEST_ERROR,
736+
status: 'error',
751737
error: action.error,
752738
requestId: action.requestId,
753739
});

‎packages/compass-collection/src/stores/collection-tab.ts‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ import reducer, {
99
analyzeCollectionSchema,
1010
} from '../modules/collection-tab';
1111
import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
12-
import { MOCK_DATA_GENERATOR_REQUEST_IDLE } from '../components/mock-data-generator-modal/types';
1312

1413
import type { Collection } from '@mongodb-js/compass-app-stores/provider';
1514
import type { ActivateHelpers } from '@mongodb-js/compass-app-registry';
@@ -94,7 +93,7 @@ export function activatePlugin(
9493
currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
9594
},
9695
fakerSchemaGeneration: {
97-
status: MOCK_DATA_GENERATOR_REQUEST_IDLE,
96+
status: 'idle',
9897
},
9998
},
10099
applyMiddleware(

0 commit comments

Comments
 (0)