[ES|QL] Follow ups on variables

src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx

@@ -406,7 +406,16 @@ export const ESQLEditor = memo(function ESQLEditor({
   const { cache: esqlFieldsCache, memoizedFieldsFromESQL } = useMemo(() => {
     // need to store the timing of the first request so we can atomically clear the cache per query
     const fn = memoize(
-      (...args: [{ esql: string }, ExpressionsStart, TimeRange, AbortController?]) => ({
+      (
+        ...args: [
+          { esql: string },
+          ExpressionsStart,
+          TimeRange,
+          AbortController?,
+          string?,
+          ESQLControlVariable[]?
+        ]
+      ) => ({
         timestamp: Date.now(),
         result: fetchFieldsFromESQL(...args),
       }),
@@ -446,7 +455,9 @@ export const ESQLEditor = memo(function ESQLEditor({
               esqlQuery,
               expressions,
               timeRange,
-              abortController
+              abortController,
+              undefined,
+              esqlVariables
             ).result;
             const columns: ESQLRealField[] =
               table?.columns.map((c) => {
@@ -478,8 +489,8 @@ export const ESQLEditor = memo(function ESQLEditor({
       },
       // @ts-expect-error To prevent circular type import, type defined here is partial of full client
       getFieldsMetadata: fieldsMetadata?.getClient(),
-      getVariablesByType: (type: ESQLVariableType) => {
-        return variablesService?.esqlVariables.filter((variable) => variable.type === type);
+      getVariables: () => {
+        return variablesService?.esqlVariables;
       },
       canSuggestVariables: () => {
         return variablesService?.areSuggestionsEnabled ?? false;
@@ -489,6 +500,7 @@ export const ESQLEditor = memo(function ESQLEditor({
     return callbacks;
   }, [
     fieldsMetadata,
+    esqlVariables,
     kibana.services?.esql?.getJoinIndicesAutocomplete,
     dataSourcesCache,
     query.esql,

src/platform/packages/private/kbn-esql-editor/src/fetch_fields_from_esql.ts

@@ -11,6 +11,7 @@ import { pluck } from 'rxjs';
 import { lastValueFrom } from 'rxjs';
 import { Query, AggregateQuery, TimeRange } from '@kbn/es-query';
 import type { ExpressionsStart } from '@kbn/expressions-plugin/public';
+import type { ESQLControlVariable } from '@kbn/esql-types';
 import type { Datatable } from '@kbn/expressions-plugin/public';
 import { textBasedQueryStateToAstWithValidation } from '@kbn/data-plugin/common';
 
@@ -26,7 +27,8 @@ export function fetchFieldsFromESQL(
   expressions: ExpressionsStart,
   time?: TimeRange,
   abortController?: AbortController,
-  timeFieldName?: string
+  timeFieldName?: string,
+  esqlVariables?: ESQLControlVariable[]
 ) {
   return textBasedQueryStateToAstWithValidation({
     query,
@@ -38,6 +40,7 @@ export function fetchFieldsFromESQL(
         const executionContract = expressions.execute(ast, null, {
           searchContext: {
             timeRange: time,
+            esqlVariables,
           },
         });
 

src/platform/packages/shared/kbn-esql-ast/index.ts

@@ -60,7 +60,7 @@ export {
   ESQLErrorListener,
 } from './src/parser';
 
-export { Walker, type WalkerOptions, walk } from './src/walker';
+export { Walker, type WalkerOptions, walk, type WalkerAstNode } from './src/walker';
 export * as synth from './src/synth';
 
 export {

src/platform/packages/shared/kbn-esql-ast/src/walker/index.ts

@@ -7,4 +7,4 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-export { Walker, type WalkerOptions, walk } from './walker';
+export { Walker, type WalkerOptions, walk, type WalkerAstNode } from './walker';

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

@@ -378,7 +378,7 @@ describe('autocomplete.suggest', () => {
           const suggestions = await suggest('FROM a | STATS /', {
             callbacks: {
               canSuggestVariables: () => true,
-              getVariablesByType: () => [],
+              getVariables: () => [],
               getColumnsFor: () => Promise.resolve([{ name: 'clientip', type: 'ip' }]),
             },
           });
@@ -399,7 +399,7 @@ describe('autocomplete.suggest', () => {
           const suggestions = await suggest('FROM a | STATS var0 = /', {
             callbacks: {
               canSuggestVariables: () => true,
-              getVariablesByType: () => [
+              getVariables: () => [
                 {
                   key: 'function',
                   value: 'avg',
@@ -426,7 +426,7 @@ describe('autocomplete.suggest', () => {
           const suggestions = await suggest('FROM a | STATS BY /', {
             callbacks: {
               canSuggestVariables: () => true,
-              getVariablesByType: () => [],
+              getVariables: () => [],
               getColumnsFor: () => Promise.resolve([{ name: 'clientip', type: 'ip' }]),
             },
           });
@@ -447,7 +447,7 @@ describe('autocomplete.suggest', () => {
           const suggestions = await suggest('FROM a | STATS BY /', {
             callbacks: {
               canSuggestVariables: () => true,
-              getVariablesByType: () => [
+              getVariables: () => [
                 {
                   key: 'field',
                   value: 'clientip',
@@ -474,7 +474,7 @@ describe('autocomplete.suggest', () => {
           const suggestions = await suggest('FROM a | STATS BY BUCKET(@timestamp, /)', {
             callbacks: {
               canSuggestVariables: () => true,
-              getVariablesByType: () => [
+              getVariables: () => [
                 {
                   key: 'interval',
                   value: '1 hour',

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

@@ -399,7 +399,7 @@ describe('WHERE <expression>', () => {
         const suggestions = await suggest('FROM a | WHERE agent.name == /', {
           callbacks: {
             canSuggestVariables: () => true,
-            getVariablesByType: () => [],
+            getVariables: () => [],
             getColumnsFor: () => Promise.resolve([{ name: 'agent.name', type: 'keyword' }]),
           },
         });
@@ -421,7 +421,7 @@ describe('WHERE <expression>', () => {
         const suggestions = await suggest('FROM a | WHERE agent.name == /', {
           callbacks: {
             canSuggestVariables: () => true,
-            getVariablesByType: () => [
+            getVariables: () => [
               {
                 key: 'value',
                 value: 'java',

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

@@ -16,7 +16,7 @@ import {
   type ESQLFunction,
   type ESQLSingleAstItem,
 } from '@kbn/esql-ast';
-import { ESQLVariableType, type ESQLControlVariable } from '@kbn/esql-types';
+import type { ESQLControlVariable } from '@kbn/esql-types';
 import { ESQL_NUMBER_TYPES, isNumericType } from '../shared/esql_types';
 import type { EditorContext, ItemKind, SuggestionRawDefinition, GetColumnsByTypeFn } from './types';
 import {
@@ -138,7 +138,7 @@ export async function suggest(
     resourceRetriever
   );
   const supportsControls = resourceRetriever?.canSuggestVariables?.() ?? false;
-  const getVariablesByType = resourceRetriever?.getVariablesByType;
+  const getVariables = resourceRetriever?.getVariables;
   const getSources = getSourcesHelper(resourceRetriever);
   const { getPolicies, getPolicyMetadata } = getPolicyRetriever(resourceRetriever);
 
@@ -188,7 +188,7 @@ export async function suggest(
       getFieldsMap,
       getPolicies,
       getPolicyMetadata,
-      getVariablesByType,
+      getVariables,
       resourceRetriever?.getPreferences,
       resourceRetriever,
       supportsControls
@@ -217,7 +217,7 @@ export async function suggest(
       getFieldsMap,
       fullText,
       offset,
-      getVariablesByType,
+      getVariables,
       supportsControls
     );
   }
@@ -239,7 +239,7 @@ export function getFieldsByTypeRetriever(
   resourceRetriever?: ESQLCallbacks
 ): { getFieldsByType: GetColumnsByTypeFn; getFieldsMap: GetFieldsMapFn } {
   const helpers = getFieldsByTypeHelper(queryString, resourceRetriever);
-  const getVariablesByType = resourceRetriever?.getVariablesByType;
+  const getVariables = resourceRetriever?.getVariables;
   const supportsControls = resourceRetriever?.canSuggestVariables?.() ?? false;
   return {
     getFieldsByType: async (
@@ -252,7 +252,7 @@ export function getFieldsByTypeRetriever(
         supportsControls,
       };
       const fields = await helpers.getFieldsByType(expectedType, ignored);
-      return buildFieldsDefinitionsWithMetadata(fields, updatedOptions, getVariablesByType);
+      return buildFieldsDefinitionsWithMetadata(fields, updatedOptions, getVariables);
     },
     getFieldsMap: helpers.getFieldsMap,
   };
@@ -355,7 +355,7 @@ async function getSuggestionsWithinCommandExpression(
   getFieldsMap: GetFieldsMapFn,
   getPolicies: GetPoliciesFn,
   getPolicyMetadata: GetPolicyMetadataFn,
-  getVariablesByType?: (type: ESQLVariableType) => ESQLControlVariable[] | undefined,
+  getVariables?: () => ESQLControlVariable[] | undefined,
   getPreferences?: () => Promise<{ histogramBarTarget: number } | undefined>,
   callbacks?: ESQLCallbacks,
   supportsControls?: boolean
@@ -396,7 +396,7 @@ async function getSuggestionsWithinCommandExpression(
       getSourcesFromQuery: (type) => getSourcesFromCommands(commands, type),
       previousCommands: commands,
       callbacks,
-      getVariablesByType,
+      getVariables,
       supportsControls,
       getPolicies,
       getPolicyMetadata,
@@ -907,7 +907,7 @@ async function getFunctionArgsSuggestions(
   getFieldsMap: GetFieldsMapFn,
   fullText: string,
   offset: number,
-  getVariablesByType?: (type: ESQLVariableType) => ESQLControlVariable[] | undefined,
+  getVariables?: () => ESQLControlVariable[] | undefined,
   supportsControls?: boolean
 ): Promise<SuggestionRawDefinition[]> {
   const fnDefinition = getFunctionDefinition(node.name);
@@ -1038,7 +1038,7 @@ async function getFunctionArgsSuggestions(
           advanceCursorAndOpenSuggestions: hasMoreMandatoryArgs,
           supportsControls,
         },
-        getVariablesByType
+        getVariables
       )
     );
 

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

@@ -59,11 +59,11 @@ const suggestFields = async (
   ]);
 
   const supportsControls = callbacks?.canSuggestVariables?.() ?? false;
-  const getVariablesByType = callbacks?.getVariablesByType;
+  const getVariables = callbacks?.getVariables;
   const joinFields = buildFieldsDefinitionsWithMetadata(
     lookupIndexFields!,
     { supportsControls },
-    getVariablesByType
+    getVariables
   );
 
   const intersection = suggestionIntersection(joinFields, sourceFields);

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

@@ -25,7 +25,7 @@ export async function suggest({
   getColumnsByType,
   getSuggestedVariableName,
   getPreferences,
-  getVariablesByType,
+  getVariables,
   supportsControls,
 }: CommandSuggestParams<'stats'>): Promise<SuggestionRawDefinition[]> {
   const pos = getPosition(innerText, command);
@@ -37,7 +37,7 @@ export async function suggest({
   const controlSuggestions = getControlSuggestionIfSupported(
     Boolean(supportsControls),
     ESQLVariableType.FUNCTIONS,
-    getVariablesByType
+    getVariables
   );
 
   switch (pos) {

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

@@ -209,7 +209,7 @@ export const buildFieldsDefinitionsWithMetadata = (
     variableType?: ESQLVariableType;
     supportsControls?: boolean;
   },
-  getVariablesByType?: (type: ESQLVariableType) => ESQLControlVariable[] | undefined
+  getVariables?: () => ESQLControlVariable[] | undefined
 ): SuggestionRawDefinition[] => {
   const fieldsSuggestions = fields.map((field) => {
     const titleCaseType = field.type.charAt(0).toUpperCase() + field.type.slice(1);
@@ -230,7 +230,7 @@ export const buildFieldsDefinitionsWithMetadata = (
   const suggestions = [...fieldsSuggestions];
   if (options?.supportsControls) {
     const variableType = options?.variableType ?? ESQLVariableType.FIELDS;
-    const variables = getVariablesByType?.(variableType) ?? [];
+    const variables = getVariables?.()?.filter((variable) => variable.type === variableType) ?? [];
 
     const controlSuggestions = fields.length
       ? getControlSuggestion(
@@ -418,7 +418,7 @@ export function getCompatibleLiterals(
     addComma?: boolean;
     supportsControls?: boolean;
   },
-  getVariablesByType?: (type: ESQLVariableType) => ESQLControlVariable[] | undefined
+  getVariables?: () => ESQLControlVariable[] | undefined
 ) {
   const suggestions: SuggestionRawDefinition[] = [];
   if (types.some(isNumericType)) {
@@ -436,7 +436,9 @@ export function getCompatibleLiterals(
       ...buildConstantsDefinitions(getUnitDuration(1), undefined, undefined, options),
     ];
     if (options?.supportsControls) {
-      const variables = getVariablesByType?.(ESQLVariableType.TIME_LITERAL) ?? [];
+      const variables =
+        getVariables?.()?.filter((variable) => variable.type === ESQLVariableType.TIME_LITERAL) ??
+        [];
       timeLiteralSuggestions.push(
         ...getControlSuggestion(
           ESQLVariableType.TIME_LITERAL,
@@ -520,13 +522,13 @@ export function getDateLiterals(options?: {
 export function getControlSuggestionIfSupported(
   supportsControls: boolean,
   type: ESQLVariableType,
-  getVariablesByType?: (type: ESQLVariableType) => ESQLControlVariable[] | undefined
+  getVariables?: () => ESQLControlVariable[] | undefined
 ) {
   if (!supportsControls) {
     return [];
   }
   const variableType = type;
-  const variables = getVariablesByType?.(variableType) ?? [];
+  const variables = getVariables?.()?.filter((variable) => variable.type === variableType) ?? [];
   const controlSuggestion = getControlSuggestion(
     variableType,
     variables?.map((v) => `?${v.key}`)

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

@@ -14,7 +14,7 @@ import type {
   ESQLMessage,
   ESQLSource,
 } from '@kbn/esql-ast';
-import { ESQLControlVariable, ESQLVariableType } from '@kbn/esql-types';
+import { ESQLControlVariable } from '@kbn/esql-types';
 import { GetColumnsByTypeFn, SuggestionRawDefinition } from '../autocomplete/types';
 import type { ESQLPolicy } from '../validation/types';
 import { ESQLCallbacks, ESQLSourceResult } from '../shared/types';
@@ -272,7 +272,7 @@ export interface CommandSuggestParams<CommandName extends string> {
    */
   previousCommands?: ESQLCommand[];
   callbacks?: ESQLCallbacks;
-  getVariablesByType?: (type: ESQLVariableType) => ESQLControlVariable[] | undefined;
+  getVariables?: () => ESQLControlVariable[] | undefined;
   supportsControls?: boolean;
 }
 

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

@@ -6,7 +6,7 @@
  * 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 { ESQLVariableType, type ESQLControlVariable } from '@kbn/esql-types';
+import type { ESQLControlVariable } from '@kbn/esql-types';
 import type { ESQLRealField, JoinIndexAutocompleteItem } from '../validation/types';
 
 /** @internal **/
@@ -46,7 +46,7 @@ export interface ESQLCallbacks {
   >;
   getPreferences?: () => Promise<{ histogramBarTarget: number }>;
   getFieldsMetadata?: Promise<PartialFieldsMetadataClient>;
-  getVariablesByType?: (type: ESQLVariableType) => ESQLControlVariable[] | undefined;
+  getVariables?: () => ESQLControlVariable[] | undefined;
   canSuggestVariables?: () => boolean;
   getJoinIndices?: () => Promise<{ indices: JoinIndexAutocompleteItem[] }>;
 }

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

@@ -1733,7 +1733,7 @@ describe('validation logic', () => {
         getColumnsFor: /Unknown column|Argument of|it is unsupported or not indexed/,
         getPreferences: /Unknown/,
         getFieldsMetadata: /Unknown/,
-        getVariablesByType: /Unknown/,
+        getVariables: /Unknown/,
         canSuggestVariables: /Unknown/,
       };
       return excludedCallback.map((callback) => (contentByCallback as any)[callback]) || [];

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

@@ -1352,7 +1352,7 @@ export const ignoreErrorsMap: Record<keyof ESQLCallbacks, ErrorTypes[]> = {
   getPolicies: ['unknownPolicy'],
   getPreferences: [],
   getFieldsMetadata: [],
-  getVariablesByType: [],
+  getVariables: [],
   canSuggestVariables: [],
   getJoinIndices: [],
 };