Add more logic for new variable suggestions

elastic · Sep 9, 2024 · aa7eb27 · aa7eb27
1 parent caa8472
commit aa7eb27
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 8 deletions.
diff --git a/...esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.suggest.eval.test.ts b/...esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.suggest.eval.test.ts
@@ -554,7 +554,7 @@ describe('autocomplete.suggest', () => {
     test('case', async () => {
       const { assertSuggestions } = await setup();
       const comparisonOperators = ['==', '!=', '>', '<', '>=', '<=']
-        .map((op) => `${op} `)
+        .map((op) => `${op}`)
         .concat(',');
 
       // case( / ) suggest any field/eval function in this position as first argument
@@ -621,7 +621,6 @@ describe('autocomplete.suggest', () => {
         // Notice no extra space after field name
         ...getFieldNamesByType('any').map((field) => `${field}`),
         ...getFunctionSignaturesByReturnType('eval', 'any', { scalar: true }, undefined, []),
-        'var0 = ',
       ]);
 
       // case( field > 0, >) suggests fields like normal

diff --git a/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts b/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
@@ -103,6 +103,7 @@ import {
   getValidFunctionSignaturesForPreviousArgs,
   strictlyGetParamAtPosition,
   isLiteralDateItem,
+  isCursorPlacementWithinFunction,
 } from './helper';
 import {
   FunctionParameter,
@@ -646,8 +647,11 @@ async function getExpressionSuggestionsByType(
   const fieldsMap: Map<string, ESQLRealField> = await (argDef ? getFieldsMap() : new Map());
   const anyVariables = collectVariables(commands, fieldsMap, innerText);
 
+  const previousWord = findPreviousWord(innerText);
   // enrich with assignment has some special rules who are handled somewhere else
-  const canHaveAssignments = ['eval', 'stats', 'row'].includes(command.name);
+  const canHaveAssignments =
+    ['eval', 'stats', 'row'].includes(command.name) &&
+    !comparisonFunctions.map((fn) => fn.name).includes(previousWord);
 
   const references = { fields: fieldsMap, variables: anyVariables };
 
@@ -1500,10 +1504,9 @@ async function getFunctionArgsSuggestions(
       suggestions.push(
         ...comparisonFunctions.map<SuggestionRawDefinition>(({ name, description }) => ({
           label: name,
-          text: name + ' ',
+          text: name,
           kind: 'Function' as ItemKind,
           detail: description,
-          command: TRIGGER_SUGGESTION_COMMAND,
         }))
       );
     }
@@ -1820,10 +1823,11 @@ async function getOptionArgsSuggestions(
             openSuggestions: true,
           }))
         );
+        const canHaveAssignment = !isCursorPlacementWithinFunction(innerText);
 
         if (option.name === 'by') {
           // Add quick snippet for for stats ... by bucket(<>)
-          if (command.name === 'stats') {
+          if (command.name === 'stats' && canHaveAssignment) {
             suggestions.push({
               label: i18n.translate(
                 'kbn-esql-validation-autocomplete.esql.autocomplete.addDateHistogram',
@@ -1854,12 +1858,13 @@ async function getOptionArgsSuggestions(
               {
                 functions: true,
                 fields: false,
-              }
+              },
+              { ignoreFn: canHaveAssignment ? [] : ['bucket', 'case'] }
             ))
           );
         }
 
-        if (command.name === 'stats' && isNewExpression) {
+        if (command.name === 'stats' && isNewExpression && canHaveAssignment) {
           suggestions.push(buildNewVarDefinition(findNewVariable(anyVariables)));
         }
       }

diff --git a/packages/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts b/packages/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
@@ -232,3 +232,22 @@ export function isLiteralDateItem(nodeArg: ESQLAstItem): boolean {
       isValidDateString(nodeArg.text))
   );
 }
+
+/**
+ * Checks if the marker editor/cursor is within an unclosed parenthesis
+ * E.g. returns false if inner text is "fn(a, b, (c + d), e" <- cursor is at end of string
+ * @param innerText
+ * @returns boolean
+ */
+export function isCursorPlacementWithinFunction(innerText: string) {
+  let openParenCount = 0;
+  for (let i = 0; i < innerText.length; i++) {
+    const char = innerText[i];
+    if (char === '(') {
+      openParenCount++;
+    } else if (char === ')') {
+      openParenCount--;
+    }
+  }
+  return openParenCount > 0;
+}