[ES|QL] Enables enrich suggestion in fork (#224992)

## Summary

Fixing this
https://github.com/elastic/kibana/issues/192255#issuecomment-2991343277

I hadn't enabled the enrich suggestion as it has some client side
validation issues. This PR is fixing them and enables the enrich
suggestions in FORK

<img width="796" alt="image"
src="https://github.com/user-attachments/assets/0b599e36-18d3-4504-b572-50575bb9e159"
/>



### Checklist

- [ ] [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

---------

Co-authored-by: Vadim Kibana <82822460+vadimkibana@users.noreply.github.com>

src/platform/packages/shared/kbn-esql-ast/index.ts

@@ -52,6 +52,7 @@ export { Builder, type AstNodeParserFields, type AstNodeTemplate } from './src/b
 export {
   createParser,
   parse,
+  Parser,
   parseErrors,
   type ParseOptions,
   type ParseResult,

src/platform/packages/shared/kbn-esql-ast/src/parser/index.ts

@@ -7,6 +7,13 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-export { createParser, parse, parseErrors, type ParseOptions, type ParseResult } from './parser';
+export {
+  createParser,
+  parse,
+  parseErrors,
+  Parser,
+  type ParseOptions,
+  type ParseResult,
+} from './parser';
 
 export { ESQLErrorListener } from './esql_error_listener';

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

@@ -23,6 +23,7 @@ import {
   getFunctionSignaturesByReturnType,
   setup,
   lookupIndexFields,
+  policies,
 } from './helpers';
 
 describe('autocomplete.suggest', () => {
@@ -62,6 +63,7 @@ describe('autocomplete.suggest', () => {
           'COMPLETION ',
           'MV_EXPAND ',
           'DROP ',
+          'ENRICH ',
           'KEEP ',
           'RENAME ',
           'SAMPLE ',
@@ -187,6 +189,18 @@ describe('autocomplete.suggest', () => {
             expect(labels).toEqual(expected);
           });
 
+          test('enrich', async () => {
+            const expectedPolicyNameSuggestions = policies
+              .map(({ name, suggestedAs }) => suggestedAs || name)
+              .map((name) => `${name} `);
+
+            await assertSuggestions(`FROM a | FORK (ENRICH /)`, expectedPolicyNameSuggestions);
+            await assertSuggestions(
+              `FROM a | FORK (ENRICH policy ON /)`,
+              getFieldNamesByType('any').map((v) => `${v} `)
+            );
+          });
+
           describe('stats', () => {
             it('suggests for empty expression', async () => {
               await assertSuggestions('FROM a | FORK (STATS /)', EXPECTED_FOR_EMPTY_EXPRESSION);

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

@@ -36,7 +36,7 @@ const FORK_AVAILABLE_COMMANDS = [
   'rename',
   'sample',
   'join',
-  // 'enrich', // not suggesting enrich for now, there are client side validation issues
+  'enrich',
 ];
 
 export async function suggest(

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

@@ -0,0 +1,34 @@
+/*
+ * 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 { Parser } from '@kbn/esql-ast';
+import { getEnrichCommands } from './helpers';
+
+describe('getEnrichCommands', () => {
+  test('should return the command in case of an enrich command', async () => {
+    const { root } = Parser.parse('FROM index | ENRICH policy ON field');
+    expect(getEnrichCommands(root.commands).length).toBe(1);
+  });
+
+  test('should return empty array if enrich is not present', async () => {
+    const { root } = Parser.parse('FROM index | STATS COUNT() BY field');
+    expect(getEnrichCommands(root.commands)).toStrictEqual([]);
+  });
+
+  test('should return the command in case of an enrich command inside a fork branch', async () => {
+    const { root } = Parser.parse(
+      'FROM index | FORK (ENRICH policy ON @timestamp WITH col0 = bikes_count) (DROP @timestamp) '
+    );
+    expect(getEnrichCommands(root.commands).length).toBe(1);
+  });
+
+  test('should return empty array in case of forck branches without enrich', async () => {
+    const { root } = Parser.parse('FROM index | FORK (STATS COUNT()) (DROP @timestamp) ');
+    expect(getEnrichCommands(root.commands)).toStrictEqual([]);
+  });
+});

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

@@ -7,15 +7,17 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import type {
+import {
-  ESQLAst,
+  type ESQLAst,
-  ESQLAstItem,
+  type ESQLAstItem,
-  ESQLAstTimeseriesCommand,
+  type ESQLAstTimeseriesCommand,
-  ESQLAstQueryExpression,
+  type ESQLAstQueryExpression,
-  ESQLColumn,
+  type ESQLColumn,
-  ESQLMessage,
+  type ESQLMessage,
-  ESQLSingleAstItem,
+  type ESQLSingleAstItem,
-  ESQLSource,
+  type ESQLSource,
+  type ESQLCommand,
+  Walker,
 } from '@kbn/esql-ast';
 import { mutate, synth } from '@kbn/esql-ast';
 import { FunctionDefinition } from '../definitions/types';
@@ -140,3 +142,12 @@ export function collapseWrongArgumentTypeMessages(
 
   return messages;
 }
+
+/**
+ * Collects all 'enrich' commands from a list of ESQL commands.
+ * @param commands - The list of ESQL commands to search through.
+ * This function traverses the provided ESQL commands and collects all commands with the name 'enrich'.
+ * @returns {ESQLCommand[]} - An array of ESQLCommand objects that represent the 'enrich' commands found in the input.
+ */
+export const getEnrichCommands = (commands: ESQLCommand[]): ESQLCommand[] =>
+  Walker.matchAll(commands, { type: 'command', name: 'enrich' }) as ESQLCommand[];

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

@@ -19,6 +19,7 @@ import {
   buildQueryForFieldsForStringSources,
   buildQueryForFieldsFromSource,
   buildQueryForFieldsInPolicies,
+  getEnrichCommands,
 } from './helpers';
 import type { ESQLFieldWithMetadata, ESQLPolicy } from './types';
 
@@ -52,7 +53,8 @@ export async function retrievePolicies(
   commands: ESQLCommand[],
   callbacks?: ESQLCallbacks
 ): Promise<Map<string, ESQLPolicy>> {
-  if (!callbacks || commands.every(({ name }) => name !== 'enrich')) {
+  const enrichCommands = getEnrichCommands(commands);
+  if (!callbacks || !enrichCommands.length) {
     return new Map();
   }
 
@@ -82,7 +84,7 @@ export async function retrievePoliciesFields(
   if (!callbacks) {
     return new Map();
   }
-  const enrichCommands = commands.filter(({ name }) => name === 'enrich');
+  const enrichCommands = getEnrichCommands(commands);
   if (!enrichCommands.length) {
     return new Map();
   }