From ce6d90e3945a323ae1306799341cf50b030a1f54 Mon Sep 17 00:00:00 2001 From: Kibana Machine <42973632+kibanamachine@users.noreply.github.com> Date: Fri, 24 Jan 2025 22:44:45 +1100 Subject: [PATCH] [8.x] [ES|QL] Load fields of indices in `JOIN` command (#207375) (#208160) # Backport This will backport the following commits from `main` to `8.x`: - [[ES|QL] Load fields of indices in `JOIN` command (#207375)](https://github.com/elastic/kibana/pull/207375) ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport) Co-authored-by: Vadim Kibana <82822460+vadimkibana@users.noreply.github.com> --- .../shared/kbn-esql-ast/src/ast/helpers.ts | 6 + .../shared/kbn-esql-ast/src/mutate/README.md | 4 + .../kbn-esql-ast/src/mutate/commands/index.ts | 3 +- .../src/mutate/commands/join/index.test.ts | 166 ++++++++++++++++++ .../src/mutate/commands/join/index.ts | 101 +++++++++++ .../kbn-esql-ast/src/parser/factories/join.ts | 2 +- .../validation/__tests__/callbacks.test.ts | 32 ++++ .../src/validation/helpers.ts | 28 ++- 8 files changed, 338 insertions(+), 4 deletions(-) create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/mutate/commands/join/index.test.ts create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/mutate/commands/join/index.ts 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 index 6668c5f82de53..35e1c3997520a 100644 --- 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 => isBinaryExpression(node) && node.name === 'where'; +export const isAsExpression = ( + node: unknown +): node is ESQLBinaryExpression => + isBinaryExpression(node) && node.name === 'as'; + export const isFieldExpression = ( node: unknown ): node is ESQLBinaryExpression => 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 index 546032d248cca..a1066dce5f312 100644 --- 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()` — Find a `STATS` command by index. - `.summarize()` — Summarize all `STATS` commands. - `.summarizeCommand()` — Summarize a specific `STATS` command. + - `.join` + - `.list()` — List all `JOIN` commands. + - `.byIndex()` — Find a `JOIN` command by index. + - `.summarize()` — 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 index 3de5669a43eaf..2347207c979fd 100644 --- 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 new file mode 100644 index 0000000000000..629e15ad76046 --- /dev/null +++ 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 new file mode 100644 index 0000000000000..d0456cd5e3c34 --- /dev/null +++ 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 => { + return generic.commands.list( + ast, + (cmd) => cmd.name === 'join' + ) as IterableIterator; +}; + +/** + * 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 index 9ff5083ac5e28..1cc0d0107de2d 100644 --- 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 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/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/callbacks.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/callbacks.test.ts index 0461036492b94..e87fbd55c2c50 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/callbacks.test.ts +++ b/src/platform/packages/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); + }); }); diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts index 1aec1a4e512d2..00bfe5275a2a3 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts @@ -11,9 +11,11 @@ import type { ESQLAst, ESQLAstItem, ESQLAstMetricsCommand, + ESQLAstQueryExpression, ESQLMessage, ESQLSingleAstItem, } from '@kbn/esql-ast'; +import { mutate } from '@kbn/esql-ast'; import { FunctionDefinition } from '../definitions/types'; import { getAllArrayTypes, getAllArrayValues } from '../shared/helpers'; import { getMessageFromId } from './errors'; @@ -22,11 +24,33 @@ import type { ESQLPolicy, ReferenceMaps } from './types'; export function buildQueryForFieldsFromSource(queryString: string, ast: ESQLAst) { const firstCommand = ast[0]; if (!firstCommand) return ''; + + let query = ''; + if (firstCommand.name === 'metrics') { const metrics = firstCommand as ESQLAstMetricsCommand; - return `FROM ${metrics.sources.map((source) => source.name).join(', ')}`; + query = `FROM ${metrics.sources.map((source) => source.name).join(', ')}`; + } else { + query = queryString.substring(0, firstCommand.location.max + 1); + } + + const joinSummary = mutate.commands.join.summarize({ + type: 'query', + commands: ast, + } as ESQLAstQueryExpression); + const joinIndices = joinSummary.map( + ({ + target: { + index: { name }, + }, + }) => name + ); + + if (joinIndices.length > 0) { + query += `, ${joinIndices.join(', ')}`; } - return queryString.substring(0, firstCommand.location.max + 1); + + return query; } export function buildQueryForFieldsInPolicies(policies: ESQLPolicy[]) {