[Discover][ES|QL] Reset selected fields when modifying the ES|QL query

.buildkite/ftr_configs.yml

@@ -115,6 +115,7 @@ enabled:
   - test/functional/apps/discover/ccs_compatibility/config.ts
   - test/functional/apps/discover/classic/config.ts
   - test/functional/apps/discover/embeddable/config.ts
+  - test/functional/apps/discover/esql/config.ts
   - test/functional/apps/discover/group1/config.ts
   - test/functional/apps/discover/group2_data_grid1/config.ts
   - test/functional/apps/discover/group2_data_grid2/config.ts

src/plugins/discover/public/application/main/hooks/use_esql_mode.test.tsx

@@ -150,10 +150,11 @@ describe('useEsqlMode', () => {
     });
   });
 
-  test('changing an ES|QL query with same result columns should not change state when loading and finished', async () => {
+  test('changing an ES|QL query with same result columns should change state when loading and finished', async () => {
     const { replaceUrlState, stateContainer } = renderHookWithContext(false);
     const documents$ = stateContainer.dataState.data$.documents$;
     stateContainer.dataState.data$.documents$.next(msgComplete);
+    replaceUrlState.mockReset();
 
     documents$.next({
       fetchStatus: FetchStatus.PARTIAL,
@@ -166,7 +167,54 @@ describe('useEsqlMode', () => {
       ],
       query: { esql: 'from the-data-view-2' },
     });
+    await waitFor(() => expect(replaceUrlState).toHaveBeenCalledTimes(1));
+
+    await waitFor(() => {
+      expect(replaceUrlState).toHaveBeenCalledWith({
+        columns: [],
+      });
+    });
+  });
+
+  test('changing a text based query with no transformational commands should not change state when loading and finished if index pattern is the same', async () => {
+    const { replaceUrlState, stateContainer } = renderHookWithContext(false);
+    const documents$ = stateContainer.dataState.data$.documents$;
+    stateContainer.dataState.data$.documents$.next(msgComplete);
+    await waitFor(() => expect(replaceUrlState).toHaveBeenCalledTimes(0));
+    replaceUrlState.mockReset();
+
+    documents$.next({
+      fetchStatus: FetchStatus.PARTIAL,
+      result: [
+        {
+          id: '1',
+          raw: { field1: 1 },
+          flattened: { field1: 1 },
+        } as unknown as DataTableRecord,
+      ],
+      // non transformational command
+      query: { esql: 'from the-data-view-title | where field1 > 0' },
+    });
     await waitFor(() => expect(replaceUrlState).toHaveBeenCalledTimes(0));
+    replaceUrlState.mockReset();
+
+    documents$.next({
+      fetchStatus: FetchStatus.PARTIAL,
+      result: [
+        {
+          id: '1',
+          raw: { field1: 1 },
+          flattened: { field1: 1 },
+        } as unknown as DataTableRecord,
+      ],
+      // non transformational command
+      query: { esql: 'from the-data-view-title2 | where field1 > 0' },
+    });
+    await waitFor(() => {
+      expect(replaceUrlState).toHaveBeenCalledWith({
+        columns: [],
+      });
+    });
   });
 
   test('only changing an ES|QL query with same result columns should not change columns', async () => {
@@ -268,7 +316,13 @@ describe('useEsqlMode', () => {
       query: { esql: 'from the-data-view-title | keep field 1 | WHERE field1=1' },
     });
 
-    expect(replaceUrlState).toHaveBeenCalledTimes(0);
+    await waitFor(() => expect(replaceUrlState).toHaveBeenCalledTimes(1));
+    await waitFor(() => {
+      expect(replaceUrlState).toHaveBeenCalledWith({
+        columns: ['field1', 'field2'],
+      });
+    });
+    replaceUrlState.mockReset();
 
     documents$.next({
       fetchStatus: FetchStatus.PARTIAL,

src/plugins/discover/public/application/main/hooks/use_esql_mode.ts

@@ -7,8 +7,8 @@
  */
 
 import { isEqual } from 'lodash';
-import { isOfAggregateQueryType, getAggregateQueryMode } from '@kbn/es-query';
-import { hasTransformationalCommand } from '@kbn/esql-utils';
+import { isOfAggregateQueryType } from '@kbn/es-query';
+import { hasTransformationalCommand, getIndexPatternFromESQLQuery } from '@kbn/esql-utils';
 import { useCallback, useEffect, useRef } from 'react';
 import type { DataViewsContract } from '@kbn/data-views-plugin/public';
 import { switchMap } from 'rxjs';
@@ -31,9 +31,9 @@ export function useEsqlMode({
 }) {
   const prev = useRef<{
     query: string;
-    columns: string[];
+    recentlyUpdatedToColumns: string[];
   }>({
-    columns: [],
+    recentlyUpdatedToColumns: [],
     query: '',
   });
   const initialFetch = useRef<boolean>(true);
@@ -43,9 +43,10 @@ export function useEsqlMode({
     if (prev.current.query) {
       // cleanup when it's not an ES|QL query
       prev.current = {
-        columns: [],
+        recentlyUpdatedToColumns: [],
         query: '',
       };
+      initialFetch.current = true;
     }
   }, []);
 
@@ -57,55 +58,61 @@ export function useEsqlMode({
           if (!query || next.fetchStatus === FetchStatus.ERROR) {
             return;
           }
+
           const sendComplete = () => {
             stateContainer.dataState.data$.documents$.next({
               ...next,
               fetchStatus: FetchStatus.COMPLETE,
             });
           };
+
           const { viewMode } = stateContainer.appState.getState();
-          let nextColumns: string[] = [];
           const isEsqlQuery = isOfAggregateQueryType(query);
-          const hasResults = Boolean(next.result?.length);
-          let queryHasTransformationalCommands = false;
-          if ('esql' in query) {
-            if (hasTransformationalCommand(query.esql)) {
-              queryHasTransformationalCommands = true;
-            }
-          }
 
           if (isEsqlQuery) {
-            const language = getAggregateQueryMode(query);
+            const hasResults = Boolean(next.result?.length);
+
             if (next.fetchStatus !== FetchStatus.PARTIAL) {
               return;
             }
 
+            let nextColumns: string[] = prev.current.recentlyUpdatedToColumns;
+
             if (hasResults) {
-              // check if state needs to contain column transformation due to a different columns in the resultset
+              // check if state needs to contain column transformation due to a different columns in the result set
               const firstRow = next.result![0];
-              const firstRowColumns = Object.keys(firstRow.raw).slice(0, MAX_NUM_OF_COLUMNS);
-              if (!queryHasTransformationalCommands) {
-                nextColumns = [];
-                initialFetch.current = false;
+              const firstRowColumns = Object.keys(firstRow.raw);
+
+              if (hasTransformationalCommand(query.esql)) {
+                nextColumns = firstRowColumns.slice(0, MAX_NUM_OF_COLUMNS);
               } else {
-                nextColumns = firstRowColumns;
-                if (initialFetch.current && !prev.current.columns.length) {
-                  prev.current.columns = firstRowColumns;
-                }
+                nextColumns = [];
               }
             }
-            const addColumnsToState = !isEqual(nextColumns, prev.current.columns);
-            const queryChanged = query[language] !== prev.current.query;
+
+            if (initialFetch.current) {
+              initialFetch.current = false;
+              prev.current.query = query.esql;
+              prev.current.recentlyUpdatedToColumns = nextColumns;
+            }
+
+            const indexPatternChanged =
+              getIndexPatternFromESQLQuery(query.esql) !==
+              getIndexPatternFromESQLQuery(prev.current.query);
+
+            const addColumnsToState =
+              indexPatternChanged || !isEqual(nextColumns, prev.current.recentlyUpdatedToColumns);
+
             const changeViewMode = viewMode !== getValidViewMode({ viewMode, isEsqlMode: true });
-            if (!queryChanged || (!addColumnsToState && !changeViewMode)) {
+
+            if (!indexPatternChanged && !addColumnsToState && !changeViewMode) {
               sendComplete();
               return;
             }
 
-            if (queryChanged) {
-              prev.current.query = query[language];
-              prev.current.columns = nextColumns;
-            }
+            prev.current.query = query.esql;
+            prev.current.recentlyUpdatedToColumns = nextColumns;
+
             // just change URL state if necessary
             if (addColumnsToState || changeViewMode) {
               const nextState = {