From 306cea73509f815d8c879152c752995d7258b528 Mon Sep 17 00:00:00 2001
From: Joey <90795735+joey-grafana@users.noreply.github.com>
Date: Thu, 18 Apr 2024 08:52:51 +0100
Subject: [PATCH] Tempo: Group by template vars (#86022)

* Add template variables to group by field

* Add test for interpolation

* Add test to allow selecting template vars

* Show custom value
---
 .../SearchTraceQLEditor/GroupByField.test.tsx | 20 +++++++++++++++++
 .../SearchTraceQLEditor/GroupByField.tsx      | 14 ++++++++----
 .../tempo/SearchTraceQLEditor/SearchField.tsx | 22 +++++++++----------
 .../datasource/tempo/datasource.test.ts       | 14 ++++++++++++
 .../plugins/datasource/tempo/datasource.ts    | 11 ++++++++++
 5 files changed, 66 insertions(+), 15 deletions(-)

diff --git a/public/app/plugins/datasource/tempo/SearchTraceQLEditor/GroupByField.test.tsx b/public/app/plugins/datasource/tempo/SearchTraceQLEditor/GroupByField.test.tsx
index 95ebbb01afc..303bf88837a 100644
--- a/public/app/plugins/datasource/tempo/SearchTraceQLEditor/GroupByField.test.tsx
+++ b/public/app/plugins/datasource/tempo/SearchTraceQLEditor/GroupByField.test.tsx
@@ -5,6 +5,7 @@ import React, { useState } from 'react';
 import { TraceqlSearchScope } from '../dataquery.gen';
 import { TempoDatasource } from '../datasource';
 import TempoLanguageProvider from '../language_provider';
+import { initTemplateSrv } from '../test_utils';
 import { TempoQuery } from '../types';
 
 import { GroupByField } from './GroupByField';
@@ -40,6 +41,8 @@ describe('GroupByField', () => {
     // Need to use delay: null here to work with fakeTimers
     // see https://github.com/testing-library/user-event/issues/833
     user = userEvent.setup({ delay: null });
+
+    initTemplateSrv([{ name: 'templateVariable1' }, { name: 'templateVariable2' }], {});
   });
 
   afterEach(() => {
@@ -132,4 +135,21 @@ describe('GroupByField', () => {
       expect(groupByFilter?.tag).toBe('http.method');
     }
   });
+
+  it('should allow selecting template variables', async () => {
+    const { container } = render(
+      <GroupByField datasource={datasource} query={query} onChange={onChange} isTagsLoading={false} />
+    );
+
+    const tagSelect = container.querySelector(`input[aria-label="Select tag for filter 1"]`);
+    expect(tagSelect).not.toBeNull();
+    expect(tagSelect).toBeInTheDocument();
+
+    if (tagSelect) {
+      await user.click(tagSelect);
+      jest.advanceTimersByTime(1000);
+      expect(await screen.findByText('$templateVariable1')).toBeInTheDocument();
+      expect(await screen.findByText('$templateVariable2')).toBeInTheDocument();
+    }
+  });
 });
diff --git a/public/app/plugins/datasource/tempo/SearchTraceQLEditor/GroupByField.tsx b/public/app/plugins/datasource/tempo/SearchTraceQLEditor/GroupByField.tsx
index e88ab32a21c..bcf53682d7e 100644
--- a/public/app/plugins/datasource/tempo/SearchTraceQLEditor/GroupByField.tsx
+++ b/public/app/plugins/datasource/tempo/SearchTraceQLEditor/GroupByField.tsx
@@ -11,6 +11,7 @@ import { TempoDatasource } from '../datasource';
 import { TempoQuery } from '../types';
 
 import InlineSearchField from './InlineSearchField';
+import { withTemplateVariableOptions } from './SearchField';
 import { replaceAt } from './utils';
 
 interface Props {
@@ -89,15 +90,20 @@ export const GroupByField = (props: Props) => {
               <Select
                 aria-label={`Select tag for filter ${i + 1}`}
                 isClearable
+                allowCustomValue
                 isLoading={isTagsLoading}
                 key={f.tag}
                 onChange={(v) => {
                   updateFilter({ ...f, tag: v?.value });
                 }}
-                options={getTags(f)?.map((t) => ({
-                  label: t,
-                  value: t,
-                }))}
+                options={withTemplateVariableOptions(
+                  getTags(f)
+                    ?.concat(f.tag !== undefined && !getTags(f)?.includes(f.tag) ? [f.tag] : [])
+                    .map((t) => ({
+                      label: t,
+                      value: t,
+                    }))
+                )}
                 placeholder="Select tag"
                 value={f.tag || ''}
               />
diff --git a/public/app/plugins/datasource/tempo/SearchTraceQLEditor/SearchField.tsx b/public/app/plugins/datasource/tempo/SearchTraceQLEditor/SearchField.tsx
index eaaeda32221..d7094d3ceec 100644
--- a/public/app/plugins/datasource/tempo/SearchTraceQLEditor/SearchField.tsx
+++ b/public/app/plugins/datasource/tempo/SearchTraceQLEditor/SearchField.tsx
@@ -126,17 +126,6 @@ const SearchField = ({
       operatorList = numberOperators;
   }
 
-  /**
-   * Add to a list of options the current template variables.
-   *
-   * @param options a list of options
-   * @returns the list of given options plus the template variables
-   */
-  const withTemplateVariableOptions = (options: SelectableValue[] | undefined) => {
-    const templateVariables = getTemplateSrv().getVariables();
-    return [...(options || []), ...templateVariables.map((v) => ({ label: `$${v.name}`, value: `$${v.name}` }))];
-  };
-
   return (
     <>
       <HorizontalGroup spacing={'none'} width={'auto'}>
@@ -220,4 +209,15 @@ const SearchField = ({
   );
 };
 
+/**
+ * Add to a list of options the current template variables.
+ *
+ * @param options a list of options
+ * @returns the list of given options plus the template variables
+ */
+export const withTemplateVariableOptions = (options: SelectableValue[] | undefined) => {
+  const templateVariables = getTemplateSrv().getVariables();
+  return [...(options || []), ...templateVariables.map((v) => ({ label: `$${v.name}`, value: `$${v.name}` }))];
+};
+
 export default SearchField;
diff --git a/public/app/plugins/datasource/tempo/datasource.test.ts b/public/app/plugins/datasource/tempo/datasource.test.ts
index ab3d90b1a9f..6b46d252b8f 100644
--- a/public/app/plugins/datasource/tempo/datasource.test.ts
+++ b/public/app/plugins/datasource/tempo/datasource.test.ts
@@ -170,6 +170,16 @@ describe('Tempo data source', () => {
             valueType: 'string',
           },
         ],
+        groupBy: [
+          {
+            id: 'groupBy1',
+            tag: '$interpolationVar',
+          },
+          {
+            id: 'groupBy2',
+            tag: '$interpolationVar',
+          },
+        ],
       };
     }
     let templateSrv: TemplateSrv;
@@ -202,6 +212,8 @@ describe('Tempo data source', () => {
       expect(queries[0].filters[0].value).toBe(textWithPipe);
       expect(queries[0].filters[1].value).toBe(text);
       expect(queries[0].filters[1].tag).toBe(text);
+      expect(queries[0].groupBy?.[0].tag).toBe(text);
+      expect(queries[0].groupBy?.[1].tag).toBe(text);
     });
 
     it('when applying template variables', async () => {
@@ -214,6 +226,8 @@ describe('Tempo data source', () => {
       expect(resp.filters[0].value).toBe(textWithPipe);
       expect(resp.filters[1].value).toBe(scopedText);
       expect(resp.filters[1].tag).toBe(scopedText);
+      expect(resp.groupBy?.[0].tag).toBe(scopedText);
+      expect(resp.groupBy?.[1].tag).toBe(scopedText);
     });
 
     it('when serviceMapQuery is an array', async () => {
diff --git a/public/app/plugins/datasource/tempo/datasource.ts b/public/app/plugins/datasource/tempo/datasource.ts
index 9b2e3a56a34..c6b4352b887 100644
--- a/public/app/plugins/datasource/tempo/datasource.ts
+++ b/public/app/plugins/datasource/tempo/datasource.ts
@@ -486,6 +486,17 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
       });
     }
 
+    if (query.groupBy) {
+      expandedQuery.groupBy = query.groupBy.map((filter) => {
+        const updatedFilter = {
+          ...filter,
+          tag: this.templateSrv.replace(filter.tag ?? '', scopedVars),
+        };
+
+        return updatedFilter;
+      });
+    }
+
     return {
       ...expandedQuery,
       query: this.templateSrv.replace(query.query ?? '', scopedVars, VariableFormatID.Pipe),