From a819d659b8d85197767ff48e288f9193bf804299 Mon Sep 17 00:00:00 2001 From: Vadim Kibana <82822460+vadimkibana@users.noreply.github.com> Date: Tue, 8 Oct 2024 13:02:23 +0200 Subject: [PATCH] [ES|QL] AST query and mutation APIs for metadata fields (#195364) ## Summary Partially addresses https://github.com/elastic/kibana/issues/191812 This PR implements some generic ES|QL AST mutation APIs and specifically APIs for working with `FROM` command `METADATA` fields: - `from.metadata.list()` — List all `METADATA` fields. - `from.metadata.find()` — Find a `METADATA` field by name. - `from.metadata.removeByPredicate()` — Remove a `METADATA` field by matching a predicate. - `from.metadata.remove()` — Remove a `METADATA` field by name. - `from.metadata.insert()` — Insert a `METADATA` field. - `from.metadata.upsert()` — Insert `METADATA` field, if it does not exist. ### 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#kibana-release-notes-process) --- packages/kbn-esql-ast/README.md | 1 + packages/kbn-esql-ast/index.ts | 2 + packages/kbn-esql-ast/src/builder/builder.ts | 32 +- packages/kbn-esql-ast/src/builder/types.ts | 2 + packages/kbn-esql-ast/src/mutate/README.md | 36 ++ .../src/mutate/commands/from/index.ts | 12 + .../src/mutate/commands/from/metadata.test.ts | 332 ++++++++++++++++++ .../src/mutate/commands/from/metadata.ts | 206 +++++++++++ .../kbn-esql-ast/src/mutate/commands/index.ts | 12 + .../kbn-esql-ast/src/mutate/generic.test.ts | 113 ++++++ packages/kbn-esql-ast/src/mutate/generic.ts | 198 +++++++++++ packages/kbn-esql-ast/src/mutate/index.ts | 15 + packages/kbn-esql-ast/src/mutate/types.ts | 10 + packages/kbn-esql-ast/src/mutate/util.ts | 52 +++ packages/kbn-esql-ast/src/visitor/visitor.ts | 1 + 15 files changed, 1022 insertions(+), 2 deletions(-) create mode 100644 packages/kbn-esql-ast/src/mutate/README.md create mode 100644 packages/kbn-esql-ast/src/mutate/commands/from/index.ts create mode 100644 packages/kbn-esql-ast/src/mutate/commands/from/metadata.test.ts create mode 100644 packages/kbn-esql-ast/src/mutate/commands/from/metadata.ts create mode 100644 packages/kbn-esql-ast/src/mutate/commands/index.ts create mode 100644 packages/kbn-esql-ast/src/mutate/generic.test.ts create mode 100644 packages/kbn-esql-ast/src/mutate/generic.ts create mode 100644 packages/kbn-esql-ast/src/mutate/index.ts create mode 100644 packages/kbn-esql-ast/src/mutate/types.ts create mode 100644 packages/kbn-esql-ast/src/mutate/util.ts diff --git a/packages/kbn-esql-ast/README.md b/packages/kbn-esql-ast/README.md index f7be5248f2ca0..dcb244af3c381 100644 --- a/packages/kbn-esql-ast/README.md +++ b/packages/kbn-esql-ast/README.md @@ -12,6 +12,7 @@ Contents of this package: - [`walker` — Contains the ES|QL AST `Walker` utility](./src/walker/README.md). - [`visitor` — Contains the ES|QL AST `Visitor` utility](./src/visitor/README.md). - [`pretty_print` — Contains code for formatting AST to text](./src/pretty_print/README.md). +- [`mutate` — Contains code for traversing and mutating the AST.](./src/mutate/README.md). ## Demo diff --git a/packages/kbn-esql-ast/index.ts b/packages/kbn-esql-ast/index.ts index 869d1aea7e0c6..1780b75f29237 100644 --- a/packages/kbn-esql-ast/index.ts +++ b/packages/kbn-esql-ast/index.ts @@ -56,3 +56,5 @@ export { } from './src/pretty_print'; export { EsqlQuery } from './src/query'; + +export * as mutate from './src/mutate'; diff --git a/packages/kbn-esql-ast/src/builder/builder.ts b/packages/kbn-esql-ast/src/builder/builder.ts index ecc375f885d8b..ece92fbcd7d5e 100644 --- a/packages/kbn-esql-ast/src/builder/builder.ts +++ b/packages/kbn-esql-ast/src/builder/builder.ts @@ -12,7 +12,9 @@ import { ESQLAstComment, ESQLAstQueryExpression, + ESQLColumn, ESQLCommand, + ESQLCommandOption, ESQLDecimalLiteral, ESQLInlineCast, ESQLIntegerLiteral, @@ -20,7 +22,7 @@ import { ESQLLocation, ESQLSource, } from '../types'; -import { AstNodeParserFields, AstNodeTemplate } from './types'; +import { AstNodeParserFields, AstNodeTemplate, PartialFields } from './types'; export namespace Builder { /** @@ -38,16 +40,29 @@ export namespace Builder { }); export const command = ( - template: AstNodeTemplate, + template: PartialFields, 'args'>, fromParser?: Partial ): ESQLCommand => { return { ...template, ...Builder.parserFields(fromParser), + args: template.args ?? [], type: 'command', }; }; + export const option = ( + template: PartialFields, 'args'>, + fromParser?: Partial + ): ESQLCommandOption => { + return { + ...template, + ...Builder.parserFields(fromParser), + args: template.args ?? [], + type: 'option', + }; + }; + export const comment = ( subtype: ESQLAstComment['subtype'], text: string, @@ -85,6 +100,19 @@ export namespace Builder { }; }; + export const column = ( + template: Omit, 'name' | 'quoted'>, + fromParser?: Partial + ): ESQLColumn => { + return { + ...template, + ...Builder.parserFields(fromParser), + quoted: false, + name: template.parts.join('.'), + type: 'column', + }; + }; + export const inlineCast = ( template: Omit, 'name'>, fromParser?: Partial diff --git a/packages/kbn-esql-ast/src/builder/types.ts b/packages/kbn-esql-ast/src/builder/types.ts index be1c1c0d6d458..2713a15fddc0f 100644 --- a/packages/kbn-esql-ast/src/builder/types.ts +++ b/packages/kbn-esql-ast/src/builder/types.ts @@ -29,3 +29,5 @@ export type AstNodeTemplate = Omit< 'type' | 'text' | 'location' | 'incomplete' > & Partial>; + +export type PartialFields = Omit & Partial>; diff --git a/packages/kbn-esql-ast/src/mutate/README.md b/packages/kbn-esql-ast/src/mutate/README.md new file mode 100644 index 0000000000000..8c38bb72ca226 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/README.md @@ -0,0 +1,36 @@ +# Mutation API + +The ES|QL mutation API provides methods to navigate and modify the AST. + + +## Usage + +For example, insert a `FROM` command `METADATA` field: + +```typescript +import { parse, mutate, BasicPrettyPrinter } from '@elastic/esql'; + +const { root } = parse('FROM index METADATA _lang'); + +console.log([...mutate.commands.from.metadata.list(root)]); // [ '_lang' ] + +mutate.commands.from.metadata.upsert(root, '_id'); + +console.log([...mutate.commands.from.metadata.list(root)]); // [ '_lang', '_id' ] + +const src = BasicPrettyPrinter.print(root); + +console.log(src); // FROM index METADATA _lang, _id +``` + + +## API + +- `.commands.from.metadata.list()` — List all `METADATA` fields. +- `.commands.from.metadata.find()` — Find a `METADATA` field by name. +- `.commands.from.metadata.removeByPredicate()` — Remove a `METADATA` + field by matching a predicate. +- `.commands.from.metadata.remove()` — Remove a `METADATA` field by name. +- `.commands.from.metadata.insert()` — Insert a `METADATA` field. +- `.commands.from.metadata.upsert()` — Insert `METADATA` field, if it does + not exist. diff --git a/packages/kbn-esql-ast/src/mutate/commands/from/index.ts b/packages/kbn-esql-ast/src/mutate/commands/from/index.ts new file mode 100644 index 0000000000000..df76e072b346e --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/from/index.ts @@ -0,0 +1,12 @@ +/* + * 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 metadata from './metadata'; + +export { metadata }; diff --git a/packages/kbn-esql-ast/src/mutate/commands/from/metadata.test.ts b/packages/kbn-esql-ast/src/mutate/commands/from/metadata.test.ts new file mode 100644 index 0000000000000..b6cb485395a6c --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/from/metadata.test.ts @@ -0,0 +1,332 @@ +/* + * 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 { parse } from '../../../parser'; +import { BasicPrettyPrinter } from '../../../pretty_print'; +import * as commands from '..'; + +describe('commands.from.metadata', () => { + describe('.list()', () => { + it('returns empty array on no metadata in query', () => { + const src = 'FROM index | WHERE a = b | LIMIT 123'; + const { root } = parse(src); + const column = [...commands.from.metadata.list(root)]; + + expect(column.length).toBe(0); + }); + + it('returns a single METADATA field', () => { + const src = 'FROM index METADATA a'; + const { root } = parse(src); + const [column] = [...commands.from.metadata.list(root)][0]; + + expect(column).toMatchObject({ + type: 'column', + parts: ['a'], + }); + }); + + it('returns all METADATA fields', () => { + const src = 'FROM index METADATA a, b, _id, _lang | STATS avg(a) as avg_a | LIMIT 88'; + const { root } = parse(src); + const columns = [...commands.from.metadata.list(root)].map(([column]) => column); + + expect(columns).toMatchObject([ + { + type: 'column', + parts: ['a'], + }, + { + type: 'column', + parts: ['b'], + }, + { + type: 'column', + parts: ['_id'], + }, + { + type: 'column', + parts: ['_lang'], + }, + ]); + }); + }); + + describe('.find()', () => { + it('returns undefined if field is not found', () => { + const src = 'FROM index | WHERE a = b | LIMIT 123'; + const { root } = parse(src); + const column = commands.from.metadata.find(root, ['a']); + + expect(column).toBe(undefined); + }); + + it('can find a single field', () => { + const src = 'FROM index METADATA a'; + const { root } = parse(src); + const [column] = commands.from.metadata.find(root, ['a'])!; + + expect(column).toMatchObject({ + type: 'column', + name: 'a', + }); + }); + + it('can find a single METADATA field', () => { + const src = 'FROM index METADATA a, b, c, _lang, _id'; + const { root } = parse(src); + const [column1] = commands.from.metadata.find(root, 'c')!; + const [column2] = commands.from.metadata.find(root, '_id')!; + + expect(column1).toMatchObject({ + type: 'column', + name: 'c', + }); + expect(column2).toMatchObject({ + type: 'column', + name: '_id', + }); + }); + }); + + describe('.remove()', () => { + it('can remove a metadata field from a list', () => { + const src1 = 'FROM index METADATA a, b, c'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index METADATA a, b, c'); + + commands.from.metadata.remove(root, 'b'); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM index METADATA a, c'); + }); + + it('does nothing if field-to-delete does not exist', () => { + const src1 = 'FROM index METADATA a, b, c'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index METADATA a, b, c'); + + commands.from.metadata.remove(root, 'd'); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM index METADATA a, b, c'); + }); + + it('can remove all metadata fields one-by-one', () => { + const src1 = 'FROM index METADATA a, b, c'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index METADATA a, b, c'); + + commands.from.metadata.remove(root, 'b'); + commands.from.metadata.remove(root, 'c'); + commands.from.metadata.remove(root, 'a'); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM index'); + }); + }); + + describe('.insert()', () => { + it('can append a METADATA field', () => { + const src1 = 'FROM index METADATA a'; + const { root } = parse(src1); + + commands.from.metadata.insert(root, 'b'); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index METADATA a, b'); + }); + + it('return inserted `column` node, and parent `option` node', () => { + const src1 = 'FROM index METADATA a'; + const { root } = parse(src1); + + const tuple = commands.from.metadata.insert(root, 'b'); + + expect(tuple).toMatchObject([ + { + type: 'column', + name: 'b', + }, + { + type: 'option', + name: 'metadata', + }, + ]); + }); + + it('can insert at specified position', () => { + const src1 = 'FROM index METADATA a1, a2, a3'; + const { root } = parse(src1); + + commands.from.metadata.insert(root, 'x', 0); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index METADATA x, a1, a2, a3'); + + commands.from.metadata.insert(root, 'y', 2); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM index METADATA x, a1, y, a2, a3'); + + commands.from.metadata.insert(root, 'z', 4); + + const src4 = BasicPrettyPrinter.print(root); + + expect(src4).toBe('FROM index METADATA x, a1, y, a2, z, a3'); + }); + + it('appends element, when insert position too high', () => { + const src1 = 'FROM index METADATA a1, a2, a3'; + const { root } = parse(src1); + + commands.from.metadata.insert(root, 'x', 999); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index METADATA a1, a2, a3, x'); + }); + + it('can insert a field when no METADATA option present', () => { + const src1 = 'FROM index'; + const { root } = parse(src1); + + commands.from.metadata.insert(root, 'x', 999); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index METADATA x'); + + commands.from.metadata.insert(root, 'y', 999); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM index METADATA x, y'); + }); + + it('can inset the same field twice', () => { + const src1 = 'FROM index'; + const { root } = parse(src1); + + commands.from.metadata.insert(root, 'x', 999); + commands.from.metadata.insert(root, 'x', 999); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index METADATA x, x'); + }); + }); + + describe('.upsert()', () => { + it('can append a METADATA field', () => { + const src1 = 'FROM index METADATA a'; + const { root } = parse(src1); + + commands.from.metadata.upsert(root, 'b'); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index METADATA a, b'); + }); + + it('return inserted `column` node, and parent `option` node', () => { + const src1 = 'FROM index METADATA a'; + const { root } = parse(src1); + + const tuple = commands.from.metadata.upsert(root, 'b'); + + expect(tuple).toMatchObject([ + { + type: 'column', + name: 'b', + }, + { + type: 'option', + name: 'metadata', + }, + ]); + }); + + it('can insert at specified position', () => { + const src1 = 'FROM index METADATA a1, a2, a3'; + const { root } = parse(src1); + + commands.from.metadata.upsert(root, 'x', 0); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index METADATA x, a1, a2, a3'); + + commands.from.metadata.upsert(root, 'y', 2); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM index METADATA x, a1, y, a2, a3'); + + commands.from.metadata.upsert(root, 'z', 4); + + const src4 = BasicPrettyPrinter.print(root); + + expect(src4).toBe('FROM index METADATA x, a1, y, a2, z, a3'); + }); + + it('appends element, when insert position too high', () => { + const src1 = 'FROM index METADATA a1, a2, a3'; + const { root } = parse(src1); + + commands.from.metadata.upsert(root, 'x', 999); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index METADATA a1, a2, a3, x'); + }); + + it('can insert a field when no METADATA option present', () => { + const src1 = 'FROM index'; + const { root } = parse(src1); + + commands.from.metadata.upsert(root, 'x', 999); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index METADATA x'); + + commands.from.metadata.upsert(root, 'y', 999); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM index METADATA x, y'); + }); + + it('does not insert a field if it is already present', () => { + const src1 = 'FROM index'; + const { root } = parse(src1); + + commands.from.metadata.upsert(root, 'x', 999); + commands.from.metadata.upsert(root, 'x', 999); + commands.from.metadata.upsert(root, 'x', 999); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index METADATA x'); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/mutate/commands/from/metadata.ts b/packages/kbn-esql-ast/src/mutate/commands/from/metadata.ts new file mode 100644 index 0000000000000..5892b028823aa --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/from/metadata.ts @@ -0,0 +1,206 @@ +/* + * 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 { ESQLAstQueryExpression, ESQLColumn, ESQLCommandOption } from '../../../types'; +import { Visitor } from '../../../visitor'; +import { cmpArr, findByPredicate } from '../../util'; +import * as generic from '../../generic'; +import { Builder } from '../../../builder'; +import type { Predicate } from '../../types'; + +/** + * Returns all METADATA field AST nodes and their corresponding parent command + * option nodes. + * + * @param ast The root AST node to search for metadata fields. + * @returns A collection of [column, option] pairs for each metadata field found. + */ +export const list = ( + ast: ESQLAstQueryExpression +): IterableIterator<[ESQLColumn, ESQLCommandOption]> => { + type ReturnExpression = IterableIterator; + type ReturnCommand = IterableIterator<[ESQLColumn, ESQLCommandOption]>; + + return new Visitor() + .on('visitExpression', function* (): ReturnExpression {}) + .on('visitColumnExpression', function* (ctx): ReturnExpression { + yield ctx.node; + }) + .on('visitCommandOption', function* (ctx): ReturnCommand { + if (ctx.node.name !== 'metadata') { + return; + } + for (const args of ctx.visitArguments()) { + for (const column of args) { + yield [column, ctx.node]; + } + } + }) + .on('visitFromCommand', function* (ctx): ReturnCommand { + for (const options of ctx.visitOptions()) { + yield* options; + } + }) + .on('visitCommand', function* (): ReturnCommand {}) + .on('visitQuery', function* (ctx): ReturnCommand { + for (const command of ctx.visitCommands()) { + yield* command; + } + }) + .visitQuery(ast); +}; + +/** + * Find a METADATA field by its name or parts. + * + * @param ast The root AST node to search for metadata fields. + * @param fieldName The name or parts of the field to find. + * @returns A 2-tuple containing the column and the option it was found in, or + * `undefined` if the field was not found. + */ +export const find = ( + ast: ESQLAstQueryExpression, + fieldName: string | string[] +): [ESQLColumn, ESQLCommandOption] | undefined => { + if (typeof fieldName === 'string') { + fieldName = [fieldName]; + } + + const predicate: Predicate<[ESQLColumn, unknown]> = ([field]) => + cmpArr(field.parts, fieldName as string[]); + + return findByPredicate(list(ast), predicate); +}; + +/** + * Removes the first found METADATA field that satisfies the predicate. + * + * @param ast The root AST node to search for metadata fields. + * @param predicate The predicate function to filter fields. + * @returns The removed column and option, if any. + */ +export const removeByPredicate = ( + ast: ESQLAstQueryExpression, + predicate: Predicate +): [column: ESQLColumn, option: ESQLCommandOption] | undefined => { + const tuple = findByPredicate(list(ast), ([field]) => predicate(field)); + + if (!tuple) { + return; + } + + const [column, option] = tuple; + const index = option.args.indexOf(column); + + if (index === -1) { + return; + } + + option.args.splice(index, 1); + + if (option.args.length === 0) { + generic.removeCommandOption(ast, option); + } + + return tuple; +}; + +/** + * Removes the first METADATA field that matches the given name and returns + * a 2-tuple (the column and the option it was removed from). + * + * @param ast The root AST node to search for metadata fields. + * @param fieldName The name or parts of the field to remove. + * @returns The removed column and option, if any. + */ +export const remove = ( + ast: ESQLAstQueryExpression, + fieldName: string | string[] +): [column: ESQLColumn, option: ESQLCommandOption] | undefined => { + if (typeof fieldName === 'string') { + fieldName = [fieldName]; + } + + return removeByPredicate(ast, (field) => cmpArr(field.parts, fieldName as string[])); +}; + +/** + * Insert into a specific position or append a `METADATA` field to the `FROM` + * command. + * + * @param ast The root AST node. + * @param fieldName Field name or parts as an array, e.g. `['a', 'b']`. + * @param index Position to insert the field at. If `-1` or not specified, the + * field will be appended. + * @returns If the field was successfully inserted, returns a 2-tuple containing + * the column and the option it was inserted into. Otherwise, returns + * `undefined`. + */ +export const insert = ( + ast: ESQLAstQueryExpression, + fieldName: string | string[], + index: number = -1 +): [column: ESQLColumn, option: ESQLCommandOption] | undefined => { + let option = generic.findCommandOptionByName(ast, 'from', 'metadata'); + + if (!option) { + const command = generic.findCommandByName(ast, 'from'); + + if (!command) { + return; + } + + option = generic.insertCommandOption(command, 'metadata'); + } + + const parts: string[] = typeof fieldName === 'string' ? [fieldName] : fieldName; + const column = Builder.expression.column({ parts }); + + if (index === -1) { + option.args.push(column); + } else { + option.args.splice(index, 0, column); + } + + return [column, option]; +}; + +/** + * The `.upsert()` method works like `.insert()`, but will not insert a field + * if it already exists. + * + * @param ast The root AST node. + * @param fieldName The field name or parts as an array, e.g. `['a', 'b']`. + * @param index Position to insert the field at. If `-1` or not specified, the + * field will be appended. + * @returns If the field was successfully inserted, returns a 2-tuple containing + * the column and the option it was inserted into. Otherwise, returns + * `undefined`. + */ +export const upsert = ( + ast: ESQLAstQueryExpression, + fieldName: string | string[], + index: number = -1 +): [column: ESQLColumn, option: ESQLCommandOption] | undefined => { + const option = generic.findCommandOptionByName(ast, 'from', 'metadata'); + + if (option) { + const parts = Array.isArray(fieldName) ? fieldName : [fieldName]; + const existing = Walker.find( + option, + (node) => node.type === 'column' && cmpArr(node.parts, parts) + ); + if (existing) { + return undefined; + } + } + + return insert(ast, fieldName, index); +}; diff --git a/packages/kbn-esql-ast/src/mutate/commands/index.ts b/packages/kbn-esql-ast/src/mutate/commands/index.ts new file mode 100644 index 0000000000000..cc3b7f446fa88 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/index.ts @@ -0,0 +1,12 @@ +/* + * 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 from from './from'; + +export { from }; diff --git a/packages/kbn-esql-ast/src/mutate/generic.test.ts b/packages/kbn-esql-ast/src/mutate/generic.test.ts new file mode 100644 index 0000000000000..14d951db1bccb --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/generic.test.ts @@ -0,0 +1,113 @@ +/* + * 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 { parse } from '../parser'; +import { BasicPrettyPrinter } from '../pretty_print'; +import * as generic from './generic'; + +describe('generic', () => { + describe('.listCommands()', () => { + it('lists all commands', () => { + const src = 'FROM index | WHERE a == b | LIMIT 123'; + const { root } = parse(src); + const commands = [...generic.listCommands(root)].map((cmd) => cmd.name); + + expect(commands).toEqual(['from', 'where', 'limit']); + }); + }); + + describe('.findCommand()', () => { + it('can the first command', () => { + const src = 'FROM index | WHERE a == b | LIMIT 123'; + const { root } = parse(src); + const command = generic.findCommand(root, (cmd) => cmd.name === 'from'); + + expect(command).toMatchObject({ + type: 'command', + name: 'from', + args: [ + { + type: 'source', + }, + ], + }); + }); + + it('can the last command', () => { + const src = 'FROM index | WHERE a == b | LIMIT 123'; + const { root } = parse(src); + const command = generic.findCommand(root, (cmd) => cmd.name === 'limit'); + + expect(command).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + }, + ], + }); + }); + + it('find the specific of multiple commands', () => { + const src = 'FROM index | WHERE a == b | LIMIT 1 | LIMIT 2 | LIMIT 3'; + const { root } = parse(src); + const command = generic.findCommand( + root, + (cmd) => cmd.name === 'limit' && (cmd.args?.[0] as any).value === 2 + ); + + expect(command).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 2, + }, + ], + }); + }); + }); + + describe('.findCommandOptionByName()', () => { + it('can the find a command option', () => { + const src = 'FROM index METADATA _score'; + const { root } = parse(src); + const option = generic.findCommandOptionByName(root, 'from', 'metadata'); + + expect(option).toMatchObject({ + type: 'option', + name: 'metadata', + }); + }); + + it('returns undefined if there is no option', () => { + const src = 'FROM index'; + const { root } = parse(src); + const option = generic.findCommandOptionByName(root, 'from', 'metadata'); + + expect(option).toBe(undefined); + }); + }); + + describe('.removeCommandOption()', () => { + it('can remove existing command option', () => { + const src = 'FROM index METADATA _score'; + const { root } = parse(src); + const option = generic.findCommandOptionByName(root, 'from', 'metadata'); + + generic.removeCommandOption(root, option!); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index'); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/mutate/generic.ts b/packages/kbn-esql-ast/src/mutate/generic.ts new file mode 100644 index 0000000000000..968eaf84f4a46 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/generic.ts @@ -0,0 +1,198 @@ +/* + * 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 { Builder } from '../builder'; +import { ESQLAstQueryExpression, ESQLCommand, ESQLCommandOption } from '../types'; +import { Visitor } from '../visitor'; +import { Predicate } from './types'; + +/** + * Returns an iterator for all command AST nodes in the query. If a predicate is + * provided, only commands that satisfy the predicate will be returned. + * + * @param ast Root AST node to search for commands. + * @param predicate Optional predicate to filter commands. + * @returns A list of commands found in the AST. + */ +export const listCommands = ( + ast: ESQLAstQueryExpression, + predicate?: Predicate +): IterableIterator => { + return new Visitor() + .on('visitQuery', function* (ctx): IterableIterator { + for (const cmd of ctx.commands()) { + if (!predicate || predicate(cmd)) { + yield cmd; + } + } + }) + .visitQuery(ast); +}; + +/** + * Returns the first command AST node at a given index in the query that + * satisfies the predicate. If no index is provided, the first command found + * will be returned. + * + * @param ast Root AST node to search for commands. + * @param predicate Optional predicate to filter commands. + * @param index The index of the command to return. + * @returns The command found in the AST, if any. + */ +export const findCommand = ( + ast: ESQLAstQueryExpression, + predicate?: Predicate, + index: number = 0 +): ESQLCommand | undefined => { + for (const cmd of listCommands(ast, predicate)) { + if (!index) { + return cmd; + } + + index--; + } + + return undefined; +}; + +/** + * Returns the first command option AST node that satisfies the predicate. + * + * @param command The command AST node to search for options. + * @param predicate The predicate to filter options. + * @returns The option found in the command, if any. + */ +export const findCommandOption = ( + command: ESQLCommand, + predicate: Predicate +): ESQLCommandOption | undefined => { + return new Visitor() + .on('visitCommand', (ctx): ESQLCommandOption | undefined => { + for (const opt of ctx.options()) { + if (predicate(opt)) { + return opt; + } + } + + return undefined; + }) + .visitCommand(command); +}; + +/** + * Returns the first command AST node at a given index with a given name in the + * query. If no index is provided, the first command found will be returned. + * + * @param ast Root AST node to search for commands. + * @param commandName The name of the command to find. + * @param index The index of the command to return. + * @returns The command found in the AST, if any. + */ +export const findCommandByName = ( + ast: ESQLAstQueryExpression, + commandName: string, + index: number = 0 +): ESQLCommand | undefined => { + return findCommand(ast, (cmd) => cmd.name === commandName, index); +}; + +/** + * Returns the first command option AST node with a given name in the query. + * + * @param ast The root AST node to search for command options. + * @param commandName Command name to search for. + * @param optionName Option name to search for. + * @returns The option found in the command, if any. + */ +export const findCommandOptionByName = ( + ast: ESQLAstQueryExpression, + commandName: string, + optionName: string +): ESQLCommandOption | undefined => { + const command = findCommand(ast, (cmd) => cmd.name === commandName); + + if (!command) { + return undefined; + } + + return findCommandOption(command, (opt) => opt.name === optionName); +}; + +/** + * Inserts a command option into the command's arguments list. The option can + * be specified as a string or an AST node. + * + * @param command The command AST node to insert the option into. + * @param option The option to insert. + * @returns The inserted option. + */ +export const insertCommandOption = ( + command: ESQLCommand, + option: string | ESQLCommandOption +): ESQLCommandOption => { + if (typeof option === 'string') { + option = Builder.option({ name: option }); + } + + command.args.push(option); + + return option; +}; + +/** + * Removes the first command option from the command's arguments list that + * satisfies the predicate. + * + * @param command The command AST node to remove the option from. + * @param predicate The predicate to filter options. + * @returns The removed option, if any. + */ +export const removeCommandOption = ( + ast: ESQLAstQueryExpression, + option: ESQLCommandOption +): boolean => { + return new Visitor() + .on('visitCommandOption', (ctx): boolean => { + return ctx.node === option; + }) + .on('visitCommand', (ctx): boolean => { + let target: undefined | ESQLCommandOption; + + for (const opt of ctx.options()) { + if (opt === option) { + target = opt; + break; + } + } + + if (!target) { + return false; + } + + const index = ctx.node.args.indexOf(target); + + if (index === -1) { + return false; + } + + ctx.node.args.splice(index, 1); + + return true; + }) + .on('visitQuery', (ctx): boolean => { + for (const success of ctx.visitCommands()) { + if (success) { + return true; + } + } + + return false; + }) + .visitQuery(ast); +}; diff --git a/packages/kbn-esql-ast/src/mutate/index.ts b/packages/kbn-esql-ast/src/mutate/index.ts new file mode 100644 index 0000000000000..da312eb79418a --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/index.ts @@ -0,0 +1,15 @@ +/* + * 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". + */ + +export * from './types'; + +import * as generic from './generic'; +import * as commands from './commands'; + +export { generic, commands }; diff --git a/packages/kbn-esql-ast/src/mutate/types.ts b/packages/kbn-esql-ast/src/mutate/types.ts new file mode 100644 index 0000000000000..14cf5d5867d38 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/types.ts @@ -0,0 +1,10 @@ +/* + * 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". + */ + +export type Predicate = (item: T) => boolean; diff --git a/packages/kbn-esql-ast/src/mutate/util.ts b/packages/kbn-esql-ast/src/mutate/util.ts new file mode 100644 index 0000000000000..24f10c1b22f93 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/util.ts @@ -0,0 +1,52 @@ +/* + * 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 { Predicate } from './types'; + +/** + * Find the first item in an iterable (such as array) that matches a predicate. + * + * @param iterable List of items to search through. + * @param predicate Function to determine if an item is the one we are looking + * for. + * @returns The first item that matches the predicate, or undefined if no item + * matches. + */ +export const findByPredicate = ( + iterable: IterableIterator, + predicate: Predicate +): T | undefined => { + for (const item of iterable) { + if (predicate(item)) { + return item; + } + } + return undefined; +}; + +/** + * Shallowly compares two arrays for equality. + * + * @param a The first array to compare. + * @param b The second array to compare. + * @returns True if the arrays are equal, false otherwise. + */ +export const cmpArr = (a: T[], b: T[]): boolean => { + const length = a.length; + if (length !== b.length) { + return false; + } + + for (let i = 0; i < length; i++) { + if (a[i] !== b[i]) { + return false; + } + } + return true; +}; diff --git a/packages/kbn-esql-ast/src/visitor/visitor.ts b/packages/kbn-esql-ast/src/visitor/visitor.ts index 523325b49457c..8fc454ea5e3f7 100644 --- a/packages/kbn-esql-ast/src/visitor/visitor.ts +++ b/packages/kbn-esql-ast/src/visitor/visitor.ts @@ -247,6 +247,7 @@ export class Visitor< ? Builder.expression.query(nodeOrCommands) : nodeOrCommands; const queryContext = new QueryVisitorContext(this.ctx, node, null); + return this.visit(queryContext, input); }