Traces: Add more template variables in Tempo & Zipkin (#52306)

* Add support for more vars in Tempo

* Tests for Tempo vars

* Tempo ds vars

* Tempo ds vars test

* Zipkin template var

* Zipkin tests

public/app/plugins/datasource/tempo/QueryEditor/NativeSearch.test.tsx

@@ -30,7 +30,17 @@ jest.mock('../language_provider', () => {
   });
 });
 
-const mockQuery = {
+jest.mock('@grafana/runtime', () => ({
+  ...jest.requireActual('@grafana/runtime'),
+  getTemplateSrv: () => ({
+    replace: jest.fn(),
+    containsTemplate: (val: string): boolean => {
+      return val.includes('$');
+    },
+  }),
+}));
+
+let mockQuery = {
   refId: 'A',
   queryType: 'nativeSearch',
   key: 'Q-595a9bbc-2a25-49a7-9249-a52a0a475d83-0',
@@ -117,7 +127,38 @@ describe('NativeSearch', () => {
     expect(option).toBeDefined();
 
     await user.type(select, 'a');
-    option = await screen.findByText('No options found');
+    option = await screen.findByText('Hit enter to add');
     expect(option).toBeDefined();
   });
+
+  it('should add variable to select menu options', async () => {
+    mockQuery = {
+      ...mockQuery,
+      refId: '121314',
+      serviceName: '$service',
+      spanName: '$span',
+    };
+
+    render(
+      <NativeSearch datasource={{} as TempoDatasource} query={mockQuery} onChange={() => {}} onRunQuery={() => {}} />
+    );
+
+    const asyncServiceSelect = screen.getByRole('combobox', { name: 'select-service-name' });
+    expect(asyncServiceSelect).toBeInTheDocument();
+    await user.click(asyncServiceSelect);
+    jest.advanceTimersByTime(3000);
+
+    await user.type(asyncServiceSelect, '$');
+    var serviceOption = await screen.findByText('$service');
+    expect(serviceOption).toBeDefined();
+
+    const asyncSpanSelect = screen.getByRole('combobox', { name: 'select-span-name' });
+    expect(asyncSpanSelect).toBeInTheDocument();
+    await user.click(asyncSpanSelect);
+    jest.advanceTimersByTime(3000);
+
+    await user.type(asyncSpanSelect, '$');
+    var operationOption = await screen.findByText('$span');
+    expect(operationOption).toBeDefined();
+  });
 });

public/app/plugins/datasource/tempo/QueryEditor/NativeSearch.tsx

@@ -3,7 +3,7 @@ import Prism from 'prismjs';
 import React, { useCallback, useState, useEffect, useMemo } from 'react';
 import { Node } from 'slate';
 
