CloudWatch/Logs: Language provider refactor and test (#24425)

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

@@ -0,0 +1,171 @@
+import { Value } from 'slate';
+import { TypeaheadOutput } from '@grafana/ui';
+import { CloudWatchDatasource } from './datasource';
+import { GetLogGroupFieldsResponse } from './types';
+import { CloudWatchLanguageProvider } from './language_provider';
+import Prism, { Token } from 'prismjs';
+import {
+  AGGREGATION_FUNCTIONS_STATS,
+  BOOLEAN_FUNCTIONS,
+  DATETIME_FUNCTIONS,
+  FUNCTIONS,
+  IP_FUNCTIONS,
+  NUMERIC_OPERATORS,
+  QUERY_COMMANDS,
+  STRING_FUNCTIONS,
+} from './syntax';
+
+const fields = ['field1', '@message'];
+
+describe('CloudWatchLanguageProvider', () => {
+  it('should suggest ', async () => {
+    await runSuggestionTest('stats count(\\)', [fields]);
+    // Make sure having a field prefix does not brake anything
+    await runSuggestionTest('stats count(@mess\\)', [fields]);
+  });
+
+  it('should suggest query commands on start of query', async () => {
+    await runSuggestionTest('\\', [QUERY_COMMANDS.map(v => v.label)]);
+  });
+
+  it('should suggest query commands after pipe', async () => {
+    await runSuggestionTest('fields f | \\', [QUERY_COMMANDS.map(v => v.label)]);
+  });
+
+  it('should suggest fields and functions after field command', async () => {
+    await runSuggestionTest('fields \\', [fields, FUNCTIONS.map(v => v.label)]);
+  });
+
+  it('should suggest fields and functions after comma', async () => {
+    await runSuggestionTest('fields field1, \\', [fields, FUNCTIONS.map(v => v.label)]);
+  });
+
+  it('should suggest fields and functions after display command', async () => {
+    await runSuggestionTest('display \\', [fields, FUNCTIONS.map(v => v.label)]);
+  });
+
+  it('should suggest functions after stats command', async () => {
+    await runSuggestionTest('stats \\', [AGGREGATION_FUNCTIONS_STATS.map(v => v.label)]);
+  });
+
+  it('should suggest fields and some functions after `by` command', async () => {
+    await runSuggestionTest('stats count(something) by \\', [
+      fields,
+      STRING_FUNCTIONS.concat(DATETIME_FUNCTIONS, IP_FUNCTIONS).map(v => v.label),
+    ]);
+  });
+
+  it('should suggest fields and some functions after comparison operator', async () => {
+    await runSuggestionTest('filter field1 >= \\', [
+      fields,
+      BOOLEAN_FUNCTIONS.map(v => v.label),
+      NUMERIC_OPERATORS.map(v => v.label),
+    ]);
+  });
+
+  it('should suggest fields directly after sort', async () => {
+    await runSuggestionTest('sort \\', [fields]);
+  });
+
+  it('should suggest fields directly after sort after a pipe', async () => {
+    await runSuggestionTest('fields field1 | sort \\', [fields]);
+  });
+
+  it('should suggest sort order after sort command and field', async () => {
+    await runSuggestionTest('sort field1 \\', [['asc', 'desc']]);
+  });
+
+  it('should suggest fields directly after parse', async () => {
+    await runSuggestionTest('parse \\', [fields]);
+  });
+
+  it('should suggest fields and bool functions after filter', async () => {
+    await runSuggestionTest('filter \\', [fields, BOOLEAN_FUNCTIONS.map(v => v.label)]);
+  });
+});
+
+async function runSuggestionTest(query: string, expectedItems: string[][]) {
+  const result = await getProvideCompletionItems(query);
+  expectedItems.forEach((items, index) => {
+    expect(result.suggestions[index].items.map(item => item.label)).toEqual(items);
+  });
+}
+
+function makeDatasource(): CloudWatchDatasource {
+  return {
+    getLogGroupFields(): Promise<GetLogGroupFieldsResponse> {
+      return Promise.resolve({ logGroupFields: [{ name: 'field1' }, { name: '@message' }] });
+    },
+  } as any;
+}
+
+/**
+ * Get suggestion items based on query. Use `\\` to mark position of the cursor.
+ */
+function getProvideCompletionItems(query: string): Promise<TypeaheadOutput> {
+  const provider = new CloudWatchLanguageProvider(makeDatasource());
+  const cursorOffset = query.indexOf('\\');
+  const queryWithoutCursor = query.replace('\\', '');
+  let tokens: Token[] = Prism.tokenize(queryWithoutCursor, provider.getSyntax()) as any;
+  tokens = addTokenMetadata(tokens);
+  const value = new ValueMock(tokens, cursorOffset);
+  return provider.provideCompletionItems(
+    {
+      value,
+    } as any,
+    { logGroupNames: ['logGroup1'] }
+  );
+}
+
+class ValueMock {
+  selection: Value['selection'];
+  data: Value['data'];
+
+  constructor(tokens: Array<string | Token>, cursorOffset: number) {
+    this.selection = {
+      start: {
+        offset: cursorOffset,
+      },
+    } as any;
+
+    this.data = {
+      get() {
+        return tokens;
+      },
+    } as any;
+  }
+}
+
+/**
+ * Adds some Slate specific metadata
+ * @param tokens
+ */
+function addTokenMetadata(tokens: Array<string | Token>): Token[] {
+  let prev = undefined as any;
+  let offset = 0;
+  return tokens.reduce((acc, token) => {
+    let newToken: any;
+    if (typeof token === 'string') {
+      newToken = {
+        content: token,
+        // Not sure what else could it be here, probably if we do not match something
+        types: ['whitespace'],
+      };
+    } else {
+      newToken = { ...token };
+      newToken.types = [token.type];
+    }
+    newToken.prev = prev;
+    if (newToken.prev) {
+      newToken.prev.next = newToken;
+    }
+    const end = offset + token.length;
+    newToken.offsets = {
+      start: offset,
+      end,
+    };
+    prev = newToken;
+    offset = end;
+    return [...acc, newToken];
+  }, [] as Token[]);
+}

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

@@ -15,20 +15,12 @@ import syntax, {
 
 // Types
 import { CloudWatchQuery } from './types';
-import { dateTime, AbsoluteTimeRange, LanguageProvider, HistoryItem } from '@grafana/data';
+import { AbsoluteTimeRange, LanguageProvider, HistoryItem } from '@grafana/data';
 
 import { CloudWatchDatasource } from './datasource';
-import { CompletionItem, TypeaheadInput, TypeaheadOutput, Token } from '@grafana/ui';
+import { TypeaheadInput, TypeaheadOutput, Token } from '@grafana/ui';
 import { Grammar } from 'prismjs';
 
-const HISTORY_ITEM_COUNT = 10;
-const HISTORY_COUNT_CUTOFF = 1000 * 60 * 60 * 24; // 24h
-const NS_IN_MS = 1000000;
-export const LABEL_REFRESH_INTERVAL = 1000 * 30; // 30sec
-
-const wrapLabel = (label: string) => ({ label });
-export const rangeToParams = (range: AbsoluteTimeRange) => ({ start: range.from * NS_IN_MS, end: range.to * NS_IN_MS });
-
 export type CloudWatchHistoryItem = HistoryItem<CloudWatchQuery>;
 
 type TypeaheadContext = {
@@ -37,26 +29,7 @@ type TypeaheadContext = {
   logGroupNames?: string[];
 };
 
-export function addHistoryMetadata(item: CompletionItem, history: CloudWatchHistoryItem[]): CompletionItem {
-  const cutoffTs = Date.now() - HISTORY_COUNT_CUTOFF;
-  const historyForItem = history.filter(h => h.ts > cutoffTs && h.query.expression === item.label);
-  let hint = `Queried ${historyForItem.length} times in the last 24h.`;
-  const recent = historyForItem[0];
-
-  if (recent) {
-    const lastQueried = dateTime(recent.ts).fromNow();
-    hint = `${hint} Last queried ${lastQueried}.`;
-  }
-
-  return {
-    ...item,
-    documentation: hint,
-  };
-}
-
 export class CloudWatchLanguageProvider extends LanguageProvider {
-  logLabelOptions: any[];
-  logLabelFetchTs?: number;
   started: boolean;
   initialRange: AbsoluteTimeRange;
   datasource: CloudWatchDatasource;
@@ -91,18 +64,6 @@ export class CloudWatchLanguageProvider extends LanguageProvider {
     return this.startTask;
   };
 
-  fetchFields = _.throttle(async (logGroups: string[]) => {
-    const results = await Promise.all(
-      logGroups.map(logGroup => this.datasource.getLogGroupFields({ logGroupName: logGroup }))
-    );
-
-    return [
-      ...new Set<string>(
-        results.reduce((acc: string[], cur) => acc.concat(cur.logGroupFields?.map(f => f.name) as string[]), [])
-      ).values(),
-    ];
-  }, 30 * 1000);
-
   /**
    * Return suggestions based on input that can be then plugged into a typeahead dropdown.
    * Keep this DOM-free for testing
@@ -112,7 +73,6 @@ export class CloudWatchLanguageProvider extends LanguageProvider {
    * @param context.history Optional used only in getEmptyCompletionItems
    */
   async provideCompletionItems(input: TypeaheadInput, context?: TypeaheadContext): Promise<TypeaheadOutput> {
-    //console.log('Providing completion items...');
     const { value } = input;
 
     // Get tokens
@@ -127,36 +87,29 @@ export class CloudWatchLanguageProvider extends LanguageProvider {
         token.offsets.start <= value!.selection?.start?.offset && token.offsets.end >= value!.selection?.start?.offset
     )[0];
 
-    const isFirstToken = curToken.prev === null || curToken.prev === undefined;
+    const isFirstToken = !curToken.prev;
     const prevToken = prevNonWhitespaceToken(curToken);
 
+    const isCommandStart = isFirstToken || (!isFirstToken && prevToken?.types.includes('command-separator'));
+    if (isCommandStart) {
+      return this.getCommandCompletionItems();
+    }
+
     if (isInsideFunctionParenthesis(curToken)) {
       return await this.getFieldCompletionItems(context?.logGroupNames ?? []);
     }
 
-    const isCommandStart = isFirstToken || (!isFirstToken && prevToken?.types.includes('command-separator'));
+    if (isAfterKeyword('by', curToken)) {
-    if (isCommandStart) {
+      return this.handleKeyword(context);
-      return this.getCommandCompletionItems();
-    } else if (!isFirstToken) {
-      if (prevToken?.types.includes('keyword')) {
-        return this.handleKeyword(prevToken, context);
     }
 
     if (prevToken?.types.includes('comparison-operator')) {
-        const suggs = await this.getFieldCompletionItems(context?.logGroupNames ?? []);
+      return this.handleComparison(context);
-        const boolFuncSuggs = this.getBoolFuncCompletionItems();
-        const numFuncSuggs = this.getNumericFuncCompletionItems();
-
-        suggs.suggestions.push(...boolFuncSuggs.suggestions, ...numFuncSuggs.suggestions);
-        return suggs;
     }
 
-      const commandToken = this.findCommandToken(curToken);
+    const commandToken = previousCommandToken(curToken);
-
+    if (commandToken) {
-      if (commandToken !== null) {
+      return await this.handleCommand(commandToken, curToken, context);
-        const typeaheadOutput = await this.handleCommand(commandToken, curToken, context);
-        return typeaheadOutput;
-      }
     }
 
     return {
@@ -164,8 +117,19 @@ export class CloudWatchLanguageProvider extends LanguageProvider {
     };
   }
 
-  handleKeyword = async (token: Token, context?: TypeaheadContext): Promise<TypeaheadOutput | null> => {
+  private fetchFields = _.throttle(async (logGroups: string[]) => {
-    if (token.content.toLowerCase() === 'by') {
+    const results = await Promise.all(
+      logGroups.map(logGroup => this.datasource.getLogGroupFields({ logGroupName: logGroup }))
+    );
+
+    return [
+      ...new Set<string>(
+        results.reduce((acc: string[], cur) => acc.concat(cur.logGroupFields?.map(f => f.name) as string[]), [])
+      ).values(),
+    ];
+  }, 30 * 1000);
+
+  private handleKeyword = async (context?: TypeaheadContext): Promise<TypeaheadOutput | null> => {
     const suggs = await this.getFieldCompletionItems(context?.logGroupNames ?? []);
     const functionSuggestions = [
       { prefixMatch: true, label: 'Functions', items: STRING_FUNCTIONS.concat(DATETIME_FUNCTIONS, IP_FUNCTIONS) },
@@ -173,24 +137,75 @@ export class CloudWatchLanguageProvider extends LanguageProvider {
     suggs.suggestions.push(...functionSuggestions);
 
     return suggs;
-    }
-
-    return null;
   };
 
-  handleCommand = async (commandToken: Token, curToken: Token, context: TypeaheadContext): Promise<TypeaheadOutput> => {
+  private handleCommand = async (
+    commandToken: Token,
+    curToken: Token,
+    context: TypeaheadContext
+  ): Promise<TypeaheadOutput> => {
     const queryCommand = commandToken.content.toLowerCase();
     const prevToken = prevNonWhitespaceToken(curToken);
     const currentTokenIsFirstArg = prevToken === commandToken;
 
-    // console.log(
-    //   `Query Command: '${queryCommand}'. Previous token: '${prevToken}'. First arg? ${currentTokenIsFirstArg}`
-    // );
-
     if (queryCommand === 'sort') {
+      return this.handleSortCommand(currentTokenIsFirstArg, curToken, context);
+    }
+
+    if (queryCommand === 'parse') {
       if (currentTokenIsFirstArg) {
         return await this.getFieldCompletionItems(context.logGroupNames ?? []);
-      } else if (prevToken?.types.includes('field-name')) {
+      }
+    }
+
+    const currentTokenIsAfterCommandAndEmpty =
+      commandToken.next?.types.includes('whitespace') && !commandToken.next.next;
+    const currentTokenIsAfterCommand =
+      currentTokenIsAfterCommandAndEmpty || nextNonWhitespaceToken(commandToken) === curToken;
+
+    const currentTokenIsComma = curToken.content === ',' && curToken.types.includes('punctuation');
+    const currentTokenIsCommaOrAfterComma =
+      currentTokenIsComma || (curToken.prev?.content === ',' && curToken.prev.types.includes('punctuation'));
+
+    // We only show suggestions if we are after a command or after a comma which is a field separator
+    if (!(currentTokenIsAfterCommand || currentTokenIsCommaOrAfterComma)) {
+      return { suggestions: [] };
+    }
+
+    if (['display', 'fields'].includes(queryCommand)) {
+      const typeaheadOutput = await this.getFieldCompletionItems(context.logGroupNames ?? []);
+      typeaheadOutput.suggestions.push(...this.getFunctionCompletionItems().suggestions);
+
+      return typeaheadOutput;
+    }
+
+    if (queryCommand === 'stats') {
+      const typeaheadOutput = this.getStatsAggCompletionItems();
+      if (currentTokenIsComma || currentTokenIsAfterCommandAndEmpty) {
+        typeaheadOutput?.suggestions.forEach(group => {
+          group.skipFilter = true;
+        });
+      }
+      return typeaheadOutput;
+    }
+
+    if (queryCommand === 'filter' && currentTokenIsFirstArg) {
+      const sugg = await this.getFieldCompletionItems(context.logGroupNames ?? []);
+      const boolFuncs = this.getBoolFuncCompletionItems();
+      sugg.suggestions.push(...boolFuncs.suggestions);
+      return sugg;
+    }
+    return { suggestions: [] };
+  };
+
+  private async handleSortCommand(
+    isFirstArgument: boolean,
+    curToken: Token,
+    context: TypeaheadContext
+  ): Promise<TypeaheadOutput> {
+    if (isFirstArgument) {
+      return await this.getFieldCompletionItems(context.logGroupNames ?? []);
+    } else if (prevNonWhitespaceToken(curToken)?.types.includes('field-name')) {
       // suggest sort options
       return {
         suggestions: [
@@ -207,119 +222,32 @@ export class CloudWatchLanguageProvider extends LanguageProvider {
         ],
       };
     }
-    }
-
-    if (queryCommand === 'parse') {
-      if (currentTokenIsFirstArg) {
-        return await this.getFieldCompletionItems(context.logGroupNames ?? []);
-      }
-    }
-
-    let typeaheadOutput: TypeaheadOutput | null = null;
-    if (
-      (commandToken.next?.types.includes('whitespace') && commandToken.next.next === null) ||
-      nextNonWhitespaceToken(commandToken) === curToken ||
-      (curToken.content === ',' && curToken.types.includes('punctuation')) ||
-      (curToken.prev?.content === ',' && curToken.prev.types.includes('punctuation'))
-    ) {
-      if (['display', 'fields'].includes(queryCommand)) {
-        // Current token comes straight after command OR after comma
-        typeaheadOutput = await this.getFieldCompletionItems(context.logGroupNames ?? []);
-        typeaheadOutput.suggestions.push(...this.getFunctionCompletionItems().suggestions);
-
-        return typeaheadOutput;
-      } else if (queryCommand === 'stats') {
-        typeaheadOutput = this.getStatsAggCompletionItems();
-      } else if (queryCommand === 'filter') {
-        if (currentTokenIsFirstArg) {
-          const sugg = await this.getFieldCompletionItems(context.logGroupNames ?? []);
-          const boolFuncs = this.getBoolFuncCompletionItems();
-          sugg.suggestions.push(...boolFuncs.suggestions);
-          return sugg;
-        }
-      }
-
-      if (
-        (curToken.content === ',' && curToken.types.includes('punctuation')) ||
-        (commandToken.next?.types.includes('whitespace') && commandToken.next.next === null)
-      ) {
-        typeaheadOutput?.suggestions.forEach(group => {
-          group.skipFilter = true;
-        });
-      }
-
-      return typeaheadOutput!;
-    }
 
     return { suggestions: [] };
+  }
+
+  private handleComparison = async (context?: TypeaheadContext) => {
+    const fieldsSuggestions = await this.getFieldCompletionItems(context?.logGroupNames ?? []);
+    const boolFuncSuggestions = this.getBoolFuncCompletionItems();
+    const numFuncSuggestions = this.getNumericFuncCompletionItems();
+
+    fieldsSuggestions.suggestions.push(...boolFuncSuggestions.suggestions, ...numFuncSuggestions.suggestions);
+    return fieldsSuggestions;
   };
 
-  findCommandToken = (startToken: Token): Token | null => {
+  private getCommandCompletionItems = (): TypeaheadOutput => {
-    let thisToken = { ...startToken };
-
-    while (thisToken.prev !== null) {
-      thisToken = thisToken.prev;
-      const isFirstCommand = thisToken.types.includes('query-command') && thisToken.prev === null;
-      if (thisToken.types.includes('command-separator') || isFirstCommand) {
-        // next token should be command
-        if (!isFirstCommand && thisToken.next?.types.includes('query-command')) {
-          return thisToken.next;
-        } else {
-          return thisToken;
-        }
-      }
-    }
-
-    return null;
-  };
-
-  getBeginningCompletionItems = (context: TypeaheadContext): TypeaheadOutput => {
-    return {
-      suggestions: [
-        ...this.getEmptyCompletionItems(context).suggestions,
-        ...this.getCommandCompletionItems().suggestions,
-      ],
-    };
-  };
-
-  getEmptyCompletionItems(context: TypeaheadContext): TypeaheadOutput {
-    const history = context?.history;
-    const suggestions = [];
-
-    if (history?.length) {
-      const historyItems = _.chain(history)
-        .map(h => h.query.expression)
-        .filter()
-        .uniq()
-        .take(HISTORY_ITEM_COUNT)
-        .map(wrapLabel)
-        .map((item: CompletionItem) => addHistoryMetadata(item, history))
-        .value();
-
-      suggestions.push({
-        prefixMatch: true,
-        skipSort: true,
-        label: 'History',
-        items: historyItems,
-      });
-    }
-
-    return { suggestions };
-  }
-
-  getCommandCompletionItems = (): TypeaheadOutput => {
     return { suggestions: [{ prefixMatch: true, label: 'Commands', items: QUERY_COMMANDS }] };
   };
 
-  getFunctionCompletionItems = (): TypeaheadOutput => {
+  private getFunctionCompletionItems = (): TypeaheadOutput => {
     return { suggestions: [{ prefixMatch: true, label: 'Functions', items: FUNCTIONS }] };
   };
 
-  getStatsAggCompletionItems = (): TypeaheadOutput => {
+  private getStatsAggCompletionItems = (): TypeaheadOutput => {
     return { suggestions: [{ prefixMatch: true, label: 'Functions', items: AGGREGATION_FUNCTIONS_STATS }] };
   };
 
-  getBoolFuncCompletionItems = (): TypeaheadOutput => {
+  private getBoolFuncCompletionItems = (): TypeaheadOutput => {
     return {
       suggestions: [
         {
@@ -331,7 +259,7 @@ export class CloudWatchLanguageProvider extends LanguageProvider {
     };
   };
 
-  getNumericFuncCompletionItems = (): TypeaheadOutput => {
+  private getNumericFuncCompletionItems = (): TypeaheadOutput => {
     return {
       suggestions: [
         {
@@ -343,11 +271,9 @@ export class CloudWatchLanguageProvider extends LanguageProvider {
     };
   };
 
-  getFieldCompletionItems = async (logGroups: string[]): Promise<TypeaheadOutput> => {
+  private getFieldCompletionItems = async (logGroups: string[]): Promise<TypeaheadOutput> => {
-    //console.log(`Fetching fields... ${logGroups}`);
     const fields = await this.fetchFields(logGroups);
 
-    //console.log(fields);
     return {
       suggestions: [
         {
@@ -391,6 +317,20 @@ function prevNonWhitespaceToken(token: Token): Token | null {
   return null;
 }
 
+function previousCommandToken(startToken: Token): Token | null {
+  let thisToken = startToken;
+  while (!!thisToken.prev) {
+    thisToken = thisToken.prev;
+    if (
+      thisToken.types.includes('query-command') &&
+      (!thisToken.prev || prevNonWhitespaceToken(thisToken)?.types.includes('command-separator'))
+    ) {
+      return thisToken;
+    }
+  }
+  return null;
+}
+
 const funcsWithFieldArgs = [
   'avg',
   'count',
@@ -439,3 +379,8 @@ function isInsideFunctionParenthesis(curToken: Token): boolean {
   }
   return false;
 }
+
+function isAfterKeyword(keyword: string, token: Token): boolean {
+  const prevToken = prevNonWhitespaceToken(token);
+  return prevToken?.types.includes('keyword') && prevToken?.content.toLowerCase() === 'by';
+}