From 17e9300a6648fb435e3b23d219d13ad76f07d7bc Mon Sep 17 00:00:00 2001 From: Vadim Kibana <82822460+vadimkibana@users.noreply.github.com> Date: Tue, 12 Nov 2024 12:59:36 +0100 Subject: [PATCH] [ES|QL] `STATS` command APIs (#199322) ## Summary Partially addresses https://github.com/elastic/kibana/issues/191812 This PR implement higher-level convenience methods for working with `STATS` commands: - `commands.stats.list()` - iterates over all `STATS` commands. - `commands.stats.byIndex()` - retrieves the Nth `STATS` command. - `commands.stats.summarize()` - returns summary about fields used in aggregates and grouping for all `STATS` commands in the query. - `commands.stats.summarizeCommand()` - same as `.summarize()`, but returns a summary only about the requested command. Usage: ```ts const query = EsqlQuery.fromSrc('FROM index | STATS a = max(b)'); const summary = commands.stats.summarize(query); // [ { aggregates: { a: { fields: ['b'] }} ] ``` ### Checklist Delete any items that are not applicable to this PR. - [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 ### For maintainers - [x] This was checked for breaking API changes and was [labeled appropriately](https://www.elastic.co/guide/en/kibana/master/contributing.html#_add_your_labels) --- .../kbn-esql-ast/src/mutate/commands/index.ts | 3 +- .../src/mutate/commands/stats/index.test.ts | 317 ++++++++++++++++++ .../src/mutate/commands/stats/index.ts | 234 +++++++++++++ packages/kbn-esql-ast/src/query/query.ts | 4 + 4 files changed, 557 insertions(+), 1 deletion(-) create mode 100644 packages/kbn-esql-ast/src/mutate/commands/stats/index.test.ts create mode 100644 packages/kbn-esql-ast/src/mutate/commands/stats/index.ts diff --git a/packages/kbn-esql-ast/src/mutate/commands/index.ts b/packages/kbn-esql-ast/src/mutate/commands/index.ts index 9e2599c49345..8068aa5d3ef9 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/index.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/index.ts @@ -10,5 +10,6 @@ import * as from from './from'; import * as limit from './limit'; import * as sort from './sort'; +import * as stats from './stats'; -export { from, limit, sort }; +export { from, limit, sort, stats }; diff --git a/packages/kbn-esql-ast/src/mutate/commands/stats/index.test.ts b/packages/kbn-esql-ast/src/mutate/commands/stats/index.test.ts new file mode 100644 index 000000000000..6fdafa45e383 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/stats/index.test.ts @@ -0,0 +1,317 @@ +/* + * 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.stats', () => { + describe('.list()', () => { + it('lists all "STATS" commands', () => { + const src = 'FROM index | LIMIT 1 | STATS agg() | LIMIT 2 | STATS max()'; + const query = EsqlQuery.fromSrc(src); + + const nodes = [...commands.stats.list(query.ast)]; + + expect(nodes).toMatchObject([ + { + type: 'command', + name: 'stats', + args: [ + { + type: 'function', + name: 'agg', + }, + ], + }, + { + type: 'command', + name: 'stats', + args: [ + { + type: 'function', + name: 'max', + }, + ], + }, + ]); + }); + }); + + describe('.byIndex()', () => { + it('retrieves the specific "STATS" command by index', () => { + const src = 'FROM index | LIMIT 1 | STATS agg() | LIMIT 2 | STATS max()'; + const query = EsqlQuery.fromSrc(src); + + const node1 = commands.stats.byIndex(query.ast, 1); + const node2 = commands.stats.byIndex(query.ast, 0); + + expect(node1).toMatchObject({ + type: 'command', + name: 'stats', + args: [ + { + type: 'function', + name: 'max', + }, + ], + }); + expect(node2).toMatchObject({ + type: 'command', + name: 'stats', + args: [ + { + type: 'function', + name: 'agg', + }, + ], + }); + }); + }); + + describe('.summarizeCommand()', () => { + it('returns summary of a simple field, defined through assignment', () => { + const src = 'FROM index | STATS foo = agg(bar)'; + const query = EsqlQuery.fromSrc(src); + + const command = commands.stats.byIndex(query.ast, 0)!; + const summary = commands.stats.summarizeCommand(query, command); + + expect(summary).toMatchObject({ + command, + aggregates: { + foo: { + arg: { + type: 'function', + name: '=', + }, + field: 'foo', + column: { + type: 'column', + name: 'foo', + }, + definition: { + type: 'function', + name: 'agg', + args: [ + { + type: 'column', + name: 'bar', + }, + ], + }, + terminals: [ + { + type: 'column', + name: 'bar', + }, + ], + fields: ['bar'], + }, + }, + }); + }); + + it('can summarize field defined without assignment', () => { + const src = 'FROM index | STATS agg( /* haha 😅 */ max(foo), bar, baz)'; + const query = EsqlQuery.fromSrc(src); + + const command = commands.stats.byIndex(query.ast, 0)!; + const summary = commands.stats.summarizeCommand(query, command); + + expect(summary).toMatchObject({ + command, + aggregates: { + '`agg( /* haha 😅 */ max(foo), bar, baz)`': { + arg: { + type: 'function', + name: 'agg', + }, + field: '`agg( /* haha 😅 */ max(foo), bar, baz)`', + column: { + type: 'column', + name: '`agg( /* haha 😅 */ max(foo), bar, baz)`', + }, + definition: { + type: 'function', + name: 'agg', + }, + terminals: [ + { + type: 'column', + name: 'foo', + }, + { + type: 'column', + name: 'bar', + }, + { + type: 'column', + name: 'baz', + }, + ], + fields: ['foo', 'bar', 'baz'], + }, + }, + }); + }); + + it('returns a map of stats about two fields', () => { + const src = 'FROM index | STATS foo = agg(f1) + agg(f2), a.b = agg(f3)'; + const query = EsqlQuery.fromSrc(src); + + const command = commands.stats.byIndex(query.ast, 0)!; + const summary = commands.stats.summarizeCommand(query, command); + + expect(summary).toMatchObject({ + aggregates: { + foo: { + field: 'foo', + fields: ['f1', 'f2'], + }, + 'a.b': { + field: 'a.b', + fields: ['f3'], + }, + }, + }); + expect(summary.fields).toEqual(new Set(['f1', 'f2', 'f3'])); + }); + + it('can get de-duplicated list of used fields', () => { + const src = 'FROM index | STATS foo = agg(f1) + agg(f2), a.b = agg(f1)'; + const query = EsqlQuery.fromSrc(src); + + const command = commands.stats.byIndex(query.ast, 0)!; + const summary = commands.stats.summarizeCommand(query, command); + + expect(summary.fields).toEqual(new Set(['f1', 'f2'])); + }); + + describe('params', () => { + it('can use params as source field names', () => { + const src = 'FROM index | STATS foo = agg(f1.?aha) + ?aha(?nested.?param), a.b = agg(f1)'; + const query = EsqlQuery.fromSrc(src); + + const command = commands.stats.byIndex(query.ast, 0)!; + const summary = commands.stats.summarizeCommand(query, command); + + expect(summary).toMatchObject({ + aggregates: { + foo: { + fields: ['f1.?aha', '?nested.?param'], + }, + 'a.b': { + fields: ['f1'], + }, + }, + }); + expect(summary.fields).toEqual(new Set(['f1.?aha', '?nested.?param', 'f1'])); + }); + + it('can use params as destination field names', () => { + const src = 'FROM index | STATS ?dest = agg(asdf) BY asdf'; + const query = EsqlQuery.fromSrc(src); + + const command = commands.stats.byIndex(query.ast, 0)!; + const summary = commands.stats.summarizeCommand(query, command); + + expect(summary).toMatchObject({ + aggregates: { + '?dest': { + fields: ['asdf'], + }, + }, + }); + expect(summary.fields).toEqual(new Set(['asdf'])); + }); + }); + + describe('BY option', () => { + it('can collect fields from the BY option', () => { + const src = 'FROM index | STATS max(1) BY abc'; + const query = EsqlQuery.fromSrc(src); + + const command = commands.stats.byIndex(query.ast, 0)!; + const summary = commands.stats.summarizeCommand(query, command); + + expect(summary.aggregates).toEqual({ + '`max(1)`': expect.any(Object), + }); + expect(summary.fields).toEqual(new Set(['abc'])); + }); + + it('returns all "grouping" fields', () => { + const src = 'FROM index | STATS max(1) BY a, b, c'; + const query = EsqlQuery.fromSrc(src); + + const command = commands.stats.byIndex(query.ast, 0)!; + const summary = commands.stats.summarizeCommand(query, command); + + expect(summary.aggregates).toEqual({ + '`max(1)`': expect.any(Object), + }); + expect(summary.grouping).toMatchObject({ + a: { type: 'column' }, + b: { type: 'column' }, + c: { type: 'column' }, + }); + }); + + it('can have params and quoted fields in grouping', () => { + const src = 'FROM index | STATS max(1) BY `a😎`, ?123, a.?b.?0.`😎`'; + const query = EsqlQuery.fromSrc(src); + + const command = commands.stats.byIndex(query.ast, 0)!; + const summary = commands.stats.summarizeCommand(query, command); + + expect(summary.aggregates).toEqual({ + '`max(1)`': expect.any(Object), + }); + expect(summary.grouping).toMatchObject({ + '`a😎`': { type: 'column' }, + // '?123': { type: 'column' }, + 'a.?b.?0.`😎`': { type: 'column' }, + }); + }); + }); + }); + + describe('.summarize()', () => { + it('can summarize multiple stats commands', () => { + const src = 'FROM index | LIMIT 1 | STATS agg() | LIMIT 2 | STATS max(a, b, c), max2(d.e)'; + const query = EsqlQuery.fromSrc(src); + const summary = commands.stats.summarize(query); + + expect(summary).toMatchObject([ + { + aggregates: { + '`agg()`': { + field: '`agg()`', + fields: [], + }, + }, + fields: new Set([]), + }, + { + aggregates: { + '`max(a, b, c)`': { + field: '`max(a, b, c)`', + fields: ['a', 'b', 'c'], + }, + '`max2(d.e)`': { + field: '`max2(d.e)`', + fields: ['d.e'], + }, + }, + fields: new Set(['a', 'b', 'c', 'd.e']), + }, + ]); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/mutate/commands/stats/index.ts b/packages/kbn-esql-ast/src/mutate/commands/stats/index.ts new file mode 100644 index 000000000000..9a07549c0e0e --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/stats/index.ts @@ -0,0 +1,234 @@ +/* + * 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 { Walker } from '../../../walker'; +import { Visitor } from '../../../visitor'; +import { LeafPrinter } from '../../../pretty_print'; +import { Builder } from '../../../builder'; +import { singleItems } from '../../../visitor/utils'; +import type { + ESQLAstQueryExpression, + ESQLColumn, + ESQLCommand, + ESQLList, + ESQLLiteral, + ESQLProperNode, + ESQLTimeInterval, +} from '../../../types'; +import * as generic from '../../generic'; +import { isColumn, isFunctionExpression } from '../../../ast/helpers'; +import type { EsqlQuery } from '../../../query'; + +/** + * Lists all "LIMIT" commands in the query AST. + * + * @param ast The root AST node to search for "LIMIT" commands. + * @returns A collection of "LIMIT" commands. + */ +export const list = (ast: ESQLAstQueryExpression): IterableIterator => { + return generic.commands.list(ast, (cmd) => cmd.name === 'stats'); +}; + +/** + * Retrieves the "LIMIT" command at the specified index in order of appearance. + * + * @param ast The root AST node to search for "LIMIT" commands. + * @param index The index of the "LIMIT" command to retrieve. + * @returns The "LIMIT" command at the specified index, if any. + */ +export const byIndex = (ast: ESQLAstQueryExpression, index: number): ESQLCommand | undefined => { + return [...list(ast)][index]; +}; + +/** + * Summary of a STATS command. + */ +export interface StatsCommandSummary { + /** + * The "STATS" command AST node from which this summary was produced. + */ + command: ESQLCommand; + + /** + * Summary of the main arguments of the "STATS" command. + */ + aggregates: Record; + + /** + * Summary of the "BY" arguments of the "STATS" command. + */ + grouping: Record; + + /** + * De-duplicated list all of ES|QL-syntax formatted field names from the + * {@link aggregates} and {@link grouping} fields. + */ + fields: Set; +} + +/** + * Summary of STATS command "aggregates" section (main arguments). + * + * STATS [ BY ] + */ +export interface StatsAggregatesSummary { + /** + * STATS command argument AST node (as was parsed). + */ + arg: ESQLProperNode; + + /** + * The field name, correctly formatted, extracted from the AST. + */ + field: string; + + /** + * A `column` AST node, which represents the field name. If no column AST node + * was found, a new one "virtual" column node is created. + */ + column: ESQLColumn; + + /** + * The definition of the field, which is the right-hand side of the `=` + * operator, or the argument itself if no `=` operator is present. + */ + definition: ESQLProperNode; + + /** + * A list of terminal nodes that were found in the definition. + */ + terminals: Array; + + /** + * Correctly formatted list of field names that were found in the {@link terminals}. + */ + fields: string[]; +} + +const summarizeArgParts = ( + query: EsqlQuery, + arg: ESQLProperNode +): [column: ESQLColumn, definition: ESQLProperNode] => { + if (isFunctionExpression(arg) && arg.name === '=' && isColumn(arg.args[0])) { + const [column, definition] = singleItems(arg.args); + + return [column as ESQLColumn, definition as ESQLProperNode]; + } + + const name = [...query.src].slice(arg.location.min, arg.location.max + 1).join(''); + const args = [Builder.identifier({ name })]; + const column = Builder.expression.column({ args }); + + return [column, arg]; +}; + +const summarizeArg = (query: EsqlQuery, arg: ESQLProperNode): StatsAggregatesSummary => { + const [column, definition] = summarizeArgParts(query, arg); + const terminals: StatsAggregatesSummary['terminals'] = []; + const fields: StatsAggregatesSummary['fields'] = []; + + Walker.walk(definition, { + visitLiteral(node) { + terminals.push(node); + }, + visitColumn(node) { + terminals.push(node); + fields.push(LeafPrinter.column(node)); + }, + visitListLiteral(node) { + terminals.push(node); + }, + visitTimeIntervalLiteral(node) { + terminals.push(node); + }, + }); + + const summary: StatsAggregatesSummary = { + arg, + field: LeafPrinter.column(column), + column, + definition, + terminals, + fields, + }; + + return summary; +}; + +/** + * Returns a summary of the STATS command. + * + * @param query Query which contains the AST and source code. + * @param command The STATS command AST node to summarize. + * @returns Summary of the STATS command. + */ +export const summarizeCommand = (query: EsqlQuery, command: ESQLCommand): StatsCommandSummary => { + const aggregates: StatsCommandSummary['aggregates'] = {}; + const grouping: StatsCommandSummary['grouping'] = {}; + const fields: StatsCommandSummary['fields'] = new Set(); + + // Process main arguments, the "aggregates" part of the command. + new Visitor() + .on('visitExpression', (ctx) => { + const summary = summarizeArg(query, ctx.node); + aggregates[summary.field] = summary; + for (const field of summary.fields) fields.add(field); + }) + .on('visitCommand', () => {}) + .on('visitStatsCommand', (ctx) => { + for (const _ of ctx.visitArguments()); + }) + .visitCommand(command); + + // Process the "BY" arguments, the "grouping" part of the command. + new Visitor() + .on('visitExpression', () => {}) + .on('visitColumnExpression', (ctx) => { + const column = ctx.node; + const formatted = LeafPrinter.column(column); + grouping[formatted] = column; + fields.add(formatted); + }) + .on('visitCommandOption', (ctx) => { + if (ctx.node.name !== 'by') return; + for (const _ of ctx.visitArguments()); + }) + .on('visitCommand', () => {}) + .on('visitStatsCommand', (ctx) => { + for (const _ of ctx.visitOptions()); + }) + .visitCommand(command); + + const summary: StatsCommandSummary = { + command, + aggregates, + grouping, + fields, + }; + + return summary; +}; + +/** + * Summarizes all STATS commands in the query. + * + * @param query Query to summarize. + * @returns Returns a list of summaries for all STATS commands in the query in + * order of appearance. + */ +export const summarize = (query: EsqlQuery): StatsCommandSummary[] => { + const summaries: StatsCommandSummary[] = []; + + for (const command of list(query.ast)) { + const summary = summarizeCommand(query, command); + summaries.push(summary); + } + + return summaries; +}; diff --git a/packages/kbn-esql-ast/src/query/query.ts b/packages/kbn-esql-ast/src/query/query.ts index 60435cc64c97..66c9fd58df08 100644 --- a/packages/kbn-esql-ast/src/query/query.ts +++ b/packages/kbn-esql-ast/src/query/query.ts @@ -15,6 +15,10 @@ import { WrappingPrettyPrinterOptions, } from '../pretty_print/wrapping_pretty_printer'; +/** + * Represents a parsed or programmatically created ES|QL query. Keeps track of + * the AST, source code, and optionally lexer tokens. + */ export class EsqlQuery { public static readonly fromSrc = (src: string, opts?: ParseOptions): EsqlQuery => { const { root, tokens } = parse(src, opts);