[ES|QL] JOIN command Traversal API and prety-printing support (#202750)

## Summary

Partially addresses https://github.com/elastic/kibana/issues/200858

- Add support for the new `JOIN` command and `AS` expression in
Traversal API: `Walker` and `Visitor`
- Adds support for the new `JOIN`command and `AS` expression in the
pretty-printer.
- Fixes some parser bugs related to the `JOIN` command.


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

packages/kbn-esql-ast/src/parser/__tests__/comments.test.ts

@@ -8,6 +8,7 @@
  */
 
 import { parse } from '..';
+import { Walker } from '../../walker';
 
 describe('Comments', () => {
   describe('can attach "top" comment(s)', () => {
@@ -442,6 +443,35 @@ FROM index`;
         ],
       });
     });
+
+    it('to an identifier', () => {
+      const text = `FROM index | LEFT JOIN
+        // comment
+        abc
+        ON a = b`;
+      const { root } = parse(text, { withFormatting: true });
+
+      expect(root.commands[1]).toMatchObject({
+        type: 'command',
+        name: 'join',
+        args: [
+          {
+            type: 'identifier',
+            name: 'abc',
+            formatting: {
+              top: [
+                {
+                  type: 'comment',
+                  subtype: 'single-line',
+                  text: ' comment',
+                },
+              ],
+            },
+          },
+          {},
+        ],
+      });
+    });
   });
 
   describe('can attach "left" comment(s)', () => {
@@ -549,6 +579,34 @@ FROM index`;
         },
       ]);
     });
