autocomplete marker free

Author: bartoval

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

@@ -7,8 +7,15 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
+import { PromQLParser } from '@elastic/esql';
 import { EDITOR_MARKER } from '../constants';
-import { correctPromqlQuerySyntax, correctQuerySyntax, getBracketsToClose } from './ast';
+import {
+  addAutocompleteMarker,
+  correctPromqlQuerySyntax,
+  correctQuerySyntax,
+  getBracketsToClose,
+  removeAutocompleteMarkers,
+} from './ast';
 
 describe('getBracketsToClose', () => {
   it('returns the number of brackets to close', () => {
@@ -38,19 +45,33 @@ describe('getBracketsToClose', () => {
   });
 });
 
-describe('correctQuerySyntax', () => {
+describe('addAutocompleteMarker', () => {
   it('appends marker after operator', () => {
     const query = 'FROM foo | EVAL field > ';
-    const result = correctQuerySyntax(query);
+    const result = addAutocompleteMarker(query);
     expect(result.endsWith(EDITOR_MARKER)).toBe(true);
   });
 
   it('appends marker after comma', () => {
     const query = 'FROM foo | STATS field1, ';
-    const result = correctQuerySyntax(query);
+    const result = addAutocompleteMarker(query);
+    expect(result.endsWith(EDITOR_MARKER)).toBe(true);
+  });
+
+  it('appends marker if all brackets are closed and ends with operator', () => {
+    const query = 'FROM index | STATS AVG(field1) != ';
+    const result = addAutocompleteMarker(query);
     expect(result.endsWith(EDITOR_MARKER)).toBe(true);
   });
 
+  it('does not append marker for inline cast type names ending with operator-like suffixes', () => {
+    const query = 'FROM index | EVAL vec = [0.1, 0.2]::dense_vector ';
+    const result = addAutocompleteMarker(query);
+    expect(result).toEqual(query);
+  });
+});
+
+describe('correctQuerySyntax', () => {
   it('closes unclosed brackets', () => {
     const query = 'FROM foo | EVAL foo(bar[baz';
     const result = correctQuerySyntax(query);
@@ -65,22 +86,11 @@ describe('correctQuerySyntax', () => {
     expect(result).toEqual(query);
   });
 
-  it('appends marker if all brackets are closed and ends with operator', () => {
-    const query = 'FROM index | STATS AVG(field1) != ';
-    const result = correctQuerySyntax(query);
-    expect(result.endsWith(EDITOR_MARKER)).toBe(true);
-  });
-
   it('handles incomplete function signature', () => {
     const query = 'FROM foo | EVAL foo(bar, ';
     const result = correctQuerySyntax(query);
-    expect(result.endsWith(`${EDITOR_MARKER})`)).toBe(true);
-  });
-
-  it('does not append marker for inline cast type names ending with operator-like suffixes', () => {
-    const query = 'FROM index | EVAL vec = [0.1, 0.2]::dense_vector ';
-    const result = correctQuerySyntax(query);
-    expect(result).toEqual(query);
+    expect(result.endsWith(')')).toBe(true);
+    expect(result).not.toContain(EDITOR_MARKER);
   });
 });
 
@@ -115,4 +125,12 @@ describe('correctPromqlQuerySyntax', () => {
     const markerMatches = result.match(new RegExp(EDITOR_MARKER, 'g')) ?? [];
     expect(markerMatches).toHaveLength(1);
   });
+
+  it('normalizes marker nodes from parsed PromQL autocomplete ASTs', () => {
+    const query = correctPromqlQuerySyntax('rate(http_requests_total, ');
+    const { root } = PromQLParser.parse(query);
+    const normalizedRoot = removeAutocompleteMarkers(root);
+
+    expect(JSON.stringify(normalizedRoot)).not.toContain(EDITOR_MARKER);
+  });
 });

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

@@ -7,30 +7,25 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { isColumn, isIdentifier, isList, isOptionNode, isSource, Walker } from '@elastic/esql';
+import { isList, isOptionNode, PromQLParser, Walker } from '@elastic/esql';
+import type { PromQLAstNode, PromQLAstQueryExpression } from '@elastic/esql';
 import type {
   ESQLFunction,
   ESQLSingleAstItem,
-  ESQLAstExpression,
   ESQLAstItem,
   ESQLCommandOption,
-  ESQLAstAllCommands,
   ESQLAstHeaderCommand,
   ESQLAstQueryExpression,
 } from '@elastic/esql/types';
 import { EDITOR_MARKER } from '../constants';
-import { endsWithComma } from './regex';
+import { endsWithComma, endsWithWhitespace } from './regex';
 
-export function isMarkerNode(node: ESQLAstItem | undefined): boolean {
-  if (Array.isArray(node)) {
-    return false;
-  }
+const ENDS_WITH_BINARY_OPERATOR_REGEX =
+  /(?:\+|\/|==|>=|>|<=|<|:|%|\*|-|!=|=|\b(?:in|like|not in|not like|not rlike|rlike|and|or|not|as)\b)\s+$/i;
+const ENDS_WITH_CASTING_OPERATOR_REGEX = /::\s*$/i;
 
-  return Boolean(
-    node &&
-      (isColumn(node) || isIdentifier(node) || isSource(node)) &&
-      node.name.endsWith(EDITOR_MARKER)
-  );
+export function isMarkerNode(node: ESQLAstItem | PromQLAstNode | undefined): boolean {
+  return Boolean(node && !Array.isArray(node) && node.name?.endsWith(EDITOR_MARKER));
 }
 
 function findCommand(ast: ESQLAstQueryExpression, offset: number) {
@@ -65,81 +60,37 @@ function findHeaderCommand(
   return targetHeader.incomplete ? targetHeader : undefined;
 }
 
-export function isNotMarkerNodeOrArray(arg: ESQLAstItem) {
-  return Array.isArray(arg) || !isMarkerNode(arg);
-}
-
-const removeMarkerNode = (node: ESQLAstExpression) => {
-  Walker.walk(node, {
-    visitAny: (current) => {
-      if ('args' in current) {
-        current.args = current.args.filter((n) => !isMarkerNode(n));
-      } else if ('values' in current) {
-        current.values = current.values.filter((n) => !isMarkerNode(n));
-      }
-    },
-  });
-};
-
-function removeMarkerArgFromArgsList<T extends ESQLSingleAstItem | ESQLAstAllCommands>(
-  node: T | undefined
-) {
-  if (!node) {
-    return;
+// Removes parser-only autocomplete markers while preserving parser locations for cursor math.
+export function removeAutocompleteMarkers<T>(value: T): T {
+  if (Array.isArray(value)) {
+    return value
+      .filter((item) => !isMarkerNode(item))
+      .map((item) => removeAutocompleteMarkers(item)) as T;
   }
-  if (node.type === 'command' || node.type === 'option' || node.type === 'function') {
-    const cleanedNode = {
-      ...node,
-      text: node.text.replace(EDITOR_MARKER, ''),
-      args: node.args.filter(isNotMarkerNodeOrArray).map(mapToNonMarkerNode),
-    };
-
-    if (cleanedNode.type === 'command' && 'expression' in cleanedNode && cleanedNode.expression) {
-      cleanedNode.expression = isMarkerNode(cleanedNode.expression)
-        ? undefined
-        : (mapToNonMarkerNode(cleanedNode.expression) as ESQLAstExpression);
-    }
 
-    return cleanedNode;
+  if (!value || typeof value !== 'object') {
+    return value;
   }
-  return node;
-}
 
-export function mapToNonMarkerNode(arg: ESQLAstItem): ESQLAstItem {
-  if (Array.isArray(arg)) {
-    return arg.filter(isNotMarkerNodeOrArray).map(mapToNonMarkerNode);
+  if (isMarkerNode(value as ESQLAstItem | PromQLAstNode)) {
+    return undefined as T;
   }
 
-  if ('args' in arg) {
-    return {
-      ...arg,
-      text: arg.text.replace(EDITOR_MARKER, ''),
-      args: arg.args.filter(isNotMarkerNodeOrArray).map(mapToNonMarkerNode) as typeof arg.args,
-    } as typeof arg;
-  }
+  const cleanedValue = { ...value };
 
-  if ('values' in arg) {
-    return {
-      ...arg,
-      text: arg.text.replace(EDITOR_MARKER, ''),
-      values: arg.values
-        .filter((value) => !isMarkerNode(value))
-        .map((value) => mapToNonMarkerNode(value) as ESQLAstExpression) as typeof arg.values,
-    } as typeof arg;
+  for (const [key, child] of Object.entries(value)) {
+    cleanedValue[key as keyof typeof cleanedValue] = removeAutocompleteMarkers(child);
   }
 
-  if ('text' in arg && typeof arg.text === 'string' && arg.text.includes(EDITOR_MARKER)) {
-    return {
-      ...arg,
-      text: arg.text.replace(EDITOR_MARKER, ''),
-    } as typeof arg;
+  if (
+    'text' in cleanedValue &&
+    typeof cleanedValue.text === 'string' &&
+    cleanedValue.text.includes(EDITOR_MARKER)
+  ) {
+    cleanedValue.text = cleanedValue.text.replace(EDITOR_MARKER, '');
   }
 
-  return arg;
-}
-
-function cleanMarkerNode(node: ESQLSingleAstItem | undefined): ESQLSingleAstItem | undefined {
-  return isMarkerNode(node) ? undefined : node;
+  return cleanedValue;
 }
 
 function findOption(nodes: ESQLAstItem[], offset: number): ESQLCommandOption | undefined {
@@ -174,7 +125,7 @@ export function findAstPosition(ast: ESQLAstQueryExpression, offset: number) {
 
   Walker.walk(command, {
     visitSource: (_node, parent, walker) => {
-      if (_node.location.max >= offset && _node.text !== EDITOR_MARKER) {
+      if (_node.location.max >= offset) {
         node = _node as ESQLSingleAstItem;
         walker.abort();
       }
@@ -188,23 +139,17 @@ export function findAstPosition(ast: ESQLAstQueryExpression, offset: number) {
         containingFunction = _node;
       }
 
-      if (
-        _node.location.max >= offset &&
-        _node.text !== EDITOR_MARKER &&
-        (!isList(_node) || _node.subtype !== 'tuple')
-      ) {
+      if (_node.location.max >= offset && (!isList(_node) || _node.subtype !== 'tuple')) {
         node = _node as ESQLSingleAstItem;
       }
     },
   });
 
-  if (node) removeMarkerNode(node);
-
   return {
-    command: removeMarkerArgFromArgsList(command)!,
-    containingFunction: removeMarkerArgFromArgsList(containingFunction),
-    option: removeMarkerArgFromArgsList(findOption(command.args, offset)),
-    node: removeMarkerArgFromArgsList(cleanMarkerNode(node)),
+    command,
+    containingFunction,
+    option: findOption(command.args, offset),
+    node,
   };
 }
 
@@ -283,6 +228,18 @@ export function getBracketsToClose(text: string) {
   return stack.reverse().map((bracket) => pairs[bracket]);
 }
 
+export function addAutocompleteMarker(query: string) {
+  if (
+    ENDS_WITH_BINARY_OPERATOR_REGEX.test(query) ||
+    ENDS_WITH_CASTING_OPERATOR_REGEX.test(query) ||
+    (endsWithComma(query) && endsWithWhitespace(query))
+  ) {
+    return `${query} ${EDITOR_MARKER}`;
+  }
+
+  return query;
+}
+
 /**
  * This function attempts to correct the syntax of a partial query to make it valid.
  *
@@ -293,27 +250,9 @@ export function getBracketsToClose(text: string) {
  * @param context
  * @returns
  */
-export function correctQuerySyntax(_query: string) {
-  let query = _query;
+export function correctQuerySyntax(query: string) {
   // check if all brackets are closed, otherwise close them
-  const bracketsToAppend = getBracketsToClose(query);
-
-  const endsWithBinaryOperatorRegex =
-    /(?:\+|\/|==|>=|>|<=|<|:|%|\*|-|!=|=|\b(?:in|like|not in|not like|not rlike|rlike|and|or|not|as)\b)\s+$/i;
-  const endsWithCastingOperatorRegex = /::\s*$/i;
-  const endsWithCommaRegex = /,\s+$/;
-
-  if (
-    endsWithBinaryOperatorRegex.test(query) ||
-    endsWithCastingOperatorRegex.test(query) ||
-    endsWithCommaRegex.test(query)
-  ) {
-    query += ` ${EDITOR_MARKER}`;
-  }
-
-  query += bracketsToAppend.join('');
-
-  return query;
+  return query + getBracketsToClose(query).join('');
 }
 
 const PROMQL_TRAILING_COLON_REGEX = /:\s*$/;
@@ -336,6 +275,19 @@ 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),
+  };
+}
+
 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

@@ -7,10 +7,10 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { PromQLParser } from '@elastic/esql';
 import type { ICommandContext, ISuggestionItem } from '../../../../registry/types';
 import { getQueryPosition } from './query_position';
 import { getPreGroupedAggregationName } from '../../promql';
+import { parsePromqlAutocompleteQuery } from '../../ast';
 import {
   buildNextActionsSuggestion,
   buildVectorSuggestions,
@@ -31,9 +31,9 @@ export function suggestForPromqlQuery(input: SuggestForPromqlQueryInput): ISugge
     return buildVectorSuggestions(columns, [], shouldWrap);
   }
 
-  const { root } = PromQLParser.parse(queryText);
-  const position = getQueryPosition(root, cursorRelative, queryText);
-  const preGroupedAgg = getPreGroupedAggregationName(queryText.slice(0, cursorRelative));
+  const { correctedQuery, root } = parsePromqlAutocompleteQuery(queryText);
+  const position = getQueryPosition(root, cursorRelative, correctedQuery);
+  const preGroupedAgg = getPreGroupedAggregationName(correctedQuery.slice(0, cursorRelative));
 
   if (cursorRelative > root.location.max && position.type === 'inside_query') {
     return buildNextActionsSuggestion({

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

@@ -26,7 +26,6 @@ import type { ESQLColumnData } from '../../registry/types';
 import { UnmappedFieldsStrategy } from '../../registry/types';
 import { inOperators } from '../all_operators';
 import { TIME_SYSTEM_PARAMS } from './literals';
-import { isMarkerNode } from './ast';
 import { getUnmappedFieldType } from './settings';
 import { getPromqlFunctionDefinition } from './promql';
 
@@ -292,15 +291,15 @@ export function getBinaryExpressionOperand(
 }
 
 /**
- * Extracts a valid expression root from an assignment, handling arrays and marker nodes.
+ * Extracts a valid expression root from an assignment, handling array operands.
  */
 export function getAssignmentExpressionRoot(
   assignment: ESQLFunction
 ): ESQLSingleAstItem | undefined {
   const rhs = getBinaryExpressionOperand(assignment, 'right');
   const root = Array.isArray(rhs) ? rhs[0] : rhs;
 
-  if (!root || isMarkerNode(root)) {
+  if (!root) {
     return undefined;
   }