[ES|QL] detect the type of COUNT(*) (#197914)

## Summary

We weren't properly detecting the type of the expression `COUNT(*)`. Now
we are!

Before:
<img width="950" alt="Screenshot 2024-10-25 at 4 38 08 PM"
src="https://github.com/user-attachments/assets/e9bd8d78-d0c8-4069-a818-5bf3486b925b">

After:
<img width="1093" alt="Screenshot 2024-10-25 at 4 35 44 PM"
src="https://github.com/user-attachments/assets/235c63dc-7d6c-49df-9adf-e225c4550a42">


### Checklist

- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

Co-authored-by: Stratoula Kalafateli <[email protected]>

packages/kbn-esql-validation-autocomplete/src/shared/helpers.test.ts

@@ -225,79 +225,47 @@ describe('getExpressionType', () => {
     });
 
     it('detects the return type of a function', () => {
-      expect(
-        getExpressionType(getASTForExpression('returns_keyword()'), new Map(), new Map())
-      ).toBe('keyword');
+      expect(getExpressionType(getASTForExpression('returns_keyword()'))).toBe('keyword');
     });
 
     it('selects the correct signature based on the arguments', () => {
-      expect(getExpressionType(getASTForExpression('test("foo")'), new Map(), new Map())).toBe(
-        'keyword'
-      );
-      expect(getExpressionType(getASTForExpression('test(1.)'), new Map(), new Map())).toBe(
-        'double'
-      );
-      expect(getExpressionType(getASTForExpression('test(1., "foo")'), new Map(), new Map())).toBe(
-        'long'
-      );
+      expect(getExpressionType(getASTForExpression('test("foo")'))).toBe('keyword');
+      expect(getExpressionType(getASTForExpression('test(1.)'))).toBe('double');
+      expect(getExpressionType(getASTForExpression('test(1., "foo")'))).toBe('long');
     });
 
     it('supports nested functions', () => {
       expect(
-        getExpressionType(
-          getASTForExpression('test(1., test(test(test(returns_keyword()))))'),
-          new Map(),
-          new Map()
-        )
+        getExpressionType(getASTForExpression('test(1., test(test(test(returns_keyword()))))'))
       ).toBe('long');
     });
 
     it('supports functions with casted results', () => {
-      expect(
-        getExpressionType(getASTForExpression('test(1.)::keyword'), new Map(), new Map())
-      ).toBe('keyword');
+      expect(getExpressionType(getASTForExpression('test(1.)::keyword'))).toBe('keyword');
     });
 
     it('handles nulls and string-date casting', () => {
-      expect(getExpressionType(getASTForExpression('test(NULL)'), new Map(), new Map())).toBe(
-        'null'
-      );
-      expect(getExpressionType(getASTForExpression('test(NULL, NULL)'), new Map(), new Map())).toBe(
-        'null'
-      );
-      expect(
-        getExpressionType(getASTForExpression('accepts_dates("", "")'), new Map(), new Map())
-      ).toBe('keyword');
+      expect(getExpressionType(getASTForExpression('test(NULL)'))).toBe('null');
+      expect(getExpressionType(getASTForExpression('test(NULL, NULL)'))).toBe('null');
+      expect(getExpressionType(getASTForExpression('accepts_dates("", "")'))).toBe('keyword');
     });
 
     it('deals with functions that do not exist', () => {
-      expect(getExpressionType(getASTForExpression('does_not_exist()'), new Map(), new Map())).toBe(
-        'unknown'
-      );
+      expect(getExpressionType(getASTForExpression('does_not_exist()'))).toBe('unknown');
     });
 
     it('deals with bad function invocations', () => {
-      expect(
-        getExpressionType(getASTForExpression('test(1., "foo", "bar")'), new Map(), new Map())
-      ).toBe('unknown');
+      expect(getExpressionType(getASTForExpression('test(1., "foo", "bar")'))).toBe('unknown');
 
-      expect(getExpressionType(getASTForExpression('test()'), new Map(), new Map())).toBe(
-        'unknown'
-      );
+      expect(getExpressionType(getASTForExpression('test()'))).toBe('unknown');
 
-      expect(getExpressionType(getASTForExpression('test("foo", 1.)'), new Map(), new Map())).toBe(
-        'unknown'
-      );
+      expect(getExpressionType(getASTForExpression('test("foo", 1.)'))).toBe('unknown');
     });
 
     it('deals with the CASE function', () => {
-      expect(getExpressionType(getASTForExpression('CASE(true, 1, 2)'), new Map(), new Map())).toBe(
-        'integer'
-      );
+      expect(getExpressionType(getASTForExpression('CASE(true, 1, 2)'))).toBe('integer');
 
-      expect(
-        getExpressionType(getASTForExpression('CASE(true, 1., true, 1., 2.)'), new Map(), new Map())
-      ).toBe('double');
+      expect(getExpressionType(getASTForExpression('CASE(true, 1., true, 1., 2.)'))).toBe('double');
 
       expect(
         getExpressionType(
@@ -306,6 +274,20 @@ describe('getExpressionType', () => {
           new Map()
         )
       ).toBe('keyword');
+
+      expect(
+        getExpressionType(
+          getASTForExpression('CASE(true, "", true, "", keywordVar)'),
+          new Map(),
+          new Map([
+            [`keywordVar`, [{ name: 'keywordVar', type: 'keyword', location: { min: 0, max: 0 } }]],
+          ])
+        )
+      ).toBe('keyword');
+    });
+
+    it('supports COUNT(*)', () => {
+      expect(getExpressionType(getASTForExpression('COUNT(*)'))).toBe<SupportedDataType>('long');
     });
   });
 

packages/kbn-esql-validation-autocomplete/src/shared/helpers.ts

@@ -815,15 +815,31 @@ export function getExpressionType(
       return 'unknown';
     }
 
+    /**
+     * Special case for COUNT(*) because
+     * the "*" column doesn't match any
+     * of COUNT's function definitions
+     */
+    if (
+      fnDefinition.name === 'count' &&
+      root.args[0] &&
+      isColumnItem(root.args[0]) &&
+      root.args[0].name === '*'
+    ) {
+      return 'long';
+    }
+
     if (fnDefinition.name === 'case' && root.args.length) {
-      // The CASE function doesn't fit our system of function definitions
-      // and needs special handling. This is imperfect, but it's a start because
-      // at least we know that the final argument to case will never be a conditional
-      // expression, always a result expression.
-      //
-      // One problem with this is that if a false case is not provided, the return type
-      // will be null, which we aren't detecting. But this is ok because we consider
-      // variables and fields to be nullable anyways and account for that during validation.
+      /**
+       * The CASE function doesn't fit our system of function definitions
+       * and needs special handling. This is imperfect, but it's a start because
+       * at least we know that the final argument to case will never be a conditional
+       * expression, always a result expression.
+       *
+       * One problem with this is that if a false case is not provided, the return type
+       * will be null, which we aren't detecting. But this is ok because we consider
+       * variables and fields to be nullable anyways and account for that during validation.
+       */
       return getExpressionType(root.args[root.args.length - 1], fields, variables);
     }