[ES|QL] Cleanup of the autocomplete and validation package

src/platform/packages/shared/kbn-esql-ast/src/__tests__/autocomplete.ts

@@ -33,7 +33,7 @@ import {
   inOperators,
   nullCheckOperators,
 } from '../definitions/all_operators';
-import { parse } from '../parser';
+import { Parser } from '../parser';
 import type { ESQLAstAllCommands } from '../types';
 import type {
   FieldType,
@@ -83,12 +83,12 @@ export const suggest = (
 ): Promise<ISuggestionItem[]> => {
   const innerText = query.substring(0, offset ?? query.length);
   const correctedQuery = correctQuerySyntax(innerText);
-  const { ast, root } = parse(correctedQuery, { withFormatting: true });
+  const { root } = Parser.parse(correctedQuery, { withFormatting: true });
   const headerConstruction = root?.header?.find((cmd) => cmd.name === commandName);
 
   const cursorPosition = offset ?? query.length;
 
-  const command = headerConstruction ?? findAstPosition(ast, cursorPosition).command;
+  const command = headerConstruction ?? findAstPosition(root, cursorPosition).command;
 
   if (!command) {
     throw new Error(`${commandName.toUpperCase()} command not found in the parsed query`);

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/autocomplete.ts

@@ -41,7 +41,7 @@ export async function autocomplete(
   // If cursor is inside a string literal, don't suggest anything
   const correctedQuery = correctQuerySyntax(innerText);
   const { root } = Parser.parse(correctedQuery, { withFormatting: true });
-  const { node } = findAstPosition(root.commands, innerText.length);
+  const { node } = findAstPosition(root, innerText.length);
 
   if (node?.type === 'literal' && node.literalType === 'keyword') {
     return [];

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/autocomplete.test.ts

@@ -25,7 +25,7 @@ import type { ICommandCallbacks } from '../../types';
 import type { FunctionReturnType, FieldType } from '../../../definitions/types';
 import { ESQL_STRING_TYPES, ESQL_NUMBER_TYPES } from '../../../definitions/types';
 import { correctQuerySyntax, findAstPosition } from '../../../definitions/utils/ast';
-import { parse } from '../../../parser';
+import { Parser } from '../../../parser';
 
 const allEvalFnsForWhere = getFunctionSignaturesByReturnType(Location.WHERE, 'any', {
   scalar: true,
@@ -415,9 +415,10 @@ describe('FORK Autocomplete', () => {
       it('suggests pipe after complete subcommands', async () => {
         const assertSuggestsPipe = async (query: string) => {
           const correctedQuery = correctQuerySyntax(query);
-          const { ast } = parse(correctedQuery, { withFormatting: true });
+          const { root } = Parser.parse(correctedQuery, { withFormatting: true });
+
           const cursorPosition = query.length;
-          const { command } = findAstPosition(ast, cursorPosition);
+          const { command } = findAstPosition(root, cursorPosition);
           if (!command) {
             throw new Error('Command not found in the parsed query');
           }

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/autocomplete.test.ts

@@ -95,9 +95,10 @@ describe('FROM Autocomplete', () => {
     test('suggests comma or pipe after complete index name', async () => {
       const suggest = async (query: string) => {
         const correctedQuery = correctQuerySyntax(query);
-        const { ast } = Parser.parse(correctedQuery, { withFormatting: true });
+        const { root } = Parser.parse(correctedQuery, { withFormatting: true });
+
         const cursorPosition = query.length;
-        const { command } = findAstPosition(ast, cursorPosition);
+        const { command } = findAstPosition(root, cursorPosition);
 
         return autocomplete(query, command!, mockCallbacks, mockContext, cursorPosition);
       };

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/autocomplete.ts

@@ -30,7 +30,7 @@ export async function autocomplete(
   // If cursor is inside a string literal, don't suggest anything
   const correctedQuery = correctQuerySyntax(innerText);
   const { root } = Parser.parse(correctedQuery, { withFormatting: true });
-  const { node } = findAstPosition(root.commands, innerText.length);
+  const { node } = findAstPosition(root, innerText.length);
 
   if (node?.type === 'literal' && node.literalType === 'keyword') {
     return [];

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/autocomplete.test.ts

@@ -34,7 +34,7 @@ import {
   ESQL_COMMON_NUMERIC_TYPES,
 } from '../../../definitions/types';
 import { correctQuerySyntax, findAstPosition } from '../../../definitions/utils/ast';
-import { parse } from '../../../parser';
+import { Parser } from '../../../parser';
 import { setTestFunctions } from '../../../definitions/utils/test_functions';
 import { getDateHistogramCompletionItem } from '../../../..';
 
@@ -126,9 +126,9 @@ describe('STATS Autocomplete', () => {
 
   const suggest = async (query: string) => {
     const correctedQuery = correctQuerySyntax(query);
-    const { ast } = parse(correctedQuery, { withFormatting: true });
+    const { root } = Parser.parse(correctedQuery, { withFormatting: true });
     const cursorPosition = query.length;
-    const { command } = findAstPosition(ast, cursorPosition);
+    const { command } = findAstPosition(root, cursorPosition);
     if (!command) {
       throw new Error('Command not found in the parsed query');
     }

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/autocomplete.ts

@@ -18,6 +18,7 @@ import type {
   ESQLAstItem,
   ESQLSingleAstItem,
   ESQLAstAllCommands,
+  ESQLAstQueryExpression,
 } from '../../../types';
 import { type ISuggestionItem, type ICommandContext } from '../../types';
 import {
@@ -469,7 +470,13 @@ function findFunctionForSuggestions(
   command: ESQLAstAllCommands,
   cursorPosition: number
 ): ESQLFunction | null {
-  const { node, containingFunction } = findAstPosition([command], cursorPosition);
+  const { node, containingFunction } = findAstPosition(
+    {
+      type: 'query' as const,
+      commands: [command],
+    } as unknown as ESQLAstQueryExpression,
+    cursorPosition
+  );
 
   if (node && node.type === 'function') {
     const fn = node as ESQLFunction;

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/autocomplete.test.ts

@@ -11,7 +11,7 @@ import { autocomplete } from './autocomplete';
 import { expectSuggestions } from '../../../__tests__/autocomplete';
 import type { ICommandCallbacks } from '../../types';
 import { correctQuerySyntax, findAstPosition } from '../../../definitions/utils/ast';
-import { parse } from '../../../parser';
+import { Parser } from '../../../parser';
 import { METADATA_FIELDS } from '../../options/metadata';
 import { getRecommendedQueriesTemplatesFromExtensions } from '../../options/recommended_queries';
 
@@ -60,9 +60,10 @@ describe('TS Autocomplete', () => {
   describe('... <sources> ...', () => {
     const suggest = async (query: string) => {
       const correctedQuery = correctQuerySyntax(query);
-      const { ast } = parse(correctedQuery, { withFormatting: true });
+      const { root } = Parser.parse(correctedQuery, { withFormatting: true });
+
       const cursorPosition = query.length;
-      const { command } = findAstPosition(ast, cursorPosition);
+      const { command } = findAstPosition(root, cursorPosition);
       if (!command) {
         throw new Error('Command not found in the parsed query');
       }

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/where/autocomplete.test.ts

@@ -27,7 +27,7 @@ import type { ICommandCallbacks } from '../../types';
 import { ESQL_COMMON_NUMERIC_TYPES } from '../../../definitions/types';
 import { getDateLiterals } from '../../../definitions/utils';
 import { correctQuerySyntax, findAstPosition } from '../../../definitions/utils/ast';
-import { parse } from '../../../parser';
+import { Parser } from '../../../parser';
 
 const allEvalFns = getFunctionSignaturesByReturnType(Location.WHERE, 'any', {
   scalar: true,
@@ -79,9 +79,10 @@ describe('WHERE Autocomplete', () => {
   describe('within the expression', () => {
     const suggest = async (query: string) => {
       const correctedQuery = correctQuerySyntax(query);
-      const { ast } = parse(correctedQuery, { withFormatting: true });
+      const { root } = Parser.parse(correctedQuery, { withFormatting: true });
+
       const cursorPosition = query.length;
-      const { command } = findAstPosition(ast, cursorPosition);
+      const { command } = findAstPosition(root, cursorPosition);
       if (!command) {
         throw new Error('Command not found in the parsed query');
       }

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/location.ts

@@ -55,7 +55,8 @@ const functionBasedLocations: Record<
  *
  * This is primarily around for backwards compatibility with the old system of command and option names.
  */
-const getLocationFromCommandOrOptionName = (name: string) => commandOptionNameToLocation[name];
+export const getLocationFromCommandOrOptionName = (name: string) =>
+  commandOptionNameToLocation[name];
 
 /**
  * Identifies the location ID at the given position

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

@@ -16,6 +16,8 @@ import type {
   ESQLAstItem,
   ESQLCommandOption,
   ESQLAstAllCommands,
+  ESQLAstHeaderCommand,
+  ESQLAstQueryExpression,
 } from '../../types';
 import { Walker } from '../../..';
 
@@ -31,16 +33,38 @@ export function isMarkerNode(node: ESQLAstItem | undefined): boolean {
   );
 }
 
-function findCommand(ast: ESQLAstAllCommands[], offset: number) {
-  const commandIndex = ast.findIndex(
+function findCommand(ast: ESQLAstQueryExpression, offset: number) {
+  const queryCommands = ast.commands;
+  const commandIndex = queryCommands.findIndex(
     ({ location }) => location.min <= offset && location.max >= offset
   );
 
-  const command = ast[commandIndex] || ast[ast.length - 1];
+  const command = queryCommands[commandIndex] || queryCommands[queryCommands.length - 1];
+
+  if (!command) {
+    return findHeaderCommand(ast, offset);
+  }
 
   return command;
 }
 
+function findHeaderCommand(
+  ast: ESQLAstQueryExpression,
+  offset: number
+): ESQLAstHeaderCommand | undefined {
+  if (!ast.header || ast.header.length === 0) {
+    return;
+  }
+
+  const commandIndex = ast.header.findIndex(
+    ({ location }) => location.min <= offset && location.max >= offset
+  );
+
+  const targetHeader = ast.header[commandIndex] || ast.header[ast.header.length - 1];
+
+  return targetHeader.incomplete ? targetHeader : undefined;
+}
+
 export function isNotMarkerNodeOrArray(arg: ESQLAstItem) {
   return Array.isArray(arg) || !isMarkerNode(arg);
 }
@@ -101,7 +125,7 @@ function findCommandSubType<T extends ESQLCommandOption>(
   }
 }
 
-export function findAstPosition(ast: ESQLAstAllCommands[], offset: number) {
+export function findAstPosition(ast: ESQLAstQueryExpression, offset: number) {
   const command = findCommand(ast, offset);
   if (!command) {
     return { command: undefined, node: undefined };

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts

@@ -35,8 +35,8 @@ import {
 } from '@kbn/esql-ast/src/definitions/all_operators';
 import type { LicenseType } from '@kbn/licensing-types';
 import type { PricingProduct } from '@kbn/core-pricing-common/src/types';
-import { NOT_SUGGESTED_TYPES } from '../../shared/resources_helpers';
-import { getLocationFromCommandOrOptionName } from '../../shared/types';
+import { getLocationFromCommandOrOptionName } from '@kbn/esql-ast/src/commands_registry/location';
+import { NOT_SUGGESTED_TYPES } from '../../query_columns_service';
 import * as autocomplete from '../autocomplete';
 import type { ESQLCallbacks } from '../../shared/types';
 import {

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts

@@ -34,7 +34,7 @@ import {
 } from './__tests__/helpers';
 import { suggest } from './autocomplete';
 import { editorExtensions } from '../__tests__/helpers';
-import { mapRecommendedQueriesFromExtensions } from './utils/recommended_queries_helpers';
+import { mapRecommendedQueriesFromExtensions } from './recommended_queries_helpers';
 
 const getRecommendedQueriesSuggestionsFromTemplates = (
   fromCommand: string,

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

@@ -31,15 +31,25 @@ import { correctQuerySyntax } from '@kbn/esql-ast/src/definitions/utils/ast';
 import { ControlTriggerSource, ESQLVariableType } from '@kbn/esql-types';
 import type { LicenseType } from '@kbn/licensing-types';
 import type { ESQLAstAllCommands } from '@kbn/esql-ast/src/types';
-import { getAstContext } from '../shared/context';
-import { isHeaderCommandSuggestion, isSourceCommandSuggestion } from '../shared/helpers';
+import { getCursorContext } from '../shared/get_cursor_context';
 import { getFromCommandHelper } from '../shared/resources_helpers';
 import type { ESQLCallbacks } from '../shared/types';
 import { getCommandContext } from './get_command_context';
-import { mapRecommendedQueriesFromExtensions } from './utils/recommended_queries_helpers';
-import { getQueryForFields } from './get_query_for_fields';
-import type { GetColumnMapFn } from '../shared/columns';
-import { getColumnsByTypeRetriever } from '../shared/columns';
+import { mapRecommendedQueriesFromExtensions } from './recommended_queries_helpers';
+import { getQueryForFields } from '../shared/get_query_for_fields';
+import type { GetColumnMapFn } from '../shared/columns_retrieval_helpers';
+import { getColumnsByTypeRetriever } from '../shared/columns_retrieval_helpers';
+
+function isSourceCommandSuggestion({ label }: { label: string }) {
+  const sourceCommands = esqlCommandRegistry
+    .getSourceCommandNames()
+    .map((cmd) => cmd.toUpperCase());
+
+  return sourceCommands.includes(label);
+}
+function isHeaderCommandSuggestion({ label }: { label: string }) {
+  return label === 'SET';
+}
 
 const orderingEngine = new SuggestionOrderingEngine();
 
@@ -52,15 +62,15 @@ export async function suggest(
   const correctedQuery = correctQuerySyntax(innerText);
   const { root } = parse(correctedQuery, { withFormatting: true });
 
-  const astContext = getAstContext(innerText, root, offset);
+  const astContext = getCursorContext(innerText, root, offset);
 
   if (astContext.type === 'comment') {
     return [];
   }
 
   // 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.isCursorInSubquery ? astContext.astForContext : root;
+  const astForFields = astContext.astForContext;
 
   const { getColumnsByType, getColumnMap } = getColumnsByTypeRetriever(
     getQueryForFields(correctedQuery, astForFields),

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/hover/get_column_hover.ts

@@ -8,7 +8,7 @@
  */
 
 import type { ESQLColumn } from '@kbn/esql-ast';
-import type { GetColumnMapFn } from '../shared/columns';
+import type { GetColumnMapFn } from '../shared/columns_retrieval_helpers';
 
 export async function getColumnHover(
   node: ESQLColumn,

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/hover/get_function_signature_hover.ts

@@ -13,7 +13,7 @@ import {
   getFormattedFunctionSignature,
 } from '@kbn/esql-ast/src/definitions/utils';
 import { fromCache, setToCache } from './hover_cache';
-import type { ColumnsMap, GetColumnMapFn } from '../shared/columns';
+import type { ColumnsMap, GetColumnMapFn } from '../shared/columns_retrieval_helpers';
 
 export async function getFunctionSignatureHover(
   fnNode: ESQLFunction,

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/hover/index.ts

@@ -14,14 +14,14 @@ import {
 } from '@kbn/esql-ast/src/types';
 
 import type { ESQLCallbacks } from '../shared/types';
-import { getColumnsByTypeRetriever } from '../shared/columns';
+import { getColumnsByTypeRetriever } from '../shared/columns_retrieval_helpers';
 import { correctQuerySyntax, getVariablesHoverContent } from './helpers';
 import { getPolicyHover } from './get_policy_hover';
 import { getFunctionSignatureHover } from './get_function_signature_hover';
-import { getQueryForFields } from '../autocomplete/get_query_for_fields';
+import { getQueryForFields } from '../shared/get_query_for_fields';
 import { getFunctionArgumentHover } from './get_function_argument_hover';
 import { getColumnHover } from './get_column_hover';
-import { getAstContext } from '../shared/context';
+import { findSubquery } from '../shared/subqueries_helpers';
 
 interface HoverContent {
   contents: Array<{ value: string }>;
@@ -71,8 +71,8 @@ export async function getHoverItem(fullText: string, offset: number, callbacks?:
     return hoverContent;
   }
 
-  const astContext = getAstContext(correctedQuery, root, offset);
-  const astForContext = astContext.astForContext;
+  const { subQuery } = findSubquery(root, offset);
+  const astForContext = subQuery ?? root;
 
   const { getColumnMap } = getColumnsByTypeRetriever(
     getQueryForFields(fullText, astForContext),

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/inline_suggestions/inline_suggest.test.ts

@@ -16,7 +16,7 @@ jest.mock('@kbn/esql-ast/src/commands_registry/options/recommended_queries', ()
   getTimeAndCategorizationFields: jest.fn(),
 }));
 
-jest.mock('../shared/columns', () => ({
+jest.mock('../shared/columns_retrieval_helpers', () => ({
   getColumnsByTypeRetriever: jest.fn(),
 }));
 
@@ -29,7 +29,7 @@ import {
   getRecommendedQueriesTemplates,
   getTimeAndCategorizationFields,
 } from '@kbn/esql-ast/src/commands_registry/options/recommended_queries';
-import { getColumnsByTypeRetriever } from '../shared/columns';
+import { getColumnsByTypeRetriever } from '../shared/columns_retrieval_helpers';
 import { setToCache } from './inline_suggestions_cache';
 
 const mockGetRecommendedQueriesTemplates = getRecommendedQueriesTemplates as jest.MockedFunction<

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/inline_suggestions/inline_suggest.ts

@@ -13,7 +13,7 @@ import {
   getTimeAndCategorizationFields,
 } from '@kbn/esql-ast/src/commands_registry/options/recommended_queries';
 import type { ESQLCallbacks } from '../shared/types';
-import { getColumnsByTypeRetriever } from '../shared/columns';
+import { getColumnsByTypeRetriever } from '../shared/columns_retrieval_helpers';
 import { getFromCommandHelper } from '../shared/resources_helpers';
 import type { InlineSuggestionItem } from './types';
 import { fromCache, setToCache } from './inline_suggestions_cache';

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts → src/platform/packages/shared/kbn-esql-validation-autocomplete/src/query_columns_service/enrich_fields_with_ecs.test.ts

@@ -9,7 +9,7 @@
 
 import type { ESQLColumnData } from '@kbn/esql-ast/src/commands_registry/types';
 import { type ESQLFieldWithMetadata } from '@kbn/esql-ast/src/commands_registry/types';
-import { type ECSMetadata, enrichFieldsWithECSInfo } from './ecs_metadata_helper';
+import { type ECSMetadata, enrichFieldsWithECSInfo } from './enrich_fields_with_ecs';
 
 describe('enrichFieldsWithECSInfo', () => {
   it('should return original columns if fieldsMetadata is not provided', async () => {