[8.x] [ES|QL] Incorrect command option location parsing (#194115) (#1…

…94147) # Backport This will backport the following commits from `main` to `8.x`: - [[ES|QL] Incorrect command option location parsing (#194115)](#194115)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Vadim Kibana <[email protected]>
elastic · Sep 26, 2024 · 0655dff · 0655dff
1 parent e72fd8f
commit 0655dff
Show file tree

Hide file tree

Showing 3 changed files with 206 additions and 11 deletions.
diff --git a/packages/kbn-esql-ast/src/parser/__tests__/command_options.test.ts b/packages/kbn-esql-ast/src/parser/__tests__/command_options.test.ts
@@ -0,0 +1,163 @@
+/*
+ * 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 '..';
+import { Walker } from '../../walker';
+
+describe('command options', () => {
+  describe('parses correctly location', () => {
+    describe('FROM', () => {
+      it('parses correctly METADATA option position', () => {
+        const query = 'FROM a METADATA b';
+        const { root } = parse(query);
+        const option = Walker.match(root, { type: 'option', name: 'metadata' });
+
+        expect(option).toMatchObject({
+          type: 'option',
+          name: 'metadata',
+          location: {
+            min: 'FROM a '.length,
+            max: 'FROM a METADATA b'.length - 1,
+          },
+        });
+      });
+
+      it('parses correctly METADATA option position position with multiple arguments', () => {
+        const query =
+          'FROM kibana_e_commerce, index_pattern METADATA _id, _index | STATS b BY c | LIMIT 123';
+        const { root } = parse(query);
+        const option = Walker.match(root, { type: 'option', name: 'metadata' });
+
+        expect(option).toMatchObject({
+          type: 'option',
+          name: 'metadata',
+          location: {
+            min: 'FROM kibana_e_commerce, index_pattern '.length,
+            max: 'FROM kibana_e_commerce, index_pattern METADATA _id, _index'.length - 1,
+          },
+        });
+      });
+    });
+
+    describe('ENRICH', () => {
+      it('parses correctly ON option position in ENRICH command', () => {
+        const query = 'FROM a | ENRICH b ON c';
+        const { root } = parse(query);
+        const option = Walker.match(root, { type: 'option', name: 'on' });
+
+        expect(option).toMatchObject({
+          type: 'option',
+          name: 'on',
+          location: {
+            min: 'FROM a | ENRICH b '.length,
+            max: 'FROM a | ENRICH b ON c'.length - 1,
+          },
+        });
+      });
+
+      it('parses correctly WITH option in ENRICH command', () => {
+        const query = 'FROM a | ENRICH b ON c WITH d';
+        const { root } = parse(query);
+        const option = Walker.match(root, { type: 'option', name: 'with' });
+
+        expect(option).toMatchObject({
+          type: 'option',
+          name: 'with',
+          location: {
+            min: 'FROM a | ENRICH b ON c '.length,
+            max: 'FROM a | ENRICH b ON c WITH d'.length - 1,
+          },
+        });
+      });
+
+      it('parses correctly WITH option with multiple arguments in ENRICH command', () => {
+        const query = 'FROM a | ENRICH b ON c WITH d, e,f | LIMIT 1000000';
+        const { root } = parse(query);
+        const option = Walker.match(root, { type: 'option', name: 'with' });
+
+        expect(option).toMatchObject({
+          type: 'option',
+          name: 'with',
+          location: {
+            min: 'FROM a | ENRICH b ON c '.length,
+            max: 'FROM a | ENRICH b ON c WITH d, e,f'.length - 1,
+          },
+        });
+      });
+
+      it('parses correctly WITH option position with assignment in ENRICH command', () => {
+        const query = 'FROM a | ENRICH b ON c WITH d, e = policy,f = something | LIMIT 1000000';
+        const { root } = parse(query);
+        const option = Walker.match(root, { type: 'option', name: 'with' });
+
+        expect(option).toMatchObject({
+          type: 'option',
+          name: 'with',
+          location: {
+            min: 'FROM a | ENRICH b ON c '.length,
+            max: 'FROM a | ENRICH b ON c WITH d, e = policy,f = something'.length - 1,
+          },
+        });
+      });
+    });
+
+    describe('STATS', () => {
+      it('parses correctly BY option in STATS command', () => {
+        const query = 'FROM a | STATS b BY c';
+        const { root } = parse(query);
+        const option = Walker.match(root, { type: 'option', name: 'by' });
+
+        expect(option).toMatchObject({
+          type: 'option',
+          name: 'by',
+          location: {
+            min: 'FROM a | STATS b '.length,
+            max: 'FROM a | STATS b BY c'.length - 1,
+          },
+        });
+      });
+
+      it('parses correctly BY option with multiple arguments in STATS command', () => {
+        const query = 'FROM a | STATS b BY c, long.field.name | LIMIT 1000000';
+        const { root } = parse(query);
+        const option = Walker.match(root, { type: 'option', name: 'by' });
+
+        expect(option).toMatchObject({
+          type: 'option',
+          name: 'by',
+          location: {
+            min: 'FROM a | STATS b '.length,
+            max: 'FROM a | STATS b BY c, long.field.name'.length - 1,
+          },
+        });
+      });
+    });
+
+    describe('RENAME', () => {
+      it('parses correctly AS option position in RENAME command', () => {
+        const query = 'FROM a | RENAME b AS c';
+        const { root } = parse(query);
+        const option = Walker.match(root, { type: 'option', name: 'as' });
+
+        expect(option).toMatchObject({
+          type: 'option',
+          name: 'as',
+          location: {
+            // The "AS" option is unusual as the it contains the argument before
+            // it, the "a" argument. It should not be the case. The "AS" option
+            // should not exist at all, should be replaced by a *rename expression*
+            // in the future: https://github.com/elastic/kibana/issues/190360
+            min: 'FROM a | RENAME '.length,
+            max: 'FROM a | RENAME b AS c'.length - 1,
+          },
+        });
+      });
+    });
+  });
+});
diff --git a/packages/kbn-esql-ast/src/parser/walkers.ts b/packages/kbn-esql-ast/src/parser/walkers.ts
@@ -101,6 +101,7 @@ import {
   ESQLNamedParamLiteral,
   ESQLOrderExpression,
 } from '../types';
+import { firstItem, lastItem } from '../visitor/utils';
 
 export function collectAllSourceIdentifiers(ctx: FromCommandContext): ESQLAstItem[] {
   const fromContexts = ctx.getTypedRuleContexts(IndexPatternContext);
@@ -167,20 +168,24 @@ export function getMatchField(ctx: EnrichCommandContext) {
   const identifier = ctx.qualifiedNamePattern();
   if (identifier) {
     const fn = createOption(ctx.ON()!.getText().toLowerCase(), ctx);
+    let max: number = ctx.ON()!.symbol.stop;
     if (textExistsAndIsValid(identifier.getText())) {
-      fn.args.push(createColumn(identifier));
+      const column = createColumn(identifier);
+      fn.args.push(column);
+      max = column.location.max;
     }
-    // overwrite the location inferring the correct position
-    fn.location = getPosition(ctx.ON()!.symbol, ctx.WITH()?.symbol);
+    fn.location.min = ctx.ON()!.symbol.start;
+    fn.location.max = max;
     return [fn];
   }
   return [];
 }
 
 export function getEnrichClauses(ctx: EnrichCommandContext) {
   const ast: ESQLCommandOption[] = [];
-  if (ctx.WITH()) {
-    const option = createOption(ctx.WITH()!.getText().toLowerCase(), ctx);
+  const withCtx = ctx.WITH();
+  if (withCtx) {
+    const option = createOption(withCtx.getText().toLowerCase(), ctx);
     ast.push(option);
     const clauses = ctx.enrichWithClause_list();
     for (const clause of clauses) {
@@ -204,8 +209,13 @@ export function getEnrichClauses(ctx: EnrichCommandContext) {
           option.args.push(fn);
         }
       }
+
+      const location = option.location;
+      const lastArg = lastItem(option.args);
+
+      location.min = withCtx.symbol.start;
+      location.max = lastArg?.location?.max ?? withCtx.symbol.stop;
     }
-    option.location = getPosition(ctx.WITH()?.symbol);
   }
 
   return ast;
@@ -436,13 +446,18 @@ export function visitRenameClauses(clausesCtx: RenameClauseContext[]): ESQLAstIt
     .map((clause) => {
       const asToken = clause.getToken(esql_parser.AS, 0);
       if (asToken && textExistsAndIsValid(asToken.getText())) {
-        const fn = createOption(asToken.getText().toLowerCase(), clause);
+        const option = createOption(asToken.getText().toLowerCase(), clause);
         for (const arg of [clause._oldName, clause._newName]) {
           if (textExistsAndIsValid(arg.getText())) {
-            fn.args.push(createColumn(arg));
+            option.args.push(createColumn(arg));
           }
         }
-        return fn;
+        const firstArg = firstItem(option.args);
+        const lastArg = lastItem(option.args);
+        const location = option.location;
+        if (firstArg) location.min = firstArg.location.min;
+        if (lastArg) location.max = lastArg.location.max;
+        return option;
       } else if (textExistsAndIsValid(clause._oldName?.getText())) {
         return createColumn(clause._oldName);
       }
@@ -600,11 +615,15 @@ export function visitByOption(
   ctx: StatsCommandContext | InlinestatsCommandContext,
   expr: FieldsContext | undefined
 ) {
-  if (!ctx.BY() || !expr) {
+  const byCtx = ctx.BY();
+  if (!byCtx || !expr) {
     return [];
   }
-  const option = createOption(ctx.BY()!.getText().toLowerCase(), ctx);
+  const option = createOption(byCtx.getText().toLowerCase(), ctx);
   option.args.push(...collectAllFields(expr));
+  option.location.min = byCtx.symbol.start;
+  const lastArg = lastItem(option.args);
+  if (lastArg) option.location.max = lastArg.location.max;
   return [option];
 }
 

diff --git a/packages/kbn-esql-ast/src/visitor/utils.ts b/packages/kbn-esql-ast/src/visitor/utils.ts
@@ -35,3 +35,16 @@ export const firstItem = (items: ESQLAstItem[]): ESQLSingleAstItem | undefined =
     return item;
   }
 };
+
+/**
+ * Returns the last normalized "single item" from the "item" list.
+ *
+ * @param items Returns the last "single item" from the "item" list.
+ * @returns A "single item", if any.
+ */
+export const lastItem = (items: ESQLAstItem[]): ESQLSingleAstItem | undefined => {
+  const last = items[items.length - 1];
+  if (!last) return undefined;
+  if (Array.isArray(last)) return lastItem(last as ESQLAstItem[]);
+  return last as ESQLSingleAstItem;
+};