From 02cb2ec6b0a978659a0ea2c2b02d02707d0daea7 Mon Sep 17 00:00:00 2001 From: Tommy Li Date: Mon, 5 Oct 2026 16:28:16 -0700 Subject: [PATCH 1/5] add query attribution for PromQL queries --- .changeset/promql-query-attribution.md | 7 + agent_docs/observability.md | 14 ++ .../__tests__/queryAttribution.int.test.ts | 123 ++++++++++++++++++ .../api/__tests__/prometheus.int.test.ts | 43 ++++++ packages/api/src/routers/api/prometheus.ts | 65 ++++++++- .../app/src/__tests__/prometheusApi.test.ts | 27 ++++ packages/app/src/api.ts | 21 ++- .../useOffsetPaginatedQuery.test.tsx | 1 + packages/app/src/hooks/useChartConfig.tsx | 9 +- .../app/src/hooks/useOffsetPaginatedQuery.tsx | 1 + .../utils/__tests__/promqlChartQuery.test.ts | 21 +++ packages/app/src/utils/promqlChartQuery.ts | 23 +++- .../clickhouse/__tests__/attribution.test.ts | 38 ++++++ .../src/clickhouse/attribution.ts | 44 +++++++ packages/common-utils/src/clickhouse/index.ts | 2 + 15 files changed, 429 insertions(+), 10 deletions(-) create mode 100644 .changeset/promql-query-attribution.md diff --git a/.changeset/promql-query-attribution.md b/.changeset/promql-query-attribution.md new file mode 100644 index 0000000000..7c0e0cedc9 --- /dev/null +++ b/.changeset/promql-query-attribution.md @@ -0,0 +1,7 @@ +--- +'@hyperdx/common-utils': minor +'@hyperdx/api': minor +'@hyperdx/app': minor +--- + +feat: Attribute PromQL queries to the dashboard, tile or search that issued them, through both the `prometheusQuery`/`prometheusQueryRange` table functions and ClickHouse's Prometheus HTTP API diff --git a/agent_docs/observability.md b/agent_docs/observability.md index 2136eb9338..cfdde77b8f 100644 --- a/agent_docs/observability.md +++ b/agent_docs/observability.md @@ -294,6 +294,20 @@ everything below it, including the materialized-view `EXPLAIN` probes. Order of precedence: the client's default, then the scope around the current request or job, then anything passed with an individual query. +**Browser to API route** — for requests that reach ClickHouse through an API +route instead of a ClickHouse client. PromQL is the case today: the browser +calls `/api/v1/prometheus/query{,_range}` over fetch. The caller sends +`buildLogComment(attribution)` in the `x-hyperdx-query-attribution` header +(`QUERY_ATTRIBUTION_HEADER`), and the router parses it with `parseLogComment` +into a `withQueryAttribution` scope layered over the auth middleware's `api` +scope. The browser's surface and ids win; `trace` and `label` stay the +server's. Queries the route runs through a ClickHouse client (the +`prometheusQuery()`/`prometheusQueryRange()` fallback) pick the scope up as +usual. Requests proxied to ClickHouse's Prometheus HTTP API on 26.6+ are tagged +by `clickhousePrometheusRequest`: `log_comment` as a URL setting, `query_id` as +the `X-ClickHouse-Query-Id` header, because the handler rejects `query_id` as an +unknown setting in the URL. + ### Rules - **Values from the browser are hints, not proof.** The client controls them. diff --git a/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts b/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts index 1fe4ae081f..be97617dd5 100644 --- a/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts +++ b/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts @@ -1,3 +1,6 @@ +import { buildLogComment } from '@hyperdx/common-utils/dist/clickhouse'; +import { Metadata } from '@hyperdx/common-utils/dist/core/metadata'; + import * as config from '@/config'; import { DEFAULT_DATABASE, @@ -9,6 +12,9 @@ import { } from '@/fixtures'; import Connection from '@/models/connection'; +// Set by jest.setup.ts before it stubs `fetch`. +const { realFetch } = globalThis as unknown as { realFetch: typeof fetch }; + // Covers the wiring the unit tests can't: auth middleware opens the // AsyncLocalStorage scope, and a query issued deep inside the route handler // still reads it. If that scope is lost, queries fall back to hdx-unknown-*. @@ -29,11 +35,13 @@ describe('request query attribution', () => { endSec: 1700003600, }, ], + withSamples: true, }); }); afterEach(async () => { await server.clearDBs(); + jest.restoreAllMocks(); }); afterAll(async () => { @@ -66,6 +74,29 @@ describe('request query attribution', () => { return []; } + // As findRouteQueries, for the PromQL evaluation a tile's request ran. + async function findTileQueries(queryText: string, tile: string) { + const client = await getTestFixtureClickHouseClient(); + for (let attempt = 0; attempt < 10; attempt++) { + await client.command({ query: 'SYSTEM FLUSH LOGS' }); + const result = await client.query({ + query: ` + SELECT query_id, log_comment + FROM system.query_log + WHERE type = 'QueryFinish' + AND position(query, {queryText:String}) > 0 + AND JSONExtractString(log_comment, 'tile') = {tile:String}`, + query_params: { queryText, tile }, + format: 'JSON', + }); + const { data }: { data: { query_id: string; log_comment: string }[] } = + await result.json(); + if (data.length > 0) return data; + await new Promise(resolve => setTimeout(resolve, 250)); + } + return []; + } + it('tags queries from an authenticated route with the api surface', async () => { const { agent, team } = await getLoggedInAgent(server); const conn = await Connection.create({ @@ -95,4 +126,96 @@ describe('request query attribution', () => { }); } }); + + it('tags PromQL table-function queries with the attribution the browser sent', async () => { + // Below 26.6 the router evaluates through prometheusQuery() rather than + // proxying to ClickHouse's Prometheus HTTP API. + jest + .spyOn(Metadata.prototype, 'getServerVersion') + .mockResolvedValue([26, 5, 0, 0]); + const { agent, team } = await getLoggedInAgent(server); + const conn = await Connection.create({ + team: team._id, + name: 'CH', + host: config.CLICKHOUSE_HOST, + username: config.CLICKHOUSE_USER, + password: config.CLICKHOUSE_PASSWORD, + }); + const tile = `tile-${Date.now()}`; + + await agent + .get('/v1/prometheus/query') + .set( + 'x-hyperdx-query-attribution', + buildLogComment({ + surface: 'dashboard', + dashboard: 'dash-1', + tile, + label: 'from-browser', + trace: 'from-browser', + })!, + ) + .query({ + query: 'attribution_metric', + time: '1700003600', + connectionId: conn._id.toString(), + database: DEFAULT_DATABASE, + table: TABLE, + }) + .expect(200); + + const logged = await findTileQueries('prometheusQuery(', tile); + + expect(logged).toHaveLength(1); + expect(logged[0].query_id).toMatch(/^hdx-dashboard-/); + expect(JSON.parse(logged[0].log_comment)).toMatchObject({ + surface: 'dashboard', + dashboard: 'dash-1', + tile, + // The server's, not the browser's. + label: '/v1/prometheus', + }); + expect(JSON.parse(logged[0].log_comment).trace).not.toBe('from-browser'); + }); + + // The CI ClickHouse serves prometheus_api_v1, which evaluates the PromQL + // itself: this pins that it honours the log_comment and query_id we send. + it('tags PromQL proxied to the prometheus_api_v1 handler', async () => { + // A spy, so the afterEach restoreAllMocks puts the stub back. + jest.spyOn(global, 'fetch').mockImplementation(realFetch); + const { agent, team } = await getLoggedInAgent(server); + const conn = await Connection.create({ + team: team._id, + name: 'CH', + host: config.CLICKHOUSE_HOST, + username: config.CLICKHOUSE_USER, + password: config.CLICKHOUSE_PASSWORD, + }); + const tile = `tile-${Date.now()}`; + + await agent + .get('/v1/prometheus/query') + .set( + 'x-hyperdx-query-attribution', + buildLogComment({ surface: 'dashboard', dashboard: 'dash-1', tile })!, + ) + .query({ + query: 'attribution_metric', + time: '1700003600', + connectionId: conn._id.toString(), + database: DEFAULT_DATABASE, + table: TABLE, + }) + .expect(200); + + const logged = await findTileQueries('prometheus_query_step', tile); + expect(logged).toHaveLength(1); + expect(logged[0].query_id).toMatch(/^hdx-dashboard-/); + expect(JSON.parse(logged[0].log_comment)).toMatchObject({ + surface: 'dashboard', + dashboard: 'dash-1', + tile, + label: '/v1/prometheus', + }); + }); }); diff --git a/packages/api/src/routers/api/__tests__/prometheus.int.test.ts b/packages/api/src/routers/api/__tests__/prometheus.int.test.ts index d7dbf94255..1033e270d8 100644 --- a/packages/api/src/routers/api/__tests__/prometheus.int.test.ts +++ b/packages/api/src/routers/api/__tests__/prometheus.int.test.ts @@ -1,3 +1,7 @@ +import { + buildLogComment, + QUERY_ATTRIBUTION_HEADER, +} from '@hyperdx/common-utils/dist/clickhouse'; import { Metadata } from '@hyperdx/common-utils/dist/core/metadata'; import { Types } from 'mongoose'; @@ -749,6 +753,45 @@ describe('prometheus router', () => { }); }); + it('tags a request proxied to prometheus_api_v1 with log_comment and query_id', async () => { + const { agent, team } = await getLoggedInAgent(server); + const conn = await seedClickHouseConnection(team._id); + mockFetch.mockResolvedValueOnce(handlerProbeResponse()); + + await agent + .get('/v1/prometheus/query_range') + .set( + QUERY_ATTRIBUTION_HEADER, + buildLogComment({ surface: 'dashboard', tile: 'tile-1' })!, + ) + .query({ + query: 'up', + start: '1700000000', + end: '1700000060', + step: '15s', + table: 'metrics_ts', + // Not caller-settable, so this never reaches the upstream. + log_comment: 'from-caller', + connectionId: conn._id.toString(), + }) + .expect(200); + + const [calledUrl, init] = mockFetch.mock.calls[1]; + const url = new URL(String(calledUrl)); + expect(url.searchParams.getAll('log_comment')).toHaveLength(1); + expect(JSON.parse(url.searchParams.get('log_comment')!)).toMatchObject({ + surface: 'dashboard', + tile: 'tile-1', + label: '/v1/prometheus', + }); + // In the URL, ClickHouse fails the request on query_id as an unknown + // setting. + expect(url.searchParams.has('query_id')).toBe(false); + expect(init?.headers).toMatchObject({ + 'X-ClickHouse-Query-Id': expect.stringMatching(/^hdx-dashboard-/), + }); + }); + it('does not let the request loosen the pinned ClickHouse limits', async () => { const { agent, team } = await getLoggedInAgent(server); const conn = await seedClickHouseConnection(team._id); diff --git a/packages/api/src/routers/api/prometheus.ts b/packages/api/src/routers/api/prometheus.ts index ca88281827..60238868e5 100644 --- a/packages/api/src/routers/api/prometheus.ts +++ b/packages/api/src/routers/api/prometheus.ts @@ -1,3 +1,13 @@ +import { + buildLogComment, + buildQueryId, + parseLogComment, + QUERY_ATTRIBUTION_HEADER, +} from '@hyperdx/common-utils/dist/clickhouse'; +import { + getCurrentQueryAttribution, + withQueryAttribution, +} from '@hyperdx/common-utils/dist/clickhouse/node'; import express from 'express'; import { performance } from 'perf_hooks'; import { Readable } from 'stream'; @@ -53,6 +63,20 @@ router.use((_req, res, next) => { next(); }); +// The browser reaches this router over fetch, not through a ClickHouse client, +// so it sends the tile, dashboard or search a query belongs to in a header. +// Its surface and ids win over the `api` scope the auth middleware opened, but +// the trace id and route label are the server's to state: a client-sent trace +// would link the query_log row to a request that never ran it. +router.use((req, _res, next) => { + const attribution = parseLogComment(req.get(QUERY_ATTRIBUTION_HEADER)); + if (!attribution) return next(); + // Dropped rather than overwritten: with no active span there is no server + // trace to overwrite with, and the browser's would survive the merge. + const { label: _label, trace: _trace, ...fromBrowser } = attribution; + withQueryAttribution(fromBrowser, () => next()); +}); + // -------------------------- // Param parsing helpers // -------------------------- @@ -210,14 +234,45 @@ export function clickhousePrometheusUpstream( { maxExecutionSec = PROMETHEUS_MAX_EXECUTION_SEC, maxResultRows = PROMETHEUS_MAX_RESULT_ROWS, + logComment, + }: { + maxExecutionSec?: number; + maxResultRows?: number; + logComment?: string; } = {}, ): string { const url = new URL(host); url.searchParams.set('max_execution_time', String(maxExecutionSec)); url.searchParams.set('max_result_rows', String(maxResultRows)); + if (logComment) url.searchParams.set('log_comment', logComment); return url.toString(); } +/** + * Where and how to proxy a PromQL request to the connection's prometheus_api_v1 + * handler, tagged with the query attribution in scope. The handler runs the + * query itself rather than through our ClickHouse client, so the tags are set + * here. `log_comment` goes in the URL like the limits. `query_id` goes in a + * header: the handler reads every URL param as a setting, and `query_id` in + * the URL fails the request as an unknown one. + */ +function clickhousePrometheusRequest(connection: { + host: string; + username: string; + password?: string; +}): { upstream: string; headers: Record } { + const attribution = getCurrentQueryAttribution(); + return { + upstream: clickhousePrometheusUpstream(connection.host, { + logComment: buildLogComment(attribution), + }), + headers: { + ...clickhouseAuthHeaders(connection), + 'X-ClickHouse-Query-Id': buildQueryId(attribution), + }, + }; +} + function newClickhouseClient(connection: { host: string; username: string; @@ -537,12 +592,13 @@ const queryRangeHandler: express.RequestHandler = async (req, res) => { // forward-compatible while TimeSeries is in preview, so it is preferred // wherever the server has it (26.6+); the table function is the fallback. if (await clickhouseServesPrometheusHttpApi(client, connection)) { + const { upstream, headers } = clickhousePrometheusRequest(connection); const status = await proxyToPrometheus( - clickhousePrometheusUpstream(connection.host), + upstream, `${CLICKHOUSE_PROMETHEUS_API_PREFIX}/query_range`, params, res, - clickhouseAuthHeaders(connection), + headers, ); recordProxyOutcome(status, 'query_range', backend); return; @@ -648,12 +704,13 @@ const queryHandler: express.RequestHandler = async (req, res) => { const client = newClickhouseClient(connection); if (await clickhouseServesPrometheusHttpApi(client, connection)) { + const { upstream, headers } = clickhousePrometheusRequest(connection); const status = await proxyToPrometheus( - clickhousePrometheusUpstream(connection.host), + upstream, `${CLICKHOUSE_PROMETHEUS_API_PREFIX}/query`, params, res, - clickhouseAuthHeaders(connection), + headers, ); recordProxyOutcome(status, 'query', backend); return; diff --git a/packages/app/src/__tests__/prometheusApi.test.ts b/packages/app/src/__tests__/prometheusApi.test.ts index 26b3da6f16..3e658b1567 100644 --- a/packages/app/src/__tests__/prometheusApi.test.ts +++ b/packages/app/src/__tests__/prometheusApi.test.ts @@ -162,6 +162,33 @@ describe('prometheusApi.query', () => { expect(params.has('table')).toBe(false); }); + it('sends the attribution as a log_comment payload in a header', async () => { + await prometheusApi.query({ + query: 'up', + time: 1, + connectionId: 'conn', + attribution: { surface: 'dashboard', dashboard: 'd1', tile: 't1' }, + }); + + const { headers } = post.mock.calls[0][1]; + expect(JSON.parse(headers['x-hyperdx-query-attribution'])).toEqual({ + v: 1, + surface: 'dashboard', + dashboard: 'd1', + tile: 't1', + }); + }); + + it('sends no attribution header when there is nothing to say', async () => { + await prometheusApi.query({ + query: 'up', + time: 1, + connectionId: 'conn', + attribution: {}, + }); + expect(post.mock.calls[0][1].headers).toBeUndefined(); + }); + it('reports the error a failed instant query carries', async () => { post.mockReturnValue({ json: () => diff --git a/packages/app/src/api.ts b/packages/app/src/api.ts index 08a83e0a32..0be4bb6ea9 100644 --- a/packages/app/src/api.ts +++ b/packages/app/src/api.ts @@ -2,6 +2,11 @@ import { useCallback } from 'react'; import Router from 'next/router'; import type { HTTPError, Options, ResponsePromise } from 'ky'; import ky from 'ky-universal'; +import { + buildLogComment, + QUERY_ATTRIBUTION_HEADER, + type QueryAttribution, +} from '@hyperdx/common-utils/dist/clickhouse'; import type { Alert, AlertApiResponse, @@ -866,10 +871,16 @@ const prometheusFetch = ( path: string, searchParams: Record, signal?: AbortSignal, -): Promise => - withPrometheusError(() => - server.post(path, { searchParams, signal }).json(), + attribution?: QueryAttribution, +): Promise => { + const logComment = buildLogComment(attribution); + const headers = logComment + ? { [QUERY_ATTRIBUTION_HEADER]: logComment } + : undefined; + return withPrometheusError(() => + server.post(path, { searchParams, signal, headers }).json(), ); +}; export const prometheusApi = { queryRange: (params: { @@ -881,6 +892,7 @@ export const prometheusApi = { database?: string; table?: string; signal?: AbortSignal; + attribution?: QueryAttribution; }): Promise => prometheusFetch( 'v1/prometheus/query_range', @@ -894,6 +906,7 @@ export const prometheusApi = { ...(params.table ? { table: params.table } : {}), }, params.signal, + params.attribution, ), query: (params: { @@ -905,6 +918,7 @@ export const prometheusApi = { /** Maximum number of series to return. */ limit?: number; signal?: AbortSignal; + attribution?: QueryAttribution; }): Promise => prometheusFetch( 'v1/prometheus/query', @@ -917,6 +931,7 @@ export const prometheusApi = { ...(params.table ? { table: params.table } : {}), }, params.signal, + params.attribution, ), labels: (params: { diff --git a/packages/app/src/hooks/__tests__/useOffsetPaginatedQuery.test.tsx b/packages/app/src/hooks/__tests__/useOffsetPaginatedQuery.test.tsx index d6389199c0..de5ba12414 100644 --- a/packages/app/src/hooks/__tests__/useOffsetPaginatedQuery.test.tsx +++ b/packages/app/src/hooks/__tests__/useOffsetPaginatedQuery.test.tsx @@ -1225,6 +1225,7 @@ describe('useOffsetPaginatedQuery', () => { promqlConfig, promqlConfig.dateRange, expect.anything(), + {}, ); expect(renderChartConfig).not.toHaveBeenCalled(); expect(mockClickhouseClient.query).not.toHaveBeenCalled(); diff --git a/packages/app/src/hooks/useChartConfig.tsx b/packages/app/src/hooks/useChartConfig.tsx index b9e8e99f88..619859c107 100644 --- a/packages/app/src/hooks/useChartConfig.tsx +++ b/packages/app/src/hooks/useChartConfig.tsx @@ -46,6 +46,7 @@ import { useClickhouseClient } from '@/clickhouse'; import { IS_MTVIEWS_ENABLED } from '@/config'; import { buildMTViewSelectQuery } from '@/hdxMTViews'; import { useMetadataWithSettings } from '@/hooks/useMetadata'; +import { useQueryAttribution } from '@/queryAttribution'; import { useSource } from '@/source'; import { ChartQueryResult } from '@/types'; import { stripClientSideConfigFields } from '@/utils/chartConfig'; @@ -358,6 +359,7 @@ export function useQueriedChartConfig( ) { const { enabled = true } = options ?? {}; const clickhouseClient = useClickhouseClient(); + const attribution = useQueryAttribution(); const queryClient = useQueryClient(); const metadata = useMetadataWithSettings(); @@ -408,7 +410,12 @@ export function useQueriedChartConfig( queryFn: async context => { // PromQL queries go through the Prometheus API route, not ClickHouse proxy if (isPromqlChartConfig(config) && config.dateRange) { - return queryPromqlChartConfig(config, config.dateRange, context.signal); + return queryPromqlChartConfig( + config, + config.dateRange, + context.signal, + attribution, + ); } const optimizedConfig = { diff --git a/packages/app/src/hooks/useOffsetPaginatedQuery.tsx b/packages/app/src/hooks/useOffsetPaginatedQuery.tsx index 567c790fcb..e8b5b573f6 100644 --- a/packages/app/src/hooks/useOffsetPaginatedQuery.tsx +++ b/packages/app/src/hooks/useOffsetPaginatedQuery.tsx @@ -232,6 +232,7 @@ const queryFn: QueryFunction = async ({ config, config.dateRange, signal, + attribution, ); return { data, diff --git a/packages/app/src/utils/__tests__/promqlChartQuery.test.ts b/packages/app/src/utils/__tests__/promqlChartQuery.test.ts index b1bcf0c624..6b9ade9d8a 100644 --- a/packages/app/src/utils/__tests__/promqlChartQuery.test.ts +++ b/packages/app/src/utils/__tests__/promqlChartQuery.test.ts @@ -271,6 +271,7 @@ describe('queryPromqlChartConfig', () => { }, dateRange, new AbortController().signal, + {}, ); expect(prometheusApi.query).toHaveBeenCalledWith( @@ -278,6 +279,25 @@ describe('queryPromqlChartConfig', () => { ); }); + it('passes the attribution through to the request', async () => { + const attribution = { surface: 'dashboard' as const, tile: 'tile-1' }; + await queryPromqlChartConfig( + { + configType: 'promql', + displayType: DisplayType.Table, + connection: 'conn', + promqlExpression: [{ expression: 'up', queryType: 'instant' }], + }, + dateRange, + new AbortController().signal, + attribution, + ); + + expect(prometheusApi.query).toHaveBeenCalledWith( + expect.objectContaining({ attribution }), + ); + }); + // A number tile warns when an expression yields several series, so it needs // to see all of them. it('leaves an instant number query uncapped', async () => { @@ -290,6 +310,7 @@ describe('queryPromqlChartConfig', () => { }, dateRange, new AbortController().signal, + {}, ); expect(prometheusApi.query).toHaveBeenCalledWith( diff --git a/packages/app/src/utils/promqlChartQuery.ts b/packages/app/src/utils/promqlChartQuery.ts index 38949028ac..f523ef0e08 100644 --- a/packages/app/src/utils/promqlChartQuery.ts +++ b/packages/app/src/utils/promqlChartQuery.ts @@ -1,3 +1,4 @@ +import type { QueryAttribution } from '@hyperdx/common-utils/dist/clickhouse'; import { displayTypeSupportsInstantQuery, getQueriedPromqlSeries, @@ -83,6 +84,7 @@ async function fetchInstantExpression( series: PromqlSeries, evalTime: Date, signal: AbortSignal, + attribution: QueryAttribution, ): Promise { const response = await prometheusApi.query({ query: series.expression, @@ -96,6 +98,7 @@ async function fetchInstantExpression( database: config.from?.databaseName, table: config.from?.tableName, signal, + attribution, }); return { @@ -110,6 +113,7 @@ async function fetchRangeExpression( series: PromqlSeries, dateRange: [Date, Date], signal: AbortSignal, + attribution: QueryAttribution, ): Promise { const [startDate, endDate] = dateRange; const response = await prometheusApi.queryRange({ @@ -121,6 +125,7 @@ async function fetchRangeExpression( database: config.from?.databaseName, table: config.from?.tableName, signal, + attribution, }); if (response.status !== 'success' || !response.data) { @@ -260,6 +265,8 @@ export async function queryPromqlChartConfig( config: PromqlChartConfig, dateRange: [Date, Date], signal: AbortSignal, + // Required so a new caller can't forget it and leave its queries untagged. + attribution: QueryAttribution, ): Promise { // Expand dashboard variables and macros before sending to Prometheus. const substituted = substitutePromqlChartConfigTemplates({ @@ -271,8 +278,20 @@ export async function queryPromqlChartConfig( getQueriedPromqlSeries(substituted).map(series => displayTypeSupportsInstantQuery(substituted) && promqlSeriesQueryType(series) === 'instant' - ? fetchInstantExpression(substituted, series, dateRange[1], signal) - : fetchRangeExpression(substituted, series, dateRange, signal), + ? fetchInstantExpression( + substituted, + series, + dateRange[1], + signal, + attribution, + ) + : fetchRangeExpression( + substituted, + series, + dateRange, + signal, + attribution, + ), ), ); diff --git a/packages/common-utils/src/clickhouse/__tests__/attribution.test.ts b/packages/common-utils/src/clickhouse/__tests__/attribution.test.ts index b389259b44..aa17112789 100644 --- a/packages/common-utils/src/clickhouse/__tests__/attribution.test.ts +++ b/packages/common-utils/src/clickhouse/__tests__/attribution.test.ts @@ -5,6 +5,7 @@ import { buildLogComment, buildQueryId, mergeQueryAttribution, + parseLogComment, QUERY_ATTRIBUTION_VERSION, QueryAttribution, QuerySurface, @@ -190,6 +191,43 @@ describe('buildLogComment', () => { }); }); +describe('parseLogComment', () => { + it('round-trips what buildLogComment emits', () => { + const attribution: QueryAttribution = { + surface: 'dashboard', + dashboard: 'dash-1', + tile: 'tile-2', + label: '/a/b', + }; + expect(parseLogComment(buildLogComment(attribution))).toEqual(attribution); + }); + + it('drops unknown keys, non-string values and unknown surfaces', () => { + expect( + parseLogComment( + JSON.stringify({ v: 1, surface: 'nope', tile: 42, extra: 'x' }), + ), + ).toBeUndefined(); + expect( + parseLogComment(JSON.stringify({ surface: 'search', tile: 42 })), + ).toEqual({ surface: 'search' }); + }); + + it('returns undefined for anything that is not a JSON object', () => { + expect(parseLogComment(undefined)).toBeUndefined(); + expect(parseLogComment('')).toBeUndefined(); + expect(parseLogComment('not json')).toBeUndefined(); + expect(parseLogComment('null')).toBeUndefined(); + expect(parseLogComment('"dashboard"')).toBeUndefined(); + }); + + it('rejects a payload over the size cap', () => { + expect( + parseLogComment(JSON.stringify({ label: 'a'.repeat(2000) })), + ).toBeUndefined(); + }); +}); + describe('buildQueryId', () => { it('prefixes with the surface', () => { expect(buildQueryId({ surface: 'alert' })).toMatch(/^hdx-alert-/); diff --git a/packages/common-utils/src/clickhouse/attribution.ts b/packages/common-utils/src/clickhouse/attribution.ts index f3b8f7a251..34b5eaa958 100644 --- a/packages/common-utils/src/clickhouse/attribution.ts +++ b/packages/common-utils/src/clickhouse/attribution.ts @@ -16,6 +16,8 @@ * check) get a `query_id` but an empty `log_comment`. */ +import { z } from 'zod'; + /** * Which part of the product sent the query. */ @@ -207,6 +209,48 @@ export function buildLogComment( return payload.size > 1 ? serialized : undefined; } +/** + * Carries a `buildLogComment` payload on requests that reach ClickHouse through + * an API route rather than a ClickHouse client, such as PromQL. + */ +export const QUERY_ATTRIBUTION_HEADER = 'x-hyperdx-query-attribution'; + +// `.catch` per field, so one bad value drops that field rather than the lot. +const optionalString = z.string().optional().catch(undefined); +const wireAttributionSchema = z.object({ + surface: z.enum(QUERY_SURFACES).optional().catch(undefined), + dashboard: optionalString, + tile: optionalString, + search: optionalString, + alert: optionalString, + source: optionalString, + trace: optionalString, + label: optionalString, +}); + +/** + * The inverse of `buildLogComment`, for a payload that arrived over the wire. + * Unknown keys and non-string values are dropped. Values are not sanitized + * here: they go back through `buildLogComment` before reaching ClickHouse. + */ +export function parseLogComment( + value: string | undefined, +): QueryAttribution | undefined { + if (!value || value.length > MAX_LOG_COMMENT_BYTES) return undefined; + + let json: unknown; + try { + json = JSON.parse(value); + } catch { + return undefined; + } + const parsed = wireAttributionSchema.safeParse(json); + if (!parsed.success) return undefined; + + const attribution = mergeQueryAttribution(parsed.data); + return Object.keys(attribution).length > 0 ? attribution : undefined; +} + function stringify(payload: Map): string { return JSON.stringify(Object.fromEntries(payload)); } diff --git a/packages/common-utils/src/clickhouse/index.ts b/packages/common-utils/src/clickhouse/index.ts index 460c997c62..d0d1974dae 100644 --- a/packages/common-utils/src/clickhouse/index.ts +++ b/packages/common-utils/src/clickhouse/index.ts @@ -54,6 +54,8 @@ export { buildLogComment, buildQueryId, mergeQueryAttribution, + parseLogComment, + QUERY_ATTRIBUTION_HEADER, QUERY_ATTRIBUTION_VERSION, QUERY_SURFACES, type QueryAttribution, From 0f8495f3f59c54166fa4f4136fe0464c4a5cceac Mon Sep 17 00:00:00 2001 From: Tommy Li Date: Mon, 5 Oct 2026 16:45:15 -0700 Subject: [PATCH 2/5] fix lint --- .../src/middleware/__tests__/queryAttribution.int.test.ts | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts b/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts index be97617dd5..be8efccb43 100644 --- a/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts +++ b/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts @@ -12,8 +12,10 @@ import { } from '@/fixtures'; import Connection from '@/models/connection'; -// Set by jest.setup.ts before it stubs `fetch`. -const { realFetch } = globalThis as unknown as { realFetch: typeof fetch }; +declare global { + // Set by jest.setup.ts before it stubs `fetch`. + var realFetch: typeof fetch; +} // Covers the wiring the unit tests can't: auth middleware opens the // AsyncLocalStorage scope, and a query issued deep inside the route handler @@ -182,7 +184,7 @@ describe('request query attribution', () => { // itself: this pins that it honours the log_comment and query_id we send. it('tags PromQL proxied to the prometheus_api_v1 handler', async () => { // A spy, so the afterEach restoreAllMocks puts the stub back. - jest.spyOn(global, 'fetch').mockImplementation(realFetch); + jest.spyOn(global, 'fetch').mockImplementation(globalThis.realFetch); const { agent, team } = await getLoggedInAgent(server); const conn = await Connection.create({ team: team._id, From 01fc10761cfceccdf86ca06eea7e0a69949c8ee0 Mon Sep 17 00:00:00 2001 From: Tommy Li Date: Mon, 5 Oct 2026 17:30:06 -0700 Subject: [PATCH 3/5] reviews --- .../__tests__/queryAttribution.int.test.ts | 17 +++++-- .../src/clickhouse/attribution.ts | 51 +++++++++++-------- 2 files changed, 43 insertions(+), 25 deletions(-) diff --git a/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts b/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts index be8efccb43..67d32a9327 100644 --- a/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts +++ b/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts @@ -41,9 +41,19 @@ describe('request query attribution', () => { }); }); + // jest.setup.ts installs `fetch` as a plain jest.fn(), not a spy, so + // restoreAllMocks can't put it back; and restoreAllMocks would also undo + // the setup's console spies. Each mock is reset by hand instead. + const fetchStub = jest.mocked(global.fetch).getMockImplementation(); + let serverVersionSpy: jest.SpiedFunction< + Metadata['getServerVersion'] + > | null = null; + afterEach(async () => { await server.clearDBs(); - jest.restoreAllMocks(); + jest.mocked(global.fetch).mockImplementation(fetchStub); + serverVersionSpy?.mockRestore(); + serverVersionSpy = null; }); afterAll(async () => { @@ -132,7 +142,7 @@ describe('request query attribution', () => { it('tags PromQL table-function queries with the attribution the browser sent', async () => { // Below 26.6 the router evaluates through prometheusQuery() rather than // proxying to ClickHouse's Prometheus HTTP API. - jest + serverVersionSpy = jest .spyOn(Metadata.prototype, 'getServerVersion') .mockResolvedValue([26, 5, 0, 0]); const { agent, team } = await getLoggedInAgent(server); @@ -183,8 +193,7 @@ describe('request query attribution', () => { // The CI ClickHouse serves prometheus_api_v1, which evaluates the PromQL // itself: this pins that it honours the log_comment and query_id we send. it('tags PromQL proxied to the prometheus_api_v1 handler', async () => { - // A spy, so the afterEach restoreAllMocks puts the stub back. - jest.spyOn(global, 'fetch').mockImplementation(globalThis.realFetch); + jest.mocked(global.fetch).mockImplementation(globalThis.realFetch); const { agent, team } = await getLoggedInAgent(server); const conn = await Connection.create({ team: team._id, diff --git a/packages/common-utils/src/clickhouse/attribution.ts b/packages/common-utils/src/clickhouse/attribution.ts index 34b5eaa958..e9e300de13 100644 --- a/packages/common-utils/src/clickhouse/attribution.ts +++ b/packages/common-utils/src/clickhouse/attribution.ts @@ -80,20 +80,22 @@ const MAX_FIELD_LENGTH = 128; * spare, so the next field added will need this. * * Pairs rather than a list of key names, so nothing below has to read a field - * out of the attribution through a variable key. + * out of the attribution through a variable key. Built from an object literal + * so the `satisfies` fails to compile until a new id field is listed here; + * `Object.entries` keeps the insertion order. */ function idFields( attribution: QueryAttribution, ): readonly (readonly [string, string | undefined])[] { - return [ - ['dashboard', attribution.dashboard], - ['tile', attribution.tile], - ['search', attribution.search], - ['alert', attribution.alert], - ['source', attribution.source], - ['trace', attribution.trace], - ['label', attribution.label], - ]; + return Object.entries({ + dashboard: attribution.dashboard, + tile: attribution.tile, + search: attribution.search, + alert: attribution.alert, + source: attribution.source, + trace: attribution.trace, + label: attribution.label, + } satisfies Record, unknown>); } /** @@ -153,19 +155,22 @@ function isQuerySurface(value: unknown): value is QuerySurface { export function mergeQueryAttribution( ...layers: (QueryAttribution | undefined)[] ): QueryAttribution { - const merged: QueryAttribution = {}; + let merged: QueryAttribution = {}; for (const layer of layers) { if (!layer) continue; // Field by field: a spread would copy an `undefined` over a real value. - if (layer.surface) merged.surface = layer.surface; - if (layer.dashboard) merged.dashboard = layer.dashboard; - if (layer.tile) merged.tile = layer.tile; - if (layer.search) merged.search = layer.search; - if (layer.alert) merged.alert = layer.alert; - if (layer.source) merged.source = layer.source; - if (layer.trace) merged.trace = layer.trace; - if (layer.label) merged.label = layer.label; + // The `satisfies` fails to compile until a new field is merged here. + merged = { + surface: layer.surface || merged.surface, + dashboard: layer.dashboard || merged.dashboard, + tile: layer.tile || merged.tile, + search: layer.search || merged.search, + alert: layer.alert || merged.alert, + source: layer.source || merged.source, + trace: layer.trace || merged.trace, + label: layer.label || merged.label, + } satisfies Record; } return merged; } @@ -216,6 +221,9 @@ export function buildLogComment( export const QUERY_ATTRIBUTION_HEADER = 'x-hyperdx-query-attribution'; // `.catch` per field, so one bad value drops that field rather than the lot. +// The `satisfies` fails to compile when this and `QueryAttribution` disagree on +// a key: zod strips keys it doesn't know, so a field missing here would vanish +// from the header without an error. const optionalString = z.string().optional().catch(undefined); const wireAttributionSchema = z.object({ surface: z.enum(QUERY_SURFACES).optional().catch(undefined), @@ -226,7 +234,7 @@ const wireAttributionSchema = z.object({ source: optionalString, trace: optionalString, label: optionalString, -}); +} satisfies Record); /** * The inverse of `buildLogComment`, for a payload that arrived over the wire. @@ -248,7 +256,8 @@ export function parseLogComment( if (!parsed.success) return undefined; const attribution = mergeQueryAttribution(parsed.data); - return Object.keys(attribution).length > 0 ? attribution : undefined; + // Every key is present after a merge, so check for a value instead. + return Object.values(attribution).some(Boolean) ? attribution : undefined; } function stringify(payload: Map): string { From 4e76cfe8dc5a0ee7272390db83ef0427877be970 Mon Sep 17 00:00:00 2001 From: Tommy Li Date: Tue, 6 Oct 2026 10:03:59 -0700 Subject: [PATCH 4/5] attribution for label queries too --- agent_docs/observability.md | 10 +++--- .../app/src/__tests__/prometheusApi.test.ts | 32 +++++++++++++++++++ packages/app/src/api.ts | 28 +++++++++++----- .../usePromqlLabelFilterValues.test.tsx | 1 + .../src/hooks/usePromqlLabelFilterValues.tsx | 3 ++ packages/app/src/hooks/usePromqlMetadata.ts | 5 +++ 6 files changed, 67 insertions(+), 12 deletions(-) diff --git a/agent_docs/observability.md b/agent_docs/observability.md index cfdde77b8f..a814e3643f 100644 --- a/agent_docs/observability.md +++ b/agent_docs/observability.md @@ -296,14 +296,16 @@ request or job, then anything passed with an individual query. **Browser to API route** — for requests that reach ClickHouse through an API route instead of a ClickHouse client. PromQL is the case today: the browser -calls `/api/v1/prometheus/query{,_range}` over fetch. The caller sends +calls `/api/v1/prometheus/query{,_range}`, and `/labels` and +`/label/:name/values` for autocomplete and dashboard filters (tagged with +`useMetadataQueryAttribution()`), over fetch. The caller sends `buildLogComment(attribution)` in the `x-hyperdx-query-attribution` header (`QUERY_ATTRIBUTION_HEADER`), and the router parses it with `parseLogComment` into a `withQueryAttribution` scope layered over the auth middleware's `api` scope. The browser's surface and ids win; `trace` and `label` stay the -server's. Queries the route runs through a ClickHouse client (the -`prometheusQuery()`/`prometheusQueryRange()` fallback) pick the scope up as -usual. Requests proxied to ClickHouse's Prometheus HTTP API on 26.6+ are tagged +server's. Queries the route runs through a ClickHouse client (label lookups, +and the `prometheusQuery()`/`prometheusQueryRange()` fallback) pick the scope +up as usual. Requests proxied to ClickHouse's Prometheus HTTP API on 26.6+ are tagged by `clickhousePrometheusRequest`: `log_comment` as a URL setting, `query_id` as the `X-ClickHouse-Query-Id` header, because the handler rejects `query_id` as an unknown setting in the URL. diff --git a/packages/app/src/__tests__/prometheusApi.test.ts b/packages/app/src/__tests__/prometheusApi.test.ts index 3e658b1567..8371e70f65 100644 --- a/packages/app/src/__tests__/prometheusApi.test.ts +++ b/packages/app/src/__tests__/prometheusApi.test.ts @@ -49,6 +49,21 @@ describe('prometheusApi.labelValues', () => { expect(params.get('end')).toBe('200'); }); + it('sends the attribution as a header', async () => { + await prometheusApi.labelValues({ + label: 'pod', + connectionId: 'conn', + attribution: { surface: 'metadata', dashboard: 'd1' }, + }); + + const { headers } = get.mock.calls[0][1]; + expect(JSON.parse(headers['x-hyperdx-query-attribution'])).toEqual({ + v: 1, + surface: 'metadata', + dashboard: 'd1', + }); + }); + it('omits the selector when there is none', async () => { await prometheusApi.labelValues({ label: 'pod', connectionId: 'conn' }); @@ -90,6 +105,23 @@ describe('prometheusApi.labels', () => { expect(resp.data).toEqual(['byoc', 'instance']); }); + it('sends the attribution as a header', async () => { + get.mockReturnValue({ + json: () => Promise.resolve({ status: 'success', data: [] }), + }); + + await prometheusApi.labels({ + connectionId: 'conn', + attribution: { surface: 'metadata' }, + }); + + const { headers } = get.mock.calls[0][1]; + expect(JSON.parse(headers['x-hyperdx-query-attribution'])).toEqual({ + v: 1, + surface: 'metadata', + }); + }); + it('leaves a response without data alone', async () => { get.mockReturnValue({ json: () => Promise.resolve({ status: 'error', error: 'nope' }), diff --git a/packages/app/src/api.ts b/packages/app/src/api.ts index 0be4bb6ea9..2f60022f92 100644 --- a/packages/app/src/api.ts +++ b/packages/app/src/api.ts @@ -867,20 +867,28 @@ const uniqueLabels = ( ): PrometheusLabelsResponse => resp.data ? { ...resp, data: [...new Set(resp.data)] } : resp; +const attributionHeaders = ( + attribution: QueryAttribution | undefined, +): Record | undefined => { + const logComment = buildLogComment(attribution); + return logComment ? { [QUERY_ATTRIBUTION_HEADER]: logComment } : undefined; +}; + const prometheusFetch = ( path: string, searchParams: Record, signal?: AbortSignal, attribution?: QueryAttribution, -): Promise => { - const logComment = buildLogComment(attribution); - const headers = logComment - ? { [QUERY_ATTRIBUTION_HEADER]: logComment } - : undefined; - return withPrometheusError(() => - server.post(path, { searchParams, signal, headers }).json(), +): Promise => + withPrometheusError(() => + server + .post(path, { + searchParams, + signal, + headers: attributionHeaders(attribution), + }) + .json(), ); -}; export const prometheusApi = { queryRange: (params: { @@ -940,10 +948,12 @@ export const prometheusApi = { table?: string; start?: number; end?: number; + attribution?: QueryAttribution; }): Promise => server .get('v1/prometheus/labels', { searchParams: labelLookupSearchParams(params), + headers: attributionHeaders(params.attribution), }) .json() .then(uniqueLabels), @@ -956,11 +966,13 @@ export const prometheusApi = { start?: number; end?: number; match?: string; + attribution?: QueryAttribution; }): Promise => withPrometheusError(() => server .get(`v1/prometheus/label/${params.label}/values`, { searchParams: labelLookupSearchParams(params), + headers: attributionHeaders(params.attribution), }) .json() .then(uniqueLabels), diff --git a/packages/app/src/hooks/__tests__/usePromqlLabelFilterValues.test.tsx b/packages/app/src/hooks/__tests__/usePromqlLabelFilterValues.test.tsx index 03e1778603..e036e2a61c 100644 --- a/packages/app/src/hooks/__tests__/usePromqlLabelFilterValues.test.tsx +++ b/packages/app/src/hooks/__tests__/usePromqlLabelFilterValues.test.tsx @@ -102,6 +102,7 @@ describe('usePromqlLabelFilterValues', () => { table: 'metrics', start: 1704067200, end: 1704153600, + attribution: expect.objectContaining({ surface: 'metadata' }), }); expect(result.current.data.get('promql1')).toEqual({ values: ['api', 'web'], diff --git a/packages/app/src/hooks/usePromqlLabelFilterValues.tsx b/packages/app/src/hooks/usePromqlLabelFilterValues.tsx index 9c954734ff..03b2a9c278 100644 --- a/packages/app/src/hooks/usePromqlLabelFilterValues.tsx +++ b/packages/app/src/hooks/usePromqlLabelFilterValues.tsx @@ -15,6 +15,7 @@ import { } from '@tanstack/react-query'; import { prometheusApi } from '@/api'; +import { useMetadataQueryAttribution } from '@/queryAttribution'; import { useSources } from '@/source'; import { mapKeyBy } from '@/utils'; @@ -87,6 +88,7 @@ export function usePromqlLabelFilterValues({ }, [filters, sourcesById, isLoadingSources, resolvedByFilterId]); const queryClient = useQueryClient(); + const attribution = useMetadataQueryAttribution(); const results: UseQueryResult[] = useQueries({ queries: calls.map(call => { @@ -126,6 +128,7 @@ export function usePromqlLabelFilterValues({ start: startSec, end: endSec, match: call.match, + attribution, }); if (resp.status === 'error') { throw new Error(resp.error ?? 'Label values query failed'); diff --git a/packages/app/src/hooks/usePromqlMetadata.ts b/packages/app/src/hooks/usePromqlMetadata.ts index c60a935bcb..351cc7632c 100644 --- a/packages/app/src/hooks/usePromqlMetadata.ts +++ b/packages/app/src/hooks/usePromqlMetadata.ts @@ -1,12 +1,14 @@ import { useQuery } from '@tanstack/react-query'; import { prometheusApi } from '@/api'; +import { useMetadataQueryAttribution } from '@/queryAttribution'; export function usePromqlMetricNames( connectionId: string | undefined, database?: string, table?: string, ) { + const attribution = useMetadataQueryAttribution(); return useQuery({ queryKey: ['promql-metric-names', connectionId, database, table], queryFn: async () => { @@ -16,6 +18,7 @@ export function usePromqlMetricNames( connectionId, database, table, + attribution, }); return resp.data ?? []; }, @@ -30,6 +33,7 @@ export function usePromqlLabelNames( database?: string, table?: string, ) { + const attribution = useMetadataQueryAttribution(); return useQuery({ queryKey: ['promql-label-names', connectionId, database, table], queryFn: async () => { @@ -38,6 +42,7 @@ export function usePromqlLabelNames( connectionId, database, table, + attribution, }); return resp.data ?? []; }, From 16096bfe96185b9ead39a71dd66ad5dd342716cd Mon Sep 17 00:00:00 2001 From: Tommy Li Date: Tue, 6 Oct 2026 10:27:08 -0700 Subject: [PATCH 5/5] test fix --- .../__tests__/PromqlLabelFilterEditForm.test.tsx | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/app/src/components/DashboardFiltersModal/__tests__/PromqlLabelFilterEditForm.test.tsx b/packages/app/src/components/DashboardFiltersModal/__tests__/PromqlLabelFilterEditForm.test.tsx index daeeb52540..ef16cc44eb 100644 --- a/packages/app/src/components/DashboardFiltersModal/__tests__/PromqlLabelFilterEditForm.test.tsx +++ b/packages/app/src/components/DashboardFiltersModal/__tests__/PromqlLabelFilterEditForm.test.tsx @@ -91,6 +91,7 @@ describe('PromqlLabelFilterEditForm', () => { connectionId: 'clickhouse-conn', database: 'telemetry', table: 'metrics', + attribution: expect.objectContaining({ surface: 'metadata' }), }); });