refactor: data trails auto query for buckets (#80170)

* refactor: data trails auto query for buckets

* refactor: vizBuilder function signature (#80178)

* fix: use closures for setting unit and title in vizBuilders

public/app/features/trails/AutomaticMetricQueries/AutoQueryEngine.test.ts

@@ -24,6 +24,7 @@ describe('getAutoQueriesForMetric', () => {
       // Bucket
       ['my_metric_bucket', 'histogram_quantile(0.99, sum by(le) (rate(...[$__rate_interval])))', 'short', 3],
       ['my_metric_seconds_bucket', 'histogram_quantile(0.99, sum by(le) (rate(...[$__rate_interval])))', 's', 3],
+      ['my_metric_bytes_bucket', 'histogram_quantile(0.99, sum by(le) (rate(...[$__rate_interval])))', 'bytes', 3],
     ])('Given metric %p expect %p with unit %p', (metric, expr, unit, queryCount) => {
       const result = getAutoQueriesForMetric(metric);
 
@@ -53,8 +54,9 @@ describe('getAutoQueriesForMetric', () => {
       ['my_metric_bytes_total', 'sum(rate(...[$__rate_interval]))', 'Bps'], // bytes/s
       ['my_metric_bytes_sum', 'avg(rate(...[$__rate_interval]))', 'Bps'],
       // Bucket
-      ['my_metric_bucket', 'histogram_quantile(0.50, sum by(le) (rate(...[$__rate_interval])))', 'short'],
+      ['my_metric_bucket', 'histogram_quantile(0.5, sum by(le) (rate(...[$__rate_interval])))', 'short'],
-      ['my_metric_seconds_bucket', 'histogram_quantile(0.50, sum by(le) (rate(...[$__rate_interval])))', 's'],
+      ['my_metric_seconds_bucket', 'histogram_quantile(0.5, sum by(le) (rate(...[$__rate_interval])))', 's'],
+      ['my_metric_bytes_bucket', 'histogram_quantile(0.5, sum by(le) (rate(...[$__rate_interval])))', 'bytes'],
     ])('Given metric %p expect %p with unit %p', (metric, expr, unit) => {
       const result = getAutoQueriesForMetric(metric);
 
@@ -86,12 +88,17 @@ describe('getAutoQueriesForMetric', () => {
       ['my_metric_bytes_total', 'sum(rate(...[$__rate_interval])) by(${groupby})', 'Bps'], // bytes/s
       ['my_metric_bytes_sum', 'avg(rate(...[$__rate_interval])) by(${groupby})', 'Bps'],
       // Bucket
-      ['my_metric_bucket', 'histogram_quantile(0.50, sum by(le, ${groupby}) (rate(...[$__rate_interval])))', 'short'],
+      ['my_metric_bucket', 'histogram_quantile(0.5, sum by(le, ${groupby}) (rate(...[$__rate_interval])))', 'short'],
       [
         'my_metric_seconds_bucket',
-        'histogram_quantile(0.50, sum by(le, ${groupby}) (rate(...[$__rate_interval])))',
+        'histogram_quantile(0.5, sum by(le, ${groupby}) (rate(...[$__rate_interval])))',
         's',
       ],
+      [
+        'my_metric_bytes_bucket',
+        'histogram_quantile(0.5, sum by(le, ${groupby}) (rate(...[$__rate_interval])))',
+        'bytes',
+      ],
     ])('Given metric %p expect %p with unit %p', (metric, expr, unit) => {
       const result = getAutoQueriesForMetric(metric);
 
@@ -125,8 +132,8 @@ describe('getAutoQueriesForMetric', () => {
             unit: 'short',
             exprs: [
               'histogram_quantile(0.99, sum by(le) (rate(${metric}{${filters}}[$__rate_interval])))',
-              'histogram_quantile(0.90, sum by(le) (rate(${metric}{${filters}}[$__rate_interval])))',
+              'histogram_quantile(0.9, sum by(le) (rate(${metric}{${filters}}[$__rate_interval])))',
-              'histogram_quantile(0.50, sum by(le) (rate(${metric}{${filters}}[$__rate_interval])))',
+              'histogram_quantile(0.5, sum by(le) (rate(${metric}{${filters}}[$__rate_interval])))',
             ],
           },
           {
@@ -144,8 +151,8 @@ describe('getAutoQueriesForMetric', () => {
             unit: 's',
             exprs: [
               'histogram_quantile(0.99, sum by(le) (rate(${metric}{${filters}}[$__rate_interval])))',
-              'histogram_quantile(0.90, sum by(le) (rate(${metric}{${filters}}[$__rate_interval])))',
+              'histogram_quantile(0.9, sum by(le) (rate(${metric}{${filters}}[$__rate_interval])))',
-              'histogram_quantile(0.50, sum by(le) (rate(${metric}{${filters}}[$__rate_interval])))',
+              'histogram_quantile(0.5, sum by(le) (rate(${metric}{${filters}}[$__rate_interval])))',
             ],
           },
           {
@@ -155,6 +162,25 @@ describe('getAutoQueriesForMetric', () => {
           },
         ],
       ],
+      [
+        'my_metric_bytes_bucket',
+        [
+          {
+            variant: 'percentiles',
+            unit: 'bytes',
+            exprs: [
+              'histogram_quantile(0.99, sum by(le) (rate(${metric}{${filters}}[$__rate_interval])))',
+              'histogram_quantile(0.9, sum by(le) (rate(${metric}{${filters}}[$__rate_interval])))',
+              'histogram_quantile(0.5, sum by(le) (rate(${metric}{${filters}}[$__rate_interval])))',
+            ],
+          },
+          {
+            variant: 'heatmap',
+            unit: 'bytes',
+            exprs: ['sum by(le) (rate(${metric}{${filters}}[$__rate_interval]))'],
+          },
+        ],
+      ],
     ])('Given metric %p should generate expected variants', (metric, expectedVariants) => {
       const defs = getAutoQueriesForMetric(metric);
 

public/app/features/trails/AutomaticMetricQueries/AutoVizPanel.tsx

@@ -51,7 +51,7 @@ export class AutoVizPanel extends SceneObjectBase<AutoVizPanelState> {
 
   private getVizPanelFor(def: AutoQueryDef) {
     return def
-      .vizBuilder(def)
+      .vizBuilder()
       .setData(
         new SceneQueryRunner({
           datasource: trailDS,

public/app/features/trails/AutomaticMetricQueries/graph-builders/heatmap.ts

@@ -1,12 +1,12 @@
 import { PanelBuilders } from '@grafana/scenes';
 import { HeatmapColorMode } from 'app/plugins/panel/heatmap/types';
 
-import { AutoQueryDef } from '../types';
+import { CommonVizParams } from './types';
 
-export function heatmapGraphBuilder(def: AutoQueryDef) {
+export function heatmapGraphBuilder({ title, unit }: CommonVizParams) {
-  return PanelBuilders.heatmap()
+  return PanelBuilders.heatmap() //
-    .setTitle(def.title)
+    .setTitle(title)
-    .setUnit(def.unit)
+    .setUnit(unit)
     .setOption('calculate', false)
     .setOption('color', {
       mode: HeatmapColorMode.Scheme,

public/app/features/trails/AutomaticMetricQueries/graph-builders/percentiles.ts

@@ -1,7 +1,10 @@
 import { PanelBuilders } from '@grafana/scenes';
 
-import { AutoQueryDef } from '../types';
+import { CommonVizParams } from './types';
 
-export function percentilesGraphBuilder(def: AutoQueryDef) {
+export function percentilesGraphBuilder({ title, unit }: CommonVizParams) {
-  return PanelBuilders.timeseries().setTitle(def.title).setUnit(def.unit).setCustomFieldConfig('fillOpacity', 9);
+  return PanelBuilders.timeseries() //
+    .setTitle(title)
+    .setUnit(unit)
+    .setCustomFieldConfig('fillOpacity', 9);
 }

public/app/features/trails/AutomaticMetricQueries/graph-builders/simple.ts

@@ -1,11 +1,11 @@
 import { PanelBuilders } from '@grafana/scenes';
 
-import { AutoQueryDef } from '../types';
+import { CommonVizParams } from './types';
 
-export function simpleGraphBuilder(def: AutoQueryDef) {
+export function simpleGraphBuilder({ title, unit }: CommonVizParams) {
-  return PanelBuilders.timeseries()
+  return PanelBuilders.timeseries() //
-    .setTitle(def.title)
+    .setTitle(title)
-    .setUnit(def.unit)
+    .setUnit(unit)
     .setOption('legend', { showLegend: false })
     .setCustomFieldConfig('fillOpacity', 9);
 }

public/app/features/trails/AutomaticMetricQueries/graph-builders/types.ts

@@ -0,0 +1,4 @@
+export type CommonVizParams = {
+  title: string;
+  unit: string;
+};

public/app/features/trails/AutomaticMetricQueries/query-generators/bucket/index.ts

@@ -1,86 +1,86 @@
+import { PromQuery } from 'app/plugins/datasource/prometheus/types';
+
 import { VAR_FILTERS_EXPR, VAR_GROUP_BY_EXP, VAR_METRIC_EXPR } from '../../../shared';
 import { heatmapGraphBuilder } from '../../graph-builders/heatmap';
 import { percentilesGraphBuilder } from '../../graph-builders/percentiles';
 import { simpleGraphBuilder } from '../../graph-builders/simple';
 import { AutoQueryDef } from '../../types';
+import { getUnit } from '../../units';
 
 function generator(metricParts: string[]) {
-  let unit = 'short';
-
   const title = `${VAR_METRIC_EXPR}`;
 
   const unitSuffix = metricParts.at(-2);
 
-  if (unitSuffix === 'seconds') {
+  const unit = getUnit(unitSuffix);
-    // TODO Map to other units
-    unit = 's';
-  }
 
-  const p50: AutoQueryDef = {
+  const common = {
     title,
-    variant: 'p50',
     unit,
-    queries: [
+  };
-      {
+
-        refId: 'A',
+  const p50: AutoQueryDef = {
-        expr: `histogram_quantile(0.50, sum by(le) (rate(${VAR_METRIC_EXPR}${VAR_FILTERS_EXPR}[$__rate_interval])))`,
+    ...common,
-      },
+    variant: 'p50',
-    ],
+    queries: [percentileQuery(50)],
-    vizBuilder: simpleGraphBuilder,
+    vizBuilder: () => simpleGraphBuilder(p50),
   };
 
   const breakdown: AutoQueryDef = {
-    title,
+    ...common,
     variant: 'p50',
-    unit,
+    queries: [percentileQuery(50, [VAR_GROUP_BY_EXP])],
-    queries: [
+    vizBuilder: () => simpleGraphBuilder(breakdown),
-      {
-        refId: 'A',
-        expr: `histogram_quantile(0.50, sum by(le, ${VAR_GROUP_BY_EXP}) (rate(${VAR_METRIC_EXPR}${VAR_FILTERS_EXPR}[$__rate_interval])))`,
-      },
-    ],
-    vizBuilder: simpleGraphBuilder,
   };
 
   const percentiles: AutoQueryDef = {
-    title,
+    ...common,
     variant: 'percentiles',
-    unit,
+    queries: [99, 90, 50].map((p) => percentileQuery(p)).map(fixRefIds),
-    queries: [
+    vizBuilder: () => percentilesGraphBuilder(percentiles),
-      {
-        refId: 'A',
-        expr: `histogram_quantile(0.99, sum by(le) (rate(${VAR_METRIC_EXPR}${VAR_FILTERS_EXPR}[$__rate_interval])))`,
-        legendFormat: '99th Percentile',
-      },
-      {
-        refId: 'B',
-        expr: `histogram_quantile(0.90, sum by(le) (rate(${VAR_METRIC_EXPR}${VAR_FILTERS_EXPR}[$__rate_interval])))`,
-        legendFormat: '90th Percentile',
-      },
-      {
-        refId: 'C',
-        expr: `histogram_quantile(0.50, sum by(le) (rate(${VAR_METRIC_EXPR}${VAR_FILTERS_EXPR}[$__rate_interval])))`,
-        legendFormat: '50th Percentile',
-      },
-    ],
-    vizBuilder: percentilesGraphBuilder,
   };
 
   const heatmap: AutoQueryDef = {
-    title,
+    ...common,
     variant: 'heatmap',
-    unit,
+    queries: [heatMapQuery()],
-    queries: [
+    vizBuilder: () => heatmapGraphBuilder(heatmap),
-      {
-        refId: 'A',
-        expr: `sum by(le) (rate(${VAR_METRIC_EXPR}${VAR_FILTERS_EXPR}[$__rate_interval]))`,
-        format: 'heatmap',
-      },
-    ],
-    vizBuilder: heatmapGraphBuilder,
   };
 
   return { preview: p50, main: percentiles, variants: [percentiles, heatmap], breakdown: breakdown };
 }
 
+function fixRefIds(queryDef: PromQuery, index: number): PromQuery {
+  // By default refIds are `"A"`
+  // This method will reassign based on `A + index` -- A, B, C, etc
+  return {
+    ...queryDef,
+    refId: String.fromCharCode('A'.charCodeAt(0) + index),
+  };
+}
+
 export default { generator };
+
+const BASE_QUERY = `rate(${VAR_METRIC_EXPR}${VAR_FILTERS_EXPR}[$__rate_interval])`;
+
+function baseQuery(groupings: string[] = []) {
+  const sumByList = ['le', ...groupings];
+  return `sum by(${sumByList.join(', ')}) (${BASE_QUERY})`;
+}
+
+function heatMapQuery(groupings: string[] = []) {
+  return {
+    refId: 'A',
+    expr: baseQuery(groupings),
+  };
+}
+
+function percentileQuery(percentile: number, groupings: string[] = []) {
+  const percent = percentile / 100;
+
+  return {
+    refId: 'A',
+    expr: `histogram_quantile(${percent}, ${baseQuery(groupings)})`,
+    legendFormat: `${percentile}th Percentile`,
+  };
+}

public/app/features/trails/AutomaticMetricQueries/query-generators/common/queries.ts

@@ -1,6 +1,6 @@
 import { VAR_FILTERS_EXPR, VAR_GROUP_BY_EXP, VAR_METRIC_EXPR } from '../../../shared';
 import { simpleGraphBuilder } from '../../graph-builders/simple';
-import { AutoQueryDef, AutoQueryInfo } from '../../types';
+import { AutoQueryInfo } from '../../types';
 
 import { AutoQueryParameters } from './types';
 
@@ -14,28 +14,20 @@ export function getGeneralBaseQuery(rate: boolean) {
 export function generateQueries({ agg, rate, unit }: AutoQueryParameters): AutoQueryInfo {
   const baseQuery = getGeneralBaseQuery(rate);
 
-  const main = createMainQuery(baseQuery, agg, unit);
+  const common = {
-
-  const breakdown = createBreakdownQuery(baseQuery, agg, unit);
-
-  return { preview: main, main: main, breakdown: breakdown, variants: [] };
-}
-
-function createMainQuery(baseQuery: string, agg: string, unit: string): AutoQueryDef {
-  return {
     title: `${VAR_METRIC_EXPR}`,
-    variant: 'graph',
     unit,
+    variant: 'graph',
+  };
+
+  const main = {
+    ...common,
     queries: [{ refId: 'A', expr: `${agg}(${baseQuery})` }],
-    vizBuilder: simpleGraphBuilder,
+    vizBuilder: () => simpleGraphBuilder(main),
   };
-}
 
-function createBreakdownQuery(baseQuery: string, agg: string, unit: string): AutoQueryDef {
+  const breakdown = {
-  return {
+    ...common,
-    title: `${VAR_METRIC_EXPR}`,
-    variant: 'graph',
-    unit,
     queries: [
       {
         refId: 'A',
@@ -43,6 +35,8 @@ function createBreakdownQuery(baseQuery: string, agg: string, unit: string): Aut
         legendFormat: `{{${VAR_GROUP_BY_EXP}}}`,
       },
     ],
-    vizBuilder: simpleGraphBuilder,
+    vizBuilder: () => simpleGraphBuilder(breakdown),
   };
+
+  return { preview: main, main: main, breakdown: breakdown, variants: [] };
 }

public/app/features/trails/AutomaticMetricQueries/types.ts

@@ -6,7 +6,7 @@ export interface AutoQueryDef {
   title: string;
   unit: string;
   queries: PromQuery[];
-  vizBuilder: (def: AutoQueryDef) => VizPanelBuilder<{}, {}>;
+  vizBuilder: VizBuilder;
 }
 
 export interface AutoQueryInfo {
@@ -15,3 +15,5 @@ export interface AutoQueryInfo {
   variants: AutoQueryDef[];
   breakdown: AutoQueryDef;
 }
+
+export type VizBuilder = () => VizPanelBuilder<{}, {}>;

public/app/features/trails/BreakdownScene.tsx

@@ -201,7 +201,6 @@ export function buildAllLayout(options: Array<SelectableValue<string>>, queryDef
       new SceneCSSGridItem({
         body: PanelBuilders.timeseries()
           .setTitle(option.label!)
-          .setUnit(queryDef.unit)
           .setData(
             new SceneQueryRunner({
               maxDataPoints: 300,
@@ -293,7 +292,7 @@ function buildNormalLayout(queryDef: AutoQueryDef) {
         getLayoutChild: (data, frame, frameIndex) => {
           return new SceneCSSGridItem({
             body: queryDef
-              .vizBuilder(queryDef)
+              .vizBuilder()
               .setTitle(getLabelValue(frame))
               .setData(new SceneDataNode({ data: { ...data, series: [frame] } }))
               .setColor({ mode: 'fixed', fixedColor: getColorByIndex(frameIndex) })
@@ -311,7 +310,7 @@ function buildNormalLayout(queryDef: AutoQueryDef) {
         getLayoutChild: (data, frame, frameIndex) => {
           return new SceneCSSGridItem({
             body: queryDef
-              .vizBuilder(queryDef)
+              .vizBuilder()
               .setTitle(getLabelValue(frame))
               .setData(new SceneDataNode({ data: { ...data, series: [frame] } }))
               .setColor({ mode: 'fixed', fixedColor: getColorByIndex(frameIndex) })

public/app/features/trails/MetricSelectScene.tsx

@@ -257,7 +257,7 @@ function getPreviewPanelFor(metric: string, index: number) {
   const autoQuery = getAutoQueriesForMetric(metric);
 
   const vizPanel = autoQuery.preview
-    .vizBuilder(autoQuery.preview)
+    .vizBuilder()
     .setColor({ mode: 'fixed', fixedColor: getColorByIndex(index) })
     .setHeaderActions(new SelectMetricAction({ metric, title: 'Select' }))
     .build();