[8.x] [ES|QL] Load fields of indices in `JOIN` command (#20…

…7375) (#208160) # Backport This will backport the following commits from `main` to `8.x`: - [[ES|QL] Load fields of indices in `JOIN` command (#207375)](#207375)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Vadim Kibana <[email protected]>
elastic · Jan 24, 2025 · ce6d90e · ce6d90e
1 parent ef88984
commit ce6d90e
Show file tree

Hide file tree

Showing 8 changed files with 338 additions and 4 deletions.
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/ast/helpers.ts b/src/platform/packages/shared/kbn-esql-ast/src/ast/helpers.ts
@@ -8,6 +8,7 @@
  */
 
 import type {
+  BinaryExpressionRenameOperator,
   BinaryExpressionWhereOperator,
   ESQLAstNode,
   ESQLBinaryExpression,
@@ -47,6 +48,11 @@ export const isWhereExpression = (
 ): node is ESQLBinaryExpression<BinaryExpressionWhereOperator> =>
   isBinaryExpression(node) && node.name === 'where';
 
+export const isAsExpression = (
+  node: unknown
+): node is ESQLBinaryExpression<BinaryExpressionRenameOperator> =>
+  isBinaryExpression(node) && node.name === 'as';
+
 export const isFieldExpression = (
   node: unknown
 ): node is ESQLBinaryExpression<BinaryExpressionWhereOperator> =>

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/mutate/README.md b/src/platform/packages/shared/kbn-esql-ast/src/mutate/README.md
@@ -65,6 +65,10 @@ console.log(src); // FROM index METADATA _lang, _id
     - `.byIndex()` &mdash; Find a `STATS` command by index.
     - `.summarize()` &mdash; Summarize all `STATS` commands.
     - `.summarizeCommand()` &mdash; Summarize a specific `STATS` command.
+  - `.join`
+    - `.list()` &mdash; List all `JOIN` commands.
+    - `.byIndex()` &mdash; Find a `JOIN` command by index.
+    - `.summarize()` &mdash; Summarize all `JOIN` commands.
 
 
 ## Examples

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/mutate/commands/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/mutate/commands/index.ts
@@ -12,5 +12,6 @@ import * as limit from './limit';
 import * as sort from './sort';
 import * as stats from './stats';
 import * as where from './where';
+import * as join from './join';
 
-export { from, limit, sort, stats, where };
+export { from, limit, sort, stats, where, join };
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/mutate/commands/join/index.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/mutate/commands/join/index.test.ts
@@ -0,0 +1,166 @@
+/*
+ * 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 * as commands from '..';
+import { EsqlQuery } from '../../../query';
+
+describe('commands.where', () => {
+  describe('.list()', () => {
+    it('lists all "JOIN" commands', () => {
+      const src =
+        'FROM index | LIMIT 1 | JOIN join_index1 ON join_field1 | WHERE b == 2 | JOIN join_index2 ON join_field2 | LIMIT 1';
+      const query = EsqlQuery.fromSrc(src);
+
+      const nodes = [...commands.join.list(query.ast)];
+
+      expect(nodes).toMatchObject([
+        {
+          type: 'command',
+          name: 'join',
+          args: [
+            {
+              type: 'identifier',
+              name: 'join_index1',
+            },
+            {},
+          ],
+        },
+        {
+          type: 'command',
+          name: 'join',
+          args: [
+            {
+              type: 'identifier',
+              name: 'join_index2',
+            },
+            {},
+          ],
+        },
+      ]);
+    });
+  });
+
+  describe('.byIndex()', () => {
+    it('retrieves the specific "WHERE" command by index', () => {
+      const src =
+        'FROM index | LIMIT 1 | JOIN join_index1 ON join_field1 | WHERE b == 2 | JOIN join_index2 ON join_field2 | LIMIT 1';
+      const query = EsqlQuery.fromSrc(src);
+
+      const node1 = commands.join.byIndex(query.ast, 1);
+      const node2 = commands.join.byIndex(query.ast, 0);
+
+      expect(node1).toMatchObject({
+        type: 'command',
+        name: 'join',
+        args: [
+          {
+            type: 'identifier',
+            name: 'join_index2',
+          },
+          {},
+        ],
+      });
+      expect(node2).toMatchObject({
+        type: 'command',
+        name: 'join',
+        args: [
+          {
+            type: 'identifier',
+            name: 'join_index1',
+          },
+          {},
+        ],
+      });
+    });
+  });
+
+  describe('.summarize', () => {
+    it('returns target index fields', () => {
+      const src =
+        'FROM index | LIMIT 1 | JOIN join_index1 ON join_field1 | WHERE b == 2 | JOIN join_index2 ON join_field2 | LIMIT 1';
+      const query = EsqlQuery.fromSrc(src);
+      const summary = commands.join.summarize(query.ast);
+
+      expect(summary).toMatchObject([
+        {
+          target: {
+            index: {
+              type: 'identifier',
+              name: 'join_index1',
+            },
+          },
+        },
+        {
+          target: {
+            index: {
+              type: 'identifier',
+              name: 'join_index2',
+            },
+          },
+        },
+      ]);
+    });
+
+    it('returns target aliases', () => {
+      const src =
+        'FROM index | LIMIT 1 | JOIN join_index1 AS a ON join_field1 | WHERE b == 2 | JOIN join_index2 AS b ON join_field2 | LIMIT 1';
+      const query = EsqlQuery.fromSrc(src);
+      const summary = commands.join.summarize(query.ast);
+
+      expect(summary).toMatchObject([
+        {
+          target: {
+            alias: {
+              type: 'identifier',
+              name: 'a',
+            },
+          },
+        },
+        {
+          target: {
+            alias: {
+              type: 'identifier',
+              name: 'b',
+            },
+          },
+        },
+      ]);
+    });
+
+    it('captures join conditions', () => {
+      const src =
+        'FROM index | LIMIT 1 | JOIN join_index1 AS a ON join_field1 | WHERE b == 2 | JOIN join_index2 AS b ON join_field2, join_field3 | LIMIT 1';
+      const query = EsqlQuery.fromSrc(src);
+      const summary = commands.join.summarize(query.ast);
+
+      expect(summary).toMatchObject([
+        {
+          conditions: [
+            {
+              type: 'column',
+              name: 'join_field1',
+            },
+          ],
+        },
+        {
+          conditions: [
+            {
+              type: 'column',
+              name: 'join_field2',
+            },
+            {
+              type: 'column',
+              name: 'join_field3',
+            },
+          ],
+        },
+      ]);
+    });
+  });
+});
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/mutate/commands/join/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/mutate/commands/join/index.ts
@@ -0,0 +1,101 @@
+/*
+ * 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 type { WalkerAstNode } from '../../../walker/walker';
+import { isAsExpression } from '../../../ast/helpers';
+import { Walker } from '../../../walker';
+import type {
+  ESQLAstExpression,
+  ESQLAstJoinCommand,
+  ESQLAstQueryExpression,
+  ESQLCommand,
+  ESQLIdentifier,
+} from '../../../types';
+import * as generic from '../../generic';
+
+/**
+ * Lists all "JOIN" commands in the query AST.
+ *
+ * @param ast The root AST node to search for "JOIN" commands.
+ * @returns A collection of "JOIN" commands.
+ */
+export const list = (ast: ESQLAstQueryExpression): IterableIterator<ESQLAstJoinCommand> => {
+  return generic.commands.list(
+    ast,
+    (cmd) => cmd.name === 'join'
+  ) as IterableIterator<ESQLAstJoinCommand>;
+};
+
+/**
+ * Retrieves the "JOIN" command at the specified index in order of appearance.
+ *
+ * @param ast The root AST node to search for "JOIN" commands.
+ * @param index The index of the "JOIN" command to retrieve.
+ * @returns The "JOIN" command at the specified index, if any.
+ */
+export const byIndex = (ast: ESQLAstQueryExpression, index: number): ESQLCommand | undefined => {
+  return [...list(ast)][index];
+};
+
+const getIdentifier = (node: WalkerAstNode): ESQLIdentifier =>
+  Walker.match(node, {
+    type: 'identifier',
+  }) as ESQLIdentifier;
+
+/**
+ * Summarizes all JOIN commands in the query.
+ *
+ * @param query Query to summarize.
+ * @returns Returns a list of summaries for all JOIN commands in the query in
+ *     order of appearance.
+ */
+export const summarize = (query: ESQLAstQueryExpression): JoinCommandSummary[] => {
+  const summaries: JoinCommandSummary[] = [];
+
+  for (const command of list(query)) {
+    const firstArg = command.args[0];
+    let index: ESQLIdentifier | undefined;
+    let alias: ESQLIdentifier | undefined;
+    const conditions: ESQLAstExpression[] = [];
+
+    if (isAsExpression(firstArg)) {
+      index = getIdentifier(firstArg.args[0]);
+      alias = getIdentifier(firstArg.args[1]);
+    } else {
+      index = getIdentifier(firstArg);
+    }
+
+    const on = generic.commands.options.find(command, ({ name }) => name === 'on');
+
+    conditions.push(...((on?.args || []) as ESQLAstExpression[]));
+
+    const target: JoinCommandTarget = {
+      index: index!,
+      alias,
+    };
+    const summary: JoinCommandSummary = {
+      target,
+      conditions,
+    };
+
+    summaries.push(summary);
+  }
+
+  return summaries;
+};
+
+export interface JoinCommandSummary {
+  target: JoinCommandTarget;
+  conditions: ESQLAstExpression[];
+}
+
+export interface JoinCommandTarget {
+  index: ESQLIdentifier;
+  alias?: ESQLIdentifier;
+}
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/parser/factories/join.ts b/src/platform/packages/shared/kbn-esql-ast/src/parser/factories/join.ts
@@ -40,7 +40,7 @@ export const createJoinCommand = (ctx: JoinCommandContext): ESQLCommand => {
   const command = createCommand('join', ctx);
 
   // Pick-up the <TYPE> of the command.
-  command.commandType = (ctx._type_.text ?? '').toLocaleLowerCase();
+  command.commandType = (ctx._type_?.text ?? 'lookup').toLocaleLowerCase();
 
   const joinTarget = createNodeFromJoinTarget(ctx.joinTarget());
   const joinCondition = ctx.joinCondition();

diff --git a/...ckages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/callbacks.test.ts b/...ckages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/callbacks.test.ts
@@ -29,4 +29,36 @@ describe('FROM', () => {
 
     expect((callbacks.getColumnsFor as any).mock.calls.length).toBe(1);
   });
+
+  test('loads fields from JOIN index', async () => {
+    const { validate, callbacks } = await setup();
+
+    await validate('FROM index1 | JOIN index2 ON field1 | LIMIT 123');
+
+    expect((callbacks.getColumnsFor as any).mock.calls.length).toBe(1);
+
+    const query = (callbacks.getColumnsFor as any).mock.calls[0][0].query as string;
+
+    expect(query.includes('index1')).toBe(true);
+    expect(query.includes('index2')).toBe(true);
+  });
+
+  test('includes all "from" and "join" index for loading fields', async () => {
+    const { validate, callbacks } = await setup();
+
+    await validate(
+      'FROM index1, index2, index3 | JOIN index4 ON field1 | KEEP abc | JOIN index5 ON field2 | LIMIT 123'
+    );
+
+    expect((callbacks.getColumnsFor as any).mock.calls.length).toBe(1);
+
+    const query = (callbacks.getColumnsFor as any).mock.calls[0][0].query as string;
+
+    expect(query.includes('index1')).toBe(true);
+    expect(query.includes('index2')).toBe(true);
+    expect(query.includes('index3')).toBe(true);
+    expect(query.includes('index4')).toBe(true);
+    expect(query.includes('index5')).toBe(true);
+    expect(query.includes('index6')).toBe(false);
+  });
 });