Fix autocomplete for by suggestion

elastic · Jul 30, 2024 · a16a9f3 · a16a9f3
1 parent 7e09343
commit a16a9f3
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 27 deletions.
diff --git a/...sql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.stats.test.ts b/...sql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.stats.test.ts
@@ -227,6 +227,11 @@ describe('autocomplete.suggest', () => {
         const { assertSuggestions } = await setup();
 
         await assertSuggestions('from a | stats avg(b) by numberField % 2 /', [',', '|']);
+
+        await assertSuggestions(
+          'from a | stats var0 = AVG(products.base_price) BY var1 = BUCKET(order_date, 1 day)/',
+          [',', '|', '+ $0', '- $0']
+        );
       });
     });
   });

diff --git a/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts b/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
@@ -1625,42 +1625,60 @@ async function getOptionArgsSuggestions(
         );
       }
     }
-  }
 
-  if (optionDef) {
-    const argDefIndex = optionDef.signature.multipleParams
-      ? 0
-      : Math.max(option.args.length - 1, 0);
-    const types = [optionDef.signature.params[argDefIndex].type].filter(nonNullable);
-    // If it's a complete expression then proposed some final suggestions
-    // A complete expression is either a function or a column: <COMMAND> <OPTION> field <here>
-    // Or an assignment complete: <COMMAND> <OPTION> field = ... <here>
+    // If it's a complete expression then propose some final suggestions
     if (
-      (!suggestions.length && option.args.length && !isNewExpression && !isAssignment(lastArg)) ||
+      (!nodeArgType &&
+        option.name === 'by' &&
+        option.args.length &&
+        !isNewExpression &&
+        !isAssignment(lastArg)) ||
       (isAssignment(lastArg) && isAssignmentComplete(lastArg))
     ) {
       suggestions.push(
         ...getFinalSuggestions({
-          comma: optionDef.signature.multipleParams,
+          comma: optionDef?.signature.multipleParams ?? option.name === 'by',
         })
       );
-    } else if (isNewExpression || (isAssignment(nodeArg) && !isAssignmentComplete(nodeArg))) {
-      // Otherwise try to complete the expression suggesting some columns
-      suggestions.push(
-        ...(await getFieldsOrFunctionsSuggestions(
-          types[0] === 'column' ? ['any'] : types,
-          command.name,
-          option.name,
-          getFieldsByType,
-          {
-            functions: option.name === 'by',
-            fields: true,
-          }
-        ))
-      );
+    }
+  }
+
+  if (optionDef) {
+    if (!suggestions.length) {
+      const argDefIndex = optionDef.signature.multipleParams
+        ? 0
+        : Math.max(option.args.length - 1, 0);
+      const types = [optionDef.signature.params[argDefIndex].type].filter(nonNullable);
+      // If it's a complete expression then proposed some final suggestions
+      // A complete expression is either a function or a column: <COMMAND> <OPTION> field <here>
+      // Or an assignment complete: <COMMAND> <OPTION> field = ... <here>
+      if (
+        (option.args.length && !isNewExpression && !isAssignment(lastArg)) ||
+        (isAssignment(lastArg) && isAssignmentComplete(lastArg))
+      ) {
+        suggestions.push(
+          ...getFinalSuggestions({
+            comma: optionDef.signature.multipleParams,
+          })
+        );
+      } else if (isNewExpression || (isAssignment(nodeArg) && !isAssignmentComplete(nodeArg))) {
+        // Otherwise try to complete the expression suggesting some columns
+        suggestions.push(
+          ...(await getFieldsOrFunctionsSuggestions(
+            types[0] === 'column' ? ['any'] : types,
+            command.name,
+            option.name,
+            getFieldsByType,
+            {
+              functions: option.name === 'by',
+              fields: true,
+            }
+          ))
+        );
 
-      if (command.name === 'stats' && isNewExpression) {
-        suggestions.push(buildNewVarDefinition(findNewVariable(anyVariables)));
+        if (command.name === 'stats' && isNewExpression) {
+          suggestions.push(buildNewVarDefinition(findNewVariable(anyVariables)));
+        }
       }
     }
   }