[ES|QL] More AST mutation APIs (elastic#196240)

## Summary

Partially addresses elastic#191812

Implements the following high-level ES|QL AST manipulation methods:

- `.generic`
- `.appendCommandArgument()` — Add a new main command argument to
a command.
- `.removeCommandArgument()` — Remove a command argument from the
AST.
- `.commands`
  - `.from`
    - `.sources`
      - `.list()` — List all `FROM` sources.
      - `.find()` — Find a source by name.
      - `.remove()` — Remove a source by name.
      - `.insert()` — Insert a source.
      - `.upsert()` — Insert a source, if it does not exist.
  - `.limit`
    - `.list()` — List all `LIMIT` commands.
    - `.byIndex()` — Find a `LIMIT` command by index.
    - `.find()` — Find a `LIMIT` command by a predicate function.
    - `.remove()` — Remove a `LIMIT` command by index.
- `.set()` — Set the limit value of a specific `LIMIT` command.
- `.upsert()` — Insert a `LIMIT` command, or update the limit
value if it already exists.

### Checklist

Delete any items that are not applicable to this PR.

- [x]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [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)

(cherry picked from commit 10364fb)

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';
+};

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>

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()` — 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.
+- `.generic`
+  - `.listCommands()` — Lists all commands. Returns an iterator.
+  - `.findCommand()` — Finds a specific command by a predicate function.
+  - `.findCommandOption()` — Finds a specific command option by a predicate function.
+  - `.findCommandByName()` — Finds a specific command by name.
+  - `.findCommandOptionByName()` — Finds a specific command option by name.
+  - `.appendCommand()` — Add a new command to the AST.
+  - `.appendCommandOption()` — Add a new command option to a command.
+  - `.appendCommandArgument()` — Add a new main command argument to a command.
+  - `.removeCommand()` — Remove a command from the AST.
+  - `.removeCommandOption()` — Remove a command option from the AST.
+  - `.removeCommandArgument()` — Remove a command argument from the AST.
+- `.commands`
+  - `.from`
+    - `.sources`
+      - `.list()` — List all `FROM` sources.
+      - `.find()` — Find a source by name.
+      - `.remove()` — Remove a source by name.
+      - `.insert()` — Insert a source.
+      - `.upsert()` — Insert a source, if it does not exist.
+    - `.metadata`
+      - `.list()` — List all `METADATA` fields.
+      - `.find()` — Find a `METADATA` field by name.
+      - `.removeByPredicate()` — Remove a `METADATA` field by matching a predicate function.
+      - `.remove()` — Remove a `METADATA` field by name.
+      - `.insert()` — Insert a `METADATA` field.
+      - `.upsert()` — Insert `METADATA` field, if it does not exist.
+  - `.limit`
+    - `.list()` — List all `LIMIT` commands.
+    - `.byIndex()` — Find a `LIMIT` command by index.
+    - `.find()` — Find a `LIMIT` command by a predicate function.
+    - `.remove()` — Remove a `LIMIT` command by index.
+    - `.set()` — Set the limit value of a specific `LIMIT` command.
+    - `.upsert()` — Insert a `LIMIT` command, or update the limit value if it already exists.

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 };

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;

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');
+    });
+  });
+});