From fc4a9904c91c85f8b8fb21a0e45494252c253b5e Mon Sep 17 00:00:00 2001
From: Andre Pereira <adrapereira@gmail.com>
Date: Thu, 20 Jun 2024 15:22:50 +0100
Subject: [PATCH] Tempo: TraceQL metrics step option (#89434)

* Add step option for metric queries

* Add support for compare metric queries

* Remove unneeded line

* Delete step if it's not defined
---
 .../dataquery/x/TempoDataQuery_types.gen.ts     |  4 ++++
 .../kinds/dataquery/types_dataquery_gen.go      |  3 +++
 .../app/plugins/datasource/tempo/dataquery.cue  |  2 ++
 .../plugins/datasource/tempo/dataquery.gen.ts   |  4 ++++
 .../app/plugins/datasource/tempo/datasource.ts  | 13 ++++++++++---
 .../tempo/traceql/TempoQueryBuilderOptions.tsx  | 17 +++++++++++++++++
 6 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/packages/grafana-schema/src/raw/composable/tempo/dataquery/x/TempoDataQuery_types.gen.ts b/packages/grafana-schema/src/raw/composable/tempo/dataquery/x/TempoDataQuery_types.gen.ts
index af517e367a2..6583afdedb2 100644
--- a/packages/grafana-schema/src/raw/composable/tempo/dataquery/x/TempoDataQuery_types.gen.ts
+++ b/packages/grafana-schema/src/raw/composable/tempo/dataquery/x/TempoDataQuery_types.gen.ts
@@ -58,6 +58,10 @@ export interface TempoQuery extends common.DataQuery {
    * Defines the maximum number of spans per spanset that are returned from Tempo
    */
   spss?: number;
+  /**
+   * For metric queries, the step size to use
+   */
+  step?: string;
   /**
    * The type of the table that is used to display the search results
    */
diff --git a/pkg/tsdb/tempo/kinds/dataquery/types_dataquery_gen.go b/pkg/tsdb/tempo/kinds/dataquery/types_dataquery_gen.go
index bfb397b514d..4b3e3510266 100644
--- a/pkg/tsdb/tempo/kinds/dataquery/types_dataquery_gen.go
+++ b/pkg/tsdb/tempo/kinds/dataquery/types_dataquery_gen.go
@@ -129,6 +129,9 @@ type TempoQuery struct {
 	// Defines the maximum number of spans per spanset that are returned from Tempo
 	Spss *int64 `json:"spss,omitempty"`
 
+	// For metric queries, the step size to use
+	Step *string `json:"step,omitempty"`
+
 	// The type of the table that is used to display the search results
 	TableType *SearchTableType `json:"tableType,omitempty"`
 }
diff --git a/public/app/plugins/datasource/tempo/dataquery.cue b/public/app/plugins/datasource/tempo/dataquery.cue
index b47078559b0..69dd0f653e2 100644
--- a/public/app/plugins/datasource/tempo/dataquery.cue
+++ b/public/app/plugins/datasource/tempo/dataquery.cue
@@ -51,6 +51,8 @@ composableKinds: DataQuery: {
 					groupBy?: [...#TraceqlFilter]
 					// The type of the table that is used to display the search results
 					tableType?: #SearchTableType
+					// For metric queries, the step size to use
+					step?: string
 				} @cuetsy(kind="interface") @grafana(TSVeneer="type")
 
 				#TempoQueryType: "traceql" | "traceqlSearch" | "serviceMap" | "upload" | "nativeSearch" | "traceId" | "clear" @cuetsy(kind="type")
diff --git a/public/app/plugins/datasource/tempo/dataquery.gen.ts b/public/app/plugins/datasource/tempo/dataquery.gen.ts
index 5e0bb9543ee..7086fcc5c63 100644
--- a/public/app/plugins/datasource/tempo/dataquery.gen.ts
+++ b/public/app/plugins/datasource/tempo/dataquery.gen.ts
@@ -56,6 +56,10 @@ export interface TempoQuery extends common.DataQuery {
    * Defines the maximum number of spans per spanset that are returned from Tempo
    */
   spss?: number;
+  /**
+   * For metric queries, the step size to use
+   */
+  step?: string;
   /**
    * The type of the table that is used to display the search results
    */
diff --git a/public/app/plugins/datasource/tempo/datasource.ts b/public/app/plugins/datasource/tempo/datasource.ts
index bcf330a29ec..b7ae4f1c77d 100644
--- a/public/app/plugins/datasource/tempo/datasource.ts
+++ b/public/app/plugins/datasource/tempo/datasource.ts
@@ -438,7 +438,7 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
   isTraceQlMetricsQuery(query: string): boolean {
     // Check whether this is a metrics query by checking if it contains a metrics function
     const metricsFnRegex =
-      /\|\s*(rate|count_over_time|avg_over_time|max_over_time|min_over_time|quantile_over_time|histogram_over_time)\s*\(/;
+      /\|\s*(rate|count_over_time|avg_over_time|max_over_time|min_over_time|quantile_over_time|histogram_over_time|compare)\s*\(/;
     return !!query.trim().match(metricsFnRegex);
   }
 
@@ -643,11 +643,18 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
     options: DataQueryRequest<TempoQuery>,
     queryValue: string
   ): Observable<DataQueryResponse> => {
-    return this._request('/api/metrics/query_range', {
+    const requestData = {
       query: queryValue,
       start: options.range.from.unix(),
       end: options.range.to.unix(),
-    }).pipe(
+      step: options.targets[0].step,
+    };
+
+    if (!requestData.step) {
+      delete requestData.step;
+    }
+
+    return this._request('/api/metrics/query_range', requestData).pipe(
       map((response) => {
         return {
           data: formatTraceQLMetrics(queryValue, response.data),
diff --git a/public/app/plugins/datasource/tempo/traceql/TempoQueryBuilderOptions.tsx b/public/app/plugins/datasource/tempo/traceql/TempoQueryBuilderOptions.tsx
index 98658945e54..d4e40701ffb 100644
--- a/public/app/plugins/datasource/tempo/traceql/TempoQueryBuilderOptions.tsx
+++ b/public/app/plugins/datasource/tempo/traceql/TempoQueryBuilderOptions.tsx
@@ -44,11 +44,15 @@ export const TempoQueryBuilderOptions = React.memo<Props>(({ onChange, query })
   const onTableTypeChange = (val: SearchTableType) => {
     onChange({ ...query, tableType: val });
   };
+  const onStepChange = (e: React.FormEvent<HTMLInputElement>) => {
+    onChange({ ...query, step: e.currentTarget.value });
+  };
 
   const collapsedInfoList = [
     `Limit: ${query.limit || DEFAULT_LIMIT}`,
     `Spans Limit: ${query.spss || DEFAULT_SPSS}`,
     `Table Format: ${query.tableType === SearchTableType.Traces ? 'Traces' : 'Spans'}`,
+    `Step: ${query.step || 'auto'}`,
   ];
 
   return (
@@ -87,6 +91,19 @@ export const TempoQueryBuilderOptions = React.memo<Props>(({ onChange, query })
               onChange={onTableTypeChange}
             />
           </EditorField>
+          <EditorField
+            label="Step"
+            tooltip="Defines the step for metric queries. Use duration notation, for example 30s or 1m"
+          >
+            <AutoSizeInput
+              className="width-4"
+              placeholder="auto"
+              type="string"
+              defaultValue={query.step}
+              onCommitChange={onStepChange}
+              value={query.step}
+            />
+          </EditorField>
         </QueryOptionGroup>
       </EditorRow>
     </>