[ES|QL] Fix some validation misconfiguration (elastic#177783)

## Summary Related issue elastic#177699 Fix variables logic for expressions at `stats by ...` level. Fix validation logic for agg functions within `eval` or `where` scope. Fix validation and autocomplete logic for nested quoted expressions * i.e. ``` from index | eval round(numberField) + 1 | eval `round(numberField) + 1` + 1 | eval ```round(numberField) + 1`` + 1` + 1 | eval ```````round(numberField) + 1```` + 1`` + 1` + 1 | eval ```````````````round(numberField) + 1```````` + 1```` + 1`` + 1` + 1 | keep ```````````````````````````````round(numberField) + 1```````````````` + 1```````` + 1```` + 1`` + 1` ``` * updated `count_distinct` agg definition to have the `precision` second optional param. ### Checklist - [x] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios --------- Co-authored-by: Stratoula Kalafateli <[email protected]> (cherry picked from commit cad276f) # Conflicts: # packages/kbn-monaco/src/esql/lib/ast/validation/esql_validation_meta_tests.json
drewdaemon · Mar 7, 2024 · 18fe06f · 18fe06f
1 parent 03636de
commit 18fe06f
Show file tree

Hide file tree

Showing 11 changed files with 8,910 additions and 97 deletions.
diff --git a/packages/kbn-monaco/src/esql/lib/ast/autocomplete/autocomplete.test.ts b/packages/kbn-monaco/src/esql/lib/ast/autocomplete/autocomplete.test.ts
@@ -562,6 +562,18 @@ describe('autocomplete', () => {
         `from a | ${command} stringField, `,
         getFieldNamesByType('any').filter((name) => name !== 'stringField')
       );
+
+      testSuggestions(
+        `from a_index | eval round(numberField) + 1 | eval \`round(numberField) + 1\` + 1 | eval \`\`\`round(numberField) + 1\`\` + 1\` + 1 | eval \`\`\`\`\`\`\`round(numberField) + 1\`\`\`\` + 1\`\` + 1\` + 1 | eval \`\`\`\`\`\`\`\`\`\`\`\`\`\`\`round(numberField) + 1\`\`\`\`\`\`\`\` + 1\`\`\`\` + 1\`\` + 1\` + 1 | ${command} `,
+        [
+          ...getFieldNamesByType('any'),
+          '`round(numberField) + 1`',
+          '```round(numberField) + 1`` + 1`',
+          '```````round(numberField) + 1```` + 1`` + 1`',
+          '```````````````round(numberField) + 1```````` + 1```` + 1`` + 1`',
+          '```````````````````````````````round(numberField) + 1```````````````` + 1```````` + 1```` + 1`` + 1`',
+        ]
+      );
     });
   }
 
@@ -927,10 +939,7 @@ describe('autocomplete', () => {
       [
         'var0 =',
         ...getFieldNamesByType('any'),
-        // @TODO: leverage the location data to get the original text
-        // For now return back the trimmed version:
-        // the ANTLR parser trims all text so that's what it's stored in the AST
-        '`abs(numberField)+1`',
+        '`abs(numberField) + 1`',
         ...getFunctionSignaturesByReturnType('eval', 'any', { evalMath: true }),
       ],
       ' '

diff --git a/packages/kbn-monaco/src/esql/lib/ast/autocomplete/autocomplete.ts b/packages/kbn-monaco/src/esql/lib/ast/autocomplete/autocomplete.ts
@@ -71,7 +71,7 @@ import {
   buildOptionDefinition,
   buildSettingDefinitions,
 } from './factories';
-import { EDITOR_MARKER } from '../shared/constants';
+import { EDITOR_MARKER, SINGLE_BACKTICK } from '../shared/constants';
 import { getAstContext, removeMarkerArgFromArgsList } from '../shared/context';
 import {
   buildQueryUntilPreviousCommand,
@@ -563,7 +563,7 @@ async function getExpressionSuggestionsByType(
 
   // collect all fields + variables to suggest
   const fieldsMap: Map<string, ESQLRealField> = await (argDef ? getFieldsMap() : new Map());
-  const anyVariables = collectVariables(commands, fieldsMap);
+  const anyVariables = collectVariables(commands, fieldsMap, innerText);
 
   // enrich with assignment has some special rules who are handled somewhere else
   const canHaveAssignments = ['eval', 'stats', 'row'].includes(command.name);
@@ -1017,13 +1017,20 @@ async function getFieldsOrFunctionsSuggestions(
     }
     // due to a bug on the ES|QL table side, filter out fields list with underscored variable names (??)
     // avg( numberField ) => avg_numberField_
+    const ALPHANUMERIC_REGEXP = /[^a-zA-Z\d]/g;
     if (
       filteredVariablesByType.length &&
-      filteredVariablesByType.some((v) => /[^a-zA-Z\d]/.test(v))
+      filteredVariablesByType.some((v) => ALPHANUMERIC_REGEXP.test(v))
     ) {
       for (const variable of filteredVariablesByType) {
-        const underscoredName = variable.replace(/[^a-zA-Z\d]/g, '_');
-        const index = filteredFieldsByType.findIndex(({ label }) => underscoredName === label);
+        // remove backticks if present
+        const sanitizedVariable = variable.startsWith(SINGLE_BACKTICK)
+          ? variable.slice(1, variable.length - 1)
+          : variable;
+        const underscoredName = sanitizedVariable.replace(ALPHANUMERIC_REGEXP, '_');
+        const index = filteredFieldsByType.findIndex(
+          ({ label }) => underscoredName === label || `_${underscoredName}_` === label
+        );
         if (index >= 0) {
           filteredFieldsByType.splice(index);
         }
@@ -1067,7 +1074,8 @@ async function getFunctionArgsSuggestions(
   const variablesExcludingCurrentCommandOnes = excludeVariablesFromCurrentCommand(
     commands,
     command,
-    fieldsMap
+    fieldsMap,
+    innerText
   );
   // pick the type of the next arg
   const shouldGetNextArgument = node.text.includes(EDITOR_MARKER);
@@ -1102,7 +1110,10 @@ async function getFunctionArgsSuggestions(
   const isUnknownColumn =
     arg &&
     isColumnItem(arg) &&
-    !columnExists(arg, { fields: fieldsMap, variables: variablesExcludingCurrentCommandOnes }).hit;
+    !columnExists(arg, {
+      fields: fieldsMap,
+      variables: variablesExcludingCurrentCommandOnes,
+    }).hit;
   if (noArgDefined || isUnknownColumn) {
     const commandArgIndex = command.args.findIndex(
       (cmdArg) => isSingleItem(cmdArg) && cmdArg.location.max >= node.location.max
@@ -1213,7 +1224,7 @@ async function getListArgsSuggestions(
   // so extract the type of the first argument and suggest fields of that type
   if (node && isFunctionItem(node)) {
     const fieldsMap: Map<string, ESQLRealField> = await getFieldsMaps();
-    const anyVariables = collectVariables(commands, fieldsMap);
+    const anyVariables = collectVariables(commands, fieldsMap, innerText);
     // extract the current node from the variables inferred
     anyVariables.forEach((values, key) => {
       if (values.some((v) => v.location === node.location)) {
@@ -1301,7 +1312,7 @@ async function getOptionArgsSuggestions(
   const isNewExpression = isRestartingExpression(innerText) || option.args.length === 0;
 
   const fieldsMap = await getFieldsMaps();
-  const anyVariables = collectVariables(commands, fieldsMap);
+  const anyVariables = collectVariables(commands, fieldsMap, innerText);
 
   const references = {
     fields: fieldsMap,
@@ -1339,7 +1350,8 @@ async function getOptionArgsSuggestions(
         const policyMetadata = await getPolicyMetadata(policyName);
         const anyEnhancedVariables = collectVariables(
           commands,
-          appendEnrichFields(fieldsMap, policyMetadata)
+          appendEnrichFields(fieldsMap, policyMetadata),
+          innerText
         );
 
         if (isNewExpression) {

diff --git a/packages/kbn-monaco/src/esql/lib/ast/autocomplete/factories.ts b/packages/kbn-monaco/src/esql/lib/ast/autocomplete/factories.ts
@@ -129,7 +129,7 @@ export const buildFieldsDefinitions = (fields: string[]): AutocompleteCommandDef
 export const buildVariablesDefinitions = (variables: string[]): AutocompleteCommandDefinition[] =>
   variables.map((label) => ({
     label,
-    insertText: getSafeInsertText(label),
+    insertText: label,
     kind: 4,
     detail: i18n.translate('monaco.esql.autocomplete.variableDefinition', {
       defaultMessage: `Variable specified by the user within the ES|QL query`,

diff --git a/packages/kbn-monaco/src/esql/lib/ast/definitions/aggs.ts b/packages/kbn-monaco/src/esql/lib/ast/definitions/aggs.ts
@@ -125,7 +125,10 @@ export const statsAggregationFunctionDefinitions: FunctionDefinition[] = [
       supportedCommands: ['stats'],
       signatures: [
         {
-          params: [{ name: 'column', type: 'any', noNestingFunctions: true }],
+          params: [
+            { name: 'column', type: 'any', noNestingFunctions: true },
+            { name: 'precision', type: 'number', noNestingFunctions: true, optional: true },
+          ],
           returnType: 'number',
           examples: [
             `from index | stats result = count_distinct(field)`,

diff --git a/packages/kbn-monaco/src/esql/lib/ast/shared/constants.ts b/packages/kbn-monaco/src/esql/lib/ast/shared/constants.ts
@@ -8,7 +8,7 @@
 
 export const EDITOR_MARKER = 'marker_esql_editor';
 
-export const TICKS_REGEX = /^(`)|(`)$/g;
+export const TICKS_REGEX = /^`{1}|`{1}$/g;
 export const DOUBLE_TICKS_REGEX = /``/g;
 export const SINGLE_TICK_REGEX = /`/g;
 export const SINGLE_BACKTICK = '`';

diff --git a/packages/kbn-monaco/src/esql/lib/ast/shared/helpers.ts b/packages/kbn-monaco/src/esql/lib/ast/shared/helpers.ts
@@ -352,7 +352,8 @@ export function isEqualType(
   item: ESQLSingleAstItem,
   argDef: SignatureArgType,
   references: ReferenceMaps,
-  parentCommand?: string
+  parentCommand?: string,
+  nameHit?: string
 ) {
   const argType = 'innerType' in argDef && argDef.innerType ? argDef.innerType : argDef.type;
   if (argType === 'any') {
@@ -375,10 +376,8 @@ export function isEqualType(
       // anything goes, so avoid any effort here
       return true;
     }
-    // perform a double check, but give priority to the non trimmed version
-    const hit = getColumnHit(item.name, references);
-    const hitTrimmed = getColumnHit(item.name.replace(/\s/g, ''), references);
-    const validHit = hit || hitTrimmed;
+    const hit = getColumnHit(nameHit ?? item.name, references);
+    const validHit = hit;
     if (!validHit) {
       return false;
     }
@@ -445,9 +444,9 @@ export function columnExists(
     return { hit: true, nameHit: column.name };
   }
   if (column.quoted) {
-    const trimmedName = column.name.replace(/`/g, '``').replace(/\s/g, '');
-    if (variables.has(trimmedName)) {
-      return { hit: true, nameHit: trimmedName };
+    const originalName = column.text;
+    if (variables.has(originalName)) {
+      return { hit: true, nameHit: originalName };
     }
   }
   if (

diff --git a/packages/kbn-monaco/src/esql/lib/ast/shared/variables.ts b/packages/kbn-monaco/src/esql/lib/ast/shared/variables.ts
@@ -6,9 +6,9 @@
  * Side Public License, v 1.
  */
 
-import type { ESQLColumn, ESQLAstItem, ESQLCommand, ESQLCommandOption } from '../types';
+import type { ESQLAstItem, ESQLCommand, ESQLCommandOption, ESQLFunction } from '../types';
 import type { ESQLVariable, ESQLRealField } from '../validation/types';
-import { DOUBLE_TICKS_REGEX, EDITOR_MARKER, SINGLE_BACKTICK, TICKS_REGEX } from './constants';
+import { DOUBLE_BACKTICK, EDITOR_MARKER, SINGLE_BACKTICK } from './constants';
 import {
   isColumnItem,
   isAssignment,
@@ -26,21 +26,6 @@ function addToVariableOccurrencies(variables: Map<string, ESQLVariable[]>, insta
   variablesOccurrencies.push(instance);
 }
 
-function replaceTrimmedVariable(
-  variables: Map<string, ESQLVariable[]>,
-  newRef: ESQLColumn,
-  oldRef: ESQLVariable[]
-) {
-  // now replace the existing trimmed version with this original one
-  addToVariableOccurrencies(variables, {
-    name: newRef.name,
-    type: oldRef[0].type,
-    location: newRef.location,
-  });
-  // remove the trimmed one
-  variables.delete(oldRef[0].name);
-}
-
 function addToVariables(
   oldArg: ESQLAstItem,
   newArg: ESQLAstItem,
@@ -55,20 +40,11 @@ function addToVariables(
     };
     // Now workout the exact type
     // it can be a rename of another variable as well
-    let oldRef = fields.get(oldArg.name) || variables.get(oldArg.name);
+    const oldRef =
+      fields.get(oldArg.name) || variables.get(oldArg.quoted ? oldArg.text : oldArg.name);
     if (oldRef) {
       addToVariableOccurrencies(variables, newVariable);
       newVariable.type = Array.isArray(oldRef) ? oldRef[0].type : oldRef.type;
-    } else if (oldArg.quoted) {
-      // a last attempt in case the user tried to rename an expression:
-      // trim every space and try a new hit
-      const expressionTrimmedRef = oldArg.name.replace(/\s/g, '');
-      oldRef = variables.get(expressionTrimmedRef);
-      if (oldRef) {
-        addToVariableOccurrencies(variables, newVariable);
-        newVariable.type = oldRef[0].type;
-        replaceTrimmedVariable(variables, oldArg, oldRef);
-      }
     }
   }
 }
@@ -99,10 +75,11 @@ function getAssignRightHandSideType(item: ESQLAstItem, fields: Map<string, ESQLR
 export function excludeVariablesFromCurrentCommand(
   commands: ESQLCommand[],
   currentCommand: ESQLCommand,
-  fieldsMap: Map<string, ESQLRealField>
+  fieldsMap: Map<string, ESQLRealField>,
+  queryString: string
 ) {
-  const anyVariables = collectVariables(commands, fieldsMap);
-  const currentCommandVariables = collectVariables([currentCommand], fieldsMap);
+  const anyVariables = collectVariables(commands, fieldsMap, queryString);
+  const currentCommandVariables = collectVariables([currentCommand], fieldsMap, queryString);
   const resultVariables = new Map<string, ESQLVariable[]>();
   anyVariables.forEach((value, key) => {
     if (!currentCommandVariables.has(key)) {
@@ -112,55 +89,78 @@ export function excludeVariablesFromCurrentCommand(
   return resultVariables;
 }
 
+function extractExpressionAsQuotedVariable(
+  originalQuery: string,
+  location: { min: number; max: number }
+) {
+  const extractExpressionText = originalQuery.substring(location.min, location.max + 1);
+  // now inject quotes and save it as variable
+  return `\`${extractExpressionText.replaceAll(SINGLE_BACKTICK, DOUBLE_BACKTICK)}\``;
+}
+
+function addVariableFromAssignment(
+  assignOperation: ESQLFunction,
+  variables: Map<string, ESQLVariable[]>,
+  fields: Map<string, ESQLRealField>
+) {
+  if (isColumnItem(assignOperation.args[0])) {
+    const rightHandSideArgType = getAssignRightHandSideType(assignOperation.args[1], fields);
+    addToVariableOccurrencies(variables, {
+      name: assignOperation.args[0].name,
+      type: rightHandSideArgType || 'number' /* fallback to number */,
+      location: assignOperation.args[0].location,
+    });
+  }
+}
+
+function addVariableFromExpression(
+  expressionOperation: ESQLFunction,
+  queryString: string,
+  variables: Map<string, ESQLVariable[]>
+) {
+  if (!expressionOperation.text.includes(EDITOR_MARKER)) {
+    // save the variable in its quoted usable way
+    // (a bit of forward thinking here to simplyfy lookups later)
+    const forwardThinkingVariableName = extractExpressionAsQuotedVariable(
+      queryString,
+      expressionOperation.location
+    );
+    const expressionType = 'number';
+    addToVariableOccurrencies(variables, {
+      name: forwardThinkingVariableName,
+      type: expressionType,
+      location: expressionOperation.location,
+    });
+  }
+}
+
 export function collectVariables(
   commands: ESQLCommand[],
-  fields: Map<string, ESQLRealField>
+  fields: Map<string, ESQLRealField>,
+  queryString: string
 ): Map<string, ESQLVariable[]> {
   const variables = new Map<string, ESQLVariable[]>();
   for (const command of commands) {
     if (['row', 'eval', 'stats'].includes(command.name)) {
-      const assignOperations = command.args.filter(isAssignment);
-      for (const assignOperation of assignOperations) {
-        if (isColumnItem(assignOperation.args[0])) {
-          const rightHandSideArgType = getAssignRightHandSideType(assignOperation.args[1], fields);
-          addToVariableOccurrencies(variables, {
-            name: assignOperation.args[0].name,
-            type: rightHandSideArgType || 'number' /* fallback to number */,
-            location: assignOperation.args[0].location,
-          });
+      for (const arg of command.args) {
+        if (isAssignment(arg)) {
+          addVariableFromAssignment(arg, variables, fields);
         }
-      }
-      const expressionOperations = command.args.filter(isExpression);
-      for (const expressionOperation of expressionOperations) {
-        if (!expressionOperation.text.includes(EDITOR_MARKER)) {
-          // just save the entire expression as variable string
-          const expressionType = 'number';
-          addToVariableOccurrencies(variables, {
-            name: expressionOperation.text
-              .replace(TICKS_REGEX, '')
-              .replace(DOUBLE_TICKS_REGEX, SINGLE_BACKTICK),
-            type: expressionType,
-            location: expressionOperation.location,
-          });
+        if (isExpression(arg)) {
+          addVariableFromExpression(arg, queryString, variables);
         }
       }
       if (command.name === 'stats') {
         const commandOptionsWithAssignment = command.args.filter(
           (arg) => isOptionItem(arg) && arg.name === 'by'
         ) as ESQLCommandOption[];
         for (const commandOption of commandOptionsWithAssignment) {
-          const optionAssignOperations = commandOption.args.filter(isAssignment);
-          for (const assignOperation of optionAssignOperations) {
-            if (isColumnItem(assignOperation.args[0])) {
-              const rightHandSideArgType = getAssignRightHandSideType(
-                assignOperation.args[1],
-                fields
-              );
-              addToVariableOccurrencies(variables, {
-                name: assignOperation.args[0].name,
-                type: rightHandSideArgType || 'number' /* fallback to number */,
-                location: assignOperation.args[0].location,
-              });
+          for (const optArg of commandOption.args) {
+            if (isAssignment(optArg)) {
+              addVariableFromAssignment(optArg, variables, fields);
+            }
+            if (isExpression(optArg)) {
+              addVariableFromExpression(optArg, queryString, variables);
             }
           }
         }
@@ -171,6 +171,7 @@ export function collectVariables(
         (arg) => isOptionItem(arg) && arg.name === 'with'
       ) as ESQLCommandOption[];
       for (const commandOption of commandOptionsWithAssignment) {
+        // Enrich assignment has some special behaviour, so do not use the version above here...
         for (const assignFn of commandOption.args) {
           if (isFunctionItem(assignFn)) {
             const [newArg, oldArg] = assignFn?.args || [];