[ES|QL] Add support for FORK wrapping pretty printing (#233888)

Partially address #231577
## Summary

Add support for FORK wrapping pretty printing

## Before

https://github.com/user-attachments/assets/785802b8-191d-438f-acbf-760c3d6b78b4
## After
<img width="463" height="364" alt="image"
src="https://github.com/user-attachments/assets/cdf22cc2-8bc5-401d-8501-547b98e816d0"
/>






### Checklist

Check the PR satisfies following conditions. 

Reviewers should verify this PR satisfies this list as well.

- [ ] Any text added follows [EUI's writing
guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses
sentence case text and includes [i18n
support](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md)
- [ ]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [ ] [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
- [ ] If a plugin configuration key changed, check if it needs to be
allowlisted in the cloud and added to the [docker
list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker)
- [ ] This was checked for breaking HTTP API changes, and any breaking
changes have been approved by the breaking-change committee. The
`release_note:breaking` label should be applied in these situations.
- [ ] [Flaky Test
Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was
used on any tests changed
- [ ] The PR description includes the appropriate Release Notes section,
and the correct `release_note:*` label is applied per the
[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)
- [ ] Review the [backport
guidelines](https://docs.google.com/document/d/1VyN5k91e5OVumlc0Gb9RPa3h1ewuPE705nRtioPiTvY/edit?usp=sharing)
and apply applicable `backport:*` labels.

### Identify risks

Does this PR introduce any risks? For example, consider risks like hard
to test bugs, performance regression, potential of data loss.

Describe the risk, its severity, and mitigation for each identified
risk. Invite stakeholders and evaluate how to proceed before merging.

- [ ] [See some risk
examples](https://github.com/elastic/kibana/blob/main/RISK_MATRIX.mdx)
- [ ] ...

Author: sddonne

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

@@ -208,6 +208,91 @@ FROM index
         WITH {"inference_id": "model"}`);
     });
   });
+
+  describe('FORK', () => {
+    test('basic fork with simple subqueries', () => {
+      const { text } = reprint('FROM index | FORK ( KEEP a ) ( KEEP b )');
+
+      expect(text).toBe('FROM index | FORK (KEEP a) (KEEP b)');
+    });
+
+    test('fork with longer subqueries', () => {
+      const { text } = reprint(
+        'FROM index | FORK ( KEEP field1, field2, field3 | WHERE x > 100 ) ( DROP field4, field5 | LIMIT 50 )',
+        { multiline: true }
+      );
+
+      expect(text).toBe(`FROM index
+  | FORK
+      (
+          KEEP field1, field2, field3
+        | WHERE x > 100
+      )
+      (
+          DROP field4, field5
+        | LIMIT 50
+      )`);
+    });
+
+    test('fork with multiple complex subqueries', () => {
+      const { text } = reprint(
+        'FROM index | FORK ( STATS count=COUNT() BY category | WHERE count > 10 | SORT count DESC ) ( KEEP name, value | WHERE value IS NOT NULL | LIMIT 100 )',
+        { multiline: true }
+      );
+
+      expect(text).toBe(`FROM index
+  | FORK
+      (
+          STATS count = COUNT()
+          BY category
+        | WHERE count > 10
+        | SORT count DESC
+      )
+      (
+          KEEP name, value
+        | WHERE value IS NOT NULL
+        | LIMIT 100
+      )`);
+    });
+
+    test('fork with commands before and after it', () => {
+      const { text } = reprint(
+        'FROM index | DROP a | FORK ( KEEP field1 | WHERE x > 100 ) ( DROP field2 | LIMIT 50 ) | LIMIT 100',
+        { multiline: true }
+      );
+
+      expect(text).toBe(`FROM index
+  | DROP a
+  | FORK
+      (
+          KEEP field1
+        | WHERE x > 100
+      )
+      (
+          DROP field2
+        | LIMIT 50
+      )
+  | LIMIT 100`);
+    });
+
+    test('fork with a very long command within', () => {
+      const { text } = reprint(
+        'FROM index | FORK ( WHERE x > 100 | KEEP field1, asd, asd, asd, asd, asd, asd, asd, asd, asd, asd, asd, asd, asd, asd, asd, asd, asd, asd) (LIMIT 10)',
+        { multiline: true }
+      );
+
+      expect(text).toBe(`FROM index
+  | FORK
+      (
+          WHERE x > 100
+        | KEEP field1, asd, asd, asd, asd, asd, asd, asd, asd, asd, asd, asd,
+            asd, asd, asd, asd, asd, asd, asd
+      )
+      (
+          LIMIT 10
+      )`);
+    });
+  });
 });
 
 describe('casing', () => {

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

@@ -25,7 +25,7 @@
  * DISSECT input "pattern"
  * ```
  */
-export const commandsWithNoCommaArgSeparator = new Set(['grok', 'dissect', 'sample']);
+export const commandsWithNoCommaArgSeparator = new Set(['grok', 'dissect', 'sample', 'fork']);
 
 /**
  * This set tracks command options which use an equals sign to separate

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

@@ -302,6 +302,10 @@ export class WrappingPrettyPrinter {
     let remainingCurrentLine = inp.remaining;
     let oneArgumentPerLine = false;
 
+    if (ctx.node.name === 'fork') {
+      oneArgumentPerLine = true;
+    }
+
     for (const child of children(ctx.node)) {
       if (getPrettyPrintStats(child).hasLineBreakingDecorations) {
         oneArgumentPerLine = true;
@@ -385,9 +389,13 @@ export class WrappingPrettyPrinter {
           remaining: this.opts.wrap - indent.length,
           suffix: isLastArg ? '' : commaBetweenArgs ? ',' : '',
         });
-        const separator = isFirstArg ? '' : '\n';
         const indentation = arg.indented ? '' : indent;
-        txt += separator + indentation + arg.txt;
+        let formattedArg = arg.txt;
+        if (args[i].type === 'query') {
+          formattedArg = `(\n${this.opts.tab.repeat(2)}${formattedArg}\n${indentation})`;
+        }
+        const separator = isFirstArg ? '' : '\n';
+        txt += separator + indentation + formattedArg;
         lines++;
       }
     }
@@ -733,9 +741,10 @@ export class WrappingPrettyPrinter {
       return { txt, lines: args.lines /* add options lines count */ };
     })
 
-    .on('visitQuery', (ctx) => {
+    .on('visitQuery', (ctx, inp: Input): Output => {
       const opts = this.opts;
-      const indent = opts.indent ?? '';
+      const indent = inp?.indent ?? opts.indent ?? '';
+      const remaining = inp?.remaining ?? opts.wrap;
       const commands = ctx.node.commands;
       const commandCount = commands.length;
 
@@ -750,20 +759,21 @@ export class WrappingPrettyPrinter {
 
       if (!multiline) {
         const oneLine = indent + BasicPrettyPrinter.print(ctx.node, opts);
-        if (oneLine.length <= opts.wrap) {
-          return oneLine;
+        if (oneLine.length <= remaining) {
+          return { txt: oneLine };
         } else {
           multiline = true;
         }
       }
 
+      // Regular top-level query formatting
       let text = indent;
       const pipedCommandIndent = `${indent}${opts.pipeTab ?? '  '}`;
       const cmdSeparator = multiline ? `${pipedCommandIndent}| ` : ' | ';
       let i = 0;
       let prevOut: Output | undefined;
 
-      for (const out of ctx.visitCommands({ indent, remaining: opts.wrap - indent.length })) {
+      for (const out of ctx.visitCommands({ indent, remaining: remaining - indent.length })) {
         const isFirstCommand = i === 0;
         const isSecondCommand = i === 1;
 
@@ -794,10 +804,10 @@ export class WrappingPrettyPrinter {
         prevOut = out;
       }
 
-      return text;
+      return { txt: text };
     });
 
   public print(query: ESQLAstQueryExpression) {
-    return this.visitor.visitQuery(query, undefined);
+    return this.visitor.visitQuery(query, undefined).txt;
   }
 }