Add codelense

Author: haohanyang

package-lock.json

@@ -69,6 +69,7 @@
     "@grafana/schema": "^10.4.0",
     "@grafana/ui": "^11.2.0",
     "@mongodb-js/mongodb-constants": "^0.11.1",
+    "jsonc-parser": "^3.3.1",
     "react": "18.2.0",
     "react-dom": "18.2.0",
     "shadowrealm-api": "^0.8.3",

package.json

@@ -3,6 +3,7 @@ import { CodeEditor, type MonacoEditor } from '@grafana/ui';
 import { useAutocomplete } from '../editor/autocomplete';
 import { useValidation } from '../editor/validation';
 import { useHover } from '../editor/hover';
+import { useCodeLens } from '../editor/codelens';
 
 interface QueryEditorRawProps {
   query: string;
@@ -20,6 +21,7 @@ export function QueryEditorRaw({ query, onBlur, language, width, height, fontSiz
   const setupAutocompleteFn = useAutocomplete();
   const setupHoverFn = useHover();
   const setupValidationFn = useValidation();
+  const setupCodeLensFn = useCodeLens();
 
   const formatQuery = useCallback(() => {
     if (monacoRef.current) {
@@ -35,6 +37,7 @@ export function QueryEditorRaw({ query, onBlur, language, width, height, fontSiz
           setupValidationFn(editor, monaco);
           setupAutocompleteFn(editor, monaco);
           setupHoverFn(editor, monaco);
+          setupCodeLensFn(editor, monaco);
         }}
         height={height || '240px'}
         width={width ? `${width - 2}px` : undefined}
@@ -43,7 +46,7 @@ export function QueryEditorRaw({ query, onBlur, language, width, height, fontSiz
         value={query}
         showMiniMap={false}
         showLineNumbers={true}
-        monacoOptions={fontSize ? { fontSize: fontSize } : undefined}
+        monacoOptions={fontSize ? { fontSize: fontSize, codeLens: true } : undefined}
       />
       {children && children({ formatQuery })}
     </div>

src/components/QueryEditorRaw.tsx

@@ -40,15 +40,18 @@ class CompletionProvider implements monacoTypes.languages.CompletionItemProvider
       endColumn: word.endColumn,
     };
 
-    const stageSuggestions: languages.CompletionItem[] = STAGE_OPERATORS.map((stage) => ({
-      label: `"${stage.name}"`,
-      kind: languages.CompletionItemKind.Function,
-      insertText: `"\\${stage.name}": ${stage.snippet}`,
-      range: range,
-      detail: stage.meta,
-      documentation: stage.description,
-      insertTextRules: languages.CompletionItemInsertTextRule.InsertAsSnippet,
-    }));
+    const stageSuggestions: languages.CompletionItem[] = STAGE_OPERATORS.map((stage) => {
+      // Add double quotation marks
+      const snippet = stage.snippet.replace(/(\s*)([a-zA-Z]+)\s*: /g, '$1"$2": ');
+      return {
+        label: `"${stage.name}"`,
+        kind: languages.CompletionItemKind.Function,
+        insertText: `"\\${stage.name}": ${snippet}`,
+        range: range,
+        detail: stage.meta,
+        documentation: stage.description,
+        insertTextRules: languages.CompletionItemInsertTextRule.InsertAsSnippet,
+  }});
 
     const expressionSuggestions: languages.CompletionItem[] = [...EXPRESSION_OPERATORS, ...ACCUMULATORS, ...CONVERSION_OPERATORS, ...QUERY_OPERATORS].map((expression) => ({
       label: `"${expression.name}"`,

src/editor/autocomplete.ts

@@ -0,0 +1,130 @@
+import { monacoTypes, type MonacoEditor, type Monaco } from '@grafana/ui';
+import { JSONPath, ParseErrorCode, type JSONVisitor, visit } from 'jsonc-parser';
+import { useEffect, useRef } from 'react';
+
+interface ParsedAggregateStages {
+  name: string;
+  startLine: number;
+  startCharacter: number;
+  length: number;
+}
+
+class CodeLensVisitor implements JSONVisitor {
+  private _currentLevel: number;
+  private _stages: ParsedAggregateStages[];
+  private _hasError: boolean;
+
+  constructor() {
+    this._hasError = false;
+    this._currentLevel = 0;
+    this._stages = [];
+  }
+
+  onObjectBegin = (
+    offset: number,
+    length: number,
+    startLine: number,
+    startCharacter: number,
+    pathSupplier: () => JSONPath,
+  ) => {
+    if (this._hasError) {
+      return false;
+    }
+    this._currentLevel += 1;
+  };
+
+  onObjectProperty = (
+    property: string,
+    offset: number,
+    length: number,
+    startLine: number,
+    startCharacter: number,
+    pathSupplier: () => JSONPath,
+  ) => {
+    if (this._currentLevel === 1) {
+      this._stages.push({
+        name: property,
+        startLine: startLine,
+        startCharacter: startCharacter,
+        length: length,
+      });
+    }
+  };
+
+  onObjectEnd = (offset: number, length: number, startLine: number, startCharacter: number) => {
+    this._currentLevel -= 1;
+  };
+
+  onError = (error: ParseErrorCode, offset: number, length: number, startLine: number, startCharacter: number) => {
+    this._hasError = true;
+  };
+
+  public get stages(): ParsedAggregateStages[] {
+    return this._stages;
+  }
+
+  public get hasError(): boolean {
+    return this._hasError;
+  }
+}
+
+class CodeLensProvider implements monacoTypes.languages.CodeLensProvider {
+  constructor(private readonly editor: MonacoEditor) {}
+
+  provideCodeLenses(
+    model: monacoTypes.editor.ITextModel,
+    _token: monacoTypes.CancellationToken,
+  ): monacoTypes.languages.ProviderResult<monacoTypes.languages.CodeLensList> {
+    if (this.editor.getModel()?.id !== model.id) {
+      return null;
+    }
+
+    const text = model.getValue();
+    const visitor = new CodeLensVisitor();
+    visit(text, visitor);
+
+    if (visitor.hasError) {
+      return null;
+    }
+
+    const stages = visitor.stages;
+    return {
+      lenses: stages.map((stage) => ({
+        range: {
+          startLineNumber: stage.startLine,
+          startColumn: stage.startCharacter,
+          endLineNumber: stage.startLine,
+          endColumn: stage.startCharacter + stage.length,
+        },
+        command: {
+          id: 'mongodb.aggregate.stage',
+          title: `Stage: ${stage.name}`,
+        },
+      })),
+      dispose: () => {},
+    };
+  }
+
+  resolveCodeLens(
+    _model: monacoTypes.editor.ITextModel,
+    codeLens: monacoTypes.languages.CodeLens,
+    _token: monacoTypes.CancellationToken,
+  ): monacoTypes.languages.ProviderResult<monacoTypes.languages.CodeLens> {
+    return codeLens;
+  }
+}
+
+export function useCodeLens() {
+  const codeLensDisposeFun = useRef<(() => void) | null>(null);
+  useEffect(() => {
+    return () => {
+      codeLensDisposeFun.current?.();
+    };
+  }, []);
+
+  return (editor: MonacoEditor, monaco: Monaco) => {
+    const provider = new CodeLensProvider(editor);
+    const { dispose } = monaco.languages.registerCodeLensProvider('json', provider);
+    codeLensDisposeFun.current = dispose;
+  };
+}