Refactor: Consolidate code for query generator (#91415)

* Refactor: Consolidate code for query generator

* Update public/app/features/trails/AutomaticMetricQueries/query-generators/default.ts

Co-authored-by: ismail simsek <ismailsimsek09@gmail.com>

* fix: rename index.ts to getQueryGeneratorFor.ts to avoid barrel files

---------

Co-authored-by: ismail simsek <ismailsimsek09@gmail.com>

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

@@ -1,4 +1,4 @@
-import { getQueryGeneratorFor } from './query-generators';
+import { getQueryGeneratorFor } from './query-generators/getQueryGeneratorFor';
 import { AutoQueryInfo } from './types';
 
 export function getAutoQueriesForMetric(metric: string): AutoQueryInfo {

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

@@ -0,0 +1,8 @@
+import { VAR_METRIC_EXPR, VAR_FILTERS_EXPR } from 'app/features/trails/shared';
+
+const GENERAL_BASE_QUERY = `${VAR_METRIC_EXPR}${VAR_FILTERS_EXPR}`;
+const GENERAL_RATE_BASE_QUERY = `rate(${GENERAL_BASE_QUERY}[$__rate_interval])`;
+
+export function getGeneralBaseQuery(rate: boolean) {
+  return rate ? GENERAL_RATE_BASE_QUERY : GENERAL_BASE_QUERY;
+}

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

@@ -1,9 +0,0 @@
-import { generateQueries } from './queries';
-import { getGeneratorParameters } from './rules';
-
-function generator(metricParts: string[]) {
-  const params = getGeneratorParameters(metricParts);
-  return generateQueries(params);
-}
-
-export default { generator };

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

@@ -1,39 +0,0 @@
-import { VAR_FILTERS_EXPR, VAR_GROUP_BY_EXP, VAR_METRIC_EXPR } from '../../../shared';
-import { AutoQueryInfo } from '../../types';
-import { generateCommonAutoQueryInfo } from '../common/generator';
-
-import { AutoQueryParameters } from './types';
-
-const GENERAL_BASE_QUERY = `${VAR_METRIC_EXPR}${VAR_FILTERS_EXPR}`;
-const GENERAL_RATE_BASE_QUERY = `rate(${GENERAL_BASE_QUERY}[$__rate_interval])`;
-
-export function getGeneralBaseQuery(rate: boolean) {
-  return rate ? GENERAL_RATE_BASE_QUERY : GENERAL_BASE_QUERY;
-}
-
-const aggLabels: Record<string, string> = {
-  avg: 'average',
-  sum: 'overall',
-};
-
-function getAggLabel(agg: string) {
-  return aggLabels[agg] || agg;
-}
-
-export function generateQueries({ agg, rate, unit }: AutoQueryParameters): AutoQueryInfo {
-  const baseQuery = getGeneralBaseQuery(rate);
-
-  const aggregationDescription = rate ? `${getAggLabel(agg)} per-second rate` : `${getAggLabel(agg)}`;
-
-  const description = `${VAR_METRIC_EXPR} (${aggregationDescription})`;
-
-  const mainQueryExpr = `${agg}(${baseQuery})`;
-  const breakdownQueryExpr = `${agg}(${baseQuery})by(${VAR_GROUP_BY_EXP})`;
-
-  return generateCommonAutoQueryInfo({
-    description,
-    mainQueryExpr,
-    breakdownQueryExpr,
-    unit,
-  });
-}

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

@@ -1,39 +0,0 @@
-import { getUnit, getPerSecondRateUnit } from '../../units';
-
-import { AutoQueryParameters } from './types';
-
-/** These suffixes will set rate to true */
-const RATE_SUFFIXES = new Set(['count', 'total']);
-
-const UNSUPPORTED_SUFFIXES = new Set(['sum', 'bucket']);
-
-/** Non-default aggregattion keyed by suffix */
-const SPECIFIC_AGGREGATIONS_FOR_SUFFIX: Record<string, string> = {
-  count: 'sum',
-  total: 'sum',
-};
-
-function checkPreviousForUnit(suffix: string) {
-  return suffix === 'total';
-}
-
-export function getGeneratorParameters(metricParts: string[]): AutoQueryParameters {
-  const suffix = metricParts.at(-1);
-
-  if (suffix == null || UNSUPPORTED_SUFFIXES.has(suffix)) {
-    throw new Error(`This function does not support a metric suffix of "${suffix}"`);
-  }
-
-  const rate = RATE_SUFFIXES.has(suffix);
-  const agg = SPECIFIC_AGGREGATIONS_FOR_SUFFIX[suffix] || 'avg';
-
-  const unitSuffix = checkPreviousForUnit(suffix) ? metricParts.at(-2) : suffix;
-
-  const unit = rate ? getPerSecondRateUnit(unitSuffix) : getUnit(unitSuffix);
-
-  return {
-    agg,
-    unit,
-    rate,
-  };
-}

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

@@ -1,5 +0,0 @@
-export type AutoQueryParameters = {
-  agg: string;
-  unit: string;
-  rate: boolean;
-};

public/app/features/trails/AutomaticMetricQueries/query-generators/default.test.ts

@@ -1,7 +1,8 @@
-import { VAR_GROUP_BY_EXP } from '../../../shared';
-import { AutoQueryDef, AutoQueryInfo } from '../../types';
+import { VAR_GROUP_BY_EXP } from '../../shared';
+import { AutoQueryDef, AutoQueryInfo } from '../types';
 
-import { generateQueries, getGeneralBaseQuery } from './queries';
+import { getGeneralBaseQuery } from './common/baseQuery';
+import { generateQueries } from './default';
 
 describe('generateQueries', () => {
   const agg = 'sum';

public/app/features/trails/AutomaticMetricQueries/query-generators/default.ts

@@ -0,0 +1,82 @@
+import { VAR_GROUP_BY_EXP, VAR_METRIC_EXPR } from 'app/features/trails/shared';
+
+import { AutoQueryInfo } from '../types';
+import { getPerSecondRateUnit, getUnit } from '../units';
+
+import { getGeneralBaseQuery } from './common/baseQuery';
+import { generateCommonAutoQueryInfo } from './common/generator';
+
+/** These suffixes will set rate to true */
+const RATE_SUFFIXES = new Set(['count', 'total']);
+
+const UNSUPPORTED_SUFFIXES = new Set(['sum', 'bucket']);
+
+/** Non-default aggregation keyed by suffix */
+const SPECIFIC_AGGREGATIONS_FOR_SUFFIX: Record<string, string> = {
+  count: 'sum',
+  total: 'sum',
+};
+
+function shouldCheckPreviousSuffixForUnit(suffix: string) {
+  return suffix === 'total';
+}
+
+const aggLabels: Record<string, string> = {
+  avg: 'average',
+  sum: 'overall',
+};
+
+function getAggLabel(agg: string) {
+  return aggLabels[agg] || agg;
+}
+
+export type AutoQueryParameters = {
+  agg: string;
+  unit: string;
+  rate: boolean;
+};
+
+export function generateQueries({ agg, rate, unit }: AutoQueryParameters): AutoQueryInfo {
+  const baseQuery = getGeneralBaseQuery(rate);
+
+  const aggregationDescription = rate ? `${getAggLabel(agg)} per-second rate` : `${getAggLabel(agg)}`;
+
+  const description = `${VAR_METRIC_EXPR} (${aggregationDescription})`;
+
+  const mainQueryExpr = `${agg}(${baseQuery})`;
+  const breakdownQueryExpr = `${agg}(${baseQuery})by(${VAR_GROUP_BY_EXP})`;
+
+  return generateCommonAutoQueryInfo({
+    description,
+    mainQueryExpr,
+    breakdownQueryExpr,
+    unit,
+  });
+}
+
+export function createDefaultMetricQueryDefs(metricParts: string[]) {
+  // Get the last part of the metric name
+  const suffix = metricParts.at(-1);
+
+  // If the suffix is null or is in the set of unsupported suffixes, throw an error because the metric should be delegated to a different generator (summary or histogram)
+  if (suffix == null || UNSUPPORTED_SUFFIXES.has(suffix)) {
+    throw new Error(`This function does not support a metric suffix of "${suffix}"`);
+  }
+
+  // Check if generating rate query and/or aggregation query
+  const rate = RATE_SUFFIXES.has(suffix);
+  const agg = SPECIFIC_AGGREGATIONS_FOR_SUFFIX[suffix] || 'avg';
+
+  // Try to find the unit in the Prometheus metric name
+  const unitSuffix = shouldCheckPreviousSuffixForUnit(suffix) ? metricParts.at(-2) : suffix;
+
+  // Get the Grafana unit or Grafana rate unit
+  const unit = rate ? getPerSecondRateUnit(unitSuffix) : getUnit(unitSuffix);
+
+  const params = {
+    agg,
+    unit,
+    rate,
+  };
+  return generateQueries(params);
+}

public/app/features/trails/AutomaticMetricQueries/query-generators/getQueryGeneratorFor.ts

@@ -0,0 +1,20 @@
+import { AutoQueryInfo } from '../types';
+
+import { createDefaultMetricQueryDefs } from './default';
+import { createHistogramMetricQueryDefs } from './histogram';
+import { createSummaryMetricQueryDefs } from './summary';
+
+// TODO: when we have a known unit parameter, use that rather than having the generator functions infer from suffix
+export type MetricQueriesGenerator = (metricParts: string[]) => AutoQueryInfo;
+
+export function getQueryGeneratorFor(suffix?: string): MetricQueriesGenerator {
+  if (suffix === 'sum') {
+    return createSummaryMetricQueryDefs;
+  }
+
+  if (suffix === 'bucket') {
+    return createHistogramMetricQueryDefs;
+  }
+
+  return createDefaultMetricQueryDefs;
+}

public/app/features/trails/AutomaticMetricQueries/query-generators/histogram.ts

@@ -7,7 +7,7 @@ import { simpleGraphBuilder } from '../graph-builders/simple';
 import { AutoQueryDef } from '../types';
 import { getUnit } from '../units';
 
-export function createHistogramQueryDefs(metricParts: string[]) {
+export function createHistogramMetricQueryDefs(metricParts: string[]) {
   const title = `${VAR_METRIC_EXPR}`;
 
   const unitSuffix = metricParts.at(-2);

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

@@ -1,13 +0,0 @@
-import general from './common';
-import { createHistogramQueryDefs } from './histogram';
-import { createSummaryQueryDefs } from './summary';
-import { MetricQueriesGenerator } from './types';
-
-const SUFFIX_TO_ALTERNATIVE_GENERATOR: Record<string, MetricQueriesGenerator> = {
-  sum: createSummaryQueryDefs,
-  bucket: createHistogramQueryDefs,
-};
-
-export function getQueryGeneratorFor(suffix?: string) {
-  return (suffix && SUFFIX_TO_ALTERNATIVE_GENERATOR[suffix]) || general.generator;
-}

public/app/features/trails/AutomaticMetricQueries/query-generators/summary.ts

@@ -2,13 +2,13 @@ import { VAR_GROUP_BY_EXP, VAR_METRIC_EXPR } from '../../shared';
 import { AutoQueryInfo } from '../types';
 import { getUnit } from '../units';
 
+import { getGeneralBaseQuery } from './common/baseQuery';
 import { generateCommonAutoQueryInfo } from './common/generator';
-import { getGeneralBaseQuery } from './common/queries';
 
-export function createSummaryQueryDefs(metricParts: string[]): AutoQueryInfo {
+export function createSummaryMetricQueryDefs(metricParts: string[]): AutoQueryInfo {
   const suffix = metricParts.at(-1);
   if (suffix !== 'sum') {
-    throw new Error('createSummaryQueryDefs is only to be used for metrics that end in "_sum"');
+    throw new Error('createSummaryMetricQueryDefs is only to be used for metrics that end in "_sum"');
   }
 
   const unitSuffix = metricParts.at(-2);

public/app/features/trails/AutomaticMetricQueries/query-generators/types.ts

@@ -1,3 +0,0 @@
-import { AutoQueryInfo } from '../types';
-
-export type MetricQueriesGenerator = (metricParts: string[]) => AutoQueryInfo;