[8.x] [ES|QL] Move to query locations (#216110)

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

@@ -18,6 +18,7 @@ import {
   FunctionReturnType,
   Signature,
   FunctionDefinitionTypes,
+  Location,
 } from '../src/definitions/types';
 import { FULL_TEXT_SEARCH_FUNCTIONS } from '../src/shared/constants';
 const aliasTable: Record<string, string[]> = {
@@ -79,14 +80,17 @@ const bucketParameterTypes: Array<
   ['long', 'integer', 'long', 'long', 'double'],
 ];
 
-const scalarSupportedCommandsAndOptions = {
-  supportedCommands: ['stats', 'inlinestats', 'metrics', 'eval', 'where', 'row', 'sort'],
-  supportedOptions: ['by'],
-};
+const defaultScalarFunctionLocations: Location[] = [
+  Location.EVAL,
+  Location.ROW,
+  Location.SORT,
+  Location.WHERE,
+  Location.STATS,
+  Location.STATS_BY,
+  Location.STATS_WHERE,
+];
 
-const aggregationSupportedCommandsAndOptions = {
-  supportedCommands: ['stats', 'inlinestats', 'metrics'],
-};
+const defaultAggFunctionLocations: Location[] = [Location.STATS];
 
 // coalesce can be removed when a test is added for version type
 // (https://github.com/elastic/elasticsearch/pull/109032#issuecomment-2150033350)
@@ -98,7 +102,7 @@ const extraFunctions: FunctionDefinition[] = [
     name: 'case',
     description:
       'Accepts pairs of conditions and values. The function returns the value that belongs to the first condition that evaluates to `true`. If the number of arguments is odd, the last argument is the default value which is returned when no condition matches.',
-    ...scalarSupportedCommandsAndOptions,
+    locationsAvailable: defaultScalarFunctionLocations,
     signatures: [
       {
         params: [
@@ -306,26 +310,20 @@ const convertDateTime = (s: string) => (s === 'datetime' ? 'date' : s);
  * @returns
  */
 function getFunctionDefinition(ESFunctionDefinition: Record<string, any>): FunctionDefinition {
-  let supportedCommandsAndOptions: Pick<
-    FunctionDefinition,
-    'supportedCommands' | 'supportedOptions'
-  > =
+  let locationsAvailable =
     ESFunctionDefinition.type === FunctionDefinitionTypes.SCALAR
-      ? scalarSupportedCommandsAndOptions
-      : aggregationSupportedCommandsAndOptions;
+      ? defaultScalarFunctionLocations
+      : defaultAggFunctionLocations;
 
   // MATCH and QSRT has limited supported for where commands only
   if (FULL_TEXT_SEARCH_FUNCTIONS.includes(ESFunctionDefinition.name)) {
-    supportedCommandsAndOptions = {
-      supportedCommands: ['where'],
-      supportedOptions: [],
-    };
+    locationsAvailable = [Location.WHERE];
   }
   const ret = {
     type: ESFunctionDefinition.type,
     name: ESFunctionDefinition.name,
     operator: ESFunctionDefinition.operator,
-    ...supportedCommandsAndOptions,
+    locationsAvailable,
     description: ESFunctionDefinition.description,
     alias: aliasTable[ESFunctionDefinition.name],
     ignoreAsSuggestion: ESFunctionDefinition.snapshot_only,
@@ -690,13 +688,13 @@ const enrichGrouping = (
       ];
       return {
         ...op,
+        locationsAvailable: [...op.locationsAvailable, Location.STATS_BY],
         signatures,
-        supportedOptions: ['by'],
       };
     }
     return {
       ...op,
-      supportedOptions: ['by'],
+      locationsAvailable: [...op.locationsAvailable, Location.STATS_BY],
     };
   });
 };
@@ -721,26 +719,32 @@ const enrichOperators = (
       // Elasticsearch docs uses lhs and rhs instead of left and right that Kibana code uses
       params: s.params.map((param) => ({ ...param, name: replaceParamName(param.name) })),
     }));
-    let supportedCommands = op.supportedCommands;
-    let supportedOptions = op.supportedOptions;
+
+    let locationsAvailable = op.locationsAvailable;
+
     if (isComparisonOperator) {
-      supportedCommands = _.uniq([...op.supportedCommands, 'eval', 'where', 'row', 'sort']);
-      supportedOptions = ['by'];
+      locationsAvailable = _.uniq([
+        ...op.locationsAvailable,
+        Location.EVAL,
+        Location.WHERE,
+        Location.ROW,
+        Location.SORT,
+        Location.STATS_BY,
+      ]);
     }
     if (isMathOperator) {
-      supportedCommands = _.uniq([
-        ...op.supportedCommands,
-        'eval',
-        'where',
-        'row',
-        'stats',
-        'metrics',
-        'sort',
+      locationsAvailable = _.uniq([
+        ...op.locationsAvailable,
+        Location.EVAL,
+        Location.WHERE,
+        Location.ROW,
+        Location.STATS,
+        Location.SORT,
+        Location.STATS_BY,
       ]);
-      supportedOptions = ['by'];
     }
     if (isInOperator || isLikeOperator || isNotOperator) {
-      supportedCommands = ['eval', 'where', 'row', 'sort'];
+      locationsAvailable = [Location.EVAL, Location.WHERE, Location.SORT, Location.ROW];
     }
     if (isInOperator) {
       // Override the signatures to be array types instead of singular
@@ -766,8 +770,7 @@ const enrichOperators = (
       signatures,
       // Elasticsearch docs does not include the full supported commands for math operators
       // so we are overriding to add proper support
-      supportedCommands,
-      supportedOptions,
+      locationsAvailable,
       type: FunctionDefinitionTypes.OPERATOR,
       validate: validators[op.name],
       ...(isNotOperator ? { ignoreAsSuggestion: true } : {}),
@@ -829,6 +832,12 @@ function printGeneratedFunctionsFile(
     if (name.toLowerCase() === 'match') {
       functionName = 'match';
     }
+
+    // Map locationsAvailable to enum names
+    const locationsAvailable = functionDefinition.locationsAvailable.map(
+      (location) => `Location.${location.toUpperCase()}`
+    );
+
     return `// Do not edit this manually... generated by scripts/generate_function_definitions.ts
     const ${getDefinitionName(name)}: FunctionDefinition = {
     type: FunctionDefinitionTypes.${type.toUpperCase()},
@@ -839,8 +848,7 @@ function printGeneratedFunctionsFile(
     preview: ${functionDefinition.preview || 'false'},
     alias: ${alias ? `['${alias.join("', '")}']` : 'undefined'},
     signatures: ${JSON.stringify(signatures, null, 2)},
-    supportedCommands: ${JSON.stringify(functionDefinition.supportedCommands)},
-    supportedOptions: ${JSON.stringify(functionDefinition.supportedOptions)},
+    locationsAvailable: [${locationsAvailable.join(', ')}],
     validate: ${functionDefinition.validate || 'undefined'},
     examples: ${JSON.stringify(functionDefinition.examples || [])},${
       customParametersSnippet
@@ -870,7 +878,7 @@ function printGeneratedFunctionsFile(
  */
 
 import { i18n } from '@kbn/i18n';
-import { type FunctionDefinition, FunctionDefinitionTypes } from '../types';
+import { type FunctionDefinition, FunctionDefinitionTypes, Location } from '../types';
 ${
   functionsType === FunctionDefinitionTypes.SCALAR
     ? `import type { ESQLFunction } from '@kbn/esql-ast';

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

@@ -14,7 +14,10 @@ import { operatorsDefinitions } from '../../definitions/all_operators';
 import { NOT_SUGGESTED_TYPES } from '../../shared/resources_helpers';
 import { aggFunctionDefinitions } from '../../definitions/generated/aggregation_functions';
 import { timeUnitsToSuggest } from '../../definitions/literals';
-import { FunctionDefinitionTypes } from '../../definitions/types';
+import {
+  FunctionDefinitionTypes,
+  getLocationFromCommandOrOptionName,
+} from '../../definitions/types';
 import { groupingFunctionDefinitions } from '../../definitions/generated/grouping_functions';
 import * as autocomplete from '../autocomplete';
 import type { ESQLCallbacks } from '../../shared/types';
@@ -175,14 +178,16 @@ export function getFunctionSignaturesByReturnType(
   const deduped = Array.from(new Set(list));
 
   const commands = Array.isArray(command) ? command : [command];
+
   return deduped
-    .filter(({ signatures, ignoreAsSuggestion, supportedCommands, supportedOptions, name }) => {
+    .filter(({ signatures, ignoreAsSuggestion, locationsAvailable }) => {
       if (ignoreAsSuggestion) {
         return false;
       }
       if (
-        !commands.some((c) => supportedCommands.includes(c)) &&
-        !supportedOptions?.includes(option || '')
+        !(option ? [...commands, option] : commands).some((name) =>
+          locationsAvailable.includes(getLocationFromCommandOrOptionName(name))
+        )
       ) {
         return false;
       }

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

@@ -8,7 +8,7 @@
  */
 
 import { setTestFunctions } from '../../shared/test_functions';
-import { FunctionDefinitionTypes } from '../../definitions/types';
+import { FunctionDefinitionTypes, Location } from '../../definitions/types';
 import { setup } from './helpers';
 
 describe('hidden commands', () => {
@@ -33,15 +33,15 @@ describe('hidden functions', () => {
         name: 'HIDDEN_FUNCTION',
         description: 'This is a hidden function',
         signatures: [{ params: [], returnType: 'text' }],
-        supportedCommands: ['eval'],
+        locationsAvailable: [Location.EVAL],
         ignoreAsSuggestion: true,
       },
       {
         type: FunctionDefinitionTypes.SCALAR,
         name: 'VISIBLE_FUNCTION',
         description: 'This is a visible function',
         signatures: [{ params: [], returnType: 'text' }],
-        supportedCommands: ['eval'],
+        locationsAvailable: [Location.EVAL],
         ignoreAsSuggestion: false,
       },
     ]);
@@ -59,15 +59,15 @@ describe('hidden functions', () => {
         name: 'HIDDEN_FUNCTION',
         description: 'This is a hidden function',
         signatures: [{ params: [], returnType: 'text' }],
-        supportedCommands: ['stats'],
+        locationsAvailable: [Location.STATS],
         ignoreAsSuggestion: true,
       },
       {
         type: FunctionDefinitionTypes.AGG,
         name: 'VISIBLE_FUNCTION',
         description: 'This is a visible function',
         signatures: [{ params: [], returnType: 'text' }],
-        supportedCommands: ['stats'],
+        locationsAvailable: [Location.STATS],
         ignoreAsSuggestion: false,
       },
     ]);
@@ -84,9 +84,14 @@ describe('hidden functions', () => {
         type: FunctionDefinitionTypes.OPERATOR,
         name: 'HIDDEN_OPERATOR',
         description: 'This is a hidden function',
-        supportedCommands: ['eval', 'where', 'row', 'sort'],
+        locationsAvailable: [
+          Location.EVAL,
+          Location.WHERE,
+          Location.ROW,
+          Location.SORT,
+          Location.STATS_BY,
+        ],
         ignoreAsSuggestion: true,
-        supportedOptions: ['by'],
         signatures: [
           {
             params: [
@@ -101,9 +106,14 @@ describe('hidden functions', () => {
         type: FunctionDefinitionTypes.OPERATOR,
         name: 'VISIBLE_OPERATOR',
         description: 'This is a visible function',
-        supportedCommands: ['eval', 'where', 'row', 'sort'],
+        locationsAvailable: [
+          Location.EVAL,
+          Location.WHERE,
+          Location.ROW,
+          Location.SORT,
+          Location.STATS_BY,
+        ],
         ignoreAsSuggestion: false,
-        supportedOptions: ['by'],
         signatures: [
           {
             params: [

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

@@ -74,6 +74,7 @@ import {
   FunctionParameter,
   FunctionDefinitionTypes,
   GetPolicyMetadataFn,
+  getLocationFromCommandOrOptionName,
 } from '../definitions/types';
 import { comparisonFunctions } from '../definitions/all_operators';
 import { getRecommendedQueriesSuggestions } from './recommended_queries/suggestions';
@@ -268,8 +269,7 @@ async function getSuggestionsWithinCommandExpression(
   ) {
     return await getSuggestionsToRightOfOperatorExpression({
       queryText: innerText,
-      commandName: astContext.command.name,
-      optionName: astContext.option?.name,
+      location: getLocationFromCommandOrOptionName(astContext.command.name),
       rootOperator: astContext.node,
       getExpressionType: (expression) =>
         getExpressionType(expression, references.fields, references.variables),
@@ -491,8 +491,7 @@ async function getFunctionArgsSuggestions(
     if (typesToSuggestNext.every((d) => !d.fieldsOnly)) {
       suggestions.push(
         ...getFunctionSuggestions({
-          command: command.name,
-          option: option?.name,
+          location: getLocationFromCommandOrOptionName(option?.name ?? command.name),
           returnTypes: canBeBooleanCondition
             ? ['any']
             : (getTypesFromParamDefs(typesToSuggestNext) as string[]),
@@ -596,8 +595,7 @@ async function getListArgsSuggestions(
         suggestions.push(
           ...(await getFieldsOrFunctionsSuggestions(
             [argType as string],
-            command.name,
-            undefined,
+            getLocationFromCommandOrOptionName(command.name),
             getFieldsByType,
             {
               functions: true,

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

@@ -14,7 +14,7 @@ import {
   isSingleItem,
   unescapeColumnName,
 } from '../../../shared/helpers';
-import { CommandSuggestParams } from '../../../definitions/types';
+import { CommandSuggestParams, Location } from '../../../definitions/types';
 import type { SuggestionRawDefinition } from '../../types';
 import {
   Position,
@@ -145,7 +145,7 @@ export async function suggest({
         return [pipeCompleteItem, { ...commaCompleteItem, command: TRIGGER_SUGGESTION_COMMAND }];
       } else {
         // not recognized as a field name, assume new user-defined column name
-        return getOperatorSuggestions({ command: 'enrich' });
+        return getOperatorSuggestions({ location: Location.ENRICH });
       }
     }
 

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

@@ -10,7 +10,7 @@
 import type { ESQLSingleAstItem } from '@kbn/esql-ast';
 import { isMarkerNode } from '../../../shared/context';
 import { isAssignment, isColumnItem } from '../../../..';
-import { CommandSuggestParams } from '../../../definitions/types';
+import { CommandSuggestParams, Location } from '../../../definitions/types';
 import type { SuggestionRawDefinition } from '../../types';
 import { getNewVariableSuggestion } from '../../factories';
 import { commaCompleteItem, pipeCompleteItem } from '../../complete_items';
@@ -40,7 +40,7 @@ export async function suggest(
   const suggestions = await suggestForExpression({
     ...params,
     expressionRoot,
-    commandName: 'eval',
+    location: Location.EVAL,
   });
 
   const positionInExpression = getExpressionPosition(params.innerText, expressionRoot);

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

@@ -8,7 +8,7 @@
  */
 
 import { isRestartingExpression } from '../../../shared/helpers';
-import { CommandSuggestParams } from '../../../definitions/types';
+import { CommandSuggestParams, Location } from '../../../definitions/types';
 
 import type { SuggestionRawDefinition } from '../../types';
 import {
@@ -25,7 +25,7 @@ export async function suggest({
 }: CommandSuggestParams<'row'>): Promise<SuggestionRawDefinition[]> {
   // ROW var0 = /
   if (/=\s*$/.test(innerText)) {
-    return getFunctionSuggestions({ command: 'row' });
+    return getFunctionSuggestions({ location: Location.ROW });
   }
 
   // ROW var0 = 23 /
@@ -40,6 +40,6 @@ export async function suggest({
   // ROW foo = "bar", /
   return [
     getNewVariableSuggestion(getSuggestedVariableName()),
-    ...getFunctionSuggestions({ command: 'row' }),
+    ...getFunctionSuggestions({ location: Location.ROW }),
   ];
 }