[ES|QL] support expressions in STATS

src/platform/packages/shared/kbn-esql-ast/src/__tests__/autocomplete.ts

@@ -57,7 +57,8 @@ export const suggest = (
   ) => Promise<ISuggestionItem[]>,
   offset?: number
 ): Promise<ISuggestionItem[]> => {
-  const correctedQuery = correctQuerySyntax(query);
+  const innerText = query.substring(0, offset ?? query.length);
+  const correctedQuery = correctQuerySyntax(innerText);
   const { ast } = parse(correctedQuery, { withFormatting: true });
   const cursorPosition = offset ?? query.length;
   const { command } = findAstPosition(ast, cursorPosition);

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/autocomplete.ts

@@ -65,7 +65,7 @@ export async function autocomplete(
         expressionRoot: undefined,
         location: Location.SORT,
         context,
-        advanceCursorAfterInitialField: false,
+        advanceCursorAfterInitialColumn: false,
         hasMinimumLicenseRequired: callbacks?.hasMinimumLicenseRequired,
       });
     }

src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/autocomplete.test.ts

@@ -116,6 +116,8 @@ describe('STATS Autocomplete', () => {
     return autocomplete(query, command, mockCallbacks, mockContext, cursorPosition);
   };
   describe('STATS ...', () => {
+    afterEach(() => setTestFunctions([]));
+
     describe('... <aggregates> ...', () => {
       test('suggestions for a fresh expression', async () => {
         const expected = EXPECTED_FOR_EMPTY_EXPRESSION;
@@ -138,7 +140,44 @@ describe('STATS Autocomplete', () => {
       });
 
       test('on space after aggregate field', async () => {
-        await statsExpectSuggestions('from a | stats a=min(b) ', ['WHERE ', 'BY ', ', ', '| ']);
+        await statsExpectSuggestions('from a | stats a=min(integerField) ', [
+          'WHERE ',
+          'BY ',
+          ', ',
+          '| ',
+          ...getFunctionSignaturesByReturnType(
+            Location.STATS,
+            'any',
+            { operators: true, skipAssign: true },
+            ['integer']
+          ),
+        ]);
+
+        await statsExpectSuggestions('FROM index1 | STATS AVG(doubleField) WHE', [
+          'WHERE ',
+          'BY ',
+          ', ',
+          '| ',
+          ...getFunctionSignaturesByReturnType(
+            Location.STATS,
+            'any',
+            { operators: true, skipAssign: true },
+            ['double']
+          ),
+        ]);
+
+        await statsExpectSuggestions('FROM index1 | STATS AVG(doubleField) B', [
+          'WHERE ',
+          'BY ',
+          ', ',
+          '| ',
+          ...getFunctionSignaturesByReturnType(
+            Location.STATS,
+            'any',
+            { operators: true, skipAssign: true },
+            ['double']
+          ),
+        ]);
       });
 
       test('on function left paren', async () => {
@@ -356,11 +395,22 @@ describe('STATS Autocomplete', () => {
 
       test('when typing right paren', async () => {
         await statsExpectSuggestions(
-          'from a | stats a = min(b) | sort b',
-          ['WHERE ', 'BY ', ', ', '| '],
+          'from a | stats a = min(integerField) | sort b',
+          [
+            'WHERE ',
+            'BY ',
+            ', ',
+            '| ',
+            ...getFunctionSignaturesByReturnType(
+              Location.STATS,
+              'any',
+              { operators: true, skipAssign: true },
+              ['integer']
+            ),
+          ],
           mockCallbacks,
           mockContext,
-          25
+          36
         );
       });
 
@@ -380,6 +430,28 @@ describe('STATS Autocomplete', () => {
         ]);
       });
 
+      test('expressions with aggregates', async () => {
+        await statsExpectSuggestions('from a | stats col0 = min(integerField) ', [
+          'BY ',
+          'WHERE ',
+          '| ',
+          ', ',
+          ...getFunctionSignaturesByReturnType(
+            Location.STATS,
+            'any',
+            { operators: true, skipAssign: true },
+            ['integer']
+          ),
+        ]);
+        await statsExpectSuggestions('from a | stats col0 = min(integerField) + ', [
+          ...getFunctionSignaturesByReturnType(Location.STATS, ['integer', 'double', 'long'], {
+            scalar: true,
+            agg: true,
+            grouping: true,
+          }),
+        ]);
+      });
+
       describe('...WHERE expression...', () => {
         it('suggests fields and functions in empty expression', async () => {
           await statsExpectSuggestions('FROM a | STATS MIN(b) WHERE ', [
@@ -473,13 +545,13 @@ describe('STATS Autocomplete', () => {
 
       test('on partial column name', async () => {
         const expected = [
-          ' = ',
-          getDateHistogramCompletionItem().text,
           ...allEvalFunctionsForStats,
           ...allGroupingFunctions,
+          getDateHistogramCompletionItem().text,
+          ' = ',
         ];
 
-        await statsExpectSuggestions('from a | stats a=max(b) BY keywor/', [
+        await statsExpectSuggestions('from a | stats a=max(b) BY keywor', [
           ...expected,
           ...getFieldNamesByType('any'),
         ]);
@@ -492,16 +564,24 @@ describe('STATS Autocomplete', () => {
 
       test('on complete column name', async () => {
         await statsExpectSuggestions('from a | stats a=max(b) by integerField', [
-          ' = ',
+          'integerField | ',
+          'integerField, ',
+        ]);
+
+        await statsExpectSuggestions('from a | stats a=max(b) by col0 = integerField', [
           'integerField | ',
           'integerField, ',
         ]);
 
         await statsExpectSuggestions('from a | stats a=max(b) by keywordField, integerField', [
-          ' = ',
           'integerField | ',
           'integerField, ',
         ]);
+
+        await statsExpectSuggestions(
+          'from a | stats a=max(b) by keywordField, col0 = integerField',
+          ['integerField | ', 'integerField, ']
+        );
       });
 
       test('attaches field range', async () => {
@@ -514,7 +594,35 @@ describe('STATS Autocomplete', () => {
       });
 
       test('on space after grouping field', async () => {
-        await statsExpectSuggestions('from a | stats a=c by d ', [', ', '| ']);
+        await statsExpectSuggestions('from a | stats a=c by keywordField ', [
+          ', ',
+          '| ',
+          ...getFunctionSignaturesByReturnType(
+            Location.STATS_BY,
+            'any',
+            {
+              operators: true,
+              skipAssign: true,
+            },
+            ['keyword']
+          ),
+        ]);
+      });
+
+      test('on space after grouping function', async () => {
+        await statsExpectSuggestions('from a | stats a=c by CATEGORIZE(keywordField) ', [
+          ', ',
+          '| ',
+          ...getFunctionSignaturesByReturnType(
+            Location.STATS_BY,
+            'any',
+            {
+              operators: true,
+              skipAssign: true,
+            },
+            ['keyword']
+          ),
+        ]);
       });
 
       test('after comma "," in grouping fields', async () => {
@@ -587,15 +695,46 @@ describe('STATS Autocomplete', () => {
       });
 
       test('on space after expression right hand side operand', async () => {
-        await statsExpectSuggestions('from a | stats avg(b) by doubleField % 2 ', [', ', '| ']);
+        await statsExpectSuggestions('from a | stats avg(b) by doubleField % 2 ', [
+          ', ',
+          '| ',
+          ...getFunctionSignaturesByReturnType(
+            Location.STATS_BY,
+            'any',
+            {
+              operators: true,
+              skipAssign: true,
+            },
+            ['double']
+          ),
+        ]);
 
         await statsExpectSuggestions(
           'from a | stats col0 = AVG(doubleField) BY col1 = BUCKET(dateField, 1 day) ',
-          [', ', '| ']
+          [
+            ', ',
+            '| ',
+            ...getFunctionSignaturesByReturnType(
+              Location.STATS_BY,
+              'any',
+              {
+                operators: true,
+                skipAssign: true,
+              },
+              ['date']
+            ),
+          ]
+        );
+      });
+
+      test('after NOT keyword', async () => {
+        await statsExpectSuggestions(
+          'FROM logs-apache_error | STATS count() by keywordField <= textField NOT ',
+          ['LIKE $0', 'RLIKE $0', 'IN $0']
         );
       });
 
-      test('on space within bucket()', async () => {
+      test('within bucket()', async () => {
         const expectedFields = getFieldNamesByType([
           'date',
           'date_nanos',
@@ -605,7 +744,7 @@ describe('STATS Autocomplete', () => {
           expectedFields.map((name) => ({ label: name, text: name }))
         );
         await statsExpectSuggestions(
-          'from a | stats avg(b) by BUCKET(, 50, ?_tstart, ?_tend)',
+          'FROM a | STATS COUNT() BY BUCKET(, 50, ?_tstart, ?_tend)',
           [
             // Note there's no space or comma in the suggested field names
             ...getFieldNamesByType(['date', 'date_nanos', ...ESQL_COMMON_NUMERIC_TYPES]),
@@ -619,10 +758,10 @@ describe('STATS Autocomplete', () => {
           ],
           mockCallbacks,
           mockContext,
-          32
+          33
         );
         await statsExpectSuggestions(
-          'from a | stats avg(b) by BUCKET( , 50, ?_tstart, ?_tend)',
+          'FROM a | STATS COUNT() BY BUCKET( , 50, ?_tstart, ?_tend)',
           [
             // Note there's no space or comma in the suggested field names
             ...getFieldNamesByType(['date', 'date_nanos', ...ESQL_COMMON_NUMERIC_TYPES]),
@@ -636,7 +775,7 @@ describe('STATS Autocomplete', () => {
           ],
           mockCallbacks,
           mockContext,
-          32
+          33
         );
         const expectedFields1 = getFieldNamesByType([
           'integer',
@@ -661,7 +800,7 @@ describe('STATS Autocomplete', () => {
           ],
           mockCallbacks,
           mockContext,
-          43
+          43 // at the second argument of the bucket function
         );
       });