[ES|QL] Add pretty-printing support for list literals (#195383)

## Summary Closes #194840 This PR add pretty-printing support for list literal expressions. For example, this query: ``` ROW ["..............................................", "..............................................", ".............................................."] ``` will be formatted as so: ``` ROW [ "..............................................", "..............................................", ".............................................."] ``` ### 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) Co-authored-by: Stratoula Kalafateli <[email protected]>
elastic · Oct 10, 2024 · 4df2d9f · 4df2d9f
1 parent 84d6899
commit 4df2d9f
Show file tree

Hide file tree

Showing 5 changed files with 132 additions and 14 deletions.
diff --git a/packages/kbn-esql-ast/src/pretty_print/__tests__/wrapping_pretty_printer.test.ts b/packages/kbn-esql-ast/src/pretty_print/__tests__/wrapping_pretty_printer.test.ts
@@ -593,6 +593,85 @@ ROW (asdf + asdf)::string, 1.2::string, "1234"::integer, (12321342134 + 23412341
 -     "aaaaaaaaaaa")::boolean`);
     });
   });
+
+  describe('list literals', () => {
+    describe('numeric', () => {
+      test('wraps long list literals one line', () => {
+        const query =
+          'ROW [1234567890, 1234567890, 1234567890, 1234567890, 1234567890, 1234567890, 1234567890, 1234567890, 1234567890]';
+        const text = reprint(query).text;
+
+        expect('\n' + text).toBe(`
+ROW
+  [1234567890, 1234567890, 1234567890, 1234567890, 1234567890, 1234567890,
+    1234567890, 1234567890, 1234567890]`);
+      });
+
+      test('wraps long list literals to multiple lines one line', () => {
+        const query = `ROW [1234567890, 1234567890, 1234567890, 1234567890, 1234567890, 1234567890,
+          1234567890, 1234567890, 1234567890, 1234567890, 1234567890, 1234567890,
+          1234567890, 1234567890, 1234567890, 1234567890, 1234567890, 1234567890,
+          1234567890, 1234567890, 1234567890]`;
+        const text = reprint(query).text;
+
+        expect('\n' + text).toBe(`
+ROW
+  [1234567890, 1234567890, 1234567890, 1234567890, 1234567890, 1234567890,
+    1234567890, 1234567890, 1234567890, 1234567890, 1234567890, 1234567890,
+    1234567890, 1234567890, 1234567890, 1234567890, 1234567890, 1234567890,
+    1234567890, 1234567890, 1234567890]`);
+      });
+
+      test('breaks very long values one-per-line', () => {
+        const query = `ROW fn1(fn2(fn3(fn4(fn5(fn6(fn7(fn8([1234567890, 1234567890, 1234567890, 1234567890, 1234567890]))))))))`;
+        const text = reprint(query, { wrap: 40 }).text;
+
+        expect('\n' + text).toBe(`
+ROW
+  FN1(
+    FN2(
+      FN3(
+        FN4(
+          FN5(
+            FN6(
+              FN7(
+                FN8(
+                  [
+                    1234567890,
+                    1234567890,
+                    1234567890,
+                    1234567890,
+                    1234567890]))))))))`);
+      });
+    });
+
+    describe('string', () => {
+      test('wraps long list literals one line', () => {
+        const query =
+          'ROW ["some text", "another text", "one more text literal", "and another one", "and one more", "and one more", "and one more", "and one more", "and one more"]';
+        const text = reprint(query).text;
+
+        expect('\n' + text).toBe(`
+ROW
+  ["some text", "another text", "one more text literal", "and another one",
+    "and one more", "and one more", "and one more", "and one more",
+    "and one more"]`);
+      });
+
+      test('can break very long strings per line', () => {
+        const query =
+          'ROW ["..............................................", "..............................................", ".............................................."]';
+        const text = reprint(query).text;
+
+        expect('\n' + text).toBe(`
+ROW
+  [
+    "..............................................",
+    "..............................................",
+    ".............................................."]`);
+      });
+    });
+  });
 });
 
 test.todo('Idempotence on multiple times pretty printing');
diff --git a/packages/kbn-esql-ast/src/pretty_print/wrapping_pretty_printer.ts b/packages/kbn-esql-ast/src/pretty_print/wrapping_pretty_printer.ts
@@ -15,9 +15,10 @@ import {
   CommandVisitorContext,
   ExpressionVisitorContext,
   FunctionCallExpressionVisitorContext,
+  ListLiteralExpressionVisitorContext,
   Visitor,
 } from '../visitor';
-import { singleItems } from '../visitor/utils';
+import { children, singleItems } from '../visitor/utils';
 import { BasicPrettyPrinter, BasicPrettyPrinterOptions } from './basic_pretty_printer';
 import { getPrettyPrintStats } from './helpers';
 import { LeafPrinter } from './leaf_printer';
@@ -235,7 +236,11 @@ export class WrappingPrettyPrinter {
   }
 
   private printArguments(
-    ctx: CommandVisitorContext | CommandOptionVisitorContext | FunctionCallExpressionVisitorContext,
+    ctx:
+      | CommandVisitorContext
+      | CommandOptionVisitorContext
+      | FunctionCallExpressionVisitorContext
+      | ListLiteralExpressionVisitorContext,
     inp: Input
   ) {
     let txt = '';
@@ -247,7 +252,7 @@ export class WrappingPrettyPrinter {
     let remainingCurrentLine = inp.remaining;
     let oneArgumentPerLine = false;
 
-    for (const child of singleItems(ctx.node.args)) {
+    for (const child of children(ctx.node)) {
       if (getPrettyPrintStats(child).hasLineBreakingDecorations) {
         oneArgumentPerLine = true;
         break;
@@ -489,13 +494,11 @@ export class WrappingPrettyPrinter {
     })
 
     .on('visitListLiteralExpression', (ctx, inp: Input): Output => {
-      let elements = '';
-
-      for (const out of ctx.visitElements(inp)) {
-        elements += (elements ? ', ' : '') + out.txt;
-      }
-
-      const formatted = `[${elements}]${inp.suffix ?? ''}`;
+      const args = this.printArguments(ctx, {
+        indent: inp.indent,
+        remaining: inp.remaining - 1,
+      });
+      const formatted = `[${args.txt}]${inp.suffix ?? ''}`;
       const { txt, indented } = this.decorateWithComments(inp.indent, ctx.node, formatted);
 
       return { txt, indented };

diff --git a/packages/kbn-esql-ast/src/types.ts b/packages/kbn-esql-ast/src/types.ts
@@ -40,6 +40,7 @@ export type ESQLAstField = ESQLFunction | ESQLColumn;
 export type ESQLAstItem = ESQLSingleAstItem | ESQLAstItem[];
 
 export type ESQLAstNodeWithArgs = ESQLCommand | ESQLCommandOption | ESQLFunction;
+export type ESQLAstNodeWithChildren = ESQLAstNodeWithArgs | ESQLList;
 
 /**
  * *Proper* are nodes which are objects with `type` property, once we get rid

diff --git a/packages/kbn-esql-ast/src/visitor/contexts.ts b/packages/kbn-esql-ast/src/visitor/contexts.ts
@@ -12,11 +12,12 @@
 // and makes it harder to understand the code structure.
 
 import { type GlobalVisitorContext, SharedData } from './global_visitor_context';
-import { firstItem, singleItems } from './utils';
+import { children, firstItem, singleItems } from './utils';
 import type {
   ESQLAstCommand,
   ESQLAstItem,
   ESQLAstNodeWithArgs,
+  ESQLAstNodeWithChildren,
   ESQLAstRenameExpression,
   ESQLColumn,
   ESQLCommandOption,
@@ -47,6 +48,11 @@ import { Builder } from '../builder';
 const isNodeWithArgs = (x: unknown): x is ESQLAstNodeWithArgs =>
   !!x && typeof x === 'object' && Array.isArray((x as any).args);
 
+const isNodeWithChildren = (x: unknown): x is ESQLAstNodeWithChildren =>
+  !!x &&
+  typeof x === 'object' &&
+  (Array.isArray((x as any).args) || Array.isArray((x as any).values));
+
 export class VisitorContext<
   Methods extends VisitorMethods = VisitorMethods,
   Data extends SharedData = SharedData,
@@ -99,13 +105,13 @@ export class VisitorContext<
   public arguments(): ESQLAstExpressionNode[] {
     const node = this.node;
 
-    if (!isNodeWithArgs(node)) {
+    if (!isNodeWithChildren(node)) {
       return [];
     }
 
     const args: ESQLAstExpressionNode[] = [];
 
-    for (const arg of singleItems(node.args)) {
+    for (const arg of children(node)) {
       args.push(arg);
     }
 

diff --git a/packages/kbn-esql-ast/src/visitor/utils.ts b/packages/kbn-esql-ast/src/visitor/utils.ts
@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { ESQLAstItem, ESQLSingleAstItem } from '../types';
+import { ESQLAstItem, ESQLProperNode, ESQLSingleAstItem } from '../types';
 
 /**
  * Normalizes AST "item" list to only contain *single* items.
@@ -48,3 +48,32 @@ export const lastItem = (items: ESQLAstItem[]): ESQLSingleAstItem | undefined =>
   if (Array.isArray(last)) return lastItem(last as ESQLAstItem[]);
   return last as ESQLSingleAstItem;
 };
+
+export function* children(node: ESQLProperNode): Iterable<ESQLSingleAstItem> {
+  switch (node.type) {
+    case 'function':
+    case 'command':
+    case 'option': {
+      for (const arg of singleItems(node.args)) {
+        yield arg;
+      }
+      break;
+    }
+    case 'list': {
+      for (const item of singleItems(node.values)) {
+        yield item;
+      }
+      break;
+    }
+    case 'inlineCast': {
+      if (Array.isArray(node.value)) {
+        for (const item of singleItems(node.value)) {
+          yield item;
+        }
+      } else {
+        yield node.value;
+      }
+      break;
+    }
+  }
+}