-import { GrafanaTheme2, isValidGoDuration, SelectableValue } from '@grafana/data';
+import { GrafanaTheme2, isValidGoDuration, SelectableValue, toOption } from '@grafana/data';
 import { FetchError, getTemplateSrv, isFetchError, TemplateSrv } from '@grafana/runtime';
 import {
   InlineFieldRow,
@@ -90,7 +90,13 @@ const NativeSearch = ({ datasource, query, onChange, onBlur, onRunQuery }: Props
     const fetchOptions = async () => {
       try {
         const [services, spans] = await Promise.all([loadOptions('serviceName'), loadOptions('spanName')]);
+        if (query.serviceName && getTemplateSrv().containsTemplate(query.serviceName)) {
+          services.push(toOption(query.serviceName));
+        }
         setServiceOptions(services);
+        if (query.spanName && getTemplateSrv().containsTemplate(query.spanName)) {
+          spans.push(toOption(query.spanName));
+        }
         setSpanOptions(spans);
       } catch (error) {
         // Display message if Tempo is connected but search 404's
@@ -102,7 +108,7 @@ const NativeSearch = ({ datasource, query, onChange, onBlur, onRunQuery }: Props
       }
     };
     fetchOptions();
-  }, [languageProvider, loadOptions]);
+  }, [languageProvider, loadOptions, query.serviceName, query.spanName]);
 
   useEffect(() => {
     const fetchTags = async () => {
@@ -161,6 +167,7 @@ const NativeSearch = ({ datasource, query, onChange, onBlur, onRunQuery }: Props
               isClearable
               onKeyDown={onKeyDown}
               aria-label={'select-service-name'}
+              allowCustomValue={true}
             />
           </InlineField>
         </InlineFieldRow>
@@ -184,6 +191,7 @@ const NativeSearch = ({ datasource, query, onChange, onBlur, onRunQuery }: Props
               isClearable
               onKeyDown={onKeyDown}
               aria-label={'select-span-name'}
+              allowCustomValue={true}
             />
           </InlineField>
         </InlineFieldRow>

public/app/plugins/datasource/tempo/datasource.test.ts

@@ -77,9 +77,11 @@ describe('Tempo data source', () => {
       const queries = ds.interpolateVariablesInQueries([getQuery()], {
         interpolationVar: { text: text, value: text },
       });
-      expect(templateSrv.replace).toBeCalledTimes(5);
+      expect(templateSrv.replace).toBeCalledTimes(7);
       expect(queries[0].linkedQuery?.expr).toBe(text);
       expect(queries[0].query).toBe(text);
+      expect(queries[0].serviceName).toBe(text);
+      expect(queries[0].spanName).toBe(text);
       expect(queries[0].search).toBe(text);
       expect(queries[0].minDuration).toBe(text);
       expect(queries[0].maxDuration).toBe(text);
@@ -94,9 +96,11 @@ describe('Tempo data source', () => {
       const resp = ds.applyTemplateVariables(getQuery(), {
         interpolationVar: { text: text, value: text },
       });
-      expect(templateSrv.replace).toBeCalledTimes(5);
+      expect(templateSrv.replace).toBeCalledTimes(7);
       expect(resp.linkedQuery?.expr).toBe(text);
       expect(resp.query).toBe(text);
+      expect(resp.serviceName).toBe(text);
+      expect(resp.spanName).toBe(text);
       expect(resp.search).toBe(text);
       expect(resp.minDuration).toBe(text);
       expect(resp.maxDuration).toBe(text);

public/app/plugins/datasource/tempo/datasource.ts

@@ -296,6 +296,8 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
     return {
       ...expandedQuery,
       query: this.templateSrv.replace(query.query ?? '', scopedVars),
+      serviceName: this.templateSrv.replace(query.serviceName ?? '', scopedVars),
+      spanName: this.templateSrv.replace(query.spanName ?? '', scopedVars),
       search: this.templateSrv.replace(query.search ?? '', scopedVars),
       minDuration: this.templateSrv.replace(query.minDuration ?? '', scopedVars),
       maxDuration: this.templateSrv.replace(query.maxDuration ?? '', scopedVars),

public/app/plugins/datasource/zipkin/datasource.test.ts

@@ -2,6 +2,7 @@ import { lastValueFrom, of } from 'rxjs';
 import { createFetchResponse } from 'test/helpers/createFetchResponse';
 
 import { DataSourceInstanceSettings, FieldType } from '@grafana/data';
+import { TemplateSrv } from '@grafana/runtime';
 import { backendSrv } from 'app/core/services/backend_srv';
 
 import { ZipkinDatasource } from './datasource';
@@ -15,16 +16,24 @@ jest.mock('@grafana/runtime', () => ({
 
 describe('ZipkinDatasource', () => {
   describe('query', () => {
+    const templateSrv: TemplateSrv = {
+      replace: jest.fn(),
+      getVariables: jest.fn(),
+      containsTemplate: jest.fn(),
+      updateTimeRange: jest.fn(),
+    };
+
     it('runs query', async () => {
       setupBackendSrv(zipkinResponse);
-      const ds = new ZipkinDatasource(defaultSettings);
+      const ds = new ZipkinDatasource(defaultSettings, templateSrv);
       await expect(ds.query({ targets: [{ query: '12345' }] } as any)).toEmitValuesWith((val) => {
         expect(val[0].data[0].fields).toMatchObject(traceFrameFields);
       });
     });
+
     it('runs query with traceId that includes special characters', async () => {
       setupBackendSrv(zipkinResponse);
-      const ds = new ZipkinDatasource(defaultSettings);
+      const ds = new ZipkinDatasource(defaultSettings, templateSrv);
       await expect(ds.query({ targets: [{ query: 'a/b' }] } as any)).toEmitValuesWith((val) => {
         expect(val[0].data[0].fields).toMatchObject(traceFrameFields);
       });

public/app/plugins/datasource/zipkin/datasource.ts

@@ -9,8 +9,9 @@ import {
   DataSourceJsonData,
   FieldType,
   MutableDataFrame,
+  ScopedVars,
 } from '@grafana/data';
-import { BackendSrvRequest, FetchResponse, getBackendSrv } from '@grafana/runtime';
+import { BackendSrvRequest, FetchResponse, getBackendSrv, getTemplateSrv, TemplateSrv } from '@grafana/runtime';
 import { SpanBarOptions } from '@jaegertracing/jaeger-ui-components';
 import { NodeGraphOptions } from 'app/core/components/NodeGraphSettings';
 
@@ -29,7 +30,10 @@ export class ZipkinDatasource extends DataSourceApi<ZipkinQuery, ZipkinJsonData>
   uploadedJson: string | ArrayBuffer | null = null;
   nodeGraph?: NodeGraphOptions;
   spanBar?: SpanBarOptions;
-  constructor(private instanceSettings: DataSourceInstanceSettings<ZipkinJsonData>) {
+  constructor(
+    private instanceSettings: DataSourceInstanceSettings<ZipkinJsonData>,
+    private readonly templateSrv: TemplateSrv = getTemplateSrv()
+  ) {
     super(instanceSettings);
     this.nodeGraph = instanceSettings.jsonData.nodeGraph;
   }
@@ -50,7 +54,8 @@ export class ZipkinDatasource extends DataSourceApi<ZipkinQuery, ZipkinJsonData>
     }
 
     if (target.query) {
-      return this.request<ZipkinSpan[]>(`${apiPrefix}/trace/${encodeURIComponent(target.query)}`).pipe(
+      const query = this.applyVariables(target, options.scopedVars);
+      return this.request<ZipkinSpan[]>(`${apiPrefix}/trace/${encodeURIComponent(query.query)}`).pipe(
         map((res) => responseToDataQueryResponse(res, this.nodeGraph?.enabled))
       );
     }
@@ -71,6 +76,29 @@ export class ZipkinDatasource extends DataSourceApi<ZipkinQuery, ZipkinJsonData>
     return query.query;
   }
 
+  interpolateVariablesInQueries(queries: ZipkinQuery[], scopedVars: ScopedVars): ZipkinQuery[] {
+    if (!queries || queries.length === 0) {
+      return [];
+    }
+
+    return queries.map((query) => {
+      return {
+        ...query,
+        datasource: this.getRef(),
+        ...this.applyVariables(query, scopedVars),
+      };
+    });
+  }
+
+  applyVariables(query: ZipkinQuery, scopedVars: ScopedVars) {
+    const expandedQuery = { ...query };
+
+    return {
+      ...expandedQuery,
+      query: this.templateSrv.replace(query.query ?? '', scopedVars),
+    };
+  }
+
   private request<T = any>(
     apiUrl: string,
     data?: any,