Address PR comments

elastic · Oct 1, 2024 · 9e731df · 9e731df
1 parent cae1002
commit 9e731df
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 23 deletions.
diff --git a/src/plugins/unified_histogram/public/chart/breakdown_field_selector.tsx b/src/plugins/unified_histogram/public/chart/breakdown_field_selector.tsx
@@ -33,11 +33,9 @@ export interface BreakdownFieldSelectorProps {
 
 const mapToDropdownFields = (dataView: DataView, esqlColumns?: DatatableColumn[]) => {
   if (esqlColumns) {
-    return esqlColumns.map((column) => ({
-      name: column.name,
-      displayName: column.name,
-      type: column.meta?.type,
-    }));
+    return esqlColumns.map(
+      (column) => new DataViewField(convertDatatableColumnToDataViewFieldSpec(column))
+    );
   }
 
   return dataView.fields.filter(fieldSupportsBreakdown);
@@ -89,20 +87,12 @@ export const BreakdownFieldSelector = ({
 
   const onChange = useCallback<NonNullable<ToolbarSelectorProps['onChange']>>(
     (chosenOption) => {
-      let breakdownField: DataViewField | undefined;
-      if (esqlColumns) {
-        const breakdownColumn = esqlColumns?.find((column) => column.name === chosenOption?.value);
-        breakdownField = breakdownColumn
-          ? new DataViewField(convertDatatableColumnToDataViewFieldSpec(breakdownColumn))
-          : undefined;
-      } else {
-        breakdownField = chosenOption?.value
-          ? dataView.fields.find((currentField) => currentField.name === chosenOption.value)
-          : undefined;
-      }
+      const breakdownField = chosenOption?.value
+        ? fields.find((currentField) => currentField.name === chosenOption.value)
+        : undefined;
       onBreakdownFieldChange?.(breakdownField);
     },
-    [dataView.fields, esqlColumns, onBreakdownFieldChange]
+    [fields, onBreakdownFieldChange]
   );
 
   return (

diff --git a/src/plugins/unified_histogram/public/services/lens_vis_service.suggestions.test.ts b/src/plugins/unified_histogram/public/services/lens_vis_service.suggestions.test.ts
@@ -242,7 +242,7 @@ describe('LensVisService suggestions', () => {
 
     const histogramQuery = {
       esql: `from the-data-view | limit 100
-| EVAL timestamp=DATE_TRUNC(30 minute, @timestamp) | stats results = count(*) by timestamp, var0 | sort var0 asc | rename timestamp as \`@timestamp every 30 minute\``,
+| EVAL timestamp=DATE_TRUNC(30 minute, @timestamp) | stats results = count(*) by timestamp, \`var0\` | sort \`var0\` asc | rename timestamp as \`@timestamp every 30 minute\``,
     };
 
     expect(lensVis.visContext?.attributes.state.query).toStrictEqual(histogramQuery);

diff --git a/src/plugins/unified_histogram/public/services/lens_vis_service.ts b/src/plugins/unified_histogram/public/services/lens_vis_service.ts
@@ -461,7 +461,9 @@ export class LensVisService {
     breakdownField?: DataViewField;
   }): Suggestion | undefined => {
     const { dataView, query, timeRange, columns } = queryParams;
-    const breakdownColumn = columns?.find((column) => column.name === breakdownField?.name);
+    const breakdownColumn = breakdownField?.name
+      ? columns?.find((column) => column.name === breakdownField.name)
+      : undefined;
     if (dataView.isTimeBased() && query && isOfAggregateQueryType(query) && timeRange) {
       const isOnHistogramMode = shouldDisplayHistogram(query);
       if (!isOnHistogramMode) return undefined;
@@ -552,7 +554,7 @@ export class LensVisService {
     const language = getAggregateQueryMode(query);
     const safeQuery = removeDropCommandsFromESQLQuery(query[language]);
     const breakdown = breakdownColumn
-      ? `, ${breakdownColumn.name} | sort ${breakdownColumn.name} asc`
+      ? `, \`${breakdownColumn.name}\` | sort \`${breakdownColumn.name}\` asc`
       : '';
     return appendToESQLQuery(
       safeQuery,
@@ -612,12 +614,17 @@ export class LensVisService {
       breakdownField: breakdownField?.name,
     };
 
-    const breakdownColumn = columns?.find((column) => column.name === breakdownField?.name);
-
     const currentQuery =
       suggestionType === UnifiedHistogramSuggestionType.histogramForESQL && isTextBased && timeRange
         ? {
-            esql: this.getESQLHistogramQuery({ dataView, query, timeRange, breakdownColumn }),
+            esql: this.getESQLHistogramQuery({
+              dataView,
+              query,
+              timeRange,
+              breakdownColumn: breakdownField?.name
+                ? columns?.find((column) => column.name === breakdownField.name)
+                : undefined,
+            }),
           }
         : query;