[ES|QL] Walker and Visitor support for SET instructions (elastic#237659)

## Summary

Partially addresses elastic#233942

Implements parsing for `SET` pseudo-command traversal in `Walker` and
`Visitor`.


### Checklist

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

Author: vadimkibana

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

@@ -134,3 +134,194 @@ test('specific commands receive specific visitor contexts', () => {
     .on('visitQuery', (ctx) => [...ctx.visitCommands()])
     .visitQuery(root);
 });
+
+describe('header commands', () => {
+  test('can visit header commands', () => {
+    const { root } = parse('SET timeout = "30s"; FROM index | LIMIT 10');
+    const headerNames: string[] = [];
+
+    new Visitor()
+      .on('visitHeaderCommand', (ctx) => {
+        headerNames.push(ctx.node.name);
+      })
+      .on('visitQuery', (ctx) => {
+        for (const _headerResult of ctx.visitHeaderCommands()) {
+          // Visit header commands
+        }
+      })
+      .visitQuery(root);
+
+    expect(headerNames).toEqual(['set']);
+  });
+
+  test('can visit multiple header commands', () => {
+    const { root } = parse('SET a = 1; SET b = 2; SET c = 3; FROM index');
+    const headerNames: string[] = [];
+
+    new Visitor()
+      .on('visitHeaderCommand', (ctx) => {
+        headerNames.push(ctx.name());
+      })
+      .on('visitQuery', (ctx) => {
+        for (const _headerResult of ctx.visitHeaderCommands()) {
+          // Visit header commands
+        }
+      })
+      .visitQuery(root);
+
+    expect(headerNames).toEqual(['SET', 'SET', 'SET']);
+  });
+
+  test('can visit header command arguments', () => {
+    const { root } = parse('SET timeout = "30s"; FROM index');
+    const identifiers: string[] = [];
+    const literals: string[] = [];
+
+    new Visitor()
+      .on('visitExpression', (ctx) => {
+        // Default expression visitor
+      })
+      .on('visitFunctionCallExpression', (ctx) => {
+        // Visit function arguments recursively
+        for (const _arg of ctx.visitArguments(undefined)) {
+          // Process arguments
+        }
+      })
+      .on('visitHeaderCommand', (ctx) => {
+        for (const _argResult of ctx.visitArgs(undefined)) {
+          // Visit arguments
+        }
+      })
+      .on('visitIdentifierExpression', (ctx) => {
+        if (ctx.node.name !== '=') {
+          identifiers.push(ctx.node.name);
+        }
+      })
+      .on('visitLiteralExpression', (ctx) => {
+        if (ctx.node.literalType === 'keyword') {
+          literals.push(ctx.node.valueUnquoted || '');
+        }
+      })
+      .on('visitQuery', (ctx) => {
+        for (const _headerResult of ctx.visitHeaderCommands()) {
+          // Visit header commands
+        }
+      })
+      .visitQuery(root);
+
+    expect(identifiers).toEqual(['timeout']);
+    expect(literals).toEqual(['30s']);
+  });
+
+  test('header commands are visited before regular commands', () => {
+    const { root } = parse('SET a = 1; FROM index | LIMIT 10');
+    const visitOrder: string[] = [];
+
+    new Visitor()
+      .on('visitHeaderCommand', (ctx) => {
+        visitOrder.push(`header:${ctx.node.name}`);
+      })
+      .on('visitCommand', (ctx) => {
+        visitOrder.push(`command:${ctx.node.name}`);
+      })
+      .on('visitQuery', (ctx) => {
+        for (const _headerResult of ctx.visitHeaderCommands()) {
+          // Visit header commands first
+        }
+        for (const _cmdResult of ctx.visitCommands()) {
+          // Then visit regular commands
+        }
+      })
+      .visitQuery(root);
+
+    expect(visitOrder).toEqual(['header:set', 'command:from', 'command:limit']);
+  });
+
+  test('can iterate through header commands', () => {
+    const { root } = parse('SET a = 1; SET b = 2; FROM index');
+    const headerCommandCount = [
+      ...new Visitor().on('visitQuery', (ctx) => ctx.headerCommands()).visitQuery(root),
+    ].length;
+
+    expect(headerCommandCount).toBe(2);
+  });
+
+  test('header command context has correct parent', () => {
+    const { root } = parse('SET timeout = "30s"; FROM index');
+
+    new Visitor()
+      .on('visitHeaderCommand', (ctx) => {
+        if (ctx.parent!.node !== root) {
+          throw new Error('Expected parent to be query node');
+        }
+      })
+      .on('visitQuery', (ctx) => {
+        for (const _headerResult of ctx.visitHeaderCommands()) {
+          // Visit header commands
+        }
+      })
+      .visitQuery(root);
+  });
+
+  test('can visit header command directly', () => {
+    const { root } = parse('SET timeout = "30s"; FROM index');
+    const headerCommand = root.header![0];
+
+    const result = new Visitor()
+      .on('visitHeaderCommand', (ctx) => {
+        return `visited:${ctx.node.name}`;
+      })
+      .visitHeaderCommand(headerCommand);
+
+    expect(result).toBe('visited:set');
+  });
+
+  test('header commands with various value types', () => {
+    const { root } = parse('SET a = 1; SET b = "value"; SET c = true; FROM index');
+    const literals: any[] = [];
+
+    new Visitor()
+      .on('visitExpression', (ctx) => {
+        // Default expression visitor
+      })
+      .on('visitFunctionCallExpression', (ctx) => {
+        // Visit function arguments recursively
+        for (const _arg of ctx.visitArguments(undefined)) {
+          // Process arguments
+        }
+      })
+      .on('visitHeaderCommand', (ctx) => {
+        for (const _argResult of ctx.visitArgs(undefined)) {
+          // Visit arguments
+        }
+      })
+      .on('visitLiteralExpression', (ctx) => {
+        literals.push(ctx.node.value);
+      })
+      .on('visitQuery', (ctx) => {
+        for (const _headerResult of ctx.visitHeaderCommands()) {
+          // Visit header commands
+        }
+      })
+      .visitQuery(root);
+
+    expect(literals).toContain(1);
+    expect(literals).toContain('"value"');
+    expect(literals).toContain('true');
+  });
+
+  test('can return values from header command visitor', () => {
+    const { root } = parse('SET a = 1; SET b = 2; FROM index');
+
+    const results = new Visitor()
+      .on('visitHeaderCommand', (ctx) => {
+        return `header-${ctx.node.name}`;
+      })
+      .on('visitQuery', (ctx) => {
+        return [...ctx.visitHeaderCommands()];
+      })
+      .visitQuery(root);
+
+    expect(results).toEqual(['header-set', 'header-set']);
+  });
+});

src/platform/packages/shared/kbn-esql-ast/src/visitor/contexts.ts

@@ -18,6 +18,7 @@ import type {
   ESQLAstChangePointCommand,
   ESQLAstCommand,
   ESQLAstExpression,
+  ESQLAstHeaderCommand,
   ESQLAstItem,
   ESQLAstJoinCommand,
   ESQLAstQueryExpression,
@@ -147,12 +148,37 @@ export class VisitorContext<
   ): ExpressionVisitorOutput<Methods> {
     return this.ctx.visitCommand(this, commandNode, input);
   }
+
+  public visitHeaderCommand(
+    headerCommandNode: ESQLAstHeaderCommand,
+    input: UndefinedToVoid<Parameters<NonNullable<Methods['visitHeaderCommand']>>[1]>
+  ): ReturnType<NonNullable<Methods['visitHeaderCommand']>> {
+    return this.ctx.visitHeaderCommand(this, headerCommandNode, input);
+  }
 }
 
 export class QueryVisitorContext<
   Methods extends VisitorMethods = VisitorMethods,
   Data extends SharedData = SharedData
 > extends VisitorContext<Methods, Data, ESQLAstQueryNode> {
+  public *headerCommands(): Iterable<ESQLAstHeaderCommand> {
+    if (this.node.header) {
+      yield* this.node.header;
+    }
+  }
+
+  public *visitHeaderCommands(
+    input: UndefinedToVoid<Parameters<NonNullable<Methods['visitHeaderCommand']>>[1]>
+  ): Iterable<ReturnType<NonNullable<Methods['visitHeaderCommand']>>> {
+    this.ctx.assertMethodExists('visitHeaderCommand');
+
+    if (this.node.header) {
+      for (const headerCmd of this.node.header) {
+        yield this.visitHeaderCommand(headerCmd, input as any);
+      }
+    }
+  }
+
   public *commands(): Iterable<ESQLAstCommand> {
     yield* this.node.commands;
   }
@@ -286,6 +312,37 @@ export class CommandVisitorContext<
   }
 }
 
+export class HeaderCommandVisitorContext<
+  Methods extends VisitorMethods = VisitorMethods,
+  Data extends SharedData = SharedData,
+  Node extends ESQLAstHeaderCommand = ESQLAstHeaderCommand
+> extends VisitorContext<Methods, Data, Node> {
+  public name(): string {
+    return this.node.name.toUpperCase();
+  }
+
+  public *args(): Iterable<ESQLAstExpression> {
+    yield* this.node.args;
+  }
+
+  public *visitArgs(
+    input:
+      | VisitorInput<Methods, 'visitExpression'>
+      | (() => VisitorInput<Methods, 'visitExpression'>)
+  ): Iterable<ExpressionVisitorOutput<Methods>> {
+    this.ctx.assertMethodExists('visitExpression');
+
+    for (const arg of this.args()) {
+      yield this.visitExpression(
+        arg,
+        typeof input === 'function'
+          ? (input as () => VisitorInput<Methods, 'visitExpression'>)()
+          : (input as VisitorInput<Methods, 'visitExpression'>)
+      );
+    }
+  }
+}
+
 export class CommandOptionVisitorContext<
   Methods extends VisitorMethods = VisitorMethods,
   Data extends SharedData = SharedData

src/platform/packages/shared/kbn-esql-ast/src/visitor/global_visitor_context.ts

@@ -12,6 +12,7 @@ import type {
   ESQLAstChangePointCommand,
   ESQLAstCommand,
   ESQLAstCompletionCommand,
+  ESQLAstHeaderCommand,
   ESQLAstJoinCommand,
   ESQLAstQueryExpression,
   ESQLAstRerankCommand,
@@ -216,6 +217,19 @@ export class GlobalVisitorContext<
     return this.visitCommandGeneric(parent, commandNode, input as any);
   }
 
+  public visitHeaderCommand(
+    parent: contexts.VisitorContext | null,
+    node: ESQLAstHeaderCommand,
+    input: types.VisitorInput<Methods, 'visitHeaderCommand'>
+  ): types.VisitorOutput<Methods, 'visitHeaderCommand'> {
+    this.assertMethodExists('visitHeaderCommand');
+
+    const context = new contexts.HeaderCommandVisitorContext(this, node, parent);
+    const output = this.methods.visitHeaderCommand!(context, input);
+
+    return output;
+  }
+
   public visitFromCommand(
     parent: contexts.VisitorContext | null,
     node: ESQLAstCommand,

src/platform/packages/shared/kbn-esql-ast/src/visitor/types.ts

@@ -149,6 +149,7 @@ export interface VisitorMethods<
 > {
   visitQuery?: Visitor<contexts.QueryVisitorContext<Visitors, Data>, any, any>;
   visitCommand?: Visitor<contexts.CommandVisitorContext<Visitors, Data>, any, any>;
+  visitHeaderCommand?: Visitor<contexts.HeaderCommandVisitorContext<Visitors, Data>, any, any>;
   visitFromCommand?: Visitor<contexts.FromCommandVisitorContext<Visitors, Data>, any, any>;
   visitLimitCommand?: Visitor<contexts.LimitCommandVisitorContext<Visitors, Data>, any, any>;
   visitExplainCommand?: Visitor<contexts.ExplainCommandVisitorContext<Visitors, Data>, any, any>;
@@ -243,6 +244,8 @@ export interface VisitorMethods<
  */
 export type AstNodeToVisitorName<Node extends VisitorAstNode> = Node extends ESQLAstQueryNode
   ? 'visitQuery'
+  : Node extends ast.ESQLAstHeaderCommand
+  ? 'visitHeaderCommand'
   : Node extends ast.ESQLCommand
   ? 'visitCommand'
   : Node extends ast.ESQLCommandOption

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

@@ -225,6 +225,12 @@ export class Visitor<
             NonNullable<Methods['visitQuery']>
           >;
         }
+        case 'header-command': {
+          this.ctx.assertMethodExists('visitHeaderCommand');
+          return this.ctx.methods.visitHeaderCommand!(ctx as any, input) as ReturnType<
+            NonNullable<Methods['visitHeaderCommand']>
+          >;
+        }
         case 'command': {
           this.ctx.assertMethodExists('visitCommand');
           return this.ctx.methods.visitCommand!(ctx as any, input) as ReturnType<
@@ -270,6 +276,21 @@ export class Visitor<
     return this.ctx.visitCommand(null, node, input);
   }
 
+  /**
+   * Traverse starting from known header command node with default context.
+   *
+   * @param node Header command node to traverse.
+   * @param input Input to pass to the first visitor.
+   * @returns The output of the visitor.
+   */
+  public visitHeaderCommand(
+    node: import('../types').ESQLAstHeaderCommand,
+    input: UndefinedToVoid<Parameters<NonNullable<Methods['visitHeaderCommand']>>[1]>
+  ) {
+    this.ctx.assertMethodExists('visitHeaderCommand');
+    return this.ctx.visitHeaderCommand(null, node, input);
+  }
+
   /**
    * Traverse starting from known expression node with default context.
    *