From a6664ae08d853c54c780678bc301e278c5dadc98 Mon Sep 17 00:00:00 2001 From: Gabriel Bernal Date: Tue, 7 Apr 2026 17:20:55 +0200 Subject: [PATCH] [ENHANCEMENT] allow variable replacement on tempo queries Signed-off-by: Gabriel Bernal --- .../tempo-trace-query/get-trace-data.test.ts | 114 +++++++++++++++++- .../tempo-trace-query/get-trace-data.ts | 14 ++- 2 files changed, 121 insertions(+), 7 deletions(-) diff --git a/tempo/src/plugins/tempo-trace-query/get-trace-data.test.ts b/tempo/src/plugins/tempo-trace-query/get-trace-data.test.ts index c0d7c05d..2062a427 100644 --- a/tempo/src/plugins/tempo-trace-query/get-trace-data.test.ts +++ b/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; + const datasource: TempoDatasourceSpec = { directUrl: '/test', }; @@ -60,6 +70,11 @@ const createStubContext = (mockClient: ReturnType): 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); + }); }); diff --git a/tempo/src/plugins/tempo-trace-query/get-trace-data.ts b/tempo/src/plugins/tempo-trace-query/get-trace-data.ts index 240b784d..c606511b 100644 --- a/tempo/src/plugins/tempo-trace-query/get-trace-data.ts +++ b/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['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['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['getTraceData'] return { searchResult, metadata: { - executedQueryString: spec.query, + executedQueryString: query, hasMoreResults, notices, },