CloudWatch Logs: Support fetching fields in monaco editor (#78244)

packages/grafana-ui/src/components/Monaco/CodeEditor.tsx

@@ -27,6 +27,8 @@ class UnthemedCodeEditor extends PureComponent<Props> {
     if (this.completionCancel) {
       this.completionCancel.dispose();
     }
+
+    this.props.onEditorWillUnmount?.();
   }
 
   componentDidUpdate(oldProps: Props) {
@@ -77,6 +79,13 @@ class UnthemedCodeEditor extends PureComponent<Props> {
     }
   };
 
+  onFocus = () => {
+    const { onFocus } = this.props;
+    if (onFocus) {
+      onFocus(this.getEditorValue());
+    }
+  };
+
   onSave = () => {
     const { onSave } = this.props;
     if (onSave) {
@@ -164,7 +173,12 @@ class UnthemedCodeEditor extends PureComponent<Props> {
     }
 
     return (
-      <div className={containerStyles} onBlur={this.onBlur} data-testid={selectors.components.CodeEditor.container}>
+      <div
+        className={containerStyles}
+        onFocus={this.onFocus}
+        onBlur={this.onBlur}
+        data-testid={selectors.components.CodeEditor.container}
+      >
         <ReactMonacoEditorLazy
           width={width}
           height={height}

packages/grafana-ui/src/components/Monaco/types.ts

@@ -39,9 +39,15 @@ export interface CodeEditorProps {
    */
   onEditorDidMount?: (editor: MonacoEditor, monaco: Monaco) => void;
 
+  /** Callback before the edior has unmounted */
+  onEditorWillUnmount?: () => void;
+
   /** Handler to be performed when editor is blurred */
   onBlur?: CodeEditorChangeHandler;
 
+  /** Handler to be performed when editor is focused */
+  onFocus?: CodeEditorChangeHandler;
+
   /** Handler to be performed whenever the text inside the editor changes */
   onChange?: CodeEditorChangeHandler;
 

public/app/plugins/datasource/cloudwatch/components/QueryEditor/LogsQueryEditor/LogsQueryField.tsx

@@ -1,5 +1,5 @@
 import type * as monacoType from 'monaco-editor/esm/vs/editor/editor.api';
-import React, { ReactNode, useCallback } from 'react';
+import React, { ReactNode, useCallback, useRef } from 'react';
 
 import { QueryEditorProps } from '@grafana/data';
 import { CodeEditor, Monaco, Themeable2, withTheme2 } from '@grafana/ui';
@@ -7,7 +7,7 @@ import { CodeEditor, Monaco, Themeable2, withTheme2 } from '@grafana/ui';
 import { CloudWatchDatasource } from '../../../datasource';
 import language from '../../../language/logs/definition';
 import { TRIGGER_SUGGEST } from '../../../language/monarch/commands';
-import { registerLanguage } from '../../../language/monarch/register';
+import { registerLanguage, reRegisterCompletionProvider } from '../../../language/monarch/register';
 import { CloudWatchJsonData, CloudWatchLogsQuery, CloudWatchQuery } from '../../../types';
 import { getStatsGroups } from '../../../utils/query/getStatsGroups';
 import { LogGroupsFieldWrapper } from '../../shared/LogGroups/LogGroupsField';
@@ -22,6 +22,27 @@ export const CloudWatchLogsQueryFieldMonaco = (props: CloudWatchLogsQueryFieldPr
   const { query, datasource, onChange, ExtraFieldElement, data } = props;
 
   const showError = data?.error?.refId === query.refId;
+  const monacoRef = useRef<Monaco>();
+  const disposalRef = useRef<monacoType.IDisposable>();
+
+  const onChangeLogs = useCallback(
+    async (query: CloudWatchLogsQuery) => {
+      onChange(query);
+    },
+    [onChange]
+  );
+
+  const onFocus = useCallback(async () => {
+    disposalRef.current = await reRegisterCompletionProvider(
+      monacoRef.current!,
+      language,
+      datasource.logsCompletionItemProviderFunc({
+        region: query.region,
+        logGroups: query.logGroups,
+      }),
+      disposalRef.current
+    );
+  }, [datasource, query.logGroups, query.region]);
 
   const onChangeQuery = useCallback(
     (value: string) => {
@@ -44,6 +65,17 @@ export const CloudWatchLogsQueryFieldMonaco = (props: CloudWatchLogsQueryFieldPr
     },
     [onChangeQuery]
   );
+  const onBeforeEditorMount = async (monaco: Monaco) => {
+    monacoRef.current = monaco;
+    disposalRef.current = await registerLanguage(
+      monaco,
+      language,
+      datasource.logsCompletionItemProviderFunc({
+        region: query.region,
+        logGroups: query.logGroups,
+      })
+    );
+  };
 
   return (
     <>
@@ -53,11 +85,11 @@ export const CloudWatchLogsQueryFieldMonaco = (props: CloudWatchLogsQueryFieldPr
         legacyLogGroupNames={query.logGroupNames}
         logGroups={query.logGroups}
         onChange={(logGroups) => {
-          onChange({ ...query, logGroups, logGroupNames: undefined });
+          onChangeLogs({ ...query, logGroups, logGroupNames: undefined });
         }}
         //legacy props
         legacyOnChange={(logGroupNames) => {
-          onChange({ ...query, logGroupNames });
+          onChangeLogs({ ...query, logGroupNames });
         }}
       />
       <div className="gf-form-inline gf-form-inline--nowrap flex-grow-1">
@@ -90,11 +122,12 @@ export const CloudWatchLogsQueryFieldMonaco = (props: CloudWatchLogsQueryFieldPr
               if (value !== query.expression) {
                 onChangeQuery(value);
               }
+              disposalRef.current?.dispose();
             }}
-            onBeforeEditorMount={(monaco: Monaco) =>
-              registerLanguage(monaco, language, datasource.logsCompletionItemProvider)
-            }
+            onFocus={onFocus}
+            onBeforeEditorMount={onBeforeEditorMount}
             onEditorDidMount={onEditorMount}
+            onEditorWillUnmount={() => disposalRef.current?.dispose()}
           />
         </div>
         {ExtraFieldElement}

public/app/plugins/datasource/cloudwatch/datasource.ts

@@ -22,7 +22,11 @@ import { DEFAULT_METRICS_QUERY, getDefaultLogsQuery } from './defaultQueries';
 import { isCloudWatchAnnotationQuery, isCloudWatchLogsQuery, isCloudWatchMetricsQuery } from './guards';
 import { CloudWatchLogsLanguageProvider } from './language/cloudwatch-logs/CloudWatchLogsLanguageProvider';
 import { SQLCompletionItemProvider } from './language/cloudwatch-sql/completion/CompletionItemProvider';
-import { LogsCompletionItemProvider } from './language/logs/completion/CompletionItemProvider';
+import {
+  LogsCompletionItemProvider,
+  LogsCompletionItemProviderFunc,
+  queryContext,
+} from './language/logs/completion/CompletionItemProvider';
 import { MetricMathCompletionItemProvider } from './language/metric-math/completion/CompletionItemProvider';
 import { CloudWatchAnnotationQueryRunner } from './query-runner/CloudWatchAnnotationQueryRunner';
 import { CloudWatchLogsQueryRunner } from './query-runner/CloudWatchLogsQueryRunner';
@@ -45,7 +49,7 @@ export class CloudWatchDatasource
   languageProvider: CloudWatchLogsLanguageProvider;
   sqlCompletionItemProvider: SQLCompletionItemProvider;
   metricMathCompletionItemProvider: MetricMathCompletionItemProvider;
-  logsCompletionItemProvider: LogsCompletionItemProvider;
+  logsCompletionItemProviderFunc: (queryContext: queryContext) => LogsCompletionItemProvider;
   defaultLogGroups?: string[];
 
   type = 'cloudwatch';
@@ -67,7 +71,7 @@ export class CloudWatchDatasource
     this.sqlCompletionItemProvider = new SQLCompletionItemProvider(this.resources, this.templateSrv);
     this.metricMathCompletionItemProvider = new MetricMathCompletionItemProvider(this.resources, this.templateSrv);
     this.metricsQueryRunner = new CloudWatchMetricsQueryRunner(instanceSettings, templateSrv, super.query.bind(this));
-    this.logsCompletionItemProvider = new LogsCompletionItemProvider(this.resources, this.templateSrv);
+    this.logsCompletionItemProviderFunc = LogsCompletionItemProviderFunc(this.resources, this.templateSrv);
     this.logsQueryRunner = new CloudWatchLogsQueryRunner(
       instanceSettings,
       templateSrv,

public/app/plugins/datasource/cloudwatch/language/logs/completion/CompletionItemProvider.test.ts

@@ -6,6 +6,8 @@ import { emptyQuery, filterQuery, newCommandQuery, sortQuery } from '../../../__
 import MonacoMock from '../../../__mocks__/monarch/Monaco';
 import TextModel from '../../../__mocks__/monarch/TextModel';
 import { ResourcesAPI } from '../../../resources/ResourcesAPI';
+import { ResourceResponse } from '../../../resources/types';
+import { LogGroup, LogGroupField } from '../../../types';
 import cloudWatchLogsLanguageDefinition from '../definition';
 import { LOGS_COMMANDS, LOGS_FUNCTION_OPERATORS, SORT_DIRECTION_KEYWORDS } from '../language';
 
@@ -18,14 +20,19 @@ jest.mock('monaco-editor/esm/vs/editor/editor.api', () => ({
 const getSuggestions = async (
   value: string,
   position: monacoTypes.IPosition,
-  variables: CustomVariableModel[] = []
+  variables: CustomVariableModel[] = [],
+  logGroups: LogGroup[] = [],
+  fields: Array<ResourceResponse<LogGroupField>> = []
 ) => {
   const setup = new LogsCompletionItemProvider(
     {
       getActualRegion: () => 'us-east-2',
     } as ResourcesAPI,
-    setupMockedTemplateService(variables)
+    setupMockedTemplateService(variables),
+    { region: 'default', logGroups }
   );
+
+  setup.resources.getLogGroupFields = jest.fn().mockResolvedValue(fields);
   const monaco = MonacoMock as Monaco;
   const provider = setup.getCompletionProvider(monaco, cloudWatchLogsLanguageDefinition);
   const { suggestions } = await provider.provideCompletionItems(
@@ -76,5 +83,17 @@ describe('LogsCompletionItemProvider', () => {
       const expectedLabels = [...LOGS_COMMANDS, expectedTemplateVariableLabel];
       expect(suggestionLabels).toEqual(expect.arrayContaining(expectedLabels));
     });
+
+    it('fetches fields when logGroups are set', async () => {
+      const suggestions = await getSuggestions(
+        sortQuery.query,
+        sortQuery.position,
+        [],
+        [{ arn: 'foo', name: 'bar' }],
+        [{ value: { name: '@field' } }]
+      );
+      const suggestionLabels = suggestions.map((s) => s.label);
+      expect(suggestionLabels).toEqual(expect.arrayContaining(['@field']));
+    });
   });
 });

public/app/plugins/datasource/cloudwatch/language/logs/completion/CompletionItemProvider.ts

@@ -1,7 +1,8 @@
-import { getTemplateSrv, type TemplateSrv } from '@grafana/runtime';
+import { getTemplateSrv, TemplateSrv } from '@grafana/runtime';
 import { Monaco, monacoTypes } from '@grafana/ui';
 
 import { type ResourcesAPI } from '../../../resources/ResourcesAPI';
+import { LogGroup } from '../../../types';
 import { CompletionItemProvider } from '../../monarch/CompletionItemProvider';
 import { LinkedToken } from '../../monarch/LinkedToken';
 import { TRIGGER_SUGGEST } from '../../monarch/commands';
@@ -12,12 +13,26 @@ import { getStatementPosition } from './statementPosition';
 import { getSuggestionKinds } from './suggestionKinds';
 import { LogsTokenTypes } from './types';
 
+export type queryContext = {
+  logGroups?: LogGroup[];
+  region: string;
+};
+
+export function LogsCompletionItemProviderFunc(resources: ResourcesAPI, templateSrv: TemplateSrv = getTemplateSrv()) {
+  return (queryContext: queryContext) => {
+    return new LogsCompletionItemProvider(resources, templateSrv, queryContext);
+  };
+}
+
 export class LogsCompletionItemProvider extends CompletionItemProvider {
-  constructor(resources: ResourcesAPI, templateSrv: TemplateSrv = getTemplateSrv()) {
+  queryContext: queryContext;
+
+  constructor(resources: ResourcesAPI, templateSrv: TemplateSrv = getTemplateSrv(), queryContext: queryContext) {
     super(resources, templateSrv);
     this.getStatementPosition = getStatementPosition;
     this.getSuggestionKinds = getSuggestionKinds;
     this.tokenTypes = LogsTokenTypes;
+    this.queryContext = queryContext;
   }
 
   async getSuggestions(
@@ -56,6 +71,7 @@ export class LogsCompletionItemProvider extends CompletionItemProvider {
               insertText: `${command} $0`,
               insertTextRules: monaco.languages.CompletionItemInsertTextRule.InsertAsSnippet,
               command: TRIGGER_SUGGEST,
+              kind: monaco.languages.CompletionItemKind.Method,
             });
           });
           break;
@@ -65,12 +81,32 @@ export class LogsCompletionItemProvider extends CompletionItemProvider {
               insertText: `${f}($0)`,
               insertTextRules: monaco.languages.CompletionItemInsertTextRule.InsertAsSnippet,
               command: TRIGGER_SUGGEST,
+              kind: monaco.languages.CompletionItemKind.Function,
             });
           });
+
+          if (this.queryContext.logGroups && this.queryContext.logGroups.length > 0) {
+            let fields = await this.fetchFields(this.queryContext.logGroups, this.queryContext.region);
+            fields.push('@log');
+            fields.forEach((field) => {
+              if (field !== '') {
+                addSuggestion(field, {
+                  range,
+                  label: field,
+                  insertText: field,
+                  kind: monaco.languages.CompletionItemKind.Field,
+                  sortText: CompletionItemPriority.High,
+                });
+              }
+            });
+          }
           break;
         case SuggestionKind.SortOrderDirectionKeyword:
           SORT_DIRECTION_KEYWORDS.forEach((direction) => {
-            addSuggestion(direction, { sortText: CompletionItemPriority.High });
+            addSuggestion(direction, {
+              sortText: CompletionItemPriority.High,
+              kind: monaco.languages.CompletionItemKind.Operator,
+            });
           });
           break;
         case SuggestionKind.InKeyword:
@@ -82,19 +118,31 @@ export class LogsCompletionItemProvider extends CompletionItemProvider {
           });
           break;
       }
+    }
 
-      this.templateSrv.getVariables().map((v) => {
-        const variable = `$${v.name}`;
-        addSuggestion(variable, {
-          range,
-          label: variable,
-          insertText: variable,
-          kind: monaco.languages.CompletionItemKind.Variable,
-          sortText: CompletionItemPriority.Low,
-        });
+    this.templateSrv.getVariables().map((v) => {
+      const variable = `$${v.name}`;
+      addSuggestion(variable, {
+        range,
+        label: variable,
+        insertText: variable,
+        kind: monaco.languages.CompletionItemKind.Variable,
+        sortText: CompletionItemPriority.Low,
       });
-    }
+    });
 
     return suggestions;
   }
+
+  private fetchFields = async (logGroups: LogGroup[], region: string): Promise<string[]> => {
+    const results = await Promise.all(
+      logGroups.map((logGroup) =>
+        this.resources
+          .getLogGroupFields({ logGroupName: logGroup.name, arn: logGroup.arn, region })
+          .then((fields) => fields.filter((f) => f).map((f) => f.value.name ?? ''))
+      )
+    );
+    // Deduplicate fields
+    return [...new Set(results.flat())];
+  };
 }

public/app/plugins/datasource/cloudwatch/language/monarch/register.ts

@@ -15,7 +15,23 @@ export type LanguageDefinition = {
   }>;
 };
 
-export const registerLanguage = (
+export const reRegisterCompletionProvider = async (
+  monaco: Monaco,
+  language: LanguageDefinition,
+  completionItemProvider: Completeable,
+  disposal?: monacoType.IDisposable
+) => {
+  const { id, loader } = language;
+  disposal?.dispose();
+  return loader().then((monarch) => {
+    return monaco.languages.registerCompletionItemProvider(
+      id,
+      completionItemProvider.getCompletionProvider(monaco, language)
+    );
+  });
+};
+
+export const registerLanguage = async (
   monaco: Monaco,
   language: LanguageDefinition,
   completionItemProvider: Completeable
@@ -28,9 +44,12 @@ export const registerLanguage = (
   }
 
   monaco.languages.register({ id });
-  loader().then((monarch) => {
+  return loader().then((monarch) => {
     monaco.languages.setMonarchTokensProvider(id, monarch.language);
     monaco.languages.setLanguageConfiguration(id, monarch.conf);
-    monaco.languages.registerCompletionItemProvider(id, completionItemProvider.getCompletionProvider(monaco, language));
+    return monaco.languages.registerCompletionItemProvider(
+      id,
+      completionItemProvider.getCompletionProvider(monaco, language)
+    );
   });
 };