[Discover] [ES|QL] fix: stats by previous aggregated field

src/platform/packages/shared/kbn-esql-utils/src/utils/cascaded_documents_helpers/index.test.ts

@@ -7,6 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
+import { EsqlQuery } from '@elastic/esql';
 import type { AggregateQuery } from '@kbn/es-query';
 import { type ESQLControlVariable, ESQLVariableType } from '@kbn/esql-types';
 import type { DataViewField } from '@kbn/data-views-plugin/common';
@@ -149,6 +150,39 @@ describe('cascaded documents helpers utils', () => {
       ]);
     });
 
+    it('should return metadata when a group field references a field that was declared as an aggregate by a preceding command', () => {
+      const queryString = `
+        FROM kibana_sample_data_logs
+        | STATS x = MAX(bytes)
+        | STATS c = COUNT(*) BY x
+      `;
+
+      const result = getESQLStatsQueryMeta(queryString);
+
+      expect(result.groupByFields).toEqual([{ field: 'x', type: 'column' }]);
+      expect(result.appliedFunctions).toEqual([{ identifier: 'c', aggregation: 'COUNT' }]);
+    });
+
+    it('should return empty metadata instead of throwing when query metadata computation fails unexpectedly', () => {
+      const fromSrcSpy = jest.spyOn(EsqlQuery, 'fromSrc').mockImplementation(() => {
+        throw new Error('unexpected parse failure');
+      });
+      const consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation(() => {});
+
+      try {
+        expect(
+          getESQLStatsQueryMeta('FROM kibana_sample_data_logs | STATS COUNT(*) BY host')
+        ).toEqual({
+          groupByFields: [],
+          appliedFunctions: [],
+        });
+        expect(consoleErrorSpy).toHaveBeenCalled();
+      } finally {
+        fromSrcSpy.mockRestore();
+        consoleErrorSpy.mockRestore();
+      }
+    });
+
     it('should return a single group by field if there is a where command following a STATS by command targeting a column specified as a grouping option in the operating stats command', () => {
       const queryString = `
      FROM kibana_sample_data_logs
@@ -367,6 +401,33 @@ describe('cascaded documents helpers utils', () => {
           );
         });
 
+        it('should construct a cascade leaf query when a later STATS groups by a prior aggregate alias', () => {
+          const editorQuery: AggregateQuery = {
+            esql: `
+              FROM kibana_sample_data_logs
+              | STATS x = MAX(bytes)
+              | STATS c = COUNT(*) BY x
+            `,
+          };
+
+          const nodePath = ['x'];
+          const nodePathMap = { x: '33' };
+
+          const cascadeQuery = constructCascadeQuery({
+            query: editorQuery,
+            dataView: dataViewMock,
+            esqlVariables: [],
+            nodeType,
+            nodePath,
+            nodePathMap,
+          });
+
+          expect(cascadeQuery).toBeDefined();
+          expect(cascadeQuery!.esql).toBe(
+            'FROM kibana_sample_data_logs | INLINE STATS x = MAX(bytes) | INLINE STATS c = COUNT(*) BY x | WHERE x == 33'
+          );
+        });
+
         it('generate a valid cascade leaf query for a valid stats command that has a parameter value for a grouping option', () => {
           const editorQuery: AggregateQuery = {
             esql: `
@@ -916,6 +977,22 @@ describe('cascaded documents helpers utils', () => {
           'FROM kibana_sample_data_logs | WHERE `agent.keyword` == "Mozilla/4.0 (compatible; MSIE 6.0; Windows NT 5.1; SV1; .NET CLR 1.1.4322)" | STATS count = COUNT(*) BY agent.keyword, extension.keyword | STATS avg = AVG(count) BY agent.keyword'
         );
       });
+
+      it('handles filtering on a group field that was declared as an aggregate by a preceding STATS', () => {
+        expect(
+          appendFilteringWhereClauseForCascadeLayout(
+            'FROM kibana_sample_data_logs | STATS x = MAX(bytes) | STATS c = COUNT(*) BY x',
+            [],
+            dataViewMock,
+            'x',
+            33,
+            '+',
+            'integer'
+          )
+        ).toBe(
+          'FROM kibana_sample_data_logs | STATS x = MAX(bytes) | WHERE x == 33 | STATS c = COUNT(*) BY x'
+        );
+      });
     });
 
     describe('handling for param fields', () => {

src/platform/packages/shared/kbn-esql-utils/src/utils/cascaded_documents_helpers/index.ts

@@ -88,12 +88,33 @@ export interface ESQLStatsQueryMeta {
   appliedFunctions: AppliedStatsFunction[];
 }
 
+const EMPTY_ESQL_STATS_QUERY_META: ESQLStatsQueryMeta = {
+  groupByFields: [],
+  appliedFunctions: [],
+};
+
+const getStatsGroupingField = (
+  summary: StatsCommandSummary,
+  field: string
+): FieldSummary | undefined => summary.grouping[field] ?? summary.grouping[`\`${field}\``];
+
 /**
  * This method is used to get the metadata on STATS command to drive the cascade experience from an ESQL query,
  * if a valid STATS command is found information about the group by fields and applied functions is returned.
  * This method will exclude queries contain commands that are not valid for the cascade experience,
  */
 export const getESQLStatsQueryMeta = (queryString: string): ESQLStatsQueryMeta => {
+  try {
+    return computeESQLStatsQueryMeta(queryString);
+  } catch (error) {
+    // Unexpected AST/parse failures must not take down Discover (or other callers).
+    // eslint-disable-next-line no-console
+    console.error('Failed to compute ES|QL stats query metadata for cascade documents', error);
+    return EMPTY_ESQL_STATS_QUERY_META;
+  }
+};
+
+const computeESQLStatsQueryMeta = (queryString: string): ESQLStatsQueryMeta => {
   const groupByFields: ESQLStatsQueryMeta['groupByFields'] = [];
   const appliedFunctions: ESQLStatsQueryMeta['appliedFunctions'] = [];
 
@@ -163,9 +184,17 @@ export const getESQLStatsQueryMeta = (queryString: string): ESQLStatsQueryMeta =
           groupDeclarationStatsCommandLookupIndex
         ))
       ) {
-        groupDeclarationStatsCommandIndex = groupDeclarationStatsCommandLookupIndex;
-        // update the group field node to it's actual definition
-        groupFieldNode = groupDeclarationCommandSummary.grouping[group.field];
+        const resolvedGroupField = getStatsGroupingField(
+          groupDeclarationCommandSummary,
+          group.field
+        );
+        if (resolvedGroupField) {
+          groupDeclarationStatsCommandIndex = groupDeclarationStatsCommandLookupIndex;
+          // update the group field node to its actual grouping definition
+          groupFieldNode = resolvedGroupField;
+        }
+        // If the preceding STATS created this field as an aggregate, keep the current
+        // STATS grouping node (a column), the same way EVAL-created fields are handled.
       }
     }
 
@@ -350,12 +379,17 @@ export const constructCascadeQuery = ({
         );
       }
 
-      fieldDeclarationCommandSummary = groupDeclarationCommandSummary
-        ? {
-            ...groupDeclarationCommandSummary,
-            index: groupDeclarationCommandIndex,
-          }
-        : fieldDeclarationCommandSummary;
+      // Only walk back to a preceding STATS when that command declared the field as a grouping
+      // key (e.g. CATEGORIZE). Aggregate aliases are just columns on the operating STATS BY clause.
+      if (
+        groupDeclarationCommandSummary &&
+        getStatsGroupingField(groupDeclarationCommandSummary, pathSegment)
+      ) {
+        fieldDeclarationCommandSummary = {
+          ...groupDeclarationCommandSummary,
+          index: groupDeclarationCommandIndex,
+        };
+      }
     }
 
     const groupValue =
@@ -708,12 +742,20 @@ export const appendFilteringWhereClauseForCascadeLayout = <
         fieldDeclarationCommandSummary.index !== groupDeclarationCommandIndex
       ) {
         filterTargetIsRuntimeField = true;
-        // update the field declaration command summary to the stats command
-        // that declared the field the filtering operation is targeting
-        fieldDeclarationCommandSummary = {
-          ...getStatsCommandAtIndexSummary(ESQLQuery, groupDeclarationCommandIndex)!,
-          index: groupDeclarationCommandIndex,
-        };
+        const declaredSummary = getStatsCommandAtIndexSummary(
+          ESQLQuery,
+          groupDeclarationCommandIndex
+        );
+        const declaredGroupingField =
+          declaredSummary && getStatsGroupingField(declaredSummary, rawFieldName);
+        if (declaredGroupingField && declaredSummary) {
+          // update the field declaration command summary to the stats command
+          // that declared the field the filtering operation is targeting
+          fieldDeclarationCommandSummary = {
+            ...declaredSummary,
+            index: groupDeclarationCommandIndex,
+          };
+        }
       }
     }