[8.x] [ES|QL] More AST mutation APIs (#196240) (#196355)

# Backport This will backport the following commits from `main` to `8.x`: - [[ES|QL] More AST mutation APIs (#196240)](#196240)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Vadim Kibana <82822460+vadimkibana@users.noreply.github.com>
elastic · Oct 15, 2024 · 79f9d05 · 79f9d05
1 parent 11f535a
commit 79f9d05
Show file tree

Hide file tree

Showing 12 changed files with 1,003 additions and 13 deletions.
diff --git a/packages/kbn-esql-ast/src/ast/util.ts b/packages/kbn-esql-ast/src/ast/util.ts
@@ -0,0 +1,14 @@
+/*
+ * 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 { ESQLAstNode, ESQLCommandOption } from '../types';
+
+export const isOptionNode = (node: ESQLAstNode): node is ESQLCommandOption => {
+  return !!node && typeof node === 'object' && !Array.isArray(node) && node.type === 'option';
+};
diff --git a/packages/kbn-esql-ast/src/builder/builder.ts b/packages/kbn-esql-ast/src/builder/builder.ts
@@ -100,6 +100,23 @@ export namespace Builder {
       };
     };
 
+    export const indexSource = (
+      index: string,
+      cluster?: string,
+      template?: Omit<AstNodeTemplate<ESQLSource>, 'name' | 'index' | 'cluster'>,
+      fromParser?: Partial<AstNodeParserFields>
+    ): ESQLSource => {
+      return {
+        ...template,
+        ...Builder.parserFields(fromParser),
+        index,
+        cluster,
+        name: (cluster ? cluster + ':' : '') + index,
+        sourceType: 'index',
+        type: 'source',
+      };
+    };
+
     export const column = (
       template: Omit<AstNodeTemplate<ESQLColumn>, 'name' | 'quoted'>,
       fromParser?: Partial<AstNodeParserFields>

diff --git a/packages/kbn-esql-ast/src/mutate/README.md b/packages/kbn-esql-ast/src/mutate/README.md
@@ -26,11 +26,37 @@ console.log(src); // FROM index METADATA _lang, _id
 
 ## API
 
-- `.commands.from.metadata.list()` &mdash; List all `METADATA` fields.
-- `.commands.from.metadata.find()` &mdash; Find a `METADATA` field by name.
-- `.commands.from.metadata.removeByPredicate()` &mdash; Remove a `METADATA`
-  field by matching a predicate.
-- `.commands.from.metadata.remove()` &mdash; Remove a `METADATA` field by name.
-- `.commands.from.metadata.insert()` &mdash; Insert a `METADATA` field.
-- `.commands.from.metadata.upsert()` &mdash; Insert `METADATA` field, if it does
-  not exist.
+- `.generic`
+  - `.listCommands()` &mdash; Lists all commands. Returns an iterator.
+  - `.findCommand()` &mdash; Finds a specific command by a predicate function.
+  - `.findCommandOption()` &mdash; Finds a specific command option by a predicate function.
+  - `.findCommandByName()` &mdash; Finds a specific command by name.
+  - `.findCommandOptionByName()` &mdash; Finds a specific command option by name.
+  - `.appendCommand()` &mdash; Add a new command to the AST.
+  - `.appendCommandOption()` &mdash; Add a new command option to a command.
+  - `.appendCommandArgument()` &mdash; Add a new main command argument to a command.
+  - `.removeCommand()` &mdash; Remove a command from the AST.
+  - `.removeCommandOption()` &mdash; Remove a command option from the AST.
+  - `.removeCommandArgument()` &mdash; Remove a command argument from the AST.
+- `.commands`
+  - `.from`
+    - `.sources`
+      - `.list()` &mdash; List all `FROM` sources.
+      - `.find()` &mdash; Find a source by name.
+      - `.remove()` &mdash; Remove a source by name.
+      - `.insert()` &mdash; Insert a source.
+      - `.upsert()` &mdash; Insert a source, if it does not exist.
+    - `.metadata`
+      - `.list()` &mdash; List all `METADATA` fields.
+      - `.find()` &mdash; Find a `METADATA` field by name.
+      - `.removeByPredicate()` &mdash; Remove a `METADATA` field by matching a predicate function.
+      - `.remove()` &mdash; Remove a `METADATA` field by name.
+      - `.insert()` &mdash; Insert a `METADATA` field.
+      - `.upsert()` &mdash; Insert `METADATA` field, if it does not exist.
+  - `.limit`
+    - `.list()` &mdash; List all `LIMIT` commands.
+    - `.byIndex()` &mdash; Find a `LIMIT` command by index.
+    - `.find()` &mdash; Find a `LIMIT` command by a predicate function.
+    - `.remove()` &mdash; Remove a `LIMIT` command by index.
+    - `.set()` &mdash; Set the limit value of a specific `LIMIT` command.
+    - `.upsert()` &mdash; Insert a `LIMIT` command, or update the limit value if it already exists.
diff --git a/packages/kbn-esql-ast/src/mutate/commands/from/index.ts b/packages/kbn-esql-ast/src/mutate/commands/from/index.ts
@@ -7,6 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
+import * as sources from './sources';
 import * as metadata from './metadata';
 
-export { metadata };
+export { sources, metadata };
diff --git a/packages/kbn-esql-ast/src/mutate/commands/from/metadata.ts b/packages/kbn-esql-ast/src/mutate/commands/from/metadata.ts
@@ -157,7 +157,7 @@ export const insert = (
       return;
     }
 
-    option = generic.insertCommandOption(command, 'metadata');
+    option = generic.appendCommandOption(command, 'metadata');
   }
 
   const parts: string[] = typeof fieldName === 'string' ? [fieldName] : fieldName;

diff --git a/packages/kbn-esql-ast/src/mutate/commands/from/sources.test.ts b/packages/kbn-esql-ast/src/mutate/commands/from/sources.test.ts
@@ -0,0 +1,246 @@
+/*
+ * 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.sources', () => {
+  describe('.list()', () => {
+    it('returns empty array, if there are no sources', () => {
+      const src = 'ROW 123';
+      const { root } = parse(src);
+      const list = [...commands.from.sources.list(root)];
+
+      expect(list.length).toBe(0);
+    });
+
+    it('returns a single source', () => {
+      const src = 'FROM index METADATA a';
+      const { root } = parse(src);
+      const list = [...commands.from.sources.list(root)];
+
+      expect(list.length).toBe(1);
+      expect(list[0]).toMatchObject({
+        type: 'source',
+      });
+    });
+
+    it('returns all source fields', () => {
+      const src = 'FROM index, index2, cl:index3 METADATA a | LIMIT 88';
+      const { root } = parse(src);
+      const list = [...commands.from.sources.list(root)];
+
+      expect(list).toMatchObject([
+        {
+          type: 'source',
+          index: 'index',
+        },
+        {
+          type: 'source',
+          index: 'index2',
+        },
+        {
+          type: 'source',
+          index: 'index3',
+          cluster: 'cl',
+        },
+      ]);
+    });
+  });
+
+  describe('.find()', () => {
+    it('returns undefined if source is not found', () => {
+      const src = 'FROM index | WHERE a = b | LIMIT 123';
+      const { root } = parse(src);
+      const source = commands.from.sources.find(root, 'abc');
+
+      expect(source).toBe(undefined);
+    });
+
+    it('can find a single source', () => {
+      const src = 'FROM index METADATA a';
+      const { root } = parse(src);
+      const source = commands.from.sources.find(root, 'index')!;
+
+      expect(source).toMatchObject({
+        type: 'source',
+        name: 'index',
+        index: 'index',
+      });
+    });
+
+    it('can find a source withing other sources', () => {
+      const src = 'FROM index, a, b, c:s1, s1, s2 METADATA a, b, c, _lang, _id';
+      const { root } = parse(src);
+      const source1 = commands.from.sources.find(root, 's2')!;
+      const source2 = commands.from.sources.find(root, 's1', 'c')!;
+
+      expect(source1).toMatchObject({
+        type: 'source',
+        name: 's2',
+        index: 's2',
+      });
+      expect(source2).toMatchObject({
+        type: 'source',
+        name: 'c:s1',
+        index: 's1',
+        cluster: 'c',
+      });
+    });
+  });
+
+  describe('.remove()', () => {
+    it('can remove a source from a list', () => {
+      const src1 = 'FROM a, b, c';
+      const { root } = parse(src1);
+      const src2 = BasicPrettyPrinter.print(root);
+
+      expect(src2).toBe('FROM a, b, c');
+
+      commands.from.sources.remove(root, 'b');
+
+      const src3 = BasicPrettyPrinter.print(root);
+
+      expect(src3).toBe('FROM a, c');
+    });
+
+    it('does nothing if source-to-delete does not exist', () => {
+      const src1 = 'FROM a, b, c';
+      const { root } = parse(src1);
+      const src2 = BasicPrettyPrinter.print(root);
+
+      expect(src2).toBe('FROM a, b, c');
+
+      commands.from.sources.remove(root, 'd');
+
+      const src3 = BasicPrettyPrinter.print(root);
+
+      expect(src3).toBe('FROM a, b, c');
+    });
+  });
+
+  describe('.insert()', () => {
+    it('can append a source', () => {
+      const src1 = 'FROM index METADATA a';
+      const { root } = parse(src1);
+
+      commands.from.sources.insert(root, 'index2');
+
+      const src2 = BasicPrettyPrinter.print(root);
+
+      expect(src2).toBe('FROM index, index2 METADATA a');
+    });
+
+    it('can insert at specified position', () => {
+      const src1 = 'FROM a1, a2, a3';
+      const { root } = parse(src1);
+
+      commands.from.sources.insert(root, 'x', '', 0);
+
+      const src2 = BasicPrettyPrinter.print(root);
+
+      expect(src2).toBe('FROM x, a1, a2, a3');
+
+      commands.from.sources.insert(root, 'y', '', 2);
+
+      const src3 = BasicPrettyPrinter.print(root);
+
+      expect(src3).toBe('FROM x, a1, y, a2, a3');
+
+      commands.from.sources.insert(root, 'z', '', 4);
+
+      const src4 = BasicPrettyPrinter.print(root);
+
+      expect(src4).toBe('FROM x, a1, y, a2, z, a3');
+    });
+
+    it('appends element, when insert position too high', () => {
+      const src1 = 'FROM a1, a2, a3';
+      const { root } = parse(src1);
+
+      commands.from.sources.insert(root, 'x', '', 999);
+
+      const src2 = BasicPrettyPrinter.print(root);
+
+      expect(src2).toBe('FROM a1, a2, a3, x');
+    });
+
+    it('can inset the same source twice', () => {
+      const src1 = 'FROM index';
+      const { root } = parse(src1);
+
+      commands.from.sources.insert(root, 'x', '', 999);
+      commands.from.sources.insert(root, 'x', '', 999);
+
+      const src2 = BasicPrettyPrinter.print(root);
+
+      expect(src2).toBe('FROM index, x, x');
+    });
+  });
+
+  describe('.upsert()', () => {
+    it('can append a source', () => {
+      const src1 = 'FROM index METADATA a';
+      const { root } = parse(src1);
+
+      commands.from.sources.upsert(root, 'index2');
+
+      const src2 = BasicPrettyPrinter.print(root);
+
+      expect(src2).toBe('FROM index, index2 METADATA a');
+    });
+
+    it('can upsert at specified position', () => {
+      const src1 = 'FROM a1, a2, a3';
+      const { root } = parse(src1);
+
+      commands.from.sources.upsert(root, 'x', '', 0);
+
+      const src2 = BasicPrettyPrinter.print(root);
+
+      expect(src2).toBe('FROM x, a1, a2, a3');
+
+      commands.from.sources.upsert(root, 'y', '', 2);
+
+      const src3 = BasicPrettyPrinter.print(root);
+
+      expect(src3).toBe('FROM x, a1, y, a2, a3');
+
+      commands.from.sources.upsert(root, 'z', '', 4);
+
+      const src4 = BasicPrettyPrinter.print(root);
+
+      expect(src4).toBe('FROM x, a1, y, a2, z, a3');
+    });
+
+    it('appends element, when upsert position too high', () => {
+      const src1 = 'FROM a1, a2, a3';
+      const { root } = parse(src1);
+
+      commands.from.sources.upsert(root, 'x', '', 999);
+
+      const src2 = BasicPrettyPrinter.print(root);
+
+      expect(src2).toBe('FROM a1, a2, a3, x');
+    });
+
+    it('inserting already existing source is a no-op', () => {
+      const src1 = 'FROM index';
+      const { root } = parse(src1);
+
+      commands.from.sources.upsert(root, 'x', '', 999);
+      commands.from.sources.upsert(root, 'x', '', 999);
+
+      const src2 = BasicPrettyPrinter.print(root);
+
+      expect(src2).toBe('FROM index, x');
+    });
+  });
+});
-Original file line number
+Diff line change
@@ Expand Up / @@ -157,7 +157,7 @@ export const insert = ( @@
           return;
         }
-        option = generic.insertCommandOption(command, 'metadata');
+        option = generic.appendCommandOption(command, 'metadata');
       }
       const parts: string[] = typeof fieldName === 'string' ? [fieldName] : fieldName;
@@ Expand Down @@