Update the query correctly

stratoula · stratoula · commit f2d5cdb91f84 · 2025-02-12T14:31:42.000+01:00
diff --git a/x-pack/platform/plugins/shared/lens/public/app_plugin/shared/edit_on_the_fly/helpers.ts b/x-pack/platform/plugins/shared/lens/public/app_plugin/shared/edit_on_the_fly/helpers.ts
@@ -11,10 +11,15 @@ import {
   formatESQLColumns,
   mapVariableToColumn,
 } from '@kbn/esql-utils';
+import { isEqual, cloneDeep } from 'lodash';
 import { type AggregateQuery, buildEsQuery } from '@kbn/es-query';
 import type { ESQLControlVariable } from '@kbn/esql-validation-autocomplete';
 import type { ESQLRow } from '@kbn/es-types';
-import { getLensAttributesFromSuggestion, mapVisToChartType } from '@kbn/visualization-utils';
+import {
+  getLensAttributesFromSuggestion,
+  mapVisToChartType,
+  getDatasourceId,
+} from '@kbn/visualization-utils';
 import type { DataViewSpec } from '@kbn/data-views-plugin/public';
 import type { DataView } from '@kbn/data-views-plugin/common';
 import type { DatatableColumn } from '@kbn/expressions-plugin/common';
@@ -136,7 +141,9 @@ export const getSuggestions = async (
         datasourceMap,
         visualizationMap,
         preferredChartType,
-        preferredVisAttributes,
+        preferredVisAttributes: preferredVisAttributes
+          ? injectESQLQueryIntoLensLayers(preferredVisAttributes, query)
+          : undefined,
       }) ?? [];
 
     // Lens might not return suggestions for some cases, i.e. in case of errors
@@ -162,3 +169,39 @@ export const getSuggestions = async (
   }
   return undefined;
 };
+
+export const injectESQLQueryIntoLensLayers = (
+  attributes: TypedLensSerializedState['attributes'],
+  query: AggregateQuery
+) => {
+  const datasourceId = getDatasourceId(attributes.state.datasourceStates);
+
+  // if the datasource is formBased, we should not fix the query
+  if (!datasourceId || datasourceId === 'formBased') {
+    return attributes;
+  }
+
+  if (!attributes.state.datasourceStates[datasourceId]) {
+    return attributes;
+  }
+
+  const datasourceState = cloneDeep(attributes.state.datasourceStates[datasourceId]);
+
+  if (datasourceState && datasourceState.layers) {
+    Object.values(datasourceState.layers).forEach((layer) => {
+      if (!isEqual(layer.query, query)) {
+        layer.query = query;
+      }
+    });
+  }
+  return {
+    ...attributes,
+    state: {
+      ...attributes.state,
+      datasourceStates: {
+        ...attributes.state.datasourceStates,
+        [datasourceId]: datasourceState,
+      },
+    },
+  };
+};