+
+    it('to an identifier', () => {
+      const text = `FROM index | LEFT JOIN
+        /* left */ abc
+        ON a = b`;
+      const { root } = parse(text, { withFormatting: true });
+
+      expect(root.commands[1]).toMatchObject({
+        type: 'command',
+        name: 'join',
+        args: [
+          {
+            type: 'identifier',
+            name: 'abc',
+            formatting: {
+              left: [
+                {
+                  type: 'comment',
+                  subtype: 'multi-line',
+                  text: ' left ',
+                },
+              ],
+            },
+          },
+          {},
+        ],
+      });
+    });
   });
 
   describe('can attach "right" comment(s)', () => {
@@ -776,6 +834,61 @@ FROM index`;
         ],
       });
     });
+
+    it('to an identifier', () => {
+      const text = `FROM index | LEFT JOIN
+        abc /* right */ // right 2
+        ON a = b`;
+      const { root } = parse(text, { withFormatting: true });
+
+      expect(root.commands[1]).toMatchObject({
+        type: 'command',
+        name: 'join',
+        args: [
+          {
+            type: 'identifier',
+            name: 'abc',
+            formatting: {
+              right: [
+                {
+                  type: 'comment',
+                  subtype: 'multi-line',
+                  text: ' right ',
+                },
+              ],
+              rightSingleLine: {
+                type: 'comment',
+                subtype: 'single-line',
+                text: ' right 2',
+              },
+            },
+          },
+          {},
+        ],
+      });
+    });
+
+    it('to a column inside ON option', () => {
+      const text = `FROM index | LEFT JOIN
+        abc
+        ON a /* right */ = b`;
+      const { root } = parse(text, { withFormatting: true });
+      const a = Walker.match(root, { type: 'column', name: 'a' });
+
+      expect(a).toMatchObject({
+        type: 'column',
+        name: 'a',
+        formatting: {
+          right: [
+            {
+              type: 'comment',
+              subtype: 'multi-line',
+              text: ' right ',
+            },
+          ],
+        },
+      });
+    });
   });
 
   describe('can attach "right end" comments', () => {

packages/kbn-esql-ast/src/parser/__tests__/join.test.ts

@@ -184,19 +184,41 @@ describe('<TYPE> JOIN command', () => {
       const node2 = Walker.match(query.ast, { type: 'identifier', name: 'alias' });
       const node3 = Walker.match(query.ast, { type: 'column', name: 'on_1' });
       const node4 = Walker.match(query.ast, { type: 'column', name: 'on_2' });
+      const node5 = Walker.match(query.ast, { type: 'function', name: 'as' });
 
       expect(query.src.slice(node1?.location.min, node1?.location.max! + 1)).toBe('index');
       expect(query.src.slice(node2?.location.min, node2?.location.max! + 1)).toBe('alias');
       expect(query.src.slice(node3?.location.min, node3?.location.max! + 1)).toBe('on_1');
       expect(query.src.slice(node4?.location.min, node4?.location.max! + 1)).toBe('on_2');
+      expect(query.src.slice(node5?.location.min, node5?.location.max! + 1)).toBe('index AS alias');
+    });
+
+    it('correctly extracts JOIN command position', () => {
+      const text = `FROM employees | LOOKUP JOIN index AS alias ON on_1, on_2 | LIMIT 1`;
+      const query = EsqlQuery.fromSrc(text);
+      const join = Walker.match(query.ast, { type: 'command', name: 'join' });
+
+      expect(query.src.slice(join?.location.min, join?.location.max! + 1)).toBe(
+        'LOOKUP JOIN index AS alias ON on_1, on_2'
+      );
+    });
+
+    it('correctly extracts ON option position', () => {
+      const text = `FROM employees | LOOKUP JOIN index AS alias ON on_1, on_2 | LIMIT 1`;
+      const query = EsqlQuery.fromSrc(text);
+      const on = Walker.match(query.ast, { type: 'option', name: 'on' });
+
+      expect(query.src.slice(on?.location.min, on?.location.max! + 1)).toBe('ON on_1, on_2');
     });
   });
 
   describe('incorrectly formatted', () => {
-    const text = `FROM employees | LOOKUP JOIN index AAS alias ON on_1, on_2 | LIMIT 1`;
-    const query = EsqlQuery.fromSrc(text);
+    it('throws error on invalid "AS" keyword', () => {
+      const text = `FROM employees | LOOKUP JOIN index AAS alias ON on_1, on_2 | LIMIT 1`;
+      const query = EsqlQuery.fromSrc(text);
 
-    expect(query.errors.length > 0).toBe(true);
-    expect(query.errors[0].message.includes('AAS')).toBe(true);
+      expect(query.errors.length > 0).toBe(true);
+      expect(query.errors[0].message.includes('AAS')).toBe(true);
+    });
   });
 });

packages/kbn-esql-ast/src/parser/factories.ts

@@ -55,6 +55,8 @@ import type {
   InlineCastingType,
   ESQLFunctionCallExpression,
   ESQLIdentifier,
+  ESQLBinaryExpression,
+  BinaryExpressionOperator,
 } from '../types';
 import { parseIdentifier, getPosition } from './helpers';
 import { Builder, type AstNodeParserFields } from '../builder';
@@ -240,6 +242,25 @@ export const createFunctionCall = (ctx: FunctionContext): ESQLFunctionCallExpres
   return node;
 };
 
+export const createBinaryExpression = (
+  operator: BinaryExpressionOperator,
+  ctx: ParserRuleContext,
+  args: ESQLBinaryExpression['args']
+): ESQLBinaryExpression => {
+  const node = Builder.expression.func.binary(
+    operator,
+    args,
+    {},
+    {
+      text: ctx.getText(),
+      location: getPosition(ctx.start, ctx.stop),
+      incomplete: Boolean(ctx.exception),
+    }
+  ) as ESQLBinaryExpression;
+
+  return node;
+};
+
 export const createIdentifierOrParam = (ctx: IdentifierOrParameterContext) => {
   const identifier = ctx.identifier();
   if (identifier) {

packages/kbn-esql-ast/src/parser/factories/join.ts

@@ -8,9 +8,13 @@
  */
 
 import { JoinCommandContext, JoinTargetContext } from '../../antlr/esql_parser';
-import { Builder } from '../../builder';
 import { ESQLAstItem, ESQLBinaryExpression, ESQLCommand, ESQLIdentifier } from '../../types';
-import { createCommand, createIdentifier } from '../factories';
+import {
+  createBinaryExpression,
+  createCommand,
+  createIdentifier,
+  createOption,
+} from '../factories';
 import { visitValueExpression } from '../walkers';
 
 const createNodeFromJoinTarget = (
@@ -24,7 +28,7 @@ const createNodeFromJoinTarget = (
   }
 
   const alias = createIdentifier(aliasCtx);
-  const renameExpression = Builder.expression.func.binary('as', [
+  const renameExpression = createBinaryExpression('as', ctx, [
     index,
     alias,
   ]) as ESQLBinaryExpression;
@@ -39,10 +43,11 @@ export const createJoinCommand = (ctx: JoinCommandContext): ESQLCommand => {
   command.commandType = (ctx._type_.text ?? '').toLocaleLowerCase();
 
   const joinTarget = createNodeFromJoinTarget(ctx.joinTarget());
-  const onOption = Builder.option({ name: 'on' });
+  const joinCondition = ctx.joinCondition();
+  const onOption = createOption('on', joinCondition);
   const joinPredicates: ESQLAstItem[] = onOption.args;
 
-  for (const joinPredicateCtx of ctx.joinCondition().joinPredicate_list()) {
+  for (const joinPredicateCtx of joinCondition.joinPredicate_list()) {
     const expression = visitValueExpression(joinPredicateCtx.valueExpression());
 
     if (expression) {
@@ -51,7 +56,10 @@ export const createJoinCommand = (ctx: JoinCommandContext): ESQLCommand => {
   }
 
   command.args.push(joinTarget);
-  command.args.push(onOption);
+
+  if (onOption.args.length) {
+    command.args.push(onOption);
+  }
 
   return command;
 };

packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.comments.test.ts

@@ -14,7 +14,7 @@ const reprint = (src: string) => {
   const { root } = parse(src, { withFormatting: true });
   const text = BasicPrettyPrinter.print(root);
 
-  // console.log(JSON.stringify(ast, null, 2));
+  // console.log(JSON.stringify(root.commands, null, 2));
 
   return { text };
 };
@@ -184,3 +184,15 @@ describe('rename expressions', () => {
     assertPrint('FROM a | RENAME /*I*/ a /*II*/ AS /*III*/ b /*IV*/, c AS d');
   });
 });
+
+describe('commands', () => {
+  describe('JOIN', () => {
+    test('around JOIN targets', () => {
+      assertPrint('FROM a | LEFT JOIN /*1*/ a /*2*/ AS /*3*/ b /*4*/ ON c');
+    });
+
+    test('around JOIN conditions', () => {
+      assertPrint('FROM a | LEFT JOIN a AS b ON /*1*/ c /*2*/, /*3*/ d /*4*/');
+    });
+  });
+});

packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.test.ts

@@ -115,6 +115,40 @@ describe('single line query', () => {
         expect(text).toBe('FROM index | DISSECT input "pattern" APPEND_SEPARATOR = "<separator>"');
       });
     });
+
+    describe('JOIN', () => {
+      test('example from docs', () => {
+        const { text } = reprint(`
+          FROM employees
+            | EVAL language_code = languages
+            | LOOKUP JOIN languages_lookup ON language_code
+            | WHERE emp_no < 500
+            | KEEP emp_no, language_name
+            | SORT emp_no
+            | LIMIT 10
+        `);
+
+        expect(text).toBe(
+          'FROM employees | EVAL language_code = languages | LOOKUP JOIN languages_lookup ON language_code | WHERE emp_no < 500 | KEEP emp_no, language_name | SORT emp_no | LIMIT 10'
+        );
+      });
+
+      test('supports aliases', () => {
+        const { text } = reprint(`
+          FROM employees | LEFT JOIN languages_lookup AS something ON language_code`);
+
+        expect(text).toBe(
+          'FROM employees | LEFT JOIN languages_lookup AS something ON language_code'
+        );
+      });
+
+      test('supports multiple conditions', () => {
+        const { text } = reprint(`
+          FROM employees | LEFT JOIN a ON b, c, d.e.f`);
+
+        expect(text).toBe('FROM employees | LEFT JOIN a ON b, c, d.e.f');
+      });
+    });
   });
 
   describe('expressions', () => {

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

@@ -497,6 +497,49 @@ ROW 1
     });
   });
 
+  describe('as-expressions', () => {
+    test('JOIN main arguments surrounded in comments', () => {
+      const query = `
+        FROM index | LEFT JOIN
+        /* 1 */
+        // 2
+        /* 3 */
+        // 4
+        /* 5 */ a /* 6 */ AS /* 7 */ b
+        ON c`;
+      const text = reprint(query).text;
+      expect('\n' + text).toBe(`
+FROM index
+  | LEFT JOIN
+      /* 1 */
+      // 2
+      /* 3 */
+      // 4
+      /* 5 */ a /* 6 */ AS
+        /* 7 */ b
+        ON c`);
+    });
+
+    test('JOIN "ON" option argument comments', () => {
+      const query = `
+        FROM index | RIGHT JOIN a AS b ON
+        // c.1
+        /* c.2 */ c /* c.3 */,
+        // d.1
+        /* d.2 */ d /* d.3 */`;
+      const text = reprint(query).text;
+      expect('\n' + text).toBe(`
+FROM index
+  | RIGHT JOIN
+      a AS b
+        ON
+          // c.1
+          /* c.2 */ c, /* c.3 */
+          // d.1
+          /* d.2 */ d /* d.3 */`);
+    });
+  });
+
   describe('function call expressions', () => {
     describe('binary expressions', () => {
       test('first operand surrounded by inline comments', () => {

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

@@ -14,12 +14,33 @@ const reprint = (src: string, opts?: WrappingPrettyPrinterOptions) => {
   const { root } = parse(src);
   const text = WrappingPrettyPrinter.print(root, opts);
 
-  // console.log(JSON.stringify(ast, null, 2));
+  // console.log(JSON.stringify(root.commands, null, 2));
 
   return { text };
 };
 
 describe('commands', () => {
+  describe('JOIN', () => {
+    test('with short identifiers', () => {
+      const { text } = reprint('FROM a | RIGHT JOIN b AS c ON d, e');
+
+      expect(text).toBe('FROM a | RIGHT JOIN b AS c ON d, e');
+    });
+
+    test('with long identifiers', () => {
+      const { text } = reprint(
+        'FROM aaaaaaaaaaaa | RIGHT JOIN bbbbbbbbbbbbbbbbb AS cccccccccccccccccccc ON dddddddddddddddddddddddddddddddddddddddd, eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee'
+      );
+
+      expect('\n' + text).toBe(`
+FROM aaaaaaaaaaaa
+  | RIGHT JOIN bbbbbbbbbbbbbbbbb AS cccccccccccccccccccc
+        ON
+          dddddddddddddddddddddddddddddddddddddddd,
+          eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee`);
+    });
+  });
+
   describe('GROK', () => {
     test('two basic arguments', () => {
       const { text } = reprint('FROM search-movies | GROK Awards "text"');

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

@@ -223,6 +223,11 @@ export class BasicPrettyPrinter {
       return '<EXPRESSION>';
     })
 
+    .on('visitIdentifierExpression', (ctx) => {
+      const formatted = LeafPrinter.identifier(ctx.node);
+      return this.decorateWithComments(ctx.node, formatted);
+    })
+
     .on('visitSourceExpression', (ctx) => {
       const formatted = LeafPrinter.source(ctx.node);
       return this.decorateWithComments(ctx.node, formatted);
@@ -383,12 +388,16 @@ export class BasicPrettyPrinter {
       const argsFormatted = args ? `${separator}${args}` : '';
       const optionFormatted = `${option}${argsFormatted}`;
 
-      return optionFormatted;
+      return this.decorateWithComments(ctx.node, optionFormatted);
     })
 
     .on('visitCommand', (ctx) => {
       const opts = this.opts;
-      const cmd = opts.lowercaseCommands ? ctx.node.name : ctx.node.name.toUpperCase();
+      const node = ctx.node;
+      const cmd = opts.lowercaseCommands ? node.name : node.name.toUpperCase();
+      const cmdType = !node.commandType
+        ? ''
+        : (opts.lowercaseCommands ? node.commandType : node.commandType.toUpperCase()) + ' ';
 
       let args = '';
       let options = '';
@@ -406,9 +415,9 @@ export class BasicPrettyPrinter {
 
       const argsFormatted = args ? ` ${args}` : '';
       const optionsFormatted = options ? ` ${options}` : '';
-      const cmdFormatted = `${cmd}${argsFormatted}${optionsFormatted}`;
+      const cmdFormatted = `${cmdType}${cmd}${argsFormatted}${optionsFormatted}`;
 
-      return cmdFormatted;
+      return this.decorateWithComments(ctx.node, cmdFormatted);
     })
 
     .on('visitQuery', (ctx) => {

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

@@ -427,12 +427,19 @@ export class WrappingPrettyPrinter {
     return { txt, indented };
   }
 
-  protected readonly visitor = new Visitor()
+  protected readonly visitor: Visitor<any> = new Visitor()
     .on('visitExpression', (ctx, inp: Input): Output => {
       const txt = ctx.node.text ?? '<EXPRESSION>';
       return { txt };
     })
 
+    .on('visitIdentifierExpression', (ctx, inp: Input) => {
+      const formatted = LeafPrinter.identifier(ctx.node);
+      const { txt, indented } = this.decorateWithComments(inp.indent, ctx.node, formatted);
+
+      return { txt, indented };
+    })
+
     .on('visitSourceExpression', (ctx, inp: Input): Output => {
       const formatted = LeafPrinter.source(ctx.node) + (inp.suffix ?? '');
       const { txt, indented } = this.decorateWithComments(inp.indent, ctx.node, formatted);
@@ -570,7 +577,14 @@ export class WrappingPrettyPrinter {
 
     .on('visitCommand', (ctx, inp: Input): Output => {
       const opts = this.opts;
-      const cmd = opts.lowercaseCommands ? ctx.node.name : ctx.node.name.toUpperCase();
+      const node = ctx.node;
+      let cmd = opts.lowercaseCommands ? node.name : node.name.toUpperCase();
+
+      if (node.commandType) {
+        const type = opts.lowercaseCommands ? node.commandType : node.commandType.toUpperCase();
+        cmd = `${type} ${cmd}`;
+      }
+
       const args = this.printArguments(ctx, {
         indent: inp.indent,
         remaining: inp.remaining - cmd.length - 1,
@@ -678,6 +692,6 @@ export class WrappingPrettyPrinter {
     });
 
   public print(query: ESQLAstQueryExpression) {
-    return this.visitor.visitQuery(query);
+    return this.visitor.visitQuery(query, undefined);
   }
 }

packages/kbn-esql-ast/src/types.ts

@@ -9,7 +9,7 @@
 
 export type ESQLAst = ESQLAstCommand[];
 
-export type ESQLAstCommand = ESQLCommand | ESQLAstMetricsCommand;
+export type ESQLAstCommand = ESQLCommand | ESQLAstMetricsCommand | ESQLAstJoinCommand;
 
 export type ESQLAstNode = ESQLAstCommand | ESQLAstExpression | ESQLAstItem;
 
@@ -92,6 +92,10 @@ export interface ESQLAstMetricsCommand extends ESQLCommand<'metrics'> {
   grouping?: ESQLAstField[];
 }
 
+export interface ESQLAstJoinCommand extends ESQLCommand<'join'> {
+  commandType: 'lookup' | 'left' | 'right';
+}
+
 export interface ESQLCommandOption extends ESQLAstBaseItem {
   type: 'option';
   args: ESQLAstItem[];
@@ -199,12 +203,14 @@ export type BinaryExpressionOperator =
   | BinaryExpressionArithmeticOperator
   | BinaryExpressionAssignmentOperator
   | BinaryExpressionComparisonOperator
-  | BinaryExpressionRegexOperator;
+  | BinaryExpressionRegexOperator
+  | BinaryExpressionRenameOperator;
 
 export type BinaryExpressionArithmeticOperator = '+' | '-' | '*' | '/' | '%';
 export type BinaryExpressionAssignmentOperator = '=';
 export type BinaryExpressionComparisonOperator = '==' | '=~' | '!=' | '<' | '<=' | '>' | '>=';
 export type BinaryExpressionRegexOperator = 'like' | 'not_like' | 'rlike' | 'not_rlike';
+export type BinaryExpressionRenameOperator = 'as';
 
 // from https://github.com/elastic/elasticsearch/blob/122e7288200ee03e9087c98dff6cebbc94e774aa/docs/reference/esql/functions/kibana/inline_cast.json
 export type InlineCastingType =

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

@@ -0,0 +1,117 @@
+/*
+ * 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 { EsqlQuery } from '../../query';
+import { Visitor } from '../visitor';
+
+test('"visitCommand" captures all non-captured commands', () => {
+  const { ast } = EsqlQuery.fromSrc(`
+    FROM index
+      | STATS 1, "str", [true], a = b BY field
+      | LIMIT 123
+  `);
+  const visitor = new Visitor()
+    .on('visitStatsCommand', (ctx) => {
+      return '<STATS>';
+    })
+    .on('visitCommand', (ctx) => {
+      return `${ctx.name()}`;
+    })
+    .on('visitQuery', (ctx) => {
+      return [...ctx.visitCommands()].join(' | ');
+    });
+  const text = visitor.visitQuery(ast);
+
+  expect(text).toBe('FROM | <STATS> | LIMIT');
+});
+
+test('can visit JOIN command', () => {
+  const { ast } = EsqlQuery.fromSrc(`
+    FROM index
+      | STATS 1, "str", [true], a = b BY field
+      | RIGHT JOIN abc ON xyz
+      | LIMIT 123
+  `);
+  const visitor = new Visitor()
+    .on('visitJoinCommand', (ctx) => {
+      return `JOIN[type = ${ctx.node.commandType}]`;
+    })
+    .on('visitCommand', (ctx) => {
+      return `${ctx.name()}`;
+    })
+    .on('visitQuery', (ctx) => {
+      return [...ctx.visitCommands()].join(' | ');
+    });
+  const text = visitor.visitQuery(ast);
+
+  expect(text).toBe('FROM | STATS | JOIN[type = right] | LIMIT');
+});
+
+test('can visit JOIN command arguments', () => {
+  const { ast } = EsqlQuery.fromSrc(`
+    FROM index
+      | STATS 1, "str", [true], a = b BY field
+      | RIGHT JOIN abc AS xxx ON xyz
+      | LIMIT 123
+  `);
+  const visitor = new Visitor()
+    .on('visitFunctionCallExpression', (ctx) => {
+      if (ctx.node.subtype === 'binary-expression') {
+        return ctx.node.name;
+      } else {
+        return null;
+      }
+    })
+    .on('visitExpression', (ctx) => {
+      return null;
+    })
+    .on('visitJoinCommand', (ctx) => {
+      return [...ctx.visitArgs()];
+    })
+    .on('visitCommand', (ctx) => {
+      return null;
+    })
+    .on('visitQuery', (ctx) => {
+      return [...ctx.visitCommands()];
+    });
+  const list = visitor.visitQuery(ast).flat().filter(Boolean);
+
+  expect(list).toMatchObject(['as']);
+});
+
+test('can visit JOIN ON option', () => {
+  const { ast } = EsqlQuery.fromSrc(`
+    FROM index
+      | STATS 1, "str", [true], a = b BY field
+      | RIGHT JOIN abc AS xxx ON xyz
+      | LIMIT 123
+  `);
+  const visitor = new Visitor()
+    .on('visitColumnExpression', (ctx) => {
+      return ctx.node.name;
+    })
+    .on('visitExpression', (ctx) => {
+      return null;
+    })
+    .on('visitCommandOption', (ctx) => {
+      return [...ctx.visitArguments()].flat();
+    })
+    .on('visitJoinCommand', (ctx) => {
+      return [...ctx.visitOptions()].flat();
+    })
+    .on('visitCommand', (ctx) => {
+      return null;
+    })
+    .on('visitQuery', (ctx) => {
+      return [...ctx.visitCommands()].flat();
+    });
+  const list = visitor.visitQuery(ast).flat().filter(Boolean);
+
+  expect(list).toMatchObject(['xyz']);
+});

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

@@ -158,3 +158,25 @@ test('"visitLiteral" takes over all literal visits', () => {
 
   expect(text).toBe('FROM E | STATS <LITERAL>, <LITERAL>, E, E | LIMIT <LITERAL>');
 });
+
+test('"visitExpression" does visit identifier nodes', () => {
+  const { ast } = parse(`
+    FROM index
+      | RIGHT JOIN a AS b ON c
+  `);
+  const expressions: string[] = [];
+  new Visitor()
+    .on('visitExpression', (ctx) => {
+      expressions.push(ctx.node.name);
+      for (const _ of ctx.visitArguments(undefined));
+    })
+    .on('visitCommand', (ctx) => {
+      for (const _ of ctx.visitArguments());
+    })
+    .on('visitQuery', (ctx) => {
+      for (const _ of ctx.visitCommands());
+    })
+    .visitQuery(ast);
+
+  expect(expressions.sort()).toEqual(['a', 'as', 'b', 'index']);
+});

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

@@ -17,6 +17,7 @@ import type {
   ESQLAstCommand,
   ESQLAstExpression,
   ESQLAstItem,
+  ESQLAstJoinCommand,
   ESQLAstNodeWithArgs,
   ESQLAstNodeWithChildren,
   ESQLAstRenameExpression,
@@ -24,6 +25,7 @@ import type {
   ESQLCommandOption,
   ESQLDecimalLiteral,
   ESQLFunction,
+  ESQLIdentifier,
   ESQLInlineCast,
   ESQLIntegerLiteral,
   ESQLList,
@@ -86,7 +88,7 @@ export class VisitorContext<
     const node = this.node;
 
     if (!isNodeWithArgs(node)) {
-      throw new Error('Node does not have arguments');
+      return;
     }
 
     for (const arg of singleItems(node.args)) {
@@ -467,6 +469,12 @@ export class MvExpandCommandVisitorContext<
   Data extends SharedData = SharedData
 > extends CommandVisitorContext<Methods, Data, ESQLAstCommand> {}
 
+// <LOOKUP | LEFT | RIGHT> JOIN <target> ON <condition>
+export class JoinCommandVisitorContext<
+  Methods extends VisitorMethods = VisitorMethods,
+  Data extends SharedData = SharedData
+> extends CommandVisitorContext<Methods, Data, ESQLAstJoinCommand> {}
+
 // Expressions -----------------------------------------------------------------
 
 export class ExpressionVisitorContext<
@@ -567,3 +575,8 @@ export class OrderExpressionVisitorContext<
   Methods extends VisitorMethods = VisitorMethods,
   Data extends SharedData = SharedData
 > extends VisitorContext<Methods, Data, ESQLOrderExpression> {}
+
+export class IdentifierExpressionVisitorContext<
+  Methods extends VisitorMethods = VisitorMethods,
+  Data extends SharedData = SharedData
+> extends VisitorContext<Methods, Data, ESQLIdentifier> {}

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

@@ -10,9 +10,11 @@
 import * as contexts from './contexts';
 import type {
   ESQLAstCommand,
+  ESQLAstJoinCommand,
   ESQLAstRenameExpression,
   ESQLColumn,
   ESQLFunction,
+  ESQLIdentifier,
   ESQLInlineCast,
   ESQLList,
   ESQLLiteral,
@@ -165,6 +167,10 @@ export class GlobalVisitorContext<
         if (!this.methods.visitMvExpandCommand) break;
         return this.visitMvExpandCommand(parent, commandNode, input as any);
       }
+      case 'join': {
+        if (!this.methods.visitJoinCommand) break;
+        return this.visitJoinCommand(parent, commandNode as ESQLAstJoinCommand, input as any);
+      }
     }
     return this.visitCommandGeneric(parent, commandNode, input as any);
   }
@@ -349,6 +355,15 @@ export class GlobalVisitorContext<
     return this.visitWithSpecificContext('visitMvExpandCommand', context, input);
   }
 
+  public visitJoinCommand(
+    parent: contexts.VisitorContext | null,
+    node: ESQLAstJoinCommand,
+    input: types.VisitorInput<Methods, 'visitJoinCommand'>
+  ): types.VisitorOutput<Methods, 'visitJoinCommand'> {
+    const context = new contexts.JoinCommandVisitorContext(this, node, parent);
+    return this.visitWithSpecificContext('visitJoinCommand', context, input);
+  }
+
   // Expression visiting -------------------------------------------------------
 
   public visitExpressionGeneric(
@@ -405,6 +420,10 @@ export class GlobalVisitorContext<
         if (!this.methods.visitOrderExpression) break;
         return this.visitOrderExpression(parent, expressionNode, input as any);
       }
+      case 'identifier': {
+        if (!this.methods.visitIdentifierExpression) break;
+        return this.visitIdentifierExpression(parent, expressionNode, input as any);
+      }
       case 'option': {
         switch (expressionNode.name) {
           case 'as': {
@@ -501,4 +520,13 @@ export class GlobalVisitorContext<
     const context = new contexts.OrderExpressionVisitorContext(this, node, parent);
     return this.visitWithSpecificContext('visitOrderExpression', context, input);
   }
+
+  public visitIdentifierExpression(
+    parent: contexts.VisitorContext | null,
+    node: ESQLIdentifier,
+    input: types.VisitorInput<Methods, 'visitIdentifierExpression'>
+  ): types.VisitorOutput<Methods, 'visitIdentifierExpression'> {
+    const context = new contexts.IdentifierExpressionVisitorContext(this, node, parent);
+    return this.visitWithSpecificContext('visitIdentifierExpression', context, input);
+  }
 }

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

@@ -61,7 +61,8 @@ export type ExpressionVisitorInput<Methods extends VisitorMethods> = AnyToVoid<
       VisitorInput<Methods, 'visitTimeIntervalLiteralExpression'> &
       VisitorInput<Methods, 'visitInlineCastExpression'> &
       VisitorInput<Methods, 'visitRenameExpression'> &
-      VisitorInput<Methods, 'visitOrderExpression'>
+      VisitorInput<Methods, 'visitOrderExpression'> &
+      VisitorInput<Methods, 'visitIdentifierExpression'>
 >;
 
 /**
@@ -77,7 +78,8 @@ export type ExpressionVisitorOutput<Methods extends VisitorMethods> =
   | VisitorOutput<Methods, 'visitTimeIntervalLiteralExpression'>
   | VisitorOutput<Methods, 'visitInlineCastExpression'>
   | VisitorOutput<Methods, 'visitRenameExpression'>
-  | VisitorOutput<Methods, 'visitOrderExpression'>;
+  | VisitorOutput<Methods, 'visitOrderExpression'>
+  | VisitorOutput<Methods, 'visitIdentifierExpression'>;
 
 /**
  * Input that satisfies any command visitor input constraints.
@@ -103,7 +105,8 @@ export type CommandVisitorInput<Methods extends VisitorMethods> = AnyToVoid<
       VisitorInput<Methods, 'visitDissectCommand'> &
       VisitorInput<Methods, 'visitGrokCommand'> &
       VisitorInput<Methods, 'visitEnrichCommand'> &
-      VisitorInput<Methods, 'visitMvExpandCommand'>
+      VisitorInput<Methods, 'visitMvExpandCommand'> &
+      VisitorInput<Methods, 'visitJoinCommand'>
 >;
 
 /**
@@ -130,7 +133,8 @@ export type CommandVisitorOutput<Methods extends VisitorMethods> =
   | VisitorOutput<Methods, 'visitDissectCommand'>
   | VisitorOutput<Methods, 'visitGrokCommand'>
   | VisitorOutput<Methods, 'visitEnrichCommand'>
-  | VisitorOutput<Methods, 'visitMvExpandCommand'>;
+  | VisitorOutput<Methods, 'visitMvExpandCommand'>
+  | VisitorOutput<Methods, 'visitJoinCommand'>;
 
 export interface VisitorMethods<
   Visitors extends VisitorMethods = any,
@@ -162,6 +166,7 @@ export interface VisitorMethods<
   visitGrokCommand?: Visitor<contexts.GrokCommandVisitorContext<Visitors, Data>, any, any>;
   visitEnrichCommand?: Visitor<contexts.EnrichCommandVisitorContext<Visitors, Data>, any, any>;
   visitMvExpandCommand?: Visitor<contexts.MvExpandCommandVisitorContext<Visitors, Data>, any, any>;
+  visitJoinCommand?: Visitor<contexts.JoinCommandVisitorContext<Visitors, Data>, any, any>;
   visitCommandOption?: Visitor<contexts.CommandOptionVisitorContext<Visitors, Data>, any, any>;
   visitExpression?: Visitor<contexts.ExpressionVisitorContext<Visitors, Data>, any, any>;
   visitSourceExpression?: Visitor<
@@ -205,6 +210,11 @@ export interface VisitorMethods<
     any
   >;
   visitOrderExpression?: Visitor<contexts.OrderExpressionVisitorContext<Visitors, Data>, any, any>;
+  visitIdentifierExpression?: Visitor<
+    contexts.IdentifierExpressionVisitorContext<Visitors, Data>,
+    any,
+    any
+  >;
 }
 
 /**
@@ -230,6 +240,8 @@ export type AstNodeToVisitorName<Node extends VisitorAstNode> = Node extends ESQ
   ? 'visitTimeIntervalLiteralExpression'
   : Node extends ast.ESQLInlineCast
   ? 'visitInlineCastExpression'
+  : Node extends ast.ESQLIdentifier
+  ? 'visitIdentifierExpression'
   : never;
 
 /**

packages/kbn-esql-ast/src/walker/helpers.ts

@@ -7,11 +7,46 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { ESQLProperNode } from '../types';
+import {
+  ESQLAstCommand,
+  ESQLAstQueryExpression,
+  ESQLColumn,
+  ESQLCommandMode,
+  ESQLCommandOption,
+  ESQLFunction,
+  ESQLIdentifier,
+  ESQLInlineCast,
+  ESQLList,
+  ESQLLiteral,
+  ESQLOrderExpression,
+  ESQLProperNode,
+  ESQLSource,
+  ESQLTimeInterval,
+  ESQLUnknownItem,
+} from '../types';
+
+export type NodeMatchKeys =
+  | keyof ESQLAstCommand
+  | keyof ESQLAstQueryExpression
+  | keyof ESQLFunction
+  | keyof ESQLCommandOption
+  | keyof ESQLSource
+  | keyof ESQLColumn
+  | keyof ESQLTimeInterval
+  | keyof ESQLList
+  | keyof ESQLLiteral
+  | keyof ESQLIdentifier
+  | keyof ESQLCommandMode
+  | keyof ESQLInlineCast
+  | keyof ESQLOrderExpression
+  | keyof ESQLUnknownItem;
 
 export type NodeMatchTemplateKey<V> = V | V[] | RegExp;
+
 export type NodeMatchTemplate = {
-  [K in keyof ESQLProperNode]?: NodeMatchTemplateKey<ESQLProperNode[K]>;
+  [K in NodeMatchKeys]?: K extends keyof ESQLProperNode
+    ? NodeMatchTemplateKey<ESQLProperNode[K]>
+    : NodeMatchTemplateKey<unknown>;
 };
 
 /**

packages/kbn-esql-ast/src/walker/walker.test.ts

@@ -20,6 +20,7 @@ import {
   ESQLTimeInterval,
   ESQLInlineCast,
   ESQLUnknownItem,
+  ESQLIdentifier,
 } from '../types';
 import { walk, Walker } from './walker';
 
@@ -82,6 +83,23 @@ describe('structurally can walk all nodes', () => {
       ]);
     });
 
+    test('can traverse JOIN command', () => {
+      const { ast } = parse('FROM index | LEFT JOIN a AS b ON c, d');
+      const commands: ESQLCommand[] = [];
+      const identifiers: ESQLIdentifier[] = [];
+      const columns: ESQLColumn[] = [];
+
+      walk(ast, {
+        visitCommand: (cmd) => commands.push(cmd),
+        visitIdentifier: (id) => identifiers.push(id),
+        visitColumn: (col) => columns.push(col),
+      });
+
+      expect(commands.map(({ name }) => name).sort()).toStrictEqual(['from', 'join']);
+      expect(identifiers.map(({ name }) => name).sort()).toStrictEqual(['a', 'as', 'b', 'c', 'd']);
+      expect(columns.map(({ name }) => name).sort()).toStrictEqual(['c', 'd']);
+    });
+
     test('"visitAny" can capture command nodes', () => {
       const { ast } = parse('FROM index | STATS a = 123 | WHERE 123 | LIMIT 10');
       const commands: ESQLCommand[] = [];
@@ -1050,6 +1068,37 @@ describe('Walker.match()', () => {
       name: 'a.b.c',
     });
   });
+
+  test('can find WHERE command by its type', () => {
+    const query = 'FROM index | LEFT JOIN a | RIGHT JOIN b';
+    const { root } = parse(query);
+
+    const join1 = Walker.match(root, {
+      type: 'command',
+      name: 'join',
+      commandType: 'left',
+    })!;
+    const identifier1 = Walker.match(join1, {
+      type: 'identifier',
+      name: 'a',
+    })!;
+    const join2 = Walker.match(root, {
+      type: 'command',
+      name: 'join',
+      commandType: 'right',
+    })!;
+    const identifier2 = Walker.match(join2, {
+      type: 'identifier',
+      name: 'b',
+    })!;
+
+    expect(identifier1).toMatchObject({
+      name: 'a',
+    });
+    expect(identifier2).toMatchObject({
+      name: 'b',
+    });
+  });
 });
 
 describe('Walker.matchAll()', () => {