chore(compass-editor): add agg autocompleter for codemirror COMPASS-6575 (#4165)

* chore(compass-editor): add agg autocompleter for codemirror

* chore(compass-editor): fix types

* chore(compass-editor): fix resolveTokenAtCursor

Author: gribnoysup

configs/mocha-config-compass/register/jsdom-global-register.js

@@ -55,3 +55,7 @@ if (!globalThis.IntersectionObserver) {
 // jsdom doesn't override classes that already exist in global scope
 // https://github.com/jsdom/jsdom/issues/3331
 globalThis.EventTarget = window.EventTarget;
+
+Range.prototype.getClientRects = function () {
+  return [];
+};

packages/compass-editor/src/codemirror/ace-compat-autocompleter.ts

@@ -90,8 +90,8 @@ export function createAceCompatAutocompleter(completers: {
 
     const opts = { context, prefix, token };
 
-    if (token.type.name === 'String') {
-      return completers.String?.(opts) ?? null;
+    if (token.type.name === 'String' && completers.String) {
+      return completers.String(opts) ?? null;
     }
 
     return completers.IdentifierLike?.(opts) ?? null;

packages/compass-editor/src/codemirror/aggregation-autocompleter.test.ts

@@ -0,0 +1,75 @@
+import { expect } from 'chai';
+import { STAGE_OPERATOR_NAMES } from '@mongodb-js/mongodb-constants';
+import { createAggregationAutocompleter } from './aggregation-autocompleter';
+import { setupCodemirrorCompleter } from '../../test/completer';
+
+describe('query autocompleter', function () {
+  const { getCompletions, applySnippet, cleanup } = setupCodemirrorCompleter(
+    createAggregationAutocompleter
+  );
+
+  after(cleanup);
+
+  context('when autocompleting outside of stage', function () {
+    context('with empty pipeline', function () {
+      it('should return stages', function () {
+        const completions = getCompletions('[{ $');
+
+        expect(
+          completions.map((completion) => completion.label).sort()
+        ).to.deep.eq([...STAGE_OPERATOR_NAMES].sort());
+      });
+    });
+
+    context('with other stages in the pipeline', function () {
+      it('should return stages', function () {
+        const completions = getCompletions('[{$match:{foo: 1}},{$');
+
+        expect(
+          completions.map((completion) => completion.label).sort()
+        ).to.deep.eq([...STAGE_OPERATOR_NAMES].sort());
+      });
+    });
+
+    context('inside block', function () {
+      it('should not suggest blocks in snippets', function () {
+        const completions = getCompletions(`[{ /** comment */ $`);
+
+        completions.forEach((completion) => {
+          const snippet = applySnippet(completion);
+          expect(snippet).to.match(
+            /^[^{]/,
+            'expected snippet NOT to start with an opening bracket'
+          );
+        });
+      });
+    });
+
+    context('outside block', function () {
+      it('should have blocks in snippets', function () {
+        const completions = getCompletions(`[{ $match: {foo: 1} }, $`);
+
+        completions.forEach((completion) => {
+          const snippet = applySnippet(completion);
+          expect(snippet).to.match(
+            /^{/,
+            'expected snippet to start with an opening bracket'
+          );
+        });
+      });
+    });
+  });
+
+  context('when autocompleting inside the stage', function () {
+    it('should return stage completer results', function () {
+      const completions = getCompletions('[{$bucket: { _id: "$', {
+        fields: [{ name: 'foo' }, { name: 'bar' }],
+      });
+
+      expect(completions.map((completion) => completion.label)).to.deep.eq([
+        '$foo',
+        '$bar',
+      ]);
+    });
+  });
+});

packages/compass-editor/src/codemirror/aggregation-autocompleter.ts

@@ -0,0 +1,106 @@
+import type { EditorState } from '@codemirror/state';
+import { STAGE_OPERATOR_NAMES } from '@mongodb-js/mongodb-constants';
+import type { CompleteOptions } from '../autocompleter';
+import { completer } from '../autocompleter';
+import type { Token } from './utils';
+import { createCompletionResultForIdPrefix } from './ace-compat-autocompleter';
+import { createAceCompatAutocompleter } from './ace-compat-autocompleter';
+import { createStageAutocompleter } from './stage-autocompleter';
+
+const StageOperatorNames = new Set(STAGE_OPERATOR_NAMES as string[]);
+
+function* parents(token: Token) {
+  let parent: Token | null = token;
+  while ((parent = parent.parent)) {
+    yield parent;
+  }
+}
+
+function removeQuotes(str: string) {
+  return str.replace(/(^('|")|('|")$)/g, '');
+}
+
+// lezer tokens are immutable, we check position in syntax tree to make sure we
+// are looking at the same token
+function isTokenEqual(a: Token, b: Token) {
+  return a.from === b.from && a.to === b.to;
+}
+
+function getPropertyNameFromPropertyToken(
+  editorState: EditorState,
+  propertyToken: Token
+): string {
+  if (!propertyToken.firstChild) {
+    return '';
+  }
+  return removeQuotes(getTokenText(editorState, propertyToken.firstChild));
+}
+
+function padLines(str: string, pad = '  ') {
+  return str
+    .split('\n')
+    .map((line) => `${pad}${line}`)
+    .join('\n');
+}
+
+function getTokenText(editorState: EditorState, token: Token) {
+  return editorState.sliceDoc(token.from, token.to);
+}
+
+function getStageNameForToken(
+  editorState: EditorState,
+  token: Token
+): string | null {
+  for (const parent of parents(token)) {
+    if (parent.name === 'Property') {
+      const propertyName = getPropertyNameFromPropertyToken(
+        editorState,
+        parent
+      );
+      if (
+        parent.firstChild &&
+        // We are inside a stage, but not right at the stage name token (we
+        // don't want to autocomplete as stage while the stage is being typed)
+        !isTokenEqual(parent.firstChild, token) &&
+        StageOperatorNames.has(propertyName)
+      ) {
+        return propertyName;
+      }
+    }
+  }
+  return null;
+}
+
+export function createAggregationAutocompleter(
+  options: Pick<CompleteOptions, 'fields' | 'serverVersion'> = {}
+) {
+  const stageAutocompletions = completer('', { ...options, meta: ['stage'] });
+
+  return createAceCompatAutocompleter({
+    IdentifierLike({ context, prefix, token }) {
+      const stageOperator = getStageNameForToken(context.state, token);
+
+      if (stageOperator) {
+        return createStageAutocompleter({ stageOperator, ...options })(context);
+      }
+
+      const isInsideBlock =
+        token.name === 'PropertyName' || token.parent?.name === 'Property';
+
+      return createCompletionResultForIdPrefix({
+        prefix,
+        completions: stageAutocompletions.map((completion) => {
+          const opName = completion.value;
+          return {
+            ...completion,
+            ...(completion.snippet && {
+              snippet: !isInsideBlock
+                ? `{\n${padLines(`${opName}: ${completion.snippet}`)}\n}`
+                : `${opName}: ${completion.snippet}`,
+            }),
+          };
+        }),
+      });
+    },
+  });
+}

packages/compass-editor/src/codemirror/utils.ts

@@ -5,7 +5,7 @@ import { syntaxTree } from '@codemirror/language';
 export const completeWordsInString = ifIn(['String'], completeAnyWord);
 
 export function resolveTokenAtCursor(context: CompletionContext) {
-  return syntaxTree(context.state).resolveInner(context.pos - 1);
+  return syntaxTree(context.state).resolveInner(context.pos, -1);
 }
 
 export type Token = ReturnType<typeof resolveTokenAtCursor>;

packages/compass-editor/src/index.ts

@@ -18,3 +18,4 @@ export { createDocumentAutocompleter } from './codemirror/document-autocompleter
 export { createValidationAutocompleter } from './codemirror/validation-autocompleter';
 export { createQueryAutocompleter } from './codemirror/query-autocompleter';
 export { createStageAutocompleter } from './codemirror/stage-autocompleter';
+export { createAggregationAutocompleter } from './codemirror/aggregation-autocompleter';

packages/compass-editor/test/completer.ts

@@ -90,5 +90,9 @@ export const setupCodemirrorCompleter = <
     editor.destroy();
     el.remove();
   };
-  return { getCompletions, cleanup };
+  const applySnippet = (completion: any) => {
+    completion.apply(editor, null, 0, editor.state.doc.length);
+    return editor.state.sliceDoc(0);
+  };
+  return { getCompletions, cleanup, applySnippet };
 };