[ES|QL] Moves fork in tech preview (#224680)

## Summary

Moves the fork command on tech preview

<img width="878" alt="image"
src="https://github.com/user-attachments/assets/d21d593d-f157-45e4-94aa-c97d5f3098ca"
/>

It also enables all the processing commands in the autocomplete except
for enrich, mostly because the client side validation complains. As
enrich is considered kinda deprecated I think it is ok to skip the
suggestion. I will create an issue to track the client side validation
problem though (Update: Added in [client side validation
bugs](https://github.com/elastic/kibana/issues/192255#issuecomment-2991343277))

### 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

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

@@ -52,14 +52,13 @@ import { createRowCommand } from './factories/row';
 import { createSortCommand } from './factories/sort';
 import { createStatsCommand } from './factories/stats';
 import { createWhereCommand } from './factories/where';
+import { createMvExpandCommand } from './factories/mv_expand';
+import { createKeepCommand } from './factories/keep';
+import { createDropCommand } from './factories/drop';
+import { createRenameCommand } from './factories/rename';
+import { createSampleCommand } from './factories/sample';
 import { getPosition } from './helpers';
-import {
-  collectAllAggFields,
-  collectAllColumnIdentifiers,
-  getConstant,
-  visitByOption,
-  visitRenameClauses,
-} from './walkers';
+import { collectAllAggFields, visitByOption } from './walkers';
 import { createTimeseriesCommand } from './factories/timeseries';
 import { createRerankCommand } from './factories/rerank';
 import { createEnrichCommand } from './factories/enrich';
@@ -216,9 +215,11 @@ export class ESQLAstBuilderListener implements ESQLParserListener {
    * @param ctx the parse tree
    */
   exitKeepCommand(ctx: KeepCommandContext) {
-    const command = createCommand('keep', ctx);
+    if (this.inFork) {
+      return;
+    }
+    const command = createKeepCommand(ctx);
     this.ast.push(command);
-    command.args.push(...collectAllColumnIdentifiers(ctx));
   }
 
   /**
@@ -226,9 +227,11 @@ export class ESQLAstBuilderListener implements ESQLParserListener {
    * @param ctx the parse tree
    */
   exitDropCommand(ctx: DropCommandContext) {
-    const command = createCommand('drop', ctx);
+    if (this.inFork) {
+      return;
+    }
+    const command = createDropCommand(ctx);
     this.ast.push(command);
-    command.args.push(...collectAllColumnIdentifiers(ctx));
   }
 
   /**
@@ -236,9 +239,11 @@ export class ESQLAstBuilderListener implements ESQLParserListener {
    * @param ctx the parse tree
    */
   exitRenameCommand(ctx: RenameCommandContext) {
-    const command = createCommand('rename', ctx);
+    if (this.inFork) {
+      return;
+    }
+    const command = createRenameCommand(ctx);
     this.ast.push(command);
-    command.args.push(...visitRenameClauses(ctx.renameClause_list()));
   }
 
   /**
@@ -270,9 +275,11 @@ export class ESQLAstBuilderListener implements ESQLParserListener {
    * @param ctx the parse tree
    */
   exitMvExpandCommand(ctx: MvExpandCommandContext) {
-    const command = createCommand('mv_expand', ctx);
+    if (this.inFork) {
+      return;
+    }
+    const command = createMvExpandCommand(ctx);
     this.ast.push(command);
-    command.args.push(...collectAllColumnIdentifiers(ctx));
   }
 
   /**
@@ -289,6 +296,9 @@ export class ESQLAstBuilderListener implements ESQLParserListener {
    * @param ctx the parse tree
    */
   exitEnrichCommand(ctx: EnrichCommandContext) {
+    if (this.inFork) {
+      return;
+    }
     const command = createEnrichCommand(ctx);
 
     this.ast.push(command);
@@ -306,6 +316,9 @@ export class ESQLAstBuilderListener implements ESQLParserListener {
    * @param ctx the parse tree
    */
   exitJoinCommand(ctx: JoinCommandContext): void {
+    if (this.inFork) {
+      return;
+    }
     const command = createJoinCommand(ctx);
 
     this.ast.push(command);
@@ -378,15 +391,11 @@ export class ESQLAstBuilderListener implements ESQLParserListener {
   }
 
   exitSampleCommand(ctx: SampleCommandContext): void {
-    const command = createCommand('sample', ctx);
-    this.ast.push(command);
-
-    if (ctx.constant()) {
-      const probability = getConstant(ctx.constant());
-      if (probability != null) {
-        command.args.push(probability);
-      }
+    if (this.inFork) {
+      return;
     }
+    const command = createSampleCommand(ctx);
+    this.ast.push(command);
   }
 
   /**

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

@@ -0,0 +1,22 @@
+/*
+ * 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 { DropCommandContext } from '../../antlr/esql_parser';
+import { ESQLCommand } from '../../types';
+import { createCommand } from '../factories';
+import { collectAllColumnIdentifiers } from '../walkers';
+
+export const createDropCommand = (ctx: DropCommandContext): ESQLCommand<'drop'> => {
+  const command = createCommand('drop', ctx);
+  const identifiers = collectAllColumnIdentifiers(ctx);
+
+  command.args.push(...identifiers);
+
+  return command;
+};

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

@@ -26,6 +26,13 @@ import { createWhereCommand } from './where';
 import { createCompletionCommand } from './completion';
 import { createChangePointCommand } from './change_point';
 import { createGrokCommand } from './grok';
+import { createKeepCommand } from './keep';
+import { createMvExpandCommand } from './mv_expand';
+import { createDropCommand } from './drop';
+import { createRenameCommand } from './rename';
+import { createEnrichCommand } from './enrich';
+import { createSampleCommand } from './sample';
+import { createJoinCommand } from './join';
 
 export const createForkCommand = (ctx: ForkCommandContext): ESQLCommand<'fork'> => {
   const command = createCommand<'fork'>('fork', ctx);
@@ -113,4 +120,39 @@ function visitForkSubQueryProcessingCommandContext(ctx: ForkSubQueryProcessingCo
   if (completionCtx) {
     return createCompletionCommand(completionCtx);
   }
+
+  const mvExpandCtx = ctx.processingCommand().mvExpandCommand();
+  if (mvExpandCtx) {
+    return createMvExpandCommand(mvExpandCtx);
+  }
+
+  const keepCtx = ctx.processingCommand().keepCommand();
+  if (keepCtx) {
+    return createKeepCommand(keepCtx);
+  }
+
+  const dropCtx = ctx.processingCommand().dropCommand();
+  if (dropCtx) {
+    return createDropCommand(dropCtx);
+  }
+
+  const renameCtx = ctx.processingCommand().renameCommand();
+  if (renameCtx) {
+    return createRenameCommand(renameCtx);
+  }
+
+  const enrichCtx = ctx.processingCommand().enrichCommand();
+  if (enrichCtx) {
+    return createEnrichCommand(enrichCtx);
+  }
+
+  const sampleCtx = ctx.processingCommand().sampleCommand();
+  if (sampleCtx) {
+    return createSampleCommand(sampleCtx);
+  }
+
+  const joinCtx = ctx.processingCommand().joinCommand();
+  if (joinCtx) {
+    return createJoinCommand(joinCtx);
+  }
 }

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

@@ -0,0 +1,22 @@
+/*
+ * 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 { KeepCommandContext } from '../../antlr/esql_parser';
+import { ESQLCommand } from '../../types';
+import { createCommand } from '../factories';
+import { collectAllColumnIdentifiers } from '../walkers';
+
+export const createKeepCommand = (ctx: KeepCommandContext): ESQLCommand<'keep'> => {
+  const command = createCommand('keep', ctx);
+  const identifiers = collectAllColumnIdentifiers(ctx);
+
+  command.args.push(...identifiers);
+
+  return command;
+};

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

@@ -0,0 +1,22 @@
+/*
+ * 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 { MvExpandCommandContext } from '../../antlr/esql_parser';
+import { ESQLCommand } from '../../types';
+import { createCommand } from '../factories';
+import { collectAllColumnIdentifiers } from '../walkers';
+
+export const createMvExpandCommand = (ctx: MvExpandCommandContext): ESQLCommand<'mv_expand'> => {
+  const command = createCommand('mv_expand', ctx);
+  const identifiers = collectAllColumnIdentifiers(ctx);
+
+  command.args.push(...identifiers);
+
+  return command;
+};

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

@@ -0,0 +1,22 @@
+/*
+ * 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 { RenameCommandContext } from '../../antlr/esql_parser';
+import { ESQLCommand } from '../../types';
+import { createCommand } from '../factories';
+import { visitRenameClauses } from '../walkers';
+
+export const createRenameCommand = (ctx: RenameCommandContext): ESQLCommand<'rename'> => {
+  const command = createCommand('rename', ctx);
+  const renameArgs = visitRenameClauses(ctx.renameClause_list());
+
+  command.args.push(...renameArgs);
+
+  return command;
+};

src/platform/packages/shared/kbn-esql-ast/src/parser/factories/sample.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 { SampleCommandContext } from '../../antlr/esql_parser';
+import { ESQLCommand } from '../../types';
+import { createCommand } from '../factories';
+import { getConstant } from '../walkers';
+
+export const createSampleCommand = (ctx: SampleCommandContext): ESQLCommand<'sample'> => {
+  const command = createCommand('sample', ctx);
+  if (ctx.constant()) {
+    const probability = getConstant(ctx.constant());
+    if (probability != null) {
+      command.args.push(probability);
+    }
+  }
+
+  return command;
+};

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

@@ -8,7 +8,7 @@
  */
 
 import { Location } from '../../definitions/types';
-import { ESQL_STRING_TYPES } from '../../shared/esql_types';
+import { ESQL_STRING_TYPES, ESQL_NUMBER_TYPES } from '../../shared/esql_types';
 import { EXPECTED_FIELD_AND_FUNCTION_SUGGESTIONS } from './autocomplete.command.sort.test';
 import { AVG_TYPES, EXPECTED_FOR_EMPTY_EXPRESSION } from './autocomplete.command.stats.test';
 import {
@@ -18,9 +18,11 @@ import {
 import {
   AssertSuggestionsFn,
   SuggestFn,
+  attachTriggerCommand,
   getFieldNamesByType,
   getFunctionSignaturesByReturnType,
   setup,
+  lookupIndexFields,
 } from './helpers';
 
 describe('autocomplete.suggest', () => {
@@ -56,6 +58,13 @@ describe('autocomplete.suggest', () => {
           'STATS ',
           'EVAL ',
           'GROK ',
+          'CHANGE_POINT ',
+          'MV_EXPAND ',
+          'DROP ',
+          'KEEP ',
+          'RENAME ',
+          'SAMPLE ',
+          'LOOKUP JOIN ',
         ];
 
         it('suggests FORK sub commands in an open branch', async () => {
@@ -108,6 +117,75 @@ describe('autocomplete.suggest', () => {
             ]);
           });
 
+          test('keep', async () => {
+            await assertSuggestions('FROM a | FORK (KEEP /)', getFieldNamesByType('any'));
+            await assertSuggestions('FROM a | FORK (KEEP integerField /)', [',', '| ']);
+          });
+
+          test('drop', async () => {
+            await assertSuggestions('FROM a | FORK (DROP /)', getFieldNamesByType('any'));
+            await assertSuggestions('FROM a | FORK (DROP integerField /)', [',', '| ']);
+          });
+
+          test('mv_expand', async () => {
+            await assertSuggestions(
+              'FROM a | FORK (MV_EXPAND /)',
+              getFieldNamesByType('any').map((name) => `${name} `)
+            );
+            await assertSuggestions('FROM a | FORK (MV_EXPAND integerField /)', ['| ']);
+          });
+
+          test('sample', async () => {
+            await assertSuggestions('FROM a | FORK (SAMPLE /)', ['.001 ', '.01 ', '.1 ']);
+            await assertSuggestions('FROM a | FORK (SAMPLE 0.01 /)', ['| ']);
+          });
+
+          test('rename', async () => {
+            await assertSuggestions('FROM a | FORK (RENAME /)', [
+              'col0 = ',
+              ...getFieldNamesByType('any').map((field) => field + ' '),
+            ]);
+            await assertSuggestions('FROM a | FORK (RENAME textField /)', ['AS ']);
+            await assertSuggestions('FROM a | FORK (RENAME field /)', ['= ']);
+          });
+
+          test('change_point', async () => {
+            await assertSuggestions(
+              `FROM a | FORK (CHANGE_POINT /`,
+              getFieldNamesByType(ESQL_NUMBER_TYPES).map((v) => `${v} `)
+            );
+            await assertSuggestions(
+              `FROM a | FORK (CHANGE_POINT value /)`,
+              ['ON ', 'AS ', '| '].map(attachTriggerCommand)
+            );
+            await assertSuggestions(
+              `FROM a | FORK (CHANGE_POINT value on /)`,
+              getFieldNamesByType('any').map((v) => `${v} `)
+            );
+          });
+
+          test('lookup join', async () => {
+            await assertSuggestions('FROM a | FORK (LOOKUP JOIN /)', [
+              'join_index ',
+              'join_index_with_alias ',
+              'lookup_index ',
+              'join_index_alias_1 $0',
+              'join_index_alias_2 $0',
+            ]);
+            const suggestions = await suggest('FROM a | FORK (LOOKUP JOIN join_index ON /)');
+            const labels = suggestions.map((s) => s.text.trim()).sort();
+            const expected = getFieldNamesByType('any')
+              .sort()
+              .map((field) => field.trim());
+
+            for (const { name } of lookupIndexFields) {
+              expected.push(name.trim());
+            }
+            expected.sort();
+
+            expect(labels).toEqual(expected);
+          });
+
           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/drop/suggest.ts

@@ -8,12 +8,7 @@
  */
 
 import { CommandSuggestParams } from '../../../definitions/types';
-import {
-  findPreviousWord,
-  getLastNonWhitespaceChar,
-  isColumnItem,
-  noCaseCompare,
-} from '../../../shared/helpers';
+import { getLastNonWhitespaceChar, isColumnItem } from '../../../shared/helpers';
 import type { SuggestionRawDefinition } from '../../types';
 import { commaCompleteItem, pipeCompleteItem } from '../../complete_items';
 import { handleFragment } from '../../helper';
@@ -28,7 +23,7 @@ export async function suggest({
   if (
     /\s/.test(innerText[innerText.length - 1]) &&
     getLastNonWhitespaceChar(innerText) !== ',' &&
-    !noCaseCompare(findPreviousWord(innerText), 'drop')
+    !/drop\s+\S*$/i.test(innerText)
   ) {
     return [pipeCompleteItem, commaCompleteItem];
   }

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

@@ -29,6 +29,14 @@ const FORK_AVAILABLE_COMMANDS = [
   'eval',
   'completion',
   'grok',
+  'change_point',
+  'mv_expand',
+  'keep',
+  'drop',
+  'rename',
+  'sample',
+  'join',
+  // 'enrich', // not suggesting enrich for now, there are client side validation issues
 ];
 
 export async function suggest(

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

@@ -8,12 +8,7 @@
  */
 
 import { CommandSuggestParams } from '../../../definitions/types';
-import {
-  findPreviousWord,
-  getLastNonWhitespaceChar,
-  isColumnItem,
-  noCaseCompare,
-} from '../../../shared/helpers';
+import { getLastNonWhitespaceChar, isColumnItem } from '../../../shared/helpers';
 import type { SuggestionRawDefinition } from '../../types';
 import { commaCompleteItem, pipeCompleteItem } from '../../complete_items';
 import { handleFragment } from '../../helper';
@@ -28,7 +23,7 @@ export async function suggest({
   if (
     /\s/.test(innerText[innerText.length - 1]) &&
     getLastNonWhitespaceChar(innerText) !== ',' &&
-    !noCaseCompare(findPreviousWord(innerText), 'keep')
+    !/keep\s+\S*$/i.test(innerText)
   ) {
     return [pipeCompleteItem, commaCompleteItem];
   }

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

@@ -680,7 +680,7 @@ export const commandDefinitions: Array<CommandDefinition<any>> = [
     fieldsSuggestionsAfter: fieldsSuggestionsAfterChangePoint,
   },
   {
-    hidden: true,
+    hidden: false,
     name: 'fork',
     preview: true,
     description: i18n.translate('kbn-esql-validation-autocomplete.esql.definitions.forkDoc', {