restore function suggestions

elastic · Sep 20, 2024 · 40c2e83 · 40c2e83
1 parent a94929a
commit 40c2e83
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 23 deletions.
diff --git a/...esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.sort.test.ts b/...esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.sort.test.ts
@@ -7,22 +7,32 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { setup, getFieldNamesByType, attachTriggerCommand } from './helpers';
+import {
+  setup,
+  getFieldNamesByType,
+  attachTriggerCommand,
+  getFunctionSignaturesByReturnType,
+} from './helpers';
 
 describe('autocomplete.suggest', () => {
   describe('SORT ( <column> [ ASC / DESC ] [ NULLS FIST / NULLS LAST ] )+', () => {
     describe('SORT <column> ...', () => {
+      const expectedFieldSuggestions = getFieldNamesByType('any').map(attachTriggerCommand);
+      const expectedFunctionSuggestions = getFunctionSignaturesByReturnType('sort', 'any', {
+        scalar: true,
+      }).map(attachTriggerCommand);
+
       test('suggests column', async () => {
         const { assertSuggestions } = await setup();
 
-        await assertSuggestions(
-          'from a | sort /',
-          [...getFieldNamesByType('any')].map(attachTriggerCommand)
-        );
-        await assertSuggestions(
-          'from a | sort keyw/',
-          [...getFieldNamesByType('any')].map(attachTriggerCommand)
-        );
+        await assertSuggestions('from a | sort /', [
+          ...expectedFieldSuggestions,
+          ...expectedFunctionSuggestions,
+        ]);
+        await assertSuggestions('from a | sort keyw/', [
+          ...expectedFieldSuggestions,
+          ...expectedFunctionSuggestions,
+        ]);
         await assertSuggestions(
           'from a | sort keywordField/',
           [
@@ -56,14 +66,14 @@ describe('autocomplete.suggest', () => {
       it('suggests subsequent column after comma', async () => {
         const { assertSuggestions } = await setup();
 
-        await assertSuggestions(
-          'from a | sort keywordField, /',
-          [...getFieldNamesByType('any')].map(attachTriggerCommand)
-        );
-        await assertSuggestions(
-          'from a | sort keywordField, doubl/',
-          [...getFieldNamesByType('any')].map(attachTriggerCommand)
-        );
+        await assertSuggestions('from a | sort keywordField, /', [
+          ...expectedFieldSuggestions,
+          ...expectedFunctionSuggestions,
+        ]);
+        await assertSuggestions('from a | sort keywordField, doubl/', [
+          ...expectedFieldSuggestions,
+          ...expectedFunctionSuggestions,
+        ]);
         await assertSuggestions(
           'from a | sort keywordField, doubleField/',
           [

diff --git a/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts b/packages/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
@@ -559,7 +559,6 @@ async function getExpressionSuggestionsByType(
   const canHaveAssignments = ['eval', 'stats', 'row'].includes(command.name);
 
   const references = { fields: fieldsMap, variables: anyVariables };
-
   if (command.name === 'sort') {
     return await suggestForSortCmd(innerText, getFieldsByType, (col) =>
       Boolean(getColumnByName(col, references))
@@ -2020,17 +2019,33 @@ export const suggestForSortCmd = async (
   const fieldSuggestions = await getFieldsByType('any', [], {
     openSuggestions: true,
   });
+  const functionSuggestions = await getFieldsOrFunctionsSuggestions(
+    ['any'],
+    'sort',
+    undefined,
+    getFieldsByType,
+    {
+      functions: true,
+      fields: false,
+    }
+  );
 
   return await handleFragment(
     innerText,
     columnExists,
     (_fragment: string, rangeToReplace?: { start: number; end: number }) => {
       // SORT fie<suggest>
-      return fieldSuggestions.map((suggestion) => ({
-        ...suggestion,
-        command: TRIGGER_SUGGESTION_COMMAND,
-        rangeToReplace,
-      }));
+      return [
+        ...pushItUpInTheList(
+          fieldSuggestions.map((suggestion) => ({
+            ...suggestion,
+            command: TRIGGER_SUGGESTION_COMMAND,
+            rangeToReplace,
+          })),
+          true
+        ),
+        ...functionSuggestions,
+      ];
     },
     (fragment: string, rangeToReplace: { start: number; end: number }) => {
       // SORT field<suggest>