mongodb-js · rubydong · Apr 29, 2025 · Apr 30, 2025 · May 1, 2025 · May 1, 2025
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -103,5 +103,9 @@
     "chai-enzyme": {
       "cheerio": "1.0.0-rc.10"
     }
+  },
+  "dependencies": {
+    "@leafygreen-ui/skeleton-loader": "^2.0.11",
+    "mongodb-mql-engines": "^0.0.4"
   }
 }
diff --git a/packages/compass-components/src/index.ts b/packages/compass-components/src/index.ts
@@ -211,3 +211,4 @@ export {
   type ItemRenderer as VirtualListItemRenderer,
 } from './components/virtual-list';
 export { SelectTable } from './components/select-table';
+export { ParagraphSkeleton } from '@leafygreen-ui/skeleton-loader';
diff --git a/packages/compass-indexes/src/components/create-index-form/mdb-code-viewer.tsx b/packages/compass-indexes/src/components/create-index-form/mdb-code-viewer.tsx
@@ -24,15 +24,15 @@ const generateCode = ({
 }: {
   dbName: string;
   collectionName: string;
-  indexNameTypeMap: { [key: string]: string };
+  indexNameTypeMap: Record<string, string | number>;
 }) => {
   let codeStr = `db.getSiblingDB("${dbName}").getCollection("${escapeText(
     collectionName
   )}").createIndex({\n`;
 
   Object.entries(indexNameTypeMap).forEach(([name, type], index) => {
     // Replacing everything inside the parenthesis i.e. (asc)
-    let parsedType = escapeText(type.replace(/\(.*?\)/g, '')).trim();
+    let parsedType = escapeText(`${type}`.replace(/\(.*?\)/g, '')).trim();
     if (!NUMERIC_INDEX_TYPES.includes(Number(parsedType))) {
       parsedType = `"${parsedType}"`;
     }
@@ -59,7 +59,7 @@ const MDBCodeViewer = ({
 }: {
   dbName: string;
   collectionName: string;
-  indexNameTypeMap: { [key: string]: string };
+  indexNameTypeMap: Record<string, string | number>;
   dataTestId?: string;
 }) => {
   const GeneratedCode = generateCode({

diff --git a/packages/compass-indexes/src/components/create-index-form/query-flow-section.tsx b/packages/compass-indexes/src/components/create-index-form/query-flow-section.tsx
@@ -4,17 +4,24 @@ import {
   Body,
   cx,
   useFocusRing,
+  ParagraphSkeleton,
 } from '@mongodb-js/compass-components';
-import React, { useMemo } from 'react';
+import React, { useMemo, useCallback } from 'react';
 import { css, spacing } from '@mongodb-js/compass-components';
 import {
   CodemirrorMultilineEditor,
   createQueryAutocompleter,
 } from '@mongodb-js/compass-editor';
 import MDBCodeViewer from './mdb-code-viewer';
+import type { RootState } from '../../modules';
+import { fetchIndexSuggestions } from '../../modules/create-index';
+import type {
+  IndexSuggestionState,
+  SuggestedIndexFetchedProps,
+} from '../../modules/create-index';
+import { connect } from 'react-redux';
 
 const inputQueryContainerStyles = css({
-  marginBottom: spacing[600],
   display: 'flex',
   flexDirection: 'column',
 });
@@ -60,16 +67,34 @@ const codeEditorStyles = css({
   },
 });
 
+const indexSuggestionsLoaderStyles = css({
+  marginBottom: spacing[600],
+  padding: spacing[600],
+  background: palette.gray.light3,
+  border: `1px solid ${palette.gray.light2}`,
+  borderRadius: editorContainerRadius,
+});
+
 const QueryFlowSection = ({
   schemaFields,
   serverVersion,
   dbName,
   collectionName,
+  onSuggestedIndexButtonClick,
+  indexSuggestions,
+  fetchingSuggestionsState,
 }: {
   schemaFields: { name: string; description?: string }[];
   serverVersion: string;
   dbName: string;
   collectionName: string;
+  onSuggestedIndexButtonClick: ({
+    dbName,
+    collectionName,
+    inputQuery,
+  }: SuggestedIndexFetchedProps) => Promise<void>;
+  indexSuggestions: Record<string, number> | null;
+  fetchingSuggestionsState: IndexSuggestionState;
 }) => {
   const [inputQuery, setInputQuery] = React.useState('');
   const completer = useMemo(
@@ -88,6 +113,18 @@ const QueryFlowSection = ({
     radius: editorContainerRadius,
   });
 
+  const handleSuggestedIndexButtonClick = useCallback(async () => {
+    const sanitizedInputQuery = inputQuery.trim();
+
+    await onSuggestedIndexButtonClick({
+      dbName,
+      collectionName,
+      inputQuery: sanitizedInputQuery,
+    });
+  }, [inputQuery, dbName, collectionName, onSuggestedIndexButtonClick]);
+
+  const isFetchingIndexSuggestions = fetchingSuggestionsState === 'fetching';
+
   return (
     <>
       <Body baseFontSize={16} weight="medium" className={headerStyles}>
@@ -117,7 +154,7 @@ const QueryFlowSection = ({
         <div className={editorActionContainerStyles}>
           <Button
             onClick={() => {
-              // TODO in CLOUDP-311786
+              void handleSuggestedIndexButtonClick();
             }}
             className={suggestedIndexButtonStyles}
             size="small"
@@ -126,20 +163,45 @@ const QueryFlowSection = ({
           </Button>
         </div>
       </div>
-      <Body baseFontSize={16} weight="medium" className={headerStyles}>
-        Suggested Index
-      </Body>{' '}
-      <div className={suggestedIndexContainerStyles}>
-        {/* TODO in CLOUDP-311786, replace hardcoded values with actual data */}
-        <MDBCodeViewer
-          dataTestId="query-flow-section-suggested-index"
-          dbName={dbName}
-          collectionName={collectionName}
-          indexNameTypeMap={{ 'awards.win': '1', 'imdb.rating': '1' }}
-        />
-      </div>
+
+      {(isFetchingIndexSuggestions || indexSuggestions) && (
+        <Body baseFontSize={16} weight="medium" className={headerStyles}>
+          Suggested Index
+        </Body>
+      )}
+
+      {isFetchingIndexSuggestions ? (
+        <ParagraphSkeleton className={indexSuggestionsLoaderStyles} />
+      ) : (
+        indexSuggestions && (
+          <>
+            <div className={suggestedIndexContainerStyles}>
+              <MDBCodeViewer
+                dataTestId="query-flow-section-suggested-index"
+                dbName={dbName}
+                collectionName={collectionName}
+                indexNameTypeMap={indexSuggestions}
+              />
+            </div>
+          </>
+        )
+      )}
     </>
   );
 };
 
-export default QueryFlowSection;
+const mapState = ({ createIndex }: RootState) => {
+  const { indexSuggestions, sampleDocs, fetchingSuggestionsState } =
+    createIndex;
+  return {
+    indexSuggestions,
+    sampleDocs,
+    fetchingSuggestionsState,
+  };
+};
+
+const mapDispatch = {
+  onSuggestedIndexButtonClick: fetchIndexSuggestions,
+};
+
+export default connect(mapState, mapDispatch)(QueryFlowSection);
diff --git a/packages/compass-indexes/src/modules/create-index.tsx b/packages/compass-indexes/src/modules/create-index.tsx
@@ -1,3 +1,4 @@
+import type { Document } from 'mongodb';
 import { EJSON, ObjectId } from 'bson';
 import type { CreateIndexesOptions, IndexDirection } from 'mongodb';
 import { isCollationValid } from 'mongodb-query-parser';
@@ -8,6 +9,7 @@ import { isAction } from '../utils/is-action';
 import type { IndexesThunkAction } from '.';
 import type { RootState } from '.';
 import { createRegularIndex } from './regular-indexes';
+import * as mql from 'mongodb-mql-engines';
 
 export enum ActionTypes {
   FieldAdded = 'compass-indexes/create-index/fields/field-added',
@@ -27,6 +29,9 @@ export enum ActionTypes {
   CreateIndexFormSubmitted = 'compass-indexes/create-index/create-index-form-submitted',
 
   TabUpdated = 'compass-indexes/create-index/tab-updated',
+
+  SuggestedIndexesRequested = 'compass-indexes/create-index/suggested-indexes-requested',
+  SuggestedIndexesFetched = 'compass-indexes/create-index/suggested-indexes-fetched',
 }
 
 // fields
@@ -274,7 +279,7 @@ const INITIAL_OPTIONS_STATE = Object.fromEntries(
   })
 ) as Options;
 
-// other
+export type IndexSuggestionState = 'initial' | 'fetching' | 'success' | 'error';
 
 export type State = {
   // A unique id assigned to the create index modal on open, will be used when
@@ -296,6 +301,18 @@ export type State = {
 
   // current tab that user is on (Query Flow or Index Flow)
   currentTab: Tab;
+
+  // state of the index suggestions
+  fetchingSuggestionsState: IndexSuggestionState;
+
+  // error specific to fetching index suggestions
+  fetchingSuggestionsError: Error | string | null;
+
+  // index suggestions in a format such as {fieldName: 1}
+  indexSuggestions: Record<string, number> | null;
+
+  // sample documents used for getting index suggestions
+  sampleDocs: Array<Document> | null;
 };
 
 export const INITIAL_STATE: State = {
@@ -305,6 +322,10 @@ export const INITIAL_STATE: State = {
   fields: INITIAL_FIELDS_STATE,
   options: INITIAL_OPTIONS_STATE,
   currentTab: 'IndexFlow',
+  fetchingSuggestionsState: 'initial',
+  fetchingSuggestionsError: null,
+  indexSuggestions: null,
+  sampleDocs: null,
 };
 
 function getInitialState(): State {
@@ -337,6 +358,104 @@ export type CreateIndexSpec = {
   [key: string]: string | number;
 };
 
+type SuggestedIndexesRequestedAction = {
+  type: ActionTypes.SuggestedIndexesRequested;
+};
+
+export type SuggestedIndexFetchedAction = {
+  type: ActionTypes.SuggestedIndexesFetched;
+  sampleDocs: Array<Document>;
+  indexSuggestions: { [key: string]: number } | null;
+  fetchingSuggestionsError: string | null;
+  indexSuggestionsState: IndexSuggestionState;
+};
+
+export type SuggestedIndexFetchedProps = {
+  dbName: string;
+  collectionName: string;
+  inputQuery: string;
+};
+
+export const fetchIndexSuggestions = ({
+  dbName,
+  collectionName,
+  inputQuery,
+}: {
+  dbName: string;
+  collectionName: string;
+  inputQuery: string;
+}): IndexesThunkAction<
+  Promise<void>,
+  SuggestedIndexFetchedAction | SuggestedIndexesRequestedAction
+> => {
+  return async (dispatch, getState, { dataService }) => {
+    dispatch({
+      type: ActionTypes.SuggestedIndexesRequested,
+    });
+    const namespace = `${dbName}.${collectionName}`;
+
+    // Get sample documents from state if it's already there, otherwise fetch it
+    let sampleDocuments: Array<Document> | null =
+      getState().createIndex.sampleDocs || null;
+
+    // If it's null, that means it has not been fetched before
+    if (sampleDocuments === null) {
+      try {
+        sampleDocuments =
+          (await dataService.sample(namespace, { size: 50 })) || [];
+      } catch (e) {
+        // Swallow the error because mql package still will work fine with empty sampleDocuments
+        sampleDocuments = [];
+      }
+    }
+
+    // Analyze namespace and fetch suggestions
+    try {
+      const analyzedNamespace = mql.analyzeNamespace(
+        { database: dbName, collection: collectionName },
+        sampleDocuments
+      );
+
+      const query = mql.parseQuery(
+        EJSON.parse(inputQuery, { relaxed: false }),
+        analyzedNamespace
+      );
+      const results = await mql.suggestIndex([query]);
+      const indexSuggestions = results?.index || null;
+
+      // TODO in CLOUDP-311787: add info banner and update the current error banner to take in fetchingSuggestionsError as well
+      if (!indexSuggestions) {
+        dispatch({
+          type: ActionTypes.SuggestedIndexesFetched,
+          sampleDocs: sampleDocuments,
+          indexSuggestions,
+          fetchingSuggestionsError:
+            'No suggested index found. Please choose Start with an Index at the top to continue.',
+          indexSuggestionsState: 'error',
+        });
+        return;
+      }
+
+      dispatch({
+        type: ActionTypes.SuggestedIndexesFetched,
+        sampleDocs: sampleDocuments,
+        indexSuggestions,
+        fetchingSuggestionsError: null,
+        indexSuggestionsState: 'success',
+      });
+    } catch (e) {
+      dispatch({
+        type: ActionTypes.SuggestedIndexesFetched,
+        sampleDocs: sampleDocuments,
+        indexSuggestions: null,
+        fetchingSuggestionsError:
+          'Error parsing query. Please follow query structure.',
+        indexSuggestionsState: 'error',
+      });
+    }
+  };
+};
+
 function isEmptyValue(value: unknown) {
   if (value === '') {
     return true;
@@ -629,6 +748,35 @@ const reducer: Reducer<State, Action> = (state = INITIAL_STATE, action) => {
     };
   }
 
+  if (
+    isAction<SuggestedIndexesRequestedAction>(
+      action,
+      ActionTypes.SuggestedIndexesRequested
+    )
+  ) {
+    return {
+      ...state,
+      fetchingSuggestionsState: 'fetching',
+      fetchingSuggestionsError: null,
+      indexSuggestions: null,
+    };
+  }
+
+  if (
+    isAction<SuggestedIndexFetchedAction>(
+      action,
+      ActionTypes.SuggestedIndexesFetched
+    )
+  ) {
+    return {
+      ...state,
+      fetchingSuggestionsState: action.indexSuggestionsState,
+      fetchingSuggestionsError: action.fetchingSuggestionsError,
+      indexSuggestions: action.indexSuggestions,
+      sampleDocs: action.sampleDocs,
+    };
+  }
+
   return state;
 };