last cleanup

Author: bartoval

src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/ast.test.ts

@@ -7,13 +7,15 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { PromQLParser } from '@elastic/esql';
+import { Parser, PromQLParser, Walker } from '@elastic/esql';
+import type { ESQLAstItem, ESQLAstQueryExpression } from '@elastic/esql/types';
 import { EDITOR_MARKER } from '../constants';
 import {
   addAutocompleteMarker,
   correctPromqlQuerySyntax,
   correctQuerySyntax,
   getBracketsToClose,
+  isMarkerNode,
   removeAutocompleteMarkers,
 } from './ast';
 
@@ -134,3 +136,37 @@ describe('correctPromqlQuerySyntax', () => {
     expect(JSON.stringify(normalizedRoot)).not.toContain(EDITOR_MARKER);
   });
 });
+
+describe('removeAutocompleteMarkers', () => {
+  const parseAutocomplete = (innerText: string): ESQLAstQueryExpression => {
+    const corrected = correctQuerySyntax(addAutocompleteMarker(innerText));
+    return Parser.parse(corrected, { withFormatting: true }).root;
+  };
+
+  const countMarkers = (node: ESQLAstQueryExpression): number => {
+    let count = 0;
+    Walker.walk(node, {
+      visitAny: (current) => {
+        if (isMarkerNode(current as ESQLAstItem)) {
+          count++;
+        }
+      },
+    });
+    return count;
+  };
+
+  it('drops marker-only nodes nested in args arrays', () => {
+    const root = parseAutocomplete('FROM index | EVAL result = ROUND(doubleField, ');
+
+    expect(countMarkers(root)).toBeGreaterThan(0);
+    expect(countMarkers(removeAutocompleteMarkers(root))).toBe(0);
+  });
+
+  it('strips the marker from the inline cast type, not only from text', () => {
+    const root = parseAutocomplete('FROM index | EVAL casted = keywordField::');
+
+    // Before cleaning, the marker leaks into inlineCast.castType (a plain string property).
+    expect(JSON.stringify(root)).toContain(EDITOR_MARKER);
+    expect(JSON.stringify(removeAutocompleteMarkers(root))).not.toContain(EDITOR_MARKER);
+  });
+});

src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/ast.ts

@@ -7,8 +7,8 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { isList, isOptionNode, PromQLParser, Walker } from '@elastic/esql';
-import type { PromQLAstNode, PromQLAstQueryExpression } from '@elastic/esql';
+import { isList, isOptionNode, Walker } from '@elastic/esql';
+import type { PromQLAstNode } from '@elastic/esql';
 import type {
   ESQLFunction,
   ESQLSingleAstItem,
@@ -68,6 +68,11 @@ export function removeAutocompleteMarkers<T>(value: T): T {
       .map((item) => removeAutocompleteMarkers(item)) as T;
   }
 
+  // Strip the marker from any string value, not only `text` (e.g. inline cast `castType`).
+  if (typeof value === 'string') {
+    return (value.includes(EDITOR_MARKER) ? value.replace(EDITOR_MARKER, '') : value) as T;
+  }
+
   if (!value || typeof value !== 'object') {
     return value;
   }
@@ -82,14 +87,6 @@ export function removeAutocompleteMarkers<T>(value: T): T {
     cleanedValue[key as keyof typeof cleanedValue] = removeAutocompleteMarkers(child);
   }
 
-  if (
-    'text' in cleanedValue &&
-    typeof cleanedValue.text === 'string' &&
-    cleanedValue.text.includes(EDITOR_MARKER)
-  ) {
-    cleanedValue.text = cleanedValue.text.replace(EDITOR_MARKER, '');
-  }
-
   return cleanedValue;
 }
 
@@ -275,19 +272,6 @@ export function correctPromqlQuerySyntax(input: string): string {
   return query + getPromqlBracketsToClose(query);
 }
 
-export function parsePromqlAutocompleteQuery(query: string): {
-  correctedQuery: string;
-  root: PromQLAstQueryExpression;
-} {
-  const correctedQuery = correctPromqlQuerySyntax(query);
-  const { root } = PromQLParser.parse(correctedQuery);
-
-  return {
-    correctedQuery,
-    root: removeAutocompleteMarkers(root),
-  };
-}
-
 export function getPromqlBracketsToClose(text: string): string {
   const esqlBrackets = getBracketsToClose(text).join('');
   const promqlBraces = getPromqlBracesToClose(text);

src/platform/packages/shared/kbn-esql-language/src/commands/definitions/utils/autocomplete/promql/suggestion_query_engine.ts

@@ -11,7 +11,7 @@ import type { ESQLControlVariable } from '@kbn/esql-types';
 import type { ICommandContext, ISuggestionItem } from '../../../../registry/types';
 import { getQueryPosition } from './query_position';
 import { getPreGroupedAggregationName } from '../../promql';
-import { parsePromqlAutocompleteQuery } from '../../ast';
+import { parsePromqlAutocompleteQuery } from '../../../../../language/shared/parse_for_autocomplete_query';
 import { buildVectorSuggestions } from './query_positions/suggestion_helpers';
 import { positionHandlers } from './query_positions/dispatcher';
 

src/platform/packages/shared/kbn-esql-language/src/language/autocomplete/autocomplete.ts

@@ -37,7 +37,6 @@ import { getColumnsByTypeRetriever } from '../shared/columns_retrieval_helpers';
 import { getUnmappedFieldsStrategy } from '../../commands/definitions/utils/settings';
 import { isTimeseriesSourceCommand } from '../../commands/definitions/utils/timeseries_check';
 import { attachReplacementRanges, ReplacementRangeStrategyKind } from './utils/prefix_range';
-import { correctQuerySyntax } from '../../commands/definitions/utils/ast';
 
 function isSourceCommandSuggestion({ label }: { label: string }) {
   const sourceCommands = esqlCommandRegistry
@@ -73,10 +72,9 @@ export async function suggest(
   // Use the appropriate AST context for field retrieval
   // When in a subquery, use the subquery's AST to get only its fields
   const astForFields = astContext.astForContext;
-  const correctedQuery = correctQuerySyntax(innerText);
 
   const { getColumnsByType, getColumnMap } = getColumnsByTypeRetriever(
-    getQueryForFields(correctedQuery, astForFields),
+    getQueryForFields(innerText, astForFields),
     innerText,
     resourceRetriever
   );

src/platform/packages/shared/kbn-esql-language/src/language/shared/get_query_for_fields.test.ts

@@ -11,7 +11,6 @@ import { BasicPrettyPrinter } from '@elastic/esql';
 import { getQueryForFields } from './get_query_for_fields';
 import { EDITOR_MARKER } from '../../commands/definitions/constants';
 import { parseAutocompleteQuery } from './parse_for_autocomplete_query';
-import { correctQuerySyntax } from '../../commands/definitions/utils/ast';
 
 describe('getQueryForFields', () => {
   const assert = (queryBeforeCursor: string, expected: string) => {
@@ -56,13 +55,6 @@ describe('getQueryForFields', () => {
   });
 
   it('should not treat a comma inside an incomplete function call as an EVAL separator', () => {
-    const query = 'FROM index | EVAL result = ROUND(field, ';
-    const { root } = parseAutocompleteQuery(query, query.length);
-
-    const result = getQueryForFields(correctQuerySyntax(query), root);
-    const printedResult = BasicPrettyPrinter.print(result);
-
-    expect(printedResult).toEqual('FROM index');
-    expect(printedResult).not.toContain(EDITOR_MARKER);
+    assert('FROM index | EVAL result = ROUND(field, ', 'FROM index');
   });
 });

src/platform/packages/shared/kbn-esql-language/src/language/shared/get_query_for_fields.ts

@@ -8,6 +8,7 @@
  */
 
 import type { ESQLAstQueryExpression, ESQLAstForkCommand } from '@elastic/esql/types';
+import { getBracketsToClose } from '../../commands/definitions/utils/ast';
 import { expandEvals } from './expand_evals';
 
 /**
@@ -21,7 +22,7 @@ import { expandEvals } from './expand_evals';
  * IMPORTANT: the AST nodes in the new query still reference locations in the original query text
  *
  * @param queryString The original query string
- * @param commands
+ * @param root
  * @returns
  */
 export function getQueryForFields(
@@ -49,8 +50,10 @@ export function getQueryForFields(
   }
 
   if (lastCommand && lastCommand.name === 'eval') {
-    const endsWithComma = queryString.trim().endsWith(',');
-    if (lastCommand.args.length > 1 || endsWithComma) {
+    const hasTrailingSeparatorComma =
+      getBracketsToClose(queryString).length === 0 && queryString.trimEnd().endsWith(',');
+
+    if (lastCommand.args.length > 1 || hasTrailingSeparatorComma) {
       /**
        * If we get here, we know that we have a multi-expression EVAL statement.
        *
@@ -64,7 +67,7 @@ export function getQueryForFields(
        * Simplified:     FROM lolz | EVAL foo = 1 | EVAL bar = foo + 1 | EVAL baz = bar + 1
        */
       const expanded = expandEvals(commands);
-      const newCommands = expanded.slice(0, endsWithComma ? undefined : -1);
+      const newCommands = expanded.slice(0, hasTrailingSeparatorComma ? undefined : -1);
       return { ...root, commands: newCommands };
     }
   }

src/platform/packages/shared/kbn-esql-language/src/language/shared/parse_for_autocomplete_query.test.ts

@@ -47,6 +47,7 @@ describe('getAutocompleteCursorContext', () => {
     ['timeseries source command after comma', 'TS timeseries_index, '],
     ['row assignment', 'ROW total = '],
     ['eval assignment', 'FROM employees | EVAL total = '],
+    ['inline cast type', 'FROM employees | EVAL casted = keywordField::'],
     ['function argument', 'FROM employees | EVAL total = ROUND(salary, '],
     ['where binary expression', 'FROM employees | WHERE age > '],
     ['where list value', 'FROM employees | WHERE age IN (1, '],
@@ -87,4 +88,15 @@ describe('getAutocompleteCursorContext', () => {
       assertNoMarker(astContext.containingFunction, query, 'containingFunction');
     }
   );
+
+  it.each([
+    'FROM employees | EVAL total = ',
+    'ROW total = ',
+    'FROM employees | EVAL total = ROUND(salary, ',
+    'FROM employees | WHERE age IN (1, ',
+  ])('resolves an expression context (not just marker-free) for %s', (query) => {
+    const { astContext } = getAutocompleteCursorContext(query, query.length);
+
+    expect(astContext.type).toBe('expression');
+  });
 });

src/platform/packages/shared/kbn-esql-language/src/language/shared/parse_for_autocomplete_query.ts

@@ -7,10 +7,12 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { Parser } from '@elastic/esql';
+import { Parser, PromQLParser } from '@elastic/esql';
+import type { PromQLAstQueryExpression } from '@elastic/esql';
 import type { ESQLAstQueryExpression } from '@elastic/esql/types';
 import {
   addAutocompleteMarker,
+  correctPromqlQuerySyntax,
   correctQuerySyntax,
   findAstPosition,
   removeAutocompleteMarkers,
@@ -19,7 +21,6 @@ import { getCursorContext } from './get_cursor_context';
 
 interface ParsedAutocompleteQuery {
   innerText: string;
-  correctedQuery: string;
   root: ESQLAstQueryExpression;
 }
 
@@ -34,6 +35,22 @@ export function parseAutocompleteQuery(fullText: string, offset: number): Parsed
 
   return {
     innerText,
+    root: removeAutocompleteMarkers(root),
+  };
+}
+
+/**
+ * PromQL counterpart of {@link parseAutocompleteQuery}: corrects partial PromQL syntax,
+ * parses it, and strips autocomplete markers from the resulting AST.
+ */
+export function parsePromqlAutocompleteQuery(query: string): {
+  correctedQuery: string;
+  root: PromQLAstQueryExpression;
+} {
+  const correctedQuery = correctPromqlQuerySyntax(query);
+  const { root } = PromQLParser.parse(correctedQuery);
+
+  return {
     correctedQuery,
     root: removeAutocompleteMarkers(root),
   };