[8.x] [ES|QL] AST node synthesizer (elastic#201814) (elastic#202409)

# Backport This will backport the following commits from `main` to `8.x`: - [[ES|QL] AST node synthesizer (elastic#201814)](elastic#201814)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Vadim Kibana <[email protected]>
kibanamachine · Dec 2, 2024 · cff07f7 · cff07f7
1 parent c3ad53e
commit cff07f7
Show file tree

Hide file tree

Showing 16 changed files with 658 additions and 12 deletions.
diff --git a/packages/kbn-esql-ast/README.md b/packages/kbn-esql-ast/README.md
@@ -5,14 +5,32 @@ building, traversal, pretty-printing, and manipulation features on top of a
 custom compact AST representation, which is designed to be resilient to many
 grammar changes.
 
-Contents of this package:
 
-- [`builder` &mdash; Contains the `Builder` class for AST node construction](./src/builder/README.md).
+### Contents of the package
+
+At the lowest level, the package provides a parser that converts ES|QL text into
+an AST representation. Or, you can use the `Builder` class to construct the AST
+manually:
+
 - [`parser` &mdash; Contains text to ES|QL AST parsing code](./src/parser/README.md).
+- [`builder` &mdash; Contains the `Builder` class for AST node construction](./src/builder/README.md).
+
+The *Traversal API* lets you walk the AST. The `Walker` class is a simple
+to use, but the `Visitor` class is more powerful and flexible:
+
 - [`walker` &mdash; Contains the ES|QL AST `Walker` utility](./src/walker/README.md).
 - [`visitor` &mdash; Contains the ES|QL AST `Visitor` utility](./src/visitor/README.md).
+
+Higher-level functionality is provided by the `mutate` and `synth` modules. They
+allow you to traverse and modify the AST, or to easily construct AST nodes:
+
+- [`mutate` &mdash; Contains code for traversing and mutating the AST](./src/mutate/README.md).
+- [`synth` &mdash; Ability to construct AST nodes from template strings](./src/synth/README.md).
+
+The *Pretty-printing API* lets you format the AST to text. There are two
+implementations &mdash; a basic pretty-printer and a wrapping pretty-printer:
+
 - [`pretty_print` &mdash; Contains code for formatting AST to text](./src/pretty_print/README.md).
-- [`mutate` &mdash; 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
@@ -45,6 +45,7 @@ export {
 } from './src/parser';
 
 export { Walker, type WalkerOptions, walk } from './src/walker';
+export * as synth from './src/synth';
 
 export {
   LeafPrinter,

diff --git a/packages/kbn-esql-ast/src/synth/README.md b/packages/kbn-esql-ast/src/synth/README.md
@@ -0,0 +1,124 @@
+# ES|QL `synth` module
+
+The `synth` module lets you easily "synthesize" AST nodes from template strings.
+This is useful when you need to construct AST nodes programmatically, but don't
+want to deal with the complexity of the `Builder` class.
+
+
+## Usage
+
+You can create an assignment expression AST node as simle as:
+
+```ts
+import { synth } from '@kbn/esql-ast';
+
+const node = synth.expr `my.field = max(10, ?my_param)`;
+// { type: 'function', name: '=', args: [ ... ]}
+```
+
+To construct an equivalent AST node using the `Builder` class, you would need to
+write the following code:
+
+```ts
+import { Builder } from '@kbn/esql-ast';
+
+const node = Builder.expression.func.binary('=', [
+  Builder.expression.column({
+    args: [Builder.identifier({ name: 'my' }), Builder.identifier({ name: 'field' })],
+  }),
+  Builder.expression.func.call('max', [
+    Builder.expression.literal.integer(10),
+    Builder.param.named({ value: 'my_param' }),
+  ]),
+]);
+// { type: 'function', name: '=', args: [ ... ]}
+```
+
+You can nest template strings to create more complex AST nodes:
+
+```ts
+const field = synth.expr `my.field`;
+const value = synth.expr `max(10, ?my_param)`;
+
+const assignment = synth.expr`${field} = ${value}`;
+// { type: 'function', name: '=', args: [ 
+//     { type: 'column', args: [ ... ] },
+//     { type: 'function', name: 'max', args: [ ... ] }
+// ]}
+```
+
+Use the `synth.cmd` method to create command nodes:
+
+```ts
+const command = synth.cmd `WHERE my.field == 10`;
+// { type: 'command', name: 'WHERE', args: [ ... ]}
+```
+
+AST nodes created by the synthesizer are pretty-printed when you coerce them to
+a string or call the `toString` method:
+
+```ts
+const command = synth.cmd ` WHERE my.field == 10 `; // { type: 'command', ... }
+String(command); // "WHERE my.field == 10"
+```
+
+
+## Reference
+
+### `synth.expr`
+
+The `synth.expr` synthesizes an expression AST nodes. (*Expressions* are
+basically any thing that can go into a `WHERE` command, like boolean expression,
+function call, literal, params, etc.)
+
+Use it as a function:
+
+```ts
+const node = synth.expr('my.field = max(10, ?my_param)');
+```
+
+Specify parser options:
+
+```ts
+const node = synth.expr('my.field = max(10, ?my_param)',
+  { withFormatting: false });
+```
+
+Use it as a template string tag:
+
+```ts
+const node = synth.expr `my.field = max(10, ?my_param)`;
+```
+
+Specify parser options, when using as a template string tag:
+
+```ts
+const node = synth.expr({ withFormatting: false }) `my.field = max(10, 20)`;
+```
+
+Combine nodes using template strings:
+
+```ts
+const field = synth.expr `my.field`;
+const node = synth.expr `${field} = max(10, 20)`;
+```
+
+Print the node as a string:
+
+```ts
+const node = synth.expr `my.field = max(10, 20)`;
+String(node); // 'my.field = max(10, 20)'
+```
+
+
+### `synth.cmd`
+
+The `synth.cmd` synthesizes a command AST node (such as `SELECT`, `WHERE`,
+etc.). You use it the same as the `synth.expr` function or template string tag.
+The only difference is that the `synth.cmd` function or tag creates a command
+AST node.
+
+```ts
+const node = synth.cmd `WHERE my.field == 10`;
+// { type: 'command', name: 'where', args: [ ... ]}
+```
diff --git a/packages/kbn-esql-ast/src/synth/__tests__/cmd.test.ts b/packages/kbn-esql-ast/src/synth/__tests__/cmd.test.ts
@@ -0,0 +1,50 @@
+/*
+ * 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 { BasicPrettyPrinter } from '../../pretty_print';
+import { cmd } from '../cmd';
+import { expr } from '../expr';
+
+test('can create a WHERE command', () => {
+  const node = cmd`WHERE coordinates.lat >= 12.123123`;
+  const text = BasicPrettyPrinter.command(node);
+
+  expect(text).toBe('WHERE coordinates.lat >= 12.123123');
+});
+
+test('can create a complex STATS command', () => {
+  const node = cmd`STATS count_last_hour = SUM(count_last_hour), total_visits = SUM(total_visits), bytes_transform = SUM(bytes_transform), bytes_transform_last_hour = SUM(bytes_transform_last_hour) BY extension.keyword`;
+  const text = BasicPrettyPrinter.command(node);
+
+  expect(text).toBe(
+    'STATS count_last_hour = SUM(count_last_hour), total_visits = SUM(total_visits), bytes_transform = SUM(bytes_transform), bytes_transform_last_hour = SUM(bytes_transform_last_hour) BY extension.keyword'
+  );
+});
+
+test('can create a FROM source command', () => {
+  const node = cmd`FROM index METADATA _id`;
+  const text = BasicPrettyPrinter.command(node);
+
+  expect(text).toBe('FROM index METADATA _id');
+});
+
+test('throws if specified source is not a command', () => {
+  expect(() => cmd`123`).toThrowError();
+});
+
+test('can compose expressions into commands', () => {
+  const field = expr`a.b.c`;
+  const cmd1 = cmd` WHERE ${field} == "asdf"`;
+  const cmd2 = cmd` DISSECT ${field} """%{date}"""`;
+  const text1 = BasicPrettyPrinter.command(cmd1);
+  const text2 = BasicPrettyPrinter.command(cmd2);
+
+  expect(text1).toBe('WHERE a.b.c == "asdf"');
+  expect(text2).toBe('DISSECT a.b.c """%{date}"""');
+});
diff --git a/packages/kbn-esql-ast/src/synth/__tests__/expr_function.test.ts b/packages/kbn-esql-ast/src/synth/__tests__/expr_function.test.ts
@@ -0,0 +1,148 @@
+/*
+ * 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 { BasicPrettyPrinter } from '../../pretty_print';
+import { ESQLProperNode } from '../../types';
+import { Walker } from '../../walker/walker';
+import { expr } from '../expr';
+
+test('can generate integer literal', () => {
+  const node = expr('42');
+
+  expect(node).toMatchObject({
+    type: 'literal',
+    literalType: 'integer',
+    name: '42',
+    value: 42,
+  });
+});
+
+test('can generate integer literal and keep comment', () => {
+  const node = expr('42 /* my 42 */');
+
+  expect(node).toMatchObject({
+    type: 'literal',
+    literalType: 'integer',
+    value: 42,
+    formatting: {
+      right: [
+        {
+          type: 'comment',
+          subtype: 'multi-line',
+          text: ' my 42 ',
+        },
+      ],
+    },
+  });
+});
+
+test('can generate a function call expression', () => {
+  const node = expr('fn(1, "test")');
+
+  expect(node).toMatchObject({
+    type: 'function',
+    name: 'fn',
+    args: [
+      {
+        type: 'literal',
+        literalType: 'integer',
+        value: 1,
+      },
+      {
+        type: 'literal',
+        literalType: 'keyword',
+        value: '"test"',
+      },
+    ],
+  });
+});
+
+test('can generate assignment expression', () => {
+  const src = 'a.b.c = AGG(123)';
+  const node = expr(src);
+  const text = BasicPrettyPrinter.expression(node);
+
+  expect(text).toBe(src);
+});
+
+test('can generate comparison expression', () => {
+  const src = 'a.b.c >= FN(123)';
+  const node = expr(src);
+  const text = BasicPrettyPrinter.expression(node);
+
+  expect(text).toBe(src);
+});
+
+describe('can generate various expression types', () => {
+  const cases: Array<[name: string, src: string]> = [
+    ['integer', '42'],
+    ['negative integer', '-24'],
+    ['zero', '0'],
+    ['float', '3.14'],
+    ['negative float', '-1.23'],
+    ['string', '"doge"'],
+    ['empty string', '""'],
+    ['integer list', '[1, 2, 3]'],
+    ['string list', '["a", "b"]'],
+    ['boolean list', '[TRUE, FALSE]'],
+    ['time interval', '1d'],
+    ['cast', '"doge"::INTEGER'],
+    ['addition', '1 + 2'],
+    ['multiplication', '2 * 2'],
+    ['parens', '2 * (2 + 3)'],
+    ['star function call', 'FN(*)'],
+    ['function call with one argument', 'FN(1)'],
+    ['nested function calls', 'FN(1, MAX("asdf"))'],
+    ['basic field', 'col'],
+    ['nested field', 'a.b.c'],
+    ['unnamed param', '?'],
+    ['named param', '?hello'],
+    ['positional param', '?123'],
+    ['param in nested field', 'a.?b.c'],
+    ['escaped field', '`😎`'],
+    ['nested escaped field', 'emoji.`😎`'],
+    ['simple assignment', 't = NOW()'],
+    ['assignment expression', 'bytes_transform = ROUND(total_bytes / 1000000.0, 1)'],
+    [
+      'assignment with time intervals',
+      'key = CASE(timestamp < (t - 1 hour) AND timestamp > (t - 2 hour), "Last hour", "Other")',
+    ],
+    [
+      'assignment with casts',
+      'total_visits = TO_DOUBLE(COALESCE(count_last_hour, 0::LONG) + COALESCE(count_rest, 0::LONG))',
+    ],
+  ];
+
+  for (const [name, src] of cases) {
+    test(name, () => {
+      const node = expr(src);
+      const text = BasicPrettyPrinter.expression(node);
+
+      expect(text).toBe(src);
+    });
+  }
+});
+
+test('parser fields are empty', () => {
+  const src = 'a.b.c >= FN(123)';
+  const ast = expr(src);
+
+  const assertParserFields = (node: ESQLProperNode) => {
+    expect(node.location.min).toBe(0);
+    expect(node.location.max).toBe(0);
+    expect(node.text).toBe('');
+    expect(node.incomplete).toBe(false);
+  };
+
+  Walker.walk(ast, {
+    visitAny: (node: ESQLProperNode) => {
+      assertParserFields(node);
+    },
+  });
+});