Promql query attribution

.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

agent_docs/observability.md

@@ -294,6 +294,22 @@ 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}`, 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 (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.
+
 ### Rules
 
 - **Values from the browser are hints, not proof.** The client controls them.

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,11 @@ import {
 } from '@/fixtures';
 import Connection from '@/models/connection';
 
+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
 // still reads it. If that scope is lost, queries fall back to hdx-unknown-*.
@@ -29,11 +37,23 @@ describe('request query attribution', () => {
           endSec: 1700003600,
         },
       ],
+      withSamples: true,
     });
   });
 
+  // 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.mocked(global.fetch).mockImplementation(fetchStub);
+    serverVersionSpy?.mockRestore();
+    serverVersionSpy = null;
   });
 
   afterAll(async () => {
@@ -66,6 +86,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 +138,95 @@ 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.
+    serverVersionSpy = 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 () => {
+    jest.mocked(global.fetch).mockImplementation(globalThis.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',
+    });
+  });
 });

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);

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<string, string> } {
+  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;