[8.x] [ES|QL] separate `KEEP`, `DROP`, and `SORT` autocomplete routines (#197744)

packages/kbn-esql-editor/src/esql_editor.tsx

@@ -336,7 +336,7 @@ export const ESQLEditor = memo(function ESQLEditor({
         const sources = await memoizedSources(dataViews, core).result;
         return sources;
       },
-      getFieldsFor: async ({ query: queryToExecute }: { query?: string } | undefined = {}) => {
+      getColumnsFor: async ({ query: queryToExecute }: { query?: string } | undefined = {}) => {
         if (queryToExecute) {
           // ES|QL with limit 0 returns only the columns and is more performant
           const esqlQuery = {

packages/kbn-esql-validation-autocomplete/src/__tests__/helpers.ts

@@ -56,7 +56,7 @@ export const policies = [
 
 export function getCallbackMocks() {
   return {
-    getFieldsFor: jest.fn(async ({ query }) => {
+    getColumnsFor: jest.fn(async ({ query }) => {
       if (/enrich/.test(query)) {
         return enrichFields;
       }

packages/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.suggest.test.ts

@@ -42,6 +42,6 @@ describe('autocomplete.suggest', () => {
     await suggest('sHoW ?');
     await suggest('row ? |');
 
-    expect(callbacks.getFieldsFor.mock.calls.length).toBe(0);
+    expect(callbacks.getColumnsFor.mock.calls.length).toBe(0);
   });
 });

packages/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts

@@ -244,7 +244,17 @@ export function getDateLiteralsByFieldType(_requestedType: FieldType | FieldType
 }
 
 export function createCustomCallbackMocks(
-  customFields?: ESQLRealField[],
+  /**
+   * Columns that will come from Elasticsearch since the last command
+   * e.g. the test case may be `FROM index | EVAL foo = 1 | KEEP /`
+   *
+   * In this case, the columns available for the KEEP command will be the ones
+   * that were available after the EVAL command
+   *
+   * `FROM index | EVAL foo = 1 | LIMIT 0` will be used to fetch columns. The response
+   * will include "foo" as a column.
+   */
+  customColumnsSinceLastCommand?: ESQLRealField[],
   customSources?: Array<{ name: string; hidden: boolean }>,
   customPolicies?: Array<{
     name: string;
@@ -253,11 +263,11 @@ export function createCustomCallbackMocks(
     enrichFields: string[];
   }>
 ) {
-  const finalFields = customFields || fields;
+  const finalColumnsSinceLastCommand = customColumnsSinceLastCommand || fields;
   const finalSources = customSources || indexes;
   const finalPolicies = customPolicies || policies;
   return {
-    getFieldsFor: jest.fn(async () => finalFields),
+    getColumnsFor: jest.fn(async () => finalColumnsSinceLastCommand),
     getSources: jest.fn(async () => finalSources),
     getPolicies: jest.fn(async () => finalPolicies),
   };

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

@@ -31,6 +31,8 @@ import {
   TIME_PICKER_SUGGESTION,
   setup,
   attachTriggerCommand,
+  SuggestOptions,
+  fields,
 } from './__tests__/helpers';
 import { METADATA_FIELDS } from '../shared/constants';
 import { ESQL_COMMON_NUMERIC_TYPES, ESQL_STRING_TYPES } from '../shared/esql_types';
@@ -385,24 +387,56 @@ describe('autocomplete', () => {
           '```````round(doubleField) + 1```` + 1`` + 1`',
           '```````````````round(doubleField) + 1```````` + 1```` + 1`` + 1`',
           '```````````````````````````````round(doubleField) + 1```````````````` + 1```````` + 1```` + 1`` + 1`',
+        ],
+        undefined,
+        [
+          [
+            ...fields,
+            // the following non-field columns will come over the wire as part of the response
+            {
+              name: 'round(doubleField) + 1',
+              type: 'double',
+            },
+            {
+              name: '`round(doubleField) + 1` + 1',
+              type: 'double',
+            },
+            {
+              name: '```round(doubleField) + 1`` + 1` + 1',
+              type: 'double',
+            },
+            {
+              name: '```````round(doubleField) + 1```` + 1`` + 1` + 1',
+              type: 'double',
+            },
+            {
+              name: '```````````````round(doubleField) + 1```````` + 1```` + 1`` + 1` + 1',
+              type: 'double',
+            },
+          ],
         ]
       );
 
       it('should not suggest already-used fields and variables', async () => {
         const { suggest: suggestTest } = await setup();
-        const getSuggestions = async (query: string) =>
-          (await suggestTest(query)).map((value) => value.text);
+        const getSuggestions = async (query: string, opts?: SuggestOptions) =>
+          (await suggestTest(query, opts)).map((value) => value.text);
 
-        expect(await getSuggestions('from a_index | EVAL foo = 1 | KEEP /')).toContain('foo');
-        expect(await getSuggestions('from a_index | EVAL foo = 1 | KEEP foo, /')).not.toContain(
-          'foo'
-        );
-        expect(await getSuggestions('from a_index | EVAL foo = 1 | KEEP /')).toContain(
+        expect(
+          await getSuggestions('from a_index | EVAL foo = 1 | KEEP /', {
+            callbacks: { getColumnsFor: () => [...fields, { name: 'foo', type: 'integer' }] },
+          })
+        ).toContain('foo');
+        expect(
+          await getSuggestions('from a_index | EVAL foo = 1 | KEEP foo, /', {
+            callbacks: { getColumnsFor: () => [...fields, { name: 'foo', type: 'integer' }] },
+          })
+        ).not.toContain('foo');
+
+        expect(await getSuggestions('from a_index | KEEP /')).toContain('doubleField');
+        expect(await getSuggestions('from a_index | KEEP doubleField, /')).not.toContain(
           'doubleField'
         );
-        expect(
-          await getSuggestions('from a_index | EVAL foo = 1 | KEEP doubleField, /')
-        ).not.toContain('doubleField');
       });
     });
   }
@@ -504,7 +538,7 @@ describe('autocomplete', () => {
   });
 
   describe('callbacks', () => {
-    it('should send the fields query without the last command', async () => {
+    it('should send the columns query without the last command', async () => {
       const callbackMocks = createCustomCallbackMocks(undefined, undefined, undefined);
       const statement = 'from a | drop keywordField | eval var0 = abs(doubleField) ';
       const triggerOffset = statement.lastIndexOf(' ');
@@ -516,7 +550,7 @@ describe('autocomplete', () => {
         async (text) => (text ? getAstAndSyntaxErrors(text) : { ast: [], errors: [] }),
         callbackMocks
       );
-      expect(callbackMocks.getFieldsFor).toHaveBeenCalledWith({
+      expect(callbackMocks.getColumnsFor).toHaveBeenCalledWith({
         query: 'from a | drop keywordField',
       });
     });
@@ -532,7 +566,7 @@ describe('autocomplete', () => {
         async (text) => (text ? getAstAndSyntaxErrors(text) : { ast: [], errors: [] }),
         callbackMocks
       );
-      expect(callbackMocks.getFieldsFor).toHaveBeenCalledWith({ query: 'from a' });
+      expect(callbackMocks.getColumnsFor).toHaveBeenCalledWith({ query: 'from a' });
     });
   });