[ES|QL] Separate LIMIT and MV_EXPAND autocomplete routines (#213500)

## Summary

Part of https://github.com/elastic/kibana/issues/195418

Gives `LIMIT` and `MV_EXPAND` autocomplete logic its own home 🏡

### Checklist

- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

### Identify risks

- [ ] As with any refactor, there's a possibility this will introduce a
regression in the behavior of commands. However, all automated tests are
passing and I have tested the behavior manually and can detect no
regression.

---------

Co-authored-by: Stratoula Kalafateli <efstratia.kalafateli@elastic.co>

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.limit.test.ts

@@ -0,0 +1,25 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { setup } from './helpers';
+
+describe('autocomplete.suggest', () => {
+  describe('LIMIT <number>', () => {
+    it('suggests numbers', async () => {
+      const { assertSuggestions } = await setup();
+      assertSuggestions('from a | limit /', ['10 ', '100 ', '1000 ']);
+      assertSuggestions('from a | limit /', ['10 ', '100 ', '1000 '], { triggerCharacter: ' ' });
+    });
+
+    it('suggests pipe after number', async () => {
+      const { assertSuggestions } = await setup();
+      assertSuggestions('from a | limit 4 /', ['| ']);
+    });
+  });
+});

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.mv_expand.test.ts

@@ -0,0 +1,46 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { getFieldNamesByType, setup } from './helpers';
+
+describe('autocomplete.suggest', () => {
+  describe('MV_EXPAND <column>', () => {
+    it('suggests columns', async () => {
+      const { assertSuggestions } = await setup();
+      assertSuggestions(
+        'from a | mv_expand /',
+        getFieldNamesByType('any').map((name) => `${name} `)
+      );
+      assertSuggestions(
+        'from a | mv_expand /',
+        getFieldNamesByType('any').map((name) => `${name} `),
+        {
+          triggerCharacter: ' ',
+        }
+      );
+    });
+
+    it('works with field name prefixes', async () => {
+      const { assertSuggestions } = await setup();
+      assertSuggestions(
+        'from a | mv_expand key/',
+        getFieldNamesByType('any').map((name) => `${name} `)
+      );
+      assertSuggestions(
+        'from a | mv_expand keywordField/',
+        getFieldNamesByType('any').map((name) => `${name} `)
+      );
+    });
+
+    it('suggests pipe after column', async () => {
+      const { assertSuggestions } = await setup();
+      assertSuggestions('from a | mv_expand a /', ['| ']);
+    });
+  });
+});

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts

@@ -112,16 +112,6 @@ describe('autocomplete', () => {
     testSuggestions('from a metadata _id | eval var0 = a | /', commands);
   });
 
-  describe('limit', () => {
-    testSuggestions('from a | limit /', ['10 ', '100 ', '1000 ']);
-    testSuggestions('from a | limit 4 /', ['| ']);
-  });
-
-  describe('mv_expand', () => {
-    testSuggestions('from a | mv_expand /', getFieldNamesByType('any'));
-    testSuggestions('from a | mv_expand a /', ['| ']);
-  });
-
   describe('rename', () => {
     testSuggestions('from a | rename /', getFieldNamesByType('any'));
     testSuggestions('from a | rename keywordField /', ['AS $0'], ' ');
@@ -406,13 +396,13 @@ describe('autocomplete', () => {
     );
 
     // LIMIT argument
-    // Here we actually test that the invoke trigger kind does NOT work
-    // the assumption is that it isn't very useful to see literal suggestions when already typing a number
-    // I'm not sure if this is true or not, but it's the current behavior
-    testSuggestions('FROM a | LIMIT 1/', ['| ']);
+    testSuggestions('FROM a | LIMIT 1/', ['10 ', '100 ', '1000 ']);
 
     // MV_EXPAND field
-    testSuggestions('FROM index1 | MV_EXPAND f/', getFieldNamesByType('any'));
+    testSuggestions(
+      'FROM index1 | MV_EXPAND f/',
+      getFieldNamesByType('any').map((name) => `${name} `)
+    );
 
     // RENAME field
     testSuggestions('FROM index1 | RENAME f/', getFieldNamesByType('any'));

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

@@ -694,7 +694,7 @@ async function getExpressionSuggestionsByType(
           );
           if (isNumericType(nodeArgType) && isLiteralItem(rightArg)) {
             // ... EVAL var = 1 <suggest>
-            suggestions.push(...getCompatibleLiterals(command.name, ['time_literal_unit']));
+            suggestions.push(...getCompatibleLiterals(['time_literal_unit']));
           }
           if (isFunctionItem(rightArg)) {
             if (rightArg.args.some(isTimeIntervalItem)) {
@@ -702,7 +702,7 @@ async function getExpressionSuggestionsByType(
               const lastFnArgType = extractTypeFromASTArg(lastFnArg, references);
               if (isNumericType(lastFnArgType) && isLiteralItem(lastFnArg))
                 // ... EVAL var = 1 year + 2 <suggest>
-                suggestions.push(...getCompatibleLiterals(command.name, ['time_literal_unit']));
+                suggestions.push(...getCompatibleLiterals(['time_literal_unit']));
             }
           }
         } else {
@@ -738,7 +738,7 @@ async function getExpressionSuggestionsByType(
                 const lastFnArgType = extractTypeFromASTArg(lastFnArg, references);
                 if (isNumericType(lastFnArgType) && isLiteralItem(lastFnArg))
                   // ... EVAL var = 1 year + 2 <suggest>
-                  suggestions.push(...getCompatibleLiterals(command.name, ['time_literal_unit']));
+                  suggestions.push(...getCompatibleLiterals(['time_literal_unit']));
               }
             }
           }
@@ -763,7 +763,7 @@ async function getExpressionSuggestionsByType(
       // it can be just literal values (i.e. "string")
       if (argDef.constantOnly) {
         // ... | <COMMAND> ... <suggest>
-        suggestions.push(...getCompatibleLiterals(command.name, [argDef.type]));
+        suggestions.push(...getCompatibleLiterals([argDef.type]));
       } else {
         // or it can be anything else as long as it is of the right type and the end (i.e. column or function)
         if (!nodeArg) {
@@ -1031,7 +1031,6 @@ async function getFunctionArgsSuggestions(
     // Literals
     suggestions.push(
       ...getCompatibleLiterals(
-        command.name,
         getTypesFromParamDefs(constantOnlyParamDefs) as string[],
         {
           addComma: shouldAddComma,
@@ -1105,7 +1104,7 @@ async function getFunctionArgsSuggestions(
       if (isLiteralItem(arg) && isNumericType(arg.literalType)) {
         // ... | EVAL fn(2 <suggest>)
         suggestions.push(
-          ...getCompatibleLiterals(command.name, ['time_literal_unit'], {
+          ...getCompatibleLiterals(['time_literal_unit'], {
             addComma: shouldAddComma,
             advanceCursorAndOpenSuggestions: hasMoreMandatoryArgs,
           })

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

@@ -0,0 +1,23 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { CommandSuggestParams } from '../../../definitions/types';
+import type { SuggestionRawDefinition } from '../../types';
+import { buildConstantsDefinitions } from '../../factories';
+import { pipeCompleteItem } from '../../complete_items';
+
+export function suggest({ innerText }: CommandSuggestParams<'limit'>): SuggestionRawDefinition[] {
+  if (/[0-9]\s+$/.test(innerText)) {
+    return [pipeCompleteItem];
+  }
+
+  return buildConstantsDefinitions(['10', '100', '1000'], '', undefined, {
+    advanceCursorAndOpenSuggestions: true,
+  });
+}

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

@@ -0,0 +1,37 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { findFinalWord } from '../../../shared/helpers';
+import { CommandSuggestParams } from '../../../definitions/types';
+import type { SuggestionRawDefinition } from '../../types';
+import { pipeCompleteItem } from '../../complete_items';
+
+export async function suggest({
+  innerText,
+  getColumnsByType,
+}: CommandSuggestParams<'limit'>): Promise<SuggestionRawDefinition[]> {
+  if (/MV_EXPAND\s+\S+\s+$/i.test(innerText)) {
+    return [pipeCompleteItem];
+  }
+
+  const columnSuggestions = await getColumnsByType('any', undefined, {
+    advanceCursor: true,
+    openSuggestions: true,
+  });
+
+  const fragment = findFinalWord(innerText);
+  columnSuggestions.forEach((suggestion) => {
+    suggestion.rangeToReplace = {
+      start: innerText.length - fragment.length + 1,
+      end: innerText.length,
+    };
+  });
+
+  return columnSuggestions;
+}

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/factories.ts

@@ -26,7 +26,6 @@ import { shouldBeQuotedSource, shouldBeQuotedText } from '../shared/helpers';
 import { buildFunctionDocumentation } from './documentation_util';
 import { DOUBLE_BACKTICK, SINGLE_TICK_REGEX } from '../shared/constants';
 import { ESQLRealField } from '../validation/types';
-import { isNumericType } from '../shared/esql_types';
 import { getTestFunctions } from '../shared/test_functions';
 import { operatorsDefinitions } from '../definitions/all_operators';
 
@@ -411,7 +410,6 @@ export function getUnitDuration(unit: number = 1) {
  * definition property...
  */
 export function getCompatibleLiterals(
-  commandName: string,
   types: string[],
   options?: {
     advanceCursorAndOpenSuggestions?: boolean;
@@ -421,16 +419,6 @@ export function getCompatibleLiterals(
   getVariables?: () => ESQLControlVariable[] | undefined
 ) {
   const suggestions: SuggestionRawDefinition[] = [];
-  if (types.some(isNumericType)) {
-    if (commandName === 'limit') {
-      // suggest 10/100/1000 for limit
-      suggestions.push(
-        ...buildConstantsDefinitions(['10', '100', '1000'], '', undefined, {
-          advanceCursorAndOpenSuggestions: true,
-        })
-      );
-    }
-  }
   if (types.includes('time_literal')) {
     const timeLiteralSuggestions = [
       ...buildConstantsDefinitions(getUnitDuration(1), undefined, undefined, options),

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts

@@ -449,7 +449,7 @@ export async function getFieldsOrFunctionsSuggestions(
     variables
       ? pushItUpInTheList(buildVariablesDefinitions(filteredVariablesByType), functions)
       : [],
-    literals ? getCompatibleLiterals(commandName, types) : []
+    literals ? getCompatibleLiterals(types) : []
   );
 
   return suggestions;

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

@@ -28,6 +28,8 @@ import {
 import { ENRICH_MODES } from './settings';
 
 import { type CommandDefinition } from './types';
+import { checkAggExistence, checkFunctionContent } from './commands_helpers';
+
 import { suggest as suggestForSort } from '../autocomplete/commands/sort';
 import { suggest as suggestForKeep } from '../autocomplete/commands/keep';
 import { suggest as suggestForDrop } from '../autocomplete/commands/drop';
@@ -40,7 +42,8 @@ import { suggest as suggestForShow } from '../autocomplete/commands/show';
 import { suggest as suggestForGrok } from '../autocomplete/commands/grok';
 import { suggest as suggestForDissect } from '../autocomplete/commands/dissect';
 import { suggest as suggestForEnrich } from '../autocomplete/commands/enrich';
-import { checkAggExistence, checkFunctionContent } from './commands_helpers';
+import { suggest as suggestForLimit } from '../autocomplete/commands/limit';
+import { suggest as suggestForMvExpand } from '../autocomplete/commands/mv_expand';
 
 const statsValidator = (command: ESQLCommand) => {
   const messages: ESQLMessage[] = [];
@@ -286,6 +289,7 @@ export const commandDefinitions: Array<CommandDefinition<any>> = [
     },
     options: [],
     modes: [],
+    suggest: suggestForLimit,
   },
   {
     name: 'keep',
@@ -438,6 +442,7 @@ export const commandDefinitions: Array<CommandDefinition<any>> = [
       multipleParams: false,
       params: [{ name: 'column', type: 'column', innerTypes: ['any'] }],
     },
+    suggest: suggestForMvExpand,
   },
   {
     name: 'enrich',

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/definitions/types.ts

@@ -278,7 +278,7 @@ export interface CommandSuggestParams<CommandName extends string> {
 
 export type CommandSuggestFunction<CommandName extends string> = (
   params: CommandSuggestParams<CommandName>
-) => Promise<SuggestionRawDefinition[]>;
+) => Promise<SuggestionRawDefinition[]> | SuggestionRawDefinition[];
 
 export interface CommandBaseDefinition<CommandName extends string> {
   name: CommandName;