[ENHANCEMENT] TempoTraceQuery: support variable interpolation

tempo/src/plugins/tempo-trace-query/get-trace-data.test.ts

@@ -12,13 +12,23 @@
 // limitations under the License.
 
 import { DatasourceSpec } from '@perses-dev/core';
-import { TraceQueryContext } from '@perses-dev/plugin-system';
+import { TraceQueryContext, replaceVariables } from '@perses-dev/plugin-system';
 import { TempoDatasourceSpec } from '../tempo-datasource-types';
 import { TempoDatasource } from '../tempo-datasource';
 import { DEFAULT_SEARCH_LIMIT, SearchResponse } from '../../model/api-types';
 import { TempoClient } from '../../model';
 import { getTraceData } from './get-trace-data';
 
+jest.mock('@perses-dev/plugin-system', () => {
+  const actual = jest.requireActual('@perses-dev/plugin-system');
+  return {
+    ...actual,
+    replaceVariables: jest.fn((query: string) => query),
+  };
+});
+
+const mockedReplaceVariables = replaceVariables as jest.MockedFunction<typeof replaceVariables>;
+
 const datasource: TempoDatasourceSpec = {
   directUrl: '/test',
 };
@@ -60,6 +70,11 @@ const createStubContext = (mockClient: ReturnType<typeof createMockClient>): Tra
 };
 
 describe('getTraceData', () => {
+  beforeEach(() => {
+    mockedReplaceVariables.mockReset();
+    mockedReplaceVariables.mockImplementation((query: string) => query);
+  });
+
   it('should fetch DEFAULT_SEARCH_LIMIT+1 results and not show notice when results are within limit', async () => {
     const mockResponse: SearchResponse = {
       traces: [
@@ -132,4 +147,101 @@ describe('getTraceData', () => {
     // Verify results are trimmed to exactly customLimit
     expect(result.searchResult).toHaveLength(customLimit);
   });
+
+  it('should call replaceVariables with the query and variableState', async () => {
+    const mockResponse: SearchResponse = {
+      traces: [],
+    };
+    const mockClient = createMockClient(mockResponse);
+    const variableState = {
+      serviceName: { value: 'frontend', loading: false },
+    };
+    const stubContext = createStubContext(mockClient);
+    stubContext.variableState = variableState;
+
+    await getTraceData({ query: '{resource.service.name="$serviceName"}' }, stubContext);
+
+    expect(mockedReplaceVariables).toHaveBeenCalledWith('{resource.service.name="$serviceName"}', variableState);
+  });
+
+  it('should use the replaced query in the search request', async () => {
+    const mockResponse: SearchResponse = {
+      traces: [],
+    };
+    const mockClient = createMockClient(mockResponse);
+    const stubContext = createStubContext(mockClient);
+    stubContext.variableState = {
+      serviceName: { value: 'frontend', loading: false },
+    };
+
+    mockedReplaceVariables.mockReturnValueOnce('{resource.service.name="frontend"}');
+
+    const result = await getTraceData({ query: '{resource.service.name="$serviceName"}' }, stubContext);
+
+    expect(mockClient.searchWithFallback).toHaveBeenCalledWith(
+      expect.objectContaining({
+        q: '{resource.service.name="frontend"}',
+      })
+    );
+    expect(result.metadata?.executedQueryString).toBe('{resource.service.name="frontend"}');
+  });
+
+  it('should use the replaced query when it resolves to a valid traceId', async () => {
+    const traceId = 'a'.repeat(32); // valid 32-char hex trace ID
+    const mockResponse = {
+      batches: [],
+    };
+    const mockClient = createMockClient({ traces: [] });
+    mockClient.query = jest.fn(async () => mockResponse);
+    const stubContext = createStubContext(mockClient);
+    stubContext.variableState = {
+      traceId: { value: traceId, loading: false },
+    };
+
+    mockedReplaceVariables.mockReturnValueOnce(traceId);
+
+    const result = await getTraceData({ query: '$traceId' }, stubContext);
+
+    expect(mockedReplaceVariables).toHaveBeenCalledWith('$traceId', stubContext.variableState);
+    expect(mockClient.query).toHaveBeenCalledWith({ traceId });
+    expect(result.trace).toBeDefined();
+  });
+
+  it('should replace multiple variables in a complex TraceQL query', async () => {
+    const mockResponse: SearchResponse = {
+      traces: [
+        {
+          traceID: 'trace1',
+          rootServiceName: 'frontend',
+          rootTraceName: 'GET /api/users',
+          startTimeUnixNano: '1718122135898442804',
+          durationMs: 350,
+        },
+      ],
+    };
+    const mockClient = createMockClient(mockResponse);
+    const stubContext = createStubContext(mockClient);
+    stubContext.variableState = {
+      serviceName: { value: 'frontend', loading: false },
+      minDuration: { value: '100ms', loading: false },
+      httpMethod: { value: 'GET', loading: false },
+    };
+
+    const rawQuery =
+      '{resource.service.name="$serviceName" && span.http.method="$httpMethod" && duration > $minDuration}';
+    const replacedQuery = '{resource.service.name="frontend" && span.http.method="GET" && duration > 100ms}';
+
+    mockedReplaceVariables.mockReturnValueOnce(replacedQuery);
+
+    const result = await getTraceData({ query: rawQuery }, stubContext);
+
+    expect(mockedReplaceVariables).toHaveBeenCalledWith(rawQuery, stubContext.variableState);
+    expect(mockClient.searchWithFallback).toHaveBeenCalledWith(
+      expect.objectContaining({
+        q: replacedQuery,
+      })
+    );
+    expect(result.metadata?.executedQueryString).toBe(replacedQuery);
+    expect(result.searchResult).toHaveLength(1);
+  });
 });

tempo/src/plugins/tempo-trace-query/get-trace-data.ts

@@ -12,7 +12,7 @@
 // limitations under the License.
 
 import { AbsoluteTimeRange, isValidTraceId, Notice, otlptracev1, TraceSearchResult } from '@perses-dev/core';
-import { datasourceSelectValueToSelector, TraceQueryPlugin } from '@perses-dev/plugin-system';
+import { datasourceSelectValueToSelector, replaceVariables, TraceQueryPlugin } from '@perses-dev/plugin-system';
 import { getUnixTime } from 'date-fns';
 import {
   TEMPO_DATASOURCE_KIND,
@@ -40,6 +40,8 @@ export const getTraceData: TraceQueryPlugin<TempoTraceQuerySpec>['getTraceData']
     return { searchResult: [] };
   }
 
+  const query = replaceVariables(spec.query, context.variableState);
+
   const defaultTempoDatasource: TempoDatasourceSelector = {
     kind: TEMPO_DATASOURCE_KIND,
   };
@@ -56,17 +58,17 @@ export const getTraceData: TraceQueryPlugin<TempoTraceQuerySpec>['getTraceData']
    * if the query is a valid traceId, fetch the trace by traceId
    * otherwise, execute a TraceQL query
    */
-  if (isValidTraceId(spec.query)) {
-    const response = await client.query({ traceId: spec.query });
+  if (isValidTraceId(query)) {
+    const response = await client.query({ traceId: query });
     return {
       trace: parseTraceResponse(response),
       metadata: {
-        executedQueryString: spec.query,
+        executedQueryString: query,
       },
     };
   } else {
     const params: SearchRequestParameters = {
-      q: spec.query,
+      q: query,
     };
 
     // handle time range selection from UI drop down (e.g. last 5 minutes, last 1 hour )
@@ -100,7 +102,7 @@ export const getTraceData: TraceQueryPlugin<TempoTraceQuerySpec>['getTraceData']
     return {
       searchResult,
       metadata: {
-        executedQueryString: spec.query,
+        executedQueryString: query,
         hasMoreResults,
         notices,
       },