[8.x] [Discover] Fix issue where KEEP columns are not applied after Elasticsearch error (#205833)

src/platform/plugins/shared/discover/public/application/main/data_fetching/fetch_all.ts

@@ -183,7 +183,10 @@ export function fetchAll(
       // Only the document query should send its errors to main$, to cause the full Discover app
       // to get into an error state. The other queries will not cause all of Discover to error out
       // but their errors will be shown in-place (e.g. of the chart).
-      .catch(sendErrorTo(dataSubjects.documents$, dataSubjects.main$));
+      .catch((e) => {
+        sendErrorMsg(dataSubjects.documents$, e, { query });
+        sendErrorMsg(dataSubjects.main$, e);
+      });
 
     // Return a promise that will resolve once all the requests have finished or failed, or no results are found
     return firstValueFrom(

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

@@ -64,7 +64,7 @@ export function useEsqlMode({
         switchMap(async (next) => {
           const { query: nextQuery } = next;
 
-          if (!nextQuery || next.fetchStatus === FetchStatus.ERROR) {
+          if (!nextQuery) {
             return;
           }
 
@@ -104,6 +104,12 @@ export function useEsqlMode({
             return;
           }
 
+          if (next.fetchStatus === FetchStatus.ERROR) {
+            // An error occurred, but it's still considered an initial fetch
+            prev.current.initialFetch = false;
+            return;
+          }
+
           if (next.fetchStatus !== FetchStatus.PARTIAL) {
             return;
           }

src/platform/plugins/shared/discover/public/application/main/hooks/use_saved_search_messages.ts

@@ -99,8 +99,12 @@ export function sendLoadingMoreFinishedMsg(
 /**
  * Send ERROR message
  */
-export function sendErrorMsg(data$: DataMain$ | DataDocuments$ | DataTotalHits$, error?: Error) {
-  data$.next({ fetchStatus: FetchStatus.ERROR, error });
+export function sendErrorMsg<T extends DataMsg>(
+  data$: DataMain$ | DataDocuments$ | DataTotalHits$,
+  error?: Error,
+  props?: Omit<T, 'fetchStatus' | 'error'>
+) {
+  data$.next({ fetchStatus: FetchStatus.ERROR, error, ...props });
 }
 
 /**

test/functional/apps/discover/esql/_esql_columns.ts

@@ -217,6 +217,23 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
       expect(await dataGrid.getHeaderFields()).to.eql(['ip', '@timestamp', 'bytes']);
     });
 
+    it('should recover from an error and reset columns correctly when a transformational query is used', async () => {
+      await monacoEditor.setCodeEditorValue('from not_an_index');
+      await testSubjects.click('querySubmitButton');
+      await header.waitUntilLoadingHasFinished();
+      await discover.showsErrorCallout();
+      await browser.refresh();
+      await header.waitUntilLoadingHasFinished();
+      await discover.showsErrorCallout();
+      await monacoEditor.setCodeEditorValue(
+        'from logstash-* | keep ip, @timestamp, bytes | limit 10'
+      );
+      await testSubjects.click('querySubmitButton');
+      await header.waitUntilLoadingHasFinished();
+      await discover.waitUntilSearchingHasFinished();
+      expect(await dataGrid.getHeaderFields()).to.eql(['ip', '@timestamp', 'bytes']);
+    });
+
     it('should restore columns correctly when switching between saved searches', async () => {
       await discover.loadSavedSearch(SAVED_SEARCH_NON_TRANSFORMATIONAL_INITIAL_COLUMNS);
       await header.waitUntilLoadingHasFinished();