[ES|QL] Wrap the fork subcommands inside the parens node

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

@@ -29,6 +29,8 @@ export type {
   ESQLInlineCast,
   ESQLAstBaseItem,
   ESQLAstChangePointCommand,
+  ESQLAstForkCommand,
+  ESQLForkParens,
 } from './src/types';
 
 export * from './src/ast/is';

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/autocomplete.ts

@@ -8,7 +8,11 @@
  */
 import { i18n } from '@kbn/i18n';
 import { withAutoSuggest } from '../../../definitions/utils/autocomplete/helpers';
-import type { ESQLAstAllCommands, ESQLCommand } from '../../../types';
+import type {
+  ESQLAstAllCommands,
+  ESQLAstForkCommand,
+  ESQLAstQueryExpression,
+} from '../../../types';
 import { pipeCompleteItem, getCommandAutocompleteDefinitions } from '../../complete_items';
 import { pipePrecedesCurrentWord } from '../../../definitions/utils/shared';
 import type { ICommandCallbacks } from '../../types';
@@ -22,20 +26,22 @@ export async function autocomplete(
   context?: ICommandContext,
   cursorPosition: number = query.length
 ): Promise<ISuggestionItem[]> {
+  const forkCommand = command as ESQLAstForkCommand;
+
   const innerText = query.substring(0, cursorPosition);
   if (/FORK\s+$/i.test(innerText)) {
     return [newBranchSuggestion];
   }
 
-  const activeBranch = getActiveBranch(command);
+  const activeBranch = getActiveBranch(forkCommand);
   const withinActiveBranch =
     activeBranch &&
     activeBranch.location.min <= innerText.length &&
     activeBranch.location.max >= innerText.length;
 
   if (!withinActiveBranch && /\)\s+$/i.test(innerText)) {
     const suggestions = [newBranchSuggestion];
-    if (command.args.length > 1) {
+    if (forkCommand.args.length > 1) {
       suggestions.push(pipeCompleteItem);
     }
     return suggestions;
@@ -58,13 +64,7 @@ export async function autocomplete(
 
   const subCommandMethods = esqlCommandRegistry.getCommandMethods(subCommand.name);
   return (
-    subCommandMethods?.autocomplete(
-      innerText,
-      subCommand as ESQLCommand,
-      callbacks,
-      context,
-      cursorPosition
-    ) || []
+    subCommandMethods?.autocomplete(innerText, subCommand, callbacks, context, cursorPosition) || []
   );
 }
 
@@ -80,13 +80,12 @@ const newBranchSuggestion: ISuggestionItem = withAutoSuggest({
   asSnippet: true,
 });
 
-const getActiveBranch = (command: ESQLAstAllCommands) => {
+const getActiveBranch = (command: ESQLAstForkCommand): ESQLAstQueryExpression | undefined => {
   const finalBranch = command.args[command.args.length - 1];
 
-  if (Array.isArray(finalBranch) || finalBranch.type !== 'query') {
-    // should never happen
+  if (!finalBranch) {
     return;
   }
 
-  return finalBranch;
+  return finalBranch.child;
 };

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts

@@ -8,18 +8,18 @@
  */
 import { uniqBy } from 'lodash';
 import { esqlCommandRegistry } from '../../../..';
-import type { ESQLAstQueryExpression } from '../../../types';
-import { type ESQLCommand } from '../../../types';
+import type { ESQLAstAllCommands, ESQLAstForkCommand } from '../../../types';
 import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
 import type { IAdditionalFields } from '../../registry';
 
 export const columnsAfter = async (
-  command: ESQLCommand,
+  command: ESQLAstAllCommands,
   previousColumns: ESQLColumnData[],
   query: string,
   additionalFields: IAdditionalFields
 ) => {
-  const branches = command.args as ESQLAstQueryExpression[];
+  const forkCommand = command as ESQLAstForkCommand;
+  const branches = forkCommand.args.map((parens) => parens.child);
 
   const columnsFromBranches = [];
 

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts

@@ -6,7 +6,7 @@
  * 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 type { ESQLAst, ESQLAstAllCommands, ESQLMessage } from '../../../types';
+import type { ESQLAst, ESQLAstAllCommands, ESQLAstForkCommand, ESQLMessage } from '../../../types';
 import { Walker } from '../../../walker';
 import type { ICommandContext, ICommandCallbacks } from '../../types';
 import { validateCommandArguments } from '../../../definitions/utils/validation';
@@ -23,24 +23,31 @@ export const validate = (
   context?: ICommandContext,
   callbacks?: ICommandCallbacks
 ): ESQLMessage[] => {
+  const forkCommand = command as ESQLAstForkCommand;
   const messages: ESQLMessage[] = [];
 
-  if (command.args.length < MIN_BRANCHES) {
-    messages.push(errors.forkTooFewBranches(command));
+  if (forkCommand.args.length < MIN_BRANCHES) {
+    messages.push(errors.forkTooFewBranches(forkCommand));
   }
 
-  if (command.args.length > MAX_BRANCHES) {
-    messages.push(errors.forkTooManyBranches(command));
+  if (forkCommand.args.length > MAX_BRANCHES) {
+    messages.push(errors.forkTooManyBranches(forkCommand));
   }
 
-  messages.push(...validateCommandArguments(command, ast, context, callbacks));
+  messages.push(...validateCommandArguments(forkCommand, ast, context, callbacks));
 
-  for (const arg of command.args.flat()) {
-    if (!Array.isArray(arg) && arg.type === 'query') {
+  for (const arg of forkCommand.args) {
+    const query = arg.child;
+
+    if (!Array.isArray(query) && query.type === 'query') {
       // all the args should be commands
-      arg.commands.forEach((subCommand) => {
+      query.commands.forEach((subCommand) => {
         const subCommandMethods = esqlCommandRegistry.getCommandMethods(subCommand.name);
-        const validationMessages = subCommandMethods?.validate?.(subCommand, arg.commands, context);
+        const validationMessages = subCommandMethods?.validate?.(
+          subCommand,
+          query.commands,
+          context
+        );
         messages.push(...(validationMessages || []));
       });
     }
@@ -58,7 +65,7 @@ export const validate = (
   const hasSubqueries = fromCommands.some((cmd) => cmd.args.some((arg) => isSubQuery(arg)));
 
   if (hasSubqueries) {
-    messages.push(errors.forkNotAllowedWithSubqueries(command));
+    messages.push(errors.forkNotAllowedWithSubqueries(forkCommand));
   }
 
   return messages;

src/platform/packages/shared/kbn-esql-ast/src/mutate/generic/commands/options/index.ts

@@ -8,7 +8,12 @@
  */
 
 import { Builder } from '../../../../builder';
-import type { ESQLAstQueryExpression, ESQLCommand, ESQLCommandOption } from '../../../../types';
+import type {
+  ESQLAstItem,
+  ESQLAstQueryExpression,
+  ESQLCommand,
+  ESQLCommandOption,
+} from '../../../../types';
 import { Visitor } from '../../../../visitor';
 import type { Predicate } from '../../../types';
 import * as commands from '..';
@@ -107,7 +112,7 @@ export const remove = (ast: ESQLAstQueryExpression, option: ESQLCommandOption):
         return false;
       }
 
-      const index = ctx.node.args.indexOf(target);
+      const index = (ctx.node.args as ESQLAstItem[]).indexOf(target);
 
       if (index === -1) {
         return false;

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

@@ -9,7 +9,7 @@
 
 import { parse } from '..';
 import { EsqlQuery } from '../../query';
-import type { ESQLAstQueryExpression } from '../../types';
+import type { ESQLForkParens } from '../../types';
 import { Walker } from '../../walker';
 
 describe('FORK', () => {
@@ -24,9 +24,9 @@ describe('FORK', () => {
 
       expect(ast[1].args).toHaveLength(3);
       expect(ast[1].args).toMatchObject([
-        { type: 'query', commands: [{ name: 'where' }] },
-        { type: 'query', commands: [{ name: 'sort' }] },
-        { type: 'query', commands: [{ name: 'limit' }] },
+        { type: 'parens', child: { type: 'query', commands: [{ name: 'where' }] } },
+        { type: 'parens', child: { type: 'query', commands: [{ name: 'sort' }] } },
+        { type: 'parens', child: { type: 'query', commands: [{ name: 'limit' }] } },
       ]);
     });
 
@@ -36,12 +36,12 @@ describe('FORK', () => {
           | FORK (WHERE bytes > 1 | LIMIT 10) (SORT bytes ASC) (LIMIT 100)`;
       const { ast } = EsqlQuery.fromSrc(text);
       const fork = Walker.match(ast, { type: 'command', name: 'fork' })!;
-      const firstQuery = Walker.match(fork, { type: 'query' })!;
+      const firstParens = Walker.match(fork, { type: 'parens' })!;
 
       expect(text.slice(fork.location.min, fork.location.max + 1)).toBe(
         'FORK (WHERE bytes > 1 | LIMIT 10) (SORT bytes ASC) (LIMIT 100)'
       );
-      expect(text.slice(firstQuery.location.min, firstQuery.location.max + 1)).toBe(
+      expect(text.slice(firstParens.location.min, firstParens.location.max + 1)).toBe(
         '(WHERE bytes > 1 | LIMIT 10)'
       );
     });
@@ -55,8 +55,17 @@ describe('FORK', () => {
 
       expect(ast[1].args).toHaveLength(2);
       expect(ast[1].args).toMatchObject([
-        { type: 'query', commands: [{ name: 'where' }, { name: 'sort' }, { name: 'limit' }] },
-        { type: 'query', commands: [{ name: 'where' }, { name: 'limit' }] },
+        {
+          type: 'parens',
+          child: {
+            type: 'query',
+            commands: [{ name: 'where' }, { name: 'sort' }, { name: 'limit' }],
+          },
+        },
+        {
+          type: 'parens',
+          child: { type: 'query', commands: [{ name: 'where' }, { name: 'limit' }] },
+        },
       ]);
     });
   });
@@ -78,7 +87,11 @@ describe('FORK', () => {
       const { root } = parse(text);
 
       expect(root.commands[1].args).toHaveLength(1);
-      expect((root.commands[1].args[0] as ESQLAstQueryExpression).commands[0]).toMatchObject({
+      const parens = root.commands[1].args[0] as ESQLForkParens;
+
+      expect(parens.type).toBe('parens');
+      expect(parens.child.type).toBe('query');
+      expect(parens.child.commands[0]).toMatchObject({
         name: type,
       });
     });

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

@@ -1724,7 +1724,7 @@ export class CstToAstConverter {
     return command;
   }
 
-  private fromForkSubQuery(ctx: cst.ForkSubQueryContext): ast.ESQLAstQueryExpression {
+  private fromForkSubQuery(ctx: cst.ForkSubQueryContext): ast.ESQLParens {
     const commands: ast.ESQLCommand[] = [];
     const collectCommand = (cmdCtx: cst.ForkSubQueryProcessingCommandContext) => {
       const processingCommandCtx = cmdCtx.processingCommand();
@@ -1752,10 +1752,25 @@ export class CstToAstConverter {
 
     commands.reverse();
 
-    const parserFields = this.getParserFields(ctx);
-    const query = Builder.expression.query(commands, parserFields);
+    const openParen = ctx.LP();
+    const closeParen = ctx.RP();
 
-    return query;
+    const closeParenText = closeParen?.getText() ?? '';
+    const hasCloseParen = closeParen && !/<missing /.test(closeParenText);
+    const incomplete = Boolean(ctx.exception) || !hasCloseParen;
+
+    const query = Builder.expression.query(commands, {
+      ...this.getParserFields(ctx),
+      incomplete,
+    });
+
+    return Builder.expression.parens(query, {
+      incomplete: incomplete || query.incomplete,
+      location: getPosition(
+        openParen?.symbol ?? ctx.start,
+        hasCloseParen ? closeParen.symbol : ctx.stop
+      ),
+    });
   }
 
   // -------------------------------------------------------------- expressions

src/platform/packages/shared/kbn-esql-ast/src/pretty_print/__tests__/wrapping_pretty_printer.comments.test.ts

@@ -970,4 +970,26 @@ FROM index`;
       assertReprint(query, expected);
     });
   });
+
+  describe('FORK', () => {
+    test('preserves comments in various positions', () => {
+      const query = `FROM index
+  /* before FORK */
+  | FORK
+      /* first branch */
+      (
+          KEEP field1, field2, field3 /* important fields */
+        | WHERE x > 100 /* filter */
+      )
+      /* second branch */
+      (
+          DROP field4, field5 /* not needed */
+        | LIMIT 50
+      )`;
+
+      const text = reprint(query, { multiline: true }).text;
+
+      expect(text).toBe(query);
+    });
+  });
 });

src/platform/packages/shared/kbn-esql-ast/src/pretty_print/basic_pretty_printer.ts

@@ -501,11 +501,15 @@ export class BasicPrettyPrinter {
       if (cmd === 'FORK') {
         const branches = node.args
           .map((branch) => {
-            if (Array.isArray(branch) || branch.type !== 'query') {
+            if (Array.isArray(branch)) {
               return undefined;
             }
 
-            return ctx.visitSubQuery(branch);
+            if (branch.type === 'parens' && branch.child.type === 'query') {
+              return ctx.visitSubQuery(branch.child);
+            }
+
+            return undefined;
           })
           .filter(Boolean) as string[];
 

src/platform/packages/shared/kbn-esql-ast/src/pretty_print/wrapping_pretty_printer.ts

@@ -404,10 +404,7 @@ export class WrappingPrettyPrinter {
           suffix: isLastArg ? '' : needsComma ? ',' : '',
         });
         const indentation = arg.indented ? '' : indent;
-        let formattedArg = arg.txt;
-        if (args[i].type === 'query') {
-          formattedArg = `(\n${this.opts.tab.repeat(2)}${formattedArg}\n${indentation})`;
-        }
+        const formattedArg = arg.txt;
         const separator = isFirstArg ? '' : '\n';
         txt += separator + indentation + formattedArg;
         lines++;
@@ -680,8 +677,37 @@ export class WrappingPrettyPrinter {
     })
 
     .on('visitParensExpression', (ctx, inp: Input): Output => {
-      const child = ctx.visitChild(inp);
-      const formatted = `(${child.txt.trimStart()})`;
+      // Check if parent is FORK command
+      const parent = ctx.parent?.node;
+      const isForkBranch =
+        !Array.isArray(parent) &&
+        parent?.type === 'command' &&
+        parent.name === 'fork' &&
+        ctx.node.child?.type === 'query';
+
+      let formatted: string;
+      if (isForkBranch) {
+        const baseIndent = inp.indent + this.opts.tab;
+        const childText = this.visitor.visitQuery(ctx.node.child as ESQLAstQueryExpression, {
+          indent: baseIndent,
+          remaining: this.opts.wrap - baseIndent.length,
+        });
+
+        const lines = childText.txt.split('\n');
+
+        for (let i = 0; i < lines.length; i++) {
+          if (i === 0) {
+            lines[i] = '  ' + lines[i];
+          } else if (lines[i].startsWith('  ')) {
+            lines[i] = lines[i].slice(2);
+          }
+        }
+
+        formatted = `(\n${lines.join('\n')}\n${inp.indent})`;
+      } else {
+        const child = ctx.visitChild(inp);
+        formatted = `(${child.txt.trimStart()})`;
+      }
 
       return this.decorateWithComments(inp, ctx.node, formatted);
     })