sharing expression suggestion routine

Author: drewdaemon

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/commands/eval/index.ts

@@ -7,159 +7,17 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { ESQLSource } from '@kbn/esql-ast';
-import {
-  findFinalWord,
-  findPreviousWord,
-  isSingleItem,
-  unescapeColumnName,
-} from '../../../shared/helpers';
+import type { ESQLSingleAstItem } from '@kbn/esql-ast';
 import { CommandSuggestParams } from '../../../definitions/types';
 import type { SuggestionRawDefinition } from '../../types';
-import {
-  Position,
-  buildMatchingFieldsDefinition,
-  getPosition,
-  modeSuggestions,
-  noPoliciesAvailableSuggestion,
-  onSuggestion,
-  withSuggestion,
-} from './util';
-import { commaCompleteItem, pipeCompleteItem } from '../../complete_items';
-import {
-  TRIGGER_SUGGESTION_COMMAND,
-  buildFieldsDefinitions,
-  getNewVariableSuggestion,
-  getOperatorSuggestions,
-} from '../../factories';
-
-export async function suggest({
-  innerText,
-  command,
-  getPolicies,
-  getPolicyMetadata,
-  getAllColumnNames,
-  getSuggestedVariableName,
-}: CommandSuggestParams<'enrich'>): Promise<SuggestionRawDefinition[]> {
-  const pos = getPosition(innerText, command);
-
-  const policyName = (
-    command.args.find((arg) => isSingleItem(arg) && arg.type === 'source') as ESQLSource | undefined
-  )?.name;
-
-  const getFieldSuggestionsForWithClause = async () => {
-    if (!policyName) {
-      return [];
-    }
-
-    const policyMetadata = await getPolicyMetadata(policyName);
-    if (!policyMetadata) {
-      return [];
-    }
-
-    const fieldSuggestions = buildFieldsDefinitions(policyMetadata.enrichFields, false);
-
-    const lastWord = findFinalWord(innerText);
-    if (lastWord) {
-      // ENRICH ... WITH a <suggest>
-      const rangeToReplace = {
-        start: innerText.length - lastWord.length + 1,
-        end: innerText.length + 1,
-      };
-      fieldSuggestions.forEach((s) => {
-        s.rangeToReplace = rangeToReplace;
-      });
-    }
-
-    return fieldSuggestions;
-  };
-
-  switch (pos) {
-    case Position.MODE:
-      return modeSuggestions;
-
-    case Position.POLICY: {
-      const policies = await getPolicies();
-      const lastWord = findFinalWord(innerText);
-      if (lastWord !== '') {
-        policies.forEach((policySuggestion) => {
-          policySuggestion.rangeToReplace = {
-            start: innerText.length - lastWord.length + 1,
-            end: innerText.length + 1,
-          };
-        });
-      }
-      return policies.length ? policies : [noPoliciesAvailableSuggestion];
-    }
-
-    case Position.AFTER_POLICY:
-      return [onSuggestion, withSuggestion, pipeCompleteItem];
-
-    case Position.MATCH_FIELD: {
-      if (!policyName) {
-        return [];
-      }
-
-      const policyMetadata = await getPolicyMetadata(policyName);
-      if (!policyMetadata) {
-        return [];
-      }
-
-      return buildMatchingFieldsDefinition(policyMetadata.matchField, getAllColumnNames());
-    }
-
-    case Position.AFTER_ON_CLAUSE:
-      return [withSuggestion, pipeCompleteItem];
-
-    case Position.WITH_NEW_CLAUSE: {
-      if (!policyName) {
-        return [];
-      }
-
-      const policyMetadata = await getPolicyMetadata(policyName);
-      if (!policyMetadata) {
-        return [];
-      }
-
-      const suggestions: SuggestionRawDefinition[] = [];
-      suggestions.push(
-        getNewVariableSuggestion(getSuggestedVariableName(policyMetadata.enrichFields))
-      );
-      suggestions.push(...(await getFieldSuggestionsForWithClause()));
-      return suggestions;
-    }
-
-    case Position.WITH_AFTER_FIRST_WORD: {
-      if (!policyName) {
-        return [];
-      }
-      const policyMetadata = await getPolicyMetadata(policyName);
-
-      if (!policyMetadata) {
-        return [];
-      }
-
-      const word = findPreviousWord(innerText);
-      if (policyMetadata.enrichFields.includes(unescapeColumnName(word))) {
-        // complete field name
-        return [pipeCompleteItem, { ...commaCompleteItem, command: TRIGGER_SUGGESTION_COMMAND }];
-      } else {
-        // not recognized as a field name, assume new user-defined column name
-        return getOperatorSuggestions({ command: 'enrich' });
-      }
-    }
-
-    case Position.WITH_AFTER_ASSIGNMENT: {
-      const suggestions: SuggestionRawDefinition[] = [];
-      suggestions.push(...(await getFieldSuggestionsForWithClause()));
-      return suggestions;
-    }
-
-    case Position.WITH_AFTER_COMPLETE_CLAUSE: {
-      return [pipeCompleteItem, { ...commaCompleteItem, command: TRIGGER_SUGGESTION_COMMAND }];
-    }
-
-    default:
-      return [];
-  }
+import { suggestForExpression } from '../where';
+
+export async function suggest(
+  params: CommandSuggestParams<'eval'>
+): Promise<SuggestionRawDefinition[]> {
+  const expressionRoot = params.command.args[0] as ESQLSingleAstItem | undefined;
+  return suggestForExpression({
+    ...params,
+    expressionRoot,
+  });
 }

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/commands/eval/util.ts

@@ -27,23 +27,31 @@ import {
   UNSUPPORTED_COMMANDS_BEFORE_QSTR,
 } from '../../../shared/constants';
 
-export async function suggest({
+export async function suggest(
+  params: CommandSuggestParams<'where'>
+): Promise<SuggestionRawDefinition[]> {
+  const expressionRoot = params.command.args[0] as ESQLSingleAstItem | undefined;
+  return suggestForExpression({
+    ...params,
+    expressionRoot,
+  });
+}
+
+export async function suggestForExpression({
+  expressionRoot,
   innerText,
-  command,
-  getColumnsByType,
   getExpressionType,
+  getColumnsByType,
   previousCommands,
-}: CommandSuggestParams<'where'>): Promise<SuggestionRawDefinition[]> {
+}: {
+  expressionRoot: ESQLSingleAstItem | undefined;
+} & Pick<
+  CommandSuggestParams<string>,
+  'innerText' | 'getExpressionType' | 'getColumnsByType' | 'previousCommands'
+>): Promise<SuggestionRawDefinition[]> {
   const suggestions: SuggestionRawDefinition[] = [];
 
-  /**
-   * The logic for WHERE suggestions is basically the logic for expression suggestions.
-   * I assume we will eventually extract much of this to be a shared function among WHERE and EVAL
-   * and anywhere else the user can enter a generic expression.
-   */
-  const expressionRoot = command.args[0] as ESQLSingleAstItem | undefined;
-
-  const position = getPosition(innerText, command);
+  const position = getPosition(innerText, expressionRoot);
   switch (position) {
     /**
      * After a column name

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/commands/where/index.ts

@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { ESQLCommand, ESQLSingleAstItem } from '@kbn/esql-ast';
+import { ESQLSingleAstItem } from '@kbn/esql-ast';
 import { isColumnItem, isFunctionItem } from '../../../shared/helpers';
 
 export type CaretPosition =
@@ -17,9 +17,10 @@ export type CaretPosition =
   | 'after_operator'
   | 'empty_expression';
 
-export const getPosition = (innerText: string, command: ESQLCommand): CaretPosition => {
-  const expressionRoot = command.args[0] as ESQLSingleAstItem | undefined;
-
+export const getPosition = (
+  innerText: string,
+  expressionRoot: ESQLSingleAstItem | undefined
+): CaretPosition => {
   const endsWithNot = / not$/i.test(innerText.trimEnd());
   if (
     endsWithNot &&