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

## Summary

This PR fixes an issue where columns are not applied correctly when
using the ES|QL `KEEP` command after an Elasticsearch error has
occurred.

Fixes #205353.

### Checklist

- [ ] Any text added follows [EUI's writing
guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses
sentence case text and includes [i18n
support](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md)
- [ ]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [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
- [ ] If a plugin configuration key changed, check if it needs to be
allowlisted in the cloud and added to the [docker
list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker)
- [ ] This was checked for breaking HTTP API changes, and any breaking
changes have been approved by the breaking-change committee. The
`release_note:breaking` label should be applied in these situations.
- [x] [Flaky Test
Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was
used on any tests changed
- [x] The PR description includes the appropriate Release Notes section,
and the correct `release_note:*` label is applied per the
[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)

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

@@ -182,7 +182,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();

test/functional/services/monaco_editor.ts

@@ -62,8 +62,6 @@ export class MonacoEditorService extends FtrService {
         nthIndex,
         value
       );
-    });
-    await this.retry.try(async () => {
       const newCodeEditorValue = await this.getCodeEditorValue(nthIndex);
       expect(newCodeEditorValue).equal(
         value,