From d7978461081fcade206678d16f2a9efce855e241 Mon Sep 17 00:00:00 2001
From: Stratoula Kalafateli <efstratia.kalafateli@elastic.co>
Date: Fri, 29 Sep 2023 08:41:08 +0300
Subject: [PATCH] [ES|QL] Resets selected columns when changing query (#167492)

## Summary

Fixes a bug in ES|QL mode where you:

- have a query of index pattern 1 and select some fields
- change the query by using a different index pattern
- the selected columns do not reset


**BUG**

![withbug](https://github.com/elastic/kibana/assets/17003240/5f3fdf89-5d17-443d-b50b-f802e44482c3)

**NOW**

![now](https://github.com/elastic/kibana/assets/17003240/0c1f8b05-0e22-4708-85e8-9387c9a5d5fa)


### Checklist

- [ ] [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
---
 .../use_test_based_query_language.test.tsx    | 32 +++++++++++++++++++
 .../hooks/use_text_based_query_language.ts    |  2 +-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/src/plugins/discover/public/application/main/hooks/use_test_based_query_language.test.tsx b/src/plugins/discover/public/application/main/hooks/use_test_based_query_language.test.tsx
index a6dcafad11757..3b14bf3b80fcb 100644
--- a/src/plugins/discover/public/application/main/hooks/use_test_based_query_language.test.tsx
+++ b/src/plugins/discover/public/application/main/hooks/use_test_based_query_language.test.tsx
@@ -106,6 +106,7 @@ describe('useTextBasedQueryLanguage', () => {
     await waitFor(() => {
       expect(replaceUrlState).toHaveBeenCalledWith({
         index: 'the-data-view-id',
+        columns: [],
       });
     });
 
@@ -123,6 +124,7 @@ describe('useTextBasedQueryLanguage', () => {
     expect(replaceUrlState).toHaveBeenCalledWith({
       index: 'the-data-view-id',
       viewMode: VIEW_MODE.DOCUMENT_LEVEL,
+      columns: [],
     });
   });
   test('changing a text based query with different result columns should change state when loading and finished', async () => {
@@ -155,6 +157,35 @@ describe('useTextBasedQueryLanguage', () => {
     });
   });
 
+  test('changing a text based 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);
+    await waitFor(() => expect(replaceUrlState).toHaveBeenCalledTimes(1));
+    replaceUrlState.mockReset();
+
+    documents$.next({
+      recordRawType: RecordRawType.PLAIN,
+      fetchStatus: FetchStatus.PARTIAL,
+      result: [
+        {
+          id: '1',
+          raw: { field1: 1 },
+          flattened: { field1: 1 },
+        } as unknown as DataTableRecord,
+      ],
+      query: { esql: 'from the-data-view-2' },
+    });
+    await waitFor(() => expect(replaceUrlState).toHaveBeenCalledTimes(1));
+
+    await waitFor(() => {
+      expect(replaceUrlState).toHaveBeenCalledWith({
+        index: 'the-data-view-id',
+        columns: [],
+      });
+    });
+  });
+
   test('changing a text based query with no transformational commands should only change dataview state when loading and finished', async () => {
     const { replaceUrlState, stateContainer } = renderHookWithContext(false);
     const documents$ = stateContainer.dataState.data$.documents$;
@@ -180,6 +211,7 @@ describe('useTextBasedQueryLanguage', () => {
     await waitFor(() => {
       expect(replaceUrlState).toHaveBeenCalledWith({
         index: 'the-data-view-id',
+        columns: [],
       });
     });
   });
diff --git a/src/plugins/discover/public/application/main/hooks/use_text_based_query_language.ts b/src/plugins/discover/public/application/main/hooks/use_text_based_query_language.ts
index fb3725d12b1e1..01f3491809fb0 100644
--- a/src/plugins/discover/public/application/main/hooks/use_text_based_query_language.ts
+++ b/src/plugins/discover/public/application/main/hooks/use_text_based_query_language.ts
@@ -119,7 +119,7 @@ export function useTextBasedQueryLanguage({
         }
         const nextState = {
           ...(addDataViewToState && { index: dataViewObj.id }),
-          ...(addColumnsToState && { columns: nextColumns }),
+          ...((addColumnsToState || queryChanged) && { columns: nextColumns }),
           ...(viewMode === VIEW_MODE.AGGREGATED_LEVEL && {
             viewMode: getValidViewMode({ viewMode, isTextBasedQueryMode: true }),
           }),