diff --git a/package-lock.json b/package-lock.json index b84f3a9c0b3..4667817c56e 100644 --- a/package-lock.json +++ b/package-lock.json @@ -47440,6 +47440,7 @@ "@mongodb-js/compass-generative-ai": "^0.51.0", "@mongodb-js/compass-logging": "^1.7.12", "@mongodb-js/compass-telemetry": "^1.14.0", + "@mongodb-js/compass-utils": "^0.9.11", "@mongodb-js/compass-workspaces": "^0.52.0", "@mongodb-js/connection-info": "^0.17.2", "@mongodb-js/mongodb-constants": "^0.14.0", @@ -60784,6 +60785,7 @@ "@mongodb-js/compass-generative-ai": "^0.51.0", "@mongodb-js/compass-logging": "^1.7.12", "@mongodb-js/compass-telemetry": "^1.14.0", + "@mongodb-js/compass-utils": "^0.9.11", "@mongodb-js/compass-workspaces": "^0.52.0", "@mongodb-js/connection-info": "^0.17.2", "@mongodb-js/eslint-config-compass": "^1.4.7", diff --git a/packages/compass-collection/package.json b/packages/compass-collection/package.json index 93e13f1641b..901bef9287e 100644 --- a/packages/compass-collection/package.json +++ b/packages/compass-collection/package.json @@ -55,9 +55,11 @@ "@mongodb-js/compass-generative-ai": "^0.51.0", "@mongodb-js/compass-logging": "^1.7.12", "@mongodb-js/compass-telemetry": "^1.14.0", + "@mongodb-js/compass-utils": "^0.9.11", "@mongodb-js/compass-workspaces": "^0.52.0", "@mongodb-js/connection-info": "^0.17.2", "@mongodb-js/mongodb-constants": "^0.14.0", + "bson": "^6.10.1", "compass-preferences-model": "^2.51.0", "hadron-document": "^8.9.6", "mongodb": "^6.19.0", @@ -67,8 +69,7 @@ "react": "^17.0.2", "react-redux": "^8.1.3", "redux": "^4.2.1", - "redux-thunk": "^2.4.2", - "bson": "^6.10.1" + "redux-thunk": "^2.4.2" }, "devDependencies": { "@mongodb-js/eslint-config-compass": "^1.4.7", diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor.tsx new file mode 100644 index 00000000000..f95105066d8 --- /dev/null +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor.tsx @@ -0,0 +1,12 @@ +import React from 'react'; + +// TODO: More to come from CLOUDP-333853, CLOUDP-333854 +const FakerSchemaEditor = () => { + return ( +
+ Schema Editor Content Placeholder +
+ ); +}; + +export default FakerSchemaEditor; 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 9caf475f8a0..9d0dd164906 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 @@ -1,75 +1,231 @@ import { expect } from 'chai'; import React from 'react'; -import { render, screen } from '@mongodb-js/testing-library-compass'; -import Sinon from 'sinon'; -import { UnconnectedMockDataGeneratorModal as MockDataGeneratorModal } from './mock-data-generator-modal'; +import { + screen, + render, + waitFor, + userEvent, +} from '@mongodb-js/testing-library-compass'; +import { Provider } from 'react-redux'; +import { createStore, applyMiddleware } from 'redux'; +import thunk from 'redux-thunk'; +import MockDataGeneratorModal from './mock-data-generator-modal'; import { MockDataGeneratorStep } from './types'; import { StepButtonLabelMap } from './constants'; +import type { CollectionState } from '../../modules/collection-tab'; +import { default as collectionTabReducer } from '../../modules/collection-tab'; describe('MockDataGeneratorModal', () => { - const sandbox = Sinon.createSandbox(); - let onClose: Sinon.SinonSpy; - - beforeEach(() => { - onClose = sandbox.spy(); - }); - - afterEach(() => { - sandbox.restore(); - }); - - const onNextStep = Sinon.stub(); - const onPreviousStep = Sinon.stub(); - function renderModal({ isOpen = true, currentStep = MockDataGeneratorStep.SCHEMA_CONFIRMATION, + mockServices = createMockServices(), } = {}) { + 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 }, + }, + fakerSchemaGeneration: { + status: 'idle', + }, + mockDataGenerator: { + isModalOpen: isOpen, + currentStep: currentStep, + }, + }; + + const store = createStore( + collectionTabReducer, + initialState, + applyMiddleware(thunk.withExtraArgument(mockServices)) + ); + return render( - + + + ); } - it('renders the modal when isOpen is true', () => { - renderModal(); + function createMockServices() { + return { + dataService: {}, + atlasAiService: { + getMockDataSchema: () => { + return Promise.resolve({ + contents: { + fields: [], + }, + }); + }, + }, + workspaces: {}, + localAppRegistry: {}, + experimentationServices: {}, + connectionInfoRef: { current: {} }, + logger: { + log: { + warn: () => {}, + error: () => {}, + }, + debug: () => {}, + mongoLogId: () => 'mock-id', + }, + preferences: { getPreferences: () => ({}) }, + fakerSchemaGenerationAbortControllerRef: { current: undefined }, + }; + } - expect(screen.getByTestId('generate-mock-data-modal')).to.exist; - }); + describe('generally', () => { + it('renders the modal when isOpen is true', () => { + renderModal(); - it('does not render the modal when isOpen is false', () => { - renderModal({ isOpen: false }); + expect(screen.getByTestId('generate-mock-data-modal')).to.exist; + }); - expect(screen.queryByTestId('generate-mock-data-modal')).to.not.exist; - }); + it('does not render the modal when isOpen is false', () => { + renderModal({ isOpen: false }); - it('calls onClose when the modal is closed', () => { - renderModal(); + expect(screen.queryByTestId('generate-mock-data-modal')).to.not.exist; + }); - screen.getByLabelText('Close modal').click(); + it('closes the modal when the close button is clicked', async () => { + renderModal(); - expect(onClose.calledOnce).to.be.true; - }); + expect(screen.getByTestId('generate-mock-data-modal')).to.exist; + userEvent.click(screen.getByLabelText('Close modal')); + await waitFor( + () => + expect(screen.queryByTestId('generate-mock-data-modal')).to.not.exist + ); + }); - it('calls onClose when the cancel button is clicked', () => { - renderModal(); + it('closes the modal when the cancel button is clicked', async () => { + renderModal(); - screen.getByText('Cancel').click(); + expect(screen.getByTestId('generate-mock-data-modal')).to.exist; + userEvent.click(screen.getByText('Cancel')); + await waitFor( + () => + expect(screen.queryByTestId('generate-mock-data-modal')).to.not.exist + ); + }); + + function createMockServicesWithSlowAiRequest() { + let abortSignalReceived = false; + let rejectPromise: (reason?: any) => void; + const rejectedPromise = new Promise((_resolve, reject) => { + rejectPromise = reject; + }); + + const baseMockServices = createMockServices(); + + const mockAiService = { + ...baseMockServices.atlasAiService, + getMockDataSchema: (request: any) => { + if (request?.signal) { + request.signal.addEventListener('abort', () => { + abortSignalReceived = true; + rejectPromise(new Error('Request aborted')); + }); + } + return rejectedPromise; + }, + getAbortSignalReceived: () => abortSignalReceived, + }; + + return { + ...baseMockServices, + atlasAiService: mockAiService, + }; + } + + it('cancels in-flight faker mapping requests when the cancel button is clicked', async () => { + const mockServices = createMockServicesWithSlowAiRequest(); + renderModal({ mockServices: mockServices as any }); + + expect(screen.getByTestId('raw-schema-confirmation')).to.exist; + userEvent.click(screen.getByText('Confirm')); + + await waitFor(() => { + expect(screen.getByTestId('faker-schema-editor')).to.exist; + }); - expect(onClose.calledOnce).to.be.true; + userEvent.click(screen.getByText('Cancel')); + + expect(mockServices.atlasAiService.getAbortSignalReceived()).to.be.true; + }); + + it('cancels in-flight faker mapping requests when the back button is clicked after schema confirmation', async () => { + const mockServices = createMockServicesWithSlowAiRequest(); + renderModal({ mockServices: mockServices as any }); + + expect(screen.getByTestId('raw-schema-confirmation')).to.exist; + userEvent.click(screen.getByText('Confirm')); + + await waitFor(() => { + expect(screen.getByTestId('faker-schema-editor')).to.exist; + }); + + userEvent.click(screen.getByText('Back')); + + expect(mockServices.atlasAiService.getAbortSignalReceived()).to.be.true; + }); }); - it('disables the Back button on the first step', () => { - renderModal(); + describe('on the schema confirmation step', () => { + it('disables the Back button', () => { + renderModal(); + + expect( + screen + .getByRole('button', { name: 'Back' }) + .getAttribute('aria-disabled') + ).to.equal('true'); + }); + + it('renders the faker schema editor when the confirm button is clicked', async () => { + renderModal(); + + expect(screen.getByTestId('raw-schema-confirmation')).to.exist; + expect(screen.queryByTestId('faker-schema-editor')).to.not.exist; + userEvent.click(screen.getByText('Confirm')); + await waitFor(() => { + expect(screen.queryByTestId('raw-schema-confirmation')).to.not.exist; + expect(screen.getByTestId('faker-schema-editor')).to.exist; + }); + }); + + it('stays on the current step when an error is encountered during faker schema generation', async () => { + const mockServices = createMockServices(); + mockServices.atlasAiService.getMockDataSchema = () => + Promise.reject('faker schema generation failed'); + renderModal({ mockServices }); + + expect(screen.getByTestId('raw-schema-confirmation')).to.exist; + expect(screen.queryByTestId('faker-schema-editor')).to.not.exist; + userEvent.click(screen.getByText('Confirm')); + await waitFor(() => { + expect(screen.getByTestId('raw-schema-confirmation')).to.exist; + expect(screen.queryByTestId('faker-schema-editor')).to.not.exist; + }); + + // todo: assert a user-friendly error is displayed (CLOUDP-333852) + }); - expect( - screen.getByRole('button', { name: 'Back' }).getAttribute('aria-disabled') - ).to.equal('true'); + // todo: assert that closing then re-opening the modal after an LLM err removes the err message }); describe('when rendering the modal in a specific step', () => { @@ -77,6 +233,7 @@ describe('MockDataGeneratorModal', () => { StepButtonLabelMap ) as unknown as MockDataGeneratorStep[]; + // note: these tests can be removed after every modal step is implemented steps.forEach((currentStep) => { it(`renders the button with the correct label when the user is in step "${currentStep}"`, () => { renderModal({ currentStep }); diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx index b47f81d8c84..d09861a6b42 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx @@ -3,25 +3,26 @@ import { connect } from 'react-redux'; import { css, + Button, + ButtonVariant, ModalBody, ModalHeader, - spacing, -} from '@mongodb-js/compass-components'; - -import { - Button, Modal, ModalFooter, - ButtonVariant, + spacing, } from '@mongodb-js/compass-components'; + import { MockDataGeneratorStep } from './types'; import { StepButtonLabelMap } from './constants'; import type { CollectionState } from '../../modules/collection-tab'; import { mockDataGeneratorModalClosed, mockDataGeneratorNextButtonClicked, + generateFakerMappings, mockDataGeneratorPreviousButtonClicked, } from '../../modules/collection-tab'; +import { default as SchemaConfirmationScreen } from './raw-schema-confirmation'; +import FakerSchemaEditor from './faker-schema-editor'; const footerStyles = css` flex-direction: row; @@ -39,6 +40,7 @@ interface Props { onClose: () => void; currentStep: MockDataGeneratorStep; onNextStep: () => void; + onConfirmSchema: () => Promise; onPreviousStep: () => void; } @@ -47,10 +49,28 @@ const MockDataGeneratorModal = ({ onClose, currentStep, onNextStep, + onConfirmSchema, onPreviousStep, }: Props) => { - const handleNextClick = - currentStep === MockDataGeneratorStep.GENERATE_DATA ? onClose : onNextStep; + const handleNextClick = () => { + if (currentStep === MockDataGeneratorStep.GENERATE_DATA) { + onClose(); + } else if (currentStep === MockDataGeneratorStep.SCHEMA_CONFIRMATION) { + void onConfirmSchema(); + } else { + onNextStep(); + } + }; + + let stepContent: React.ReactNode; + + if (currentStep === MockDataGeneratorStep.SCHEMA_CONFIRMATION) { + stepContent = ; + } + + if (currentStep === MockDataGeneratorStep.SCHEMA_EDITOR) { + stepContent = ; + } return ( - {/* TODO: Render actual step content here based on currentStep. (CLOUDP-333851) */} + {stepContent}
@@ -97,8 +117,8 @@ const mapStateToProps = (state: CollectionState) => ({ const ConnectedMockDataGeneratorModal = connect(mapStateToProps, { onClose: mockDataGeneratorModalClosed, onNextStep: mockDataGeneratorNextButtonClicked, + onConfirmSchema: generateFakerMappings, onPreviousStep: mockDataGeneratorPreviousButtonClicked, })(MockDataGeneratorModal); export default ConnectedMockDataGeneratorModal; -export { MockDataGeneratorModal as UnconnectedMockDataGeneratorModal }; diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/raw-schema-confirmation.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/raw-schema-confirmation.tsx new file mode 100644 index 00000000000..6848d5b8c97 --- /dev/null +++ b/packages/compass-collection/src/components/mock-data-generator-modal/raw-schema-confirmation.tsx @@ -0,0 +1,52 @@ +import React from 'react'; +import { connect } from 'react-redux'; + +import { Code, Body, Subtitle } from '@mongodb-js/compass-components'; + +import type { CollectionState } from '../../modules/collection-tab'; +import type { FieldInfo } from '../../schema-analysis-types'; + +interface RawSchemaConfirmationProps { + schemaContent: Record | null; + namespace: string; +} + +// Note: Currently a placeholder. The final contents will be addressed by CLOUDP-333852 +const RawSchemaConfirmation = (props: RawSchemaConfirmationProps) => { + // this will change + const codeContent = props.schemaContent + ? JSON.stringify(props.schemaContent, null, 4) + : 'No schema data available'; + + return ( +
+ {props.namespace} + Document Schema Identified + + We have identified the following schema from your documents. This schema + will be sent to an LLM for processing. + + + {codeContent} + +
+ ); +}; + +const mapStateToProps = (state: CollectionState) => { + const schemaContent = + state.schemaAnalysis.status === 'complete' + ? state.schemaAnalysis.processedSchema + : null; + return { + schemaContent, + namespace: state.namespace, + }; +}; + +const ConnectedRawSchemaConfirmation = connect( + mapStateToProps, + {} +)(RawSchemaConfirmation); + +export default ConnectedRawSchemaConfirmation; 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 c5a45271504..5812f3693a4 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,3 +1,5 @@ +import type { MockDataSchemaResponse } from '@mongodb-js/compass-generative-ai'; + export enum MockDataGeneratorStep { SCHEMA_CONFIRMATION = 'SCHEMA_CONFIRMATION', SCHEMA_EDITOR = 'SCHEMA_EDITOR', @@ -5,3 +7,30 @@ export enum MockDataGeneratorStep { PREVIEW_DATA = 'PREVIEW_DATA', GENERATE_DATA = 'GENERATE_DATA', } + +type MockDataGeneratorIdleState = { + status: 'idle'; +}; + +type MockDataGeneratorInProgressState = { + status: 'in-progress'; + requestId: string; +}; + +type MockDataGeneratorCompletedState = { + status: 'completed'; + fakerSchema: MockDataSchemaResponse; + requestId: string; +}; + +type MockDataGeneratorErrorState = { + status: 'error'; + error: unknown; + requestId: string; +}; + +export type MockDataGeneratorState = + | MockDataGeneratorIdleState + | MockDataGeneratorInProgressState + | MockDataGeneratorCompletedState + | MockDataGeneratorErrorState; diff --git a/packages/compass-collection/src/index.ts b/packages/compass-collection/src/index.ts index 5ba6370bc36..05a1680409b 100644 --- a/packages/compass-collection/src/index.ts +++ b/packages/compass-collection/src/index.ts @@ -14,11 +14,11 @@ import { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provide import { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider'; import { createLoggerLocator } from '@mongodb-js/compass-logging/provider'; import { preferencesLocator } from 'compass-preferences-model/provider'; +import { atlasAiServiceLocator } from '@mongodb-js/compass-generative-ai/provider'; import { CollectionWorkspaceTitle, CollectionPluginTitleComponent, } from './plugin-tab-title'; -import { atlasAiServiceLocator } from '@mongodb-js/compass-generative-ai/provider'; export const WorkspaceTab: WorkspacePlugin = { name: CollectionWorkspaceTitle, @@ -33,12 +33,12 @@ export const WorkspaceTab: WorkspacePlugin = { { dataService: dataServiceLocator as DataServiceLocator, collection: collectionModelLocator, + atlasAiService: atlasAiServiceLocator, workspaces: workspacesServiceLocator, experimentationServices: experimentationServiceLocator, connectionInfoRef: connectionInfoRefLocator, logger: createLoggerLocator('COMPASS-COLLECTION'), preferences: preferencesLocator, - atlasAiService: atlasAiServiceLocator, } ), content: CollectionTab, diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts index 746632abe9e..7a001624dd4 100644 --- a/packages/compass-collection/src/modules/collection-tab.ts +++ b/packages/compass-collection/src/modules/collection-tab.ts @@ -1,17 +1,25 @@ import type { Reducer, AnyAction, Action } from 'redux'; import { analyzeDocuments } from 'mongodb-schema'; +import { UUID } from 'bson'; +import { isCancelError } from '@mongodb-js/compass-utils'; import type { CollectionMetadata } from 'mongodb-collection-model'; import type { ThunkAction } from 'redux-thunk'; import type AppRegistry from '@mongodb-js/compass-app-registry'; import type { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provider'; +import type { + ConnectionInfoRef, + DataService, +} from '@mongodb-js/compass-connections/provider'; import type { CollectionSubtab } from '@mongodb-js/compass-workspaces'; -import type { DataService } from '@mongodb-js/compass-connections/provider'; +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 { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider'; - +import type { + MockDataSchemaRequest, + MockDataSchemaResponse, +} from '@mongodb-js/compass-generative-ai'; import { isInternalFieldPath } from 'hadron-document'; import toNS from 'mongodb-ns'; import { @@ -27,6 +35,7 @@ import { calculateSchemaDepth } from '../calculate-schema-depth'; import { processSchema } from '../transform-schema-to-field-info'; import type { Document, MongoError } from 'mongodb'; import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types'; +import type { MockDataGeneratorState } from '../components/mock-data-generator-modal/types'; const DEFAULT_SAMPLE_SIZE = 100; @@ -63,11 +72,13 @@ type CollectionThunkAction = ThunkAction< { localAppRegistry: AppRegistry; dataService: DataService; + atlasAiService: AtlasAiService; workspaces: ReturnType; experimentationServices: ReturnType; logger: Logger; preferences: PreferencesAccess; - atlasAiService: AtlasAiService; + connectionInfoRef: ConnectionInfoRef; + fakerSchemaGenerationAbortControllerRef: { current?: AbortController }; }, A >; @@ -82,9 +93,10 @@ export type CollectionState = { isModalOpen: boolean; currentStep: MockDataGeneratorStep; }; + fakerSchemaGeneration: MockDataGeneratorState; }; -enum CollectionActions { +export enum CollectionActions { CollectionMetadataFetched = 'compass-collection/CollectionMetadataFetched', SchemaAnalysisStarted = 'compass-collection/SchemaAnalysisStarted', SchemaAnalysisFinished = 'compass-collection/SchemaAnalysisFinished', @@ -94,6 +106,9 @@ enum CollectionActions { MockDataGeneratorModalClosed = 'compass-collection/MockDataGeneratorModalClosed', MockDataGeneratorNextButtonClicked = 'compass-collection/MockDataGeneratorNextButtonClicked', MockDataGeneratorPreviousButtonClicked = 'compass-collection/MockDataGeneratorPreviousButtonClicked', + FakerMappingGenerationStarted = 'compass-collection/FakerMappingGenerationStarted', + FakerMappingGenerationCompleted = 'compass-collection/FakerMappingGenerationCompleted', + FakerMappingGenerationFailed = 'compass-collection/FakerMappingGenerationFailed', } interface CollectionMetadataFetchedAction { @@ -140,6 +155,23 @@ interface MockDataGeneratorPreviousButtonClickedAction { type: CollectionActions.MockDataGeneratorPreviousButtonClicked; } +export interface FakerMappingGenerationStartedAction { + type: CollectionActions.FakerMappingGenerationStarted; + requestId: string; +} + +export interface FakerMappingGenerationCompletedAction { + type: CollectionActions.FakerMappingGenerationCompleted; + fakerSchema: MockDataSchemaResponse; + requestId: string; +} + +export interface FakerMappingGenerationFailedAction { + type: CollectionActions.FakerMappingGenerationFailed; + error: string; + requestId: string; +} + const reducer: Reducer = ( state = { // TODO(COMPASS-7782): use hook to get the workspace tab id instead @@ -153,6 +185,9 @@ const reducer: Reducer = ( isModalOpen: false, currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION, }, + fakerSchemaGeneration: { + status: 'idle', + }, }, action ) => { @@ -256,6 +291,9 @@ const reducer: Reducer = ( ...state.mockDataGenerator, isModalOpen: false, }, + fakerSchemaGeneration: { + status: 'idle', + }, }; } @@ -269,9 +307,6 @@ const reducer: Reducer = ( let nextStep: MockDataGeneratorStep; switch (currentStep) { - case MockDataGeneratorStep.SCHEMA_CONFIRMATION: - nextStep = MockDataGeneratorStep.SCHEMA_EDITOR; - break; case MockDataGeneratorStep.SCHEMA_EDITOR: nextStep = MockDataGeneratorStep.DOCUMENT_COUNT; break; @@ -309,8 +344,16 @@ const reducer: Reducer = ( previousStep = MockDataGeneratorStep.SCHEMA_CONFIRMATION; break; case MockDataGeneratorStep.SCHEMA_EDITOR: - previousStep = MockDataGeneratorStep.SCHEMA_CONFIRMATION; - break; + return { + ...state, + fakerSchemaGeneration: { + status: 'idle', + }, + mockDataGenerator: { + ...state.mockDataGenerator, + currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION, + }, + }; case MockDataGeneratorStep.DOCUMENT_COUNT: previousStep = MockDataGeneratorStep.SCHEMA_EDITOR; break; @@ -333,6 +376,78 @@ const reducer: Reducer = ( }; } + if ( + isAction( + action, + CollectionActions.FakerMappingGenerationStarted + ) + ) { + if ( + state.mockDataGenerator.currentStep !== + MockDataGeneratorStep.SCHEMA_CONFIRMATION || + state.fakerSchemaGeneration.status === 'in-progress' || + state.fakerSchemaGeneration.status === 'completed' + ) { + return state; + } + + return { + ...state, + mockDataGenerator: { + ...state.mockDataGenerator, + currentStep: MockDataGeneratorStep.SCHEMA_EDITOR, + }, + fakerSchemaGeneration: { + status: 'in-progress', + requestId: action.requestId, + }, + }; + } + + if ( + isAction( + action, + CollectionActions.FakerMappingGenerationCompleted + ) + ) { + if (state.fakerSchemaGeneration.status !== 'in-progress') { + return state; + } + + return { + ...state, + fakerSchemaGeneration: { + status: 'completed', + fakerSchema: action.fakerSchema, + requestId: action.requestId, + }, + }; + } + + if ( + isAction( + action, + CollectionActions.FakerMappingGenerationFailed + ) + ) { + if (state.fakerSchemaGeneration.status !== 'in-progress') { + return state; + } + + return { + ...state, + fakerSchemaGeneration: { + status: 'error', + error: action.error, + requestId: action.requestId, + }, + mockDataGenerator: { + ...state.mockDataGenerator, + currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION, + }, + }; + } + return state; }; @@ -347,20 +462,32 @@ export const mockDataGeneratorModalOpened = return { type: CollectionActions.MockDataGeneratorModalOpened }; }; -export const mockDataGeneratorModalClosed = - (): MockDataGeneratorModalClosedAction => { - return { type: CollectionActions.MockDataGeneratorModalClosed }; +export const mockDataGeneratorModalClosed = (): CollectionThunkAction< + void, + MockDataGeneratorModalClosedAction +> => { + return (dispatch, _getState, { fakerSchemaGenerationAbortControllerRef }) => { + fakerSchemaGenerationAbortControllerRef.current?.abort(); + dispatch({ type: CollectionActions.MockDataGeneratorModalClosed }); }; +}; export const mockDataGeneratorNextButtonClicked = (): MockDataGeneratorNextButtonClickedAction => { return { type: CollectionActions.MockDataGeneratorNextButtonClicked }; }; -export const mockDataGeneratorPreviousButtonClicked = - (): MockDataGeneratorPreviousButtonClickedAction => { - return { type: CollectionActions.MockDataGeneratorPreviousButtonClicked }; +export const mockDataGeneratorPreviousButtonClicked = (): CollectionThunkAction< + void, + MockDataGeneratorPreviousButtonClickedAction +> => { + return (dispatch, _getState, { fakerSchemaGenerationAbortControllerRef }) => { + fakerSchemaGenerationAbortControllerRef.current?.abort(); + dispatch({ + type: CollectionActions.MockDataGeneratorPreviousButtonClicked, + }); }; +}; export const selectTab = ( tabName: CollectionSubtab @@ -480,6 +607,104 @@ export const analyzeCollectionSchema = (): CollectionThunkAction< }; }; +export const generateFakerMappings = (): CollectionThunkAction< + Promise +> => { + return async ( + dispatch, + getState, + { + logger, + atlasAiService, + preferences, + connectionInfoRef, + fakerSchemaGenerationAbortControllerRef, + } + ) => { + const { schemaAnalysis, fakerSchemaGeneration, namespace } = getState(); + if (schemaAnalysis.status !== SCHEMA_ANALYSIS_STATE_COMPLETE) { + logger.log.warn( + mongoLogId(1_001_000_305), + 'Collection', + 'Cannot call `generateFakeMappings` unless schema analysis is complete' + ); + return; + } + + if (fakerSchemaGeneration.status === 'in-progress') { + logger.debug( + 'Faker mapping generation is already in progress, skipping new generation.' + ); + return; + } + + const requestId = new UUID().toString(); + + const includeSampleValues = + preferences.getPreferences().enableGenAISampleDocumentPassing; + + try { + logger.debug('Generating faker mappings'); + + const { database, collection } = toNS(namespace); + + dispatch({ + type: CollectionActions.FakerMappingGenerationStarted, + requestId: requestId, + }); + + fakerSchemaGenerationAbortControllerRef.current?.abort(); + fakerSchemaGenerationAbortControllerRef.current = new AbortController(); + const abortSignal = + fakerSchemaGenerationAbortControllerRef.current.signal; + + const mockDataSchemaRequest: MockDataSchemaRequest = { + databaseName: database, + collectionName: collection, + schema: schemaAnalysis.processedSchema, + validationRules: schemaAnalysis.schemaMetadata.validationRules, + includeSampleValues, + requestId, + signal: abortSignal, + }; + + const response = await atlasAiService.getMockDataSchema( + mockDataSchemaRequest, + connectionInfoRef.current + ); + + fakerSchemaGenerationAbortControllerRef.current = undefined; + dispatch({ + type: CollectionActions.FakerMappingGenerationCompleted, + fakerSchema: response, + requestId: requestId, + }); + } catch (e) { + if (isCancelError(e)) { + // abort errors should not produce error logs + return; + } + + const errorMessage = e instanceof Error ? e.stack : String(e); + + logger.log.error( + mongoLogId(1_001_000_312), + 'Collection', + 'Failed to generate faker.js mappings', + { + message: errorMessage, + namespace, + } + ); + dispatch({ + type: CollectionActions.FakerMappingGenerationFailed, + error: 'faker mapping request failed', + requestId, + }); + } + }; +}; + export type CollectionTabPluginMetadata = CollectionMetadata & { /** * Initial query for the query bar diff --git a/packages/compass-collection/src/stores/collection-tab.spec.ts b/packages/compass-collection/src/stores/collection-tab.spec.ts index 4813d13bc66..6de48f4426f 100644 --- a/packages/compass-collection/src/stores/collection-tab.spec.ts +++ b/packages/compass-collection/src/stores/collection-tab.spec.ts @@ -62,6 +62,7 @@ describe('Collection Tab Content store', function () { const analyzeCollectionSchemaStub = sandbox .stub(collectionTabModule, 'analyzeCollectionSchema') .returns(async () => {}); + const dataService = {} as any; const atlasAiService = {} as any; let store: ReturnType['store']; @@ -100,6 +101,7 @@ describe('Collection Tab Content store', function () { }, { dataService, + atlasAiService, localAppRegistry, collection: mockCollection as any, workspaces: workspaces as any, @@ -107,7 +109,6 @@ describe('Collection Tab Content store', function () { connectionInfoRef: connectionInfoRef as any, logger, preferences, - atlasAiService, }, { on() {}, cleanup() {} } as any )); diff --git a/packages/compass-collection/src/stores/collection-tab.ts b/packages/compass-collection/src/stores/collection-tab.ts index db718545ec9..48b4240ff63 100644 --- a/packages/compass-collection/src/stores/collection-tab.ts +++ b/packages/compass-collection/src/stores/collection-tab.ts @@ -16,13 +16,13 @@ import type { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/pr import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider'; import type { connectionInfoRefLocator } from '@mongodb-js/compass-connections/provider'; import type { Logger } from '@mongodb-js/compass-logging/provider'; +import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider'; import { isAIFeatureEnabled, type PreferencesAccess, } from 'compass-preferences-model/provider'; import { ExperimentTestName } from '@mongodb-js/compass-telemetry/provider'; import { SCHEMA_ANALYSIS_STATE_INITIAL } from '../schema-analysis-types'; -import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider'; export type CollectionTabOptions = { /** @@ -44,12 +44,12 @@ export type CollectionTabServices = { dataService: DataService; collection: Collection; localAppRegistry: AppRegistry; + atlasAiService: AtlasAiService; workspaces: ReturnType; experimentationServices: ReturnType; connectionInfoRef: ReturnType; logger: Logger; preferences: PreferencesAccess; - atlasAiService: AtlasAiService; }; export function activatePlugin( @@ -64,12 +64,12 @@ export function activatePlugin( dataService, collection: collectionModel, localAppRegistry, + atlasAiService, workspaces, experimentationServices, connectionInfoRef, logger, preferences, - atlasAiService, } = services; if (!collectionModel) { @@ -78,6 +78,9 @@ export function activatePlugin( ); } + const fakerSchemaGenerationAbortControllerRef = { + current: undefined, + }; const store = createStore( reducer, { @@ -92,16 +95,21 @@ export function activatePlugin( isModalOpen: false, currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION, }, + fakerSchemaGeneration: { + status: 'idle', + }, }, applyMiddleware( thunk.withExtraArgument({ dataService, + atlasAiService, workspaces, localAppRegistry, experimentationServices, + connectionInfoRef, logger, preferences, - atlasAiService, + fakerSchemaGenerationAbortControllerRef, }) ) ); 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 0c9f6aa2684..e1051c1c333 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts @@ -402,6 +402,8 @@ describe('AtlasAiService', function () { }, }, includeSampleValues: false, + requestId: 'test-request-id', + signal: new AbortController().signal, }; if (apiURLPreset === 'admin-api') { @@ -460,7 +462,7 @@ describe('AtlasAiService', function () { expect(fetchStub).to.have.been.calledOnce; const { args } = fetchStub.firstCall; expect(args[0]).to.eq( - '/cloud/ai/v1/groups/testProject/mock-data-schema' + '/cloud/ai/v1/groups/testProject/mock-data-schema?request_id=test-request-id' ); expect(result).to.deep.equal(mockResponse); }); diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index 3dd8d727b47..66595ff40dc 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.ts @@ -227,6 +227,8 @@ export interface MockDataSchemaRequest { schema: Record; validationRules?: Record | null; includeSampleValues?: boolean; + requestId: string; + signal: AbortSignal; } export const MockDataSchemaResponseShape = z.object({ @@ -461,7 +463,10 @@ export class AtlasAiService { const { collectionName, databaseName } = input; let schema = input.schema; - const url = this.getUrlForEndpoint('mock-data-schema', connectionInfo); + const url = `${this.getUrlForEndpoint( + 'mock-data-schema', + connectionInfo + )}?request_id=${encodeURIComponent(input.requestId)}`; if (!input.includeSampleValues) { const newSchema: Record< @@ -485,6 +490,7 @@ export class AtlasAiService { 'Content-Type': 'application/json', Accept: 'application/json', }, + signal: input.signal, }); try {