From aa68d4177a3393f313c715a67b42b35cebba543e Mon Sep 17 00:00:00 2001
From: Quynh Nguyen <quynh.nguyen@elastic.co>
Date: Mon, 29 Jul 2024 15:34:43 -0500
Subject: [PATCH 1/3] Update push getFinalSuggestions to suggestions for stats
 by

---
 .../src/autocomplete/autocomplete.ts                  | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts b/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
index 45c0563ac833ed..25a15bb5efc908 100644
--- a/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
+++ b/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
@@ -1565,7 +1565,7 @@ async function getOptionArgsSuggestions(
   if (command.name === 'stats') {
     suggestions.push(
       ...(await getFieldsOrFunctionsSuggestions(
-        ['column'],
+        ['any'],
         command.name,
         option.name,
         getFieldsByType,
@@ -1575,7 +1575,6 @@ async function getOptionArgsSuggestions(
         }
       ))
     );
-
     const argDef = optionDef?.signature.params[argIndex];
 
     const nodeArgType = extractFinalTypeFromArg(nodeArg, references);
@@ -1598,6 +1597,14 @@ async function getOptionArgsSuggestions(
         );
       }
     }
+
+    if (option.name === 'by' && !option.incomplete) {
+      suggestions.push(
+        ...getFinalSuggestions({
+          comma: optionDef?.signature.multipleParams,
+        })
+      );
+    }
   }
 
   if (optionDef) {

From 7e09343964938d8eaf2be7c7b2b0a1ca4d8c7d3a Mon Sep 17 00:00:00 2001
From: Quynh Nguyen <quynh.nguyen@elastic.co>
Date: Mon, 29 Jul 2024 22:25:59 -0500
Subject: [PATCH 2/3] Update

---
 .../src/autocomplete/autocomplete.ts          | 73 ++++++++-----------
 .../esql_validation_meta_tests.json           | 15 ----
 2 files changed, 32 insertions(+), 56 deletions(-)

diff --git a/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts b/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
index a5783d33406ae9..5b0637d849096f 100644
--- a/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
+++ b/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
@@ -1592,7 +1592,7 @@ async function getOptionArgsSuggestions(
   if (command.name === 'stats') {
     suggestions.push(
       ...(await getFieldsOrFunctionsSuggestions(
-        ['any'],
+        ['column'],
         command.name,
         option.name,
         getFieldsByType,
@@ -1602,6 +1602,7 @@ async function getOptionArgsSuggestions(
         }
       ))
     );
+
     const argDef = optionDef?.signature.params[argIndex];
 
     const nodeArgType = extractFinalTypeFromArg(nodeArg, references);
@@ -1624,52 +1625,42 @@ async function getOptionArgsSuggestions(
         );
       }
     }
+  }
 
-    if (option.name === 'by' && !option.incomplete) {
+  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 (
+      (!suggestions.length && option.args.length && !isNewExpression && !isAssignment(lastArg)) ||
+      (isAssignment(lastArg) && isAssignmentComplete(lastArg))
+    ) {
       suggestions.push(
         ...getFinalSuggestions({
-          comma: optionDef?.signature.multipleParams,
+          comma: optionDef.signature.multipleParams,
         })
       );
-    }
-  }
-
-  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,
-            }
-          ))
-        );
+    } 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)));
       }
     }
   }
diff --git a/packages/kbn-esql-validation-autocomplete/src/validation/esql_validation_meta_tests.json b/packages/kbn-esql-validation-autocomplete/src/validation/esql_validation_meta_tests.json
index d230f9e2facc9d..b8b91d665663ff 100644
--- a/packages/kbn-esql-validation-autocomplete/src/validation/esql_validation_meta_tests.json
+++ b/packages/kbn-esql-validation-autocomplete/src/validation/esql_validation_meta_tests.json
@@ -2099,21 +2099,6 @@
       ],
       "warning": []
     },
-    {
-      "query": "show functions",
-      "error": [
-        "SyntaxError: token recognition error at: 'f'",
-        "SyntaxError: token recognition error at: 'u'",
-        "SyntaxError: token recognition error at: 'n'",
-        "SyntaxError: token recognition error at: 'c'",
-        "SyntaxError: token recognition error at: 't'",
-        "SyntaxError: token recognition error at: 'io'",
-        "SyntaxError: token recognition error at: 'n'",
-        "SyntaxError: token recognition error at: 's'",
-        "SyntaxError: missing 'info' at '<EOF>'"
-      ],
-      "warning": []
-    },
     {
       "query": "show info",
       "error": [],

From a16a9f305a707a16987cf76af40449c1db7c62e5 Mon Sep 17 00:00:00 2001
From: Quynh Nguyen <quynh.nguyen@elastic.co>
Date: Tue, 30 Jul 2024 11:02:47 -0500
Subject: [PATCH 3/3] Fix autocomplete for by suggestion

---
 .../autocomplete.command.stats.test.ts        |  5 ++
 .../src/autocomplete/autocomplete.ts          | 72 ++++++++++++-------
 2 files changed, 50 insertions(+), 27 deletions(-)

diff --git a/packages/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.stats.test.ts b/packages/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.stats.test.ts
index 7378760781d976..58d46989740899 100644
--- a/packages/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.stats.test.ts
+++ b/packages/kbn-esql-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
index 5b0637d849096f..42e82662e71101 100644
--- 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)));
+        }
       }
     }
   }