diff --git a/.changeset/query-log-attribution.md b/.changeset/query-log-attribution.md new file mode 100644 index 0000000000..f11020165d --- /dev/null +++ b/.changeset/query-log-attribution.md @@ -0,0 +1,40 @@ +--- +'@hyperdx/common-utils': minor +'@hyperdx/api': minor +'@hyperdx/app': minor +--- + +feat: attribute ClickHouse queries to the dashboard, tile, search or alert that +issued them + +A row in `system.query_log` previously couldn't be attributed to the part of +HyperDX that produced it. Now every query HyperDX runs has a small JSON +`log_comment`, plus a `query_id` starting with `hdx-`. + +Attribute load after the fact: + +```sql +SELECT + JSONExtractString(log_comment, 'surface') AS surface, + JSONExtractString(log_comment, 'dashboard') AS dashboard, + JSONExtractString(log_comment, 'tile') AS tile, + count(), + sum(read_rows) +FROM system.query_log +WHERE type = 'QueryFinish' AND log_comment != '' AND query_id like 'hdx-%' +GROUP BY ALL +ORDER BY sum(read_rows) DESC +``` + +Or spot a query while it is still running, in `system.processes`, where the +`query_id` names the surface with no JSON parsing needed. + +The payload records which part of the product asked (a dashboard, a search, an +alert, an MCP tool, the chart explorer, session replay, field lookups, and so +on), the dashboard and tile or saved search id, the source id, and on the server +the trace id of the request. Field lookups and autocomplete are labelled too, so +they can be told apart from a user's chart queries. + +The exception is queries that skip settings processing: the `system.settings`, +server version and Cloud-detection probes, and the onboarding connection check. +They get a `query_id` but an empty `log_comment`. diff --git a/agent_docs/observability.md b/agent_docs/observability.md index d2e5d855c2..2136eb9338 100644 --- a/agent_docs/observability.md +++ b/agent_docs/observability.md @@ -35,6 +35,9 @@ aligned with the 5. **Use the shared helpers** in [`packages/api/src/utils/instrumentation.ts`](../packages/api/src/utils/instrumentation.ts). Don't hand-roll tracer/meter lifecycle. +6. **A new surface that queries ClickHouse declares its query attribution.** + One provider or one `withQueryAttribution` scope at the entry point, not a + change at each query. See [Query attribution](#query-attribution). ## The shared helper library @@ -232,6 +235,84 @@ for this on functionality with real failure modes worth a target (external dependencies, query proxies, AI calls) — not on thin CRUD that the HTTP auto-instrumentation and error middleware already cover. +## Query attribution + +Every ClickHouse query HyperDX runs is tagged with what asked for it, so a row +in the customer's `system.query_log` points back to a dashboard tile, a saved +search, an alert, or an MCP tool. This is the one piece of our instrumentation +that lands in *their* data rather than ours, so it stays small, carries no +identity, and must never be able to fail a query. + +Two channels carry it, set by `BaseClickhouseClient.query`: + +| Where | Shape | What it answers | +| --- | --- | --- | +| `log_comment` | JSON payload, e.g. `{"v":1,"surface":"dashboard","dashboard":"…","tile":"…"}` | After the fact: `GROUP BY JSONExtractString(log_comment, 'tile')` over `system.query_log` | +| `query_id` | `hdx--` | Right now: reading `system.processes` while a query is still running | + +The vocabulary lives in +[`packages/common-utils/src/clickhouse/attribution.ts`](../packages/common-utils/src/clickhouse/attribution.ts). +`QUERY_SURFACES` is a fixed, short list, because it goes in the `query_id` and +people group by it. Anything naming one specific dashboard, tile or search goes +in the id fields instead. + +### Declaring it + +Not at the call site that runs the query: there are roughly sixty of those, and +`Metadata` accounts for two dozen on its own. Set it once where the reason for +the query is known, and everything underneath picks it up. + +**Browser** — React context, read by `useClickhouseClient`: + +```tsx + + {children} + +``` + +Providers nest, and the innermost value wins, so a dashboard sets its id once +and each tile adds only its own. A page that needs nothing finer than "which +page" can use `withAppNavForSurface('…')` as its `getLayout` instead. + +Do **not** put attribution on the chart config: the config is the React Query +cache key, so two surfaces asking the identical question would stop sharing a +cache entry and query volume would double. + +**Server** — a scope around the request or job (`AsyncLocalStorage`, node only): + +```ts +await withQueryAttribution({ surface: 'alert', alert: alert.id }, async () => + processAlert(...), +); +``` + +Wrap the whole request or job, not the query. One call at the top covers +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. + +### Rules + +- **Values from the browser are hints, not proof.** The client controls them. + `sanitizeField` drops control characters and caps the length, and the payload + only ever becomes JSON strings, never SQL. +- **It must never throw.** A tag describes the work; it does not gate it. Read + fields with `?.` even where the type says they are required, and keep helpers + like `getActiveTraceId` from throwing. A half-populated object must not be + able to stop an alert from running. +- **Set it where the reason is known, not just on the page.** A page-level + surface is a fallback. A modal or panel rendered outside a tile only gets the + page's surface, so give it its own (`chart-preview`) when it has ids to + report. +- **Keep the payload small.** It rides in the URL query string on every browser + query. The cap is 1 KB; real payloads are 100 to 200 bytes. Today's fields + cannot fill the budget even at full length, so nothing is ever shed. If you + add a field, check that still holds: over budget, fields are dropped whole + from the end rather than truncated, so the value stays parseable JSON. + ## Where to look for examples - Generic span + status handling: `packages/api/src/utils/instrumentation.ts` @@ -260,3 +341,9 @@ auto-instrumentation and error middleware already cover. - End-to-end + sub-operation SLO metrics on a background job (alert evaluation and its data fetch): `packages/api/src/tasks/checkAlerts/index.ts` (`alerts.evaluate`, `alerts.query`) +- Query attribution scope at the auth boundary: + `packages/api/src/middleware/auth.ts` (`nextWithQueryAttribution`) +- Query attribution around one alert evaluation: + `packages/api/src/tasks/checkAlerts/index.ts` (`alertQueryAttribution`) +- Query attribution per React subtree: `packages/app/src/queryAttribution.tsx`, + mounted per tile in `packages/app/src/DBDashboardPage.tsx` diff --git a/packages/api/src/mcp/utils/tracing.ts b/packages/api/src/mcp/utils/tracing.ts index 6c0b13e640..592379e9df 100644 --- a/packages/api/src/mcp/utils/tracing.ts +++ b/packages/api/src/mcp/utils/tracing.ts @@ -1,3 +1,4 @@ +import { withQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'; import { recordOnboardingTaskCompletion } from '@/controllers/user'; @@ -5,6 +6,7 @@ import type { McpContext, ToolResult } from '@/mcp/tools/types'; import type { McpErrorCategory, McpErrorResult } from '@/mcp/utils/errors'; import { getErrorCategory } from '@/mcp/utils/errors'; import { + getActiveTraceId, getCounter, getHistogram, SpanStatusCode, @@ -88,7 +90,16 @@ export function withToolTracing( logger.info(logContext, `MCP tool invoked: ${toolName}`); try { - const result = await handler(args); + // Around the whole handler, so every query the tool makes carries + // the tool name and who called it. + const result = await withQueryAttribution( + { + surface: 'mcp', + label: toolName, + trace: getActiveTraceId(), + }, + async () => handler(args), + ); const durationMs = Date.now() - startTime; if (result.isError) { diff --git a/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts b/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts new file mode 100644 index 0000000000..1fe4ae081f --- /dev/null +++ b/packages/api/src/middleware/__tests__/queryAttribution.int.test.ts @@ -0,0 +1,98 @@ +import * as config from '@/config'; +import { + DEFAULT_DATABASE, + dropTimeSeriesTable, + getLoggedInAgent, + getServer, + getTestFixtureClickHouseClient, + seedTimeSeriesTagsTable, +} from '@/fixtures'; +import Connection from '@/models/connection'; + +// 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-*. +describe('request query attribution', () => { + const server = getServer(); + // Unique per run so query_log rows from earlier runs can't match. + const TABLE = `query_attribution_${Date.now()}`; + + beforeAll(async () => { + await server.start(); + await seedTimeSeriesTagsTable({ + table: TABLE, + series: [ + { + metricName: 'attribution_metric', + tags: { job: 'api' }, + startSec: 1700000000, + endSec: 1700003600, + }, + ], + }); + }); + + afterEach(async () => { + await server.clearDBs(); + }); + + afterAll(async () => { + await dropTimeSeriesTable({ table: TABLE }); + await server.stop(); + }); + + // query_log is written asynchronously, so flush and poll. + async function findRouteQueries() { + 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 query_kind = 'Select' + AND position(query, {table:String}) > 0`, + // Match on the text: timeSeriesTags() reads an inner table, so + // `tables` lists `.inner_id.tags.`, not TABLE. + query_params: { table: TABLE }, + 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({ + team: team._id, + name: 'CH', + host: config.CLICKHOUSE_HOST, + username: config.CLICKHOUSE_USER, + password: config.CLICKHOUSE_PASSWORD, + }); + + await agent + .get('/v1/prometheus/labels') + .query({ + connectionId: conn._id.toString(), + database: DEFAULT_DATABASE, + table: TABLE, + }) + .expect(200); + + const logged = await findRouteQueries(); + expect(logged.length).toBeGreaterThan(0); + for (const row of logged) { + expect(row.query_id).toMatch(/^hdx-api-/); + expect(JSON.parse(row.log_comment)).toMatchObject({ + surface: 'api', + label: '/v1/prometheus', + }); + } + }); +}); diff --git a/packages/api/src/middleware/auth.ts b/packages/api/src/middleware/auth.ts index 4db9840162..5be32c9798 100644 --- a/packages/api/src/middleware/auth.ts +++ b/packages/api/src/middleware/auth.ts @@ -1,3 +1,4 @@ +import { withQueryAttribution } from '@hyperdx/common-utils/dist/clickhouse/node'; import { Connection } from '@hyperdx/common-utils/dist/types'; import type { NextFunction, Request, Response } from 'express'; import { serializeError } from 'serialize-error'; @@ -6,6 +7,7 @@ import * as config from '@/config'; import { findUserByAccessKey } from '@/controllers/user'; import type { UserDocument } from '@/models/user'; import { + getActiveTraceId, getStaticFeatureFlags, setBusinessContext, } from '@/utils/instrumentation'; @@ -109,7 +111,23 @@ export async function validateUserAccessKey( ...getStaticFeatureFlags(), }); - next(); + nextWithQueryAttribution(req, next); +} + +/** + * Continue the request with its route recorded on any ClickHouse queries it + * makes. Here because every authenticated route passes through this point. + */ +function nextWithQueryAttribution(req: Request, next: NextFunction): void { + withQueryAttribution( + { + surface: 'api', + // Mount path, not the full URL, to keep request ids out of the log. + label: req.baseUrl || undefined, + trace: getActiveTraceId(), + }, + () => next(), + ); } export function isUserAuthenticated( @@ -133,7 +151,7 @@ export function isUserAuthenticated( 'hyperdx.local_mode': true, ...getStaticFeatureFlags(), }); - return next(); + return nextWithQueryAttribution(req, next); } if (req.isAuthenticated()) { @@ -145,7 +163,7 @@ export function isUserAuthenticated( ...getStaticFeatureFlags(), }); - return next(); + return nextWithQueryAttribution(req, next); } res.sendStatus(401); } diff --git a/packages/api/src/routers/api/__tests__/clickhouseProxy.int.test.ts b/packages/api/src/routers/api/__tests__/clickhouseProxy.int.test.ts index b658edfd85..23d5c2e5ca 100644 --- a/packages/api/src/routers/api/__tests__/clickhouseProxy.int.test.ts +++ b/packages/api/src/routers/api/__tests__/clickhouseProxy.int.test.ts @@ -1,3 +1,4 @@ +import { buildLogComment } from '@hyperdx/common-utils/dist/clickhouse'; import http from 'http'; import { getLoggedInAgent, getServer } from '@/fixtures'; @@ -15,15 +16,16 @@ type CapturedRequest = { }; /** - * Pins /clickhouse-proxy body forwarding: text/plain re-injection and the + * Pins what /clickhouse-proxy forwards upstream: text/plain re-injection, the * raw-stream passthrough of multipart/form-data bodies, which * `@clickhouse/client-web` emits when query params exceed its URL budget - * (regression context: ClickHouse support-escalation #8482). + * (regression context: ClickHouse support-escalation #8482), and the + * decode/re-encode the query string goes through in `pathRewrite`. * * Known-latent cases (charset-suffixed JSON, urlencoded) are not covered; * tracked in https://github.com/hyperdxio/hyperdx/issues/2942. */ -describe('clickhouse-proxy body forwarding', () => { +describe('clickhouse-proxy forwarding', () => { const server = getServer(); let upstream: http.Server; @@ -139,6 +141,48 @@ describe('clickhouse-proxy body forwarding', () => { ); }); + // The browser client ships `log_comment` as a URL query param, and + // pathRewrite fully decodes the query string (sanitizeUrl) before + // re-encoding it with `searchParams.toString()`. Breaking that round trip + // fails every browser query, not just the attribution tag - which is why + // `buildLogComment` allowlists the characters it emits. + it('preserves the attribution query params through the path rewrite', async () => { + const { agent, team } = await getLoggedInAgent(server); + const connectionId = await createConnection(team._id.toString()); + + const logComment = buildLogComment({ + surface: 'dashboard', + dashboard: 'my dashboard 1', + tile: 'tile-42', + label: 'a/b:c_d.e', + }); + if (logComment == null) throw new Error('Expected a log_comment payload'); + + const queryId = 'hdx-dashboard-0e0d1a2b-3c4d-4e5f-8a9b-0c1d2e3f4a5b'; + await agent + .post( + `/clickhouse-proxy/?query_id=${queryId}&log_comment=${encodeURIComponent(logComment)}`, + ) + .set('x-hyperdx-connection-id', connectionId) + .set('Content-Type', 'text/plain') + .send('SELECT 1 FORMAT JSON') + .expect(200); + + expect(captured).toHaveLength(1); + const forwarded = new URL(captured[0].url ?? '', 'http://localhost'); + expect(forwarded.searchParams.get('query_id')).toBe(queryId); + expect(forwarded.searchParams.get('log_comment')).toBe(logComment); + expect(JSON.parse(forwarded.searchParams.get('log_comment') ?? '')).toEqual( + { + v: 1, + surface: 'dashboard', + dashboard: 'my dashboard 1', + tile: 'tile-42', + label: 'a/b:c_d.e', + }, + ); + }); + it('rejects requests without a connection id header', async () => { const { agent } = await getLoggedInAgent(server); diff --git a/packages/api/src/tasks/checkAlerts/__tests__/alertQueryAttribution.test.ts b/packages/api/src/tasks/checkAlerts/__tests__/alertQueryAttribution.test.ts new file mode 100644 index 0000000000..96f5d9a165 --- /dev/null +++ b/packages/api/src/tasks/checkAlerts/__tests__/alertQueryAttribution.test.ts @@ -0,0 +1,90 @@ +import { Types } from 'mongoose'; + +import { alertQueryAttribution } from '@/tasks/checkAlerts'; +import { AlertTaskType } from '@/tasks/checkAlerts/providers'; + +// Real AlertDetails comes from mongoose documents. These tests only need the +// few fields the tag reads. + +const details = (overrides: Record): any => ({ + alert: { id: 'alert-1' }, + ...overrides, +}); + +describe('alertQueryAttribution', () => { + it('reports the saved search for a saved-search alert', () => { + expect( + alertQueryAttribution( + details({ + taskType: AlertTaskType.SAVED_SEARCH, + savedSearch: { id: 'search-1' }, + source: { id: 'source-1' }, + }), + ), + ).toMatchObject({ + surface: 'alert', + alert: 'alert-1', + search: 'search-1', + source: 'source-1', + }); + }); + + it('reports the dashboard and tile for a tile alert', () => { + expect( + alertQueryAttribution( + details({ + taskType: AlertTaskType.TILE, + dashboard: { id: 'dash-1' }, + tile: { id: 'tile-1' }, + source: { id: 'source-1' }, + }), + ), + ).toMatchObject({ dashboard: 'dash-1', tile: 'tile-1' }); + }); + + // The tile path goes through `toObject()`, which leaves out the `id` + // virtual, so the source arrives with only `_id`. This used to report no + // source at all. + it('falls back to _id when the source has no id virtual', () => { + const _id = new Types.ObjectId(); + expect( + alertQueryAttribution( + details({ + taskType: AlertTaskType.TILE, + dashboard: { id: 'dash-1' }, + tile: { id: 'tile-1' }, + source: { _id }, + }), + ).source, + ).toBe(_id.toString()); + }); + + // A bare ObjectId has an `id` of its own, a byte array, which used to + // satisfy the `??` and then get thrown away as a non-string. + it('falls back to _id when the source is an unpopulated ref', () => { + const _id = new Types.ObjectId(); + expect( + alertQueryAttribution( + details({ + taskType: AlertTaskType.TILE, + dashboard: { id: 'dash-1' }, + tile: { id: 'tile-1' }, + source: _id, + }), + ).source, + ).toBe(_id.toString()); + }); + + it('omits the source when the alert has none', () => { + expect( + alertQueryAttribution( + details({ taskType: AlertTaskType.INLINE, chartConfig: {} }), + ), + ).toMatchObject({ surface: 'alert', source: undefined }); + }); + + it('does not throw on a partially populated alert', () => { + const partial: any = { taskType: AlertTaskType.SAVED_SEARCH }; + expect(() => alertQueryAttribution(partial)).not.toThrow(); + }); +}); diff --git a/packages/api/src/tasks/checkAlerts/index.ts b/packages/api/src/tasks/checkAlerts/index.ts index cb8d1a1d33..73e8cace9a 100644 --- a/packages/api/src/tasks/checkAlerts/index.ts +++ b/packages/api/src/tasks/checkAlerts/index.ts @@ -5,9 +5,13 @@ import PQueue from '@esm2cjs/p-queue'; import * as clickhouse from '@hyperdx/common-utils/dist/clickhouse'; import { chSqlToAliasMap, + type QueryAttribution, ResponseJSON, } from '@hyperdx/common-utils/dist/clickhouse'; -import { ClickhouseClient } from '@hyperdx/common-utils/dist/clickhouse/node'; +import { + ClickhouseClient, + withQueryAttribution, +} from '@hyperdx/common-utils/dist/clickhouse/node'; import { tryOptimizeConfigWithMaterializedView } from '@hyperdx/common-utils/dist/core/materializedViews'; import { getMetadata, @@ -97,6 +101,7 @@ import { } from '@/tasks/util'; import { alertConfigHasGroupBy, isPopulatedRef } from '@/utils/alerts'; import { + getActiveTraceId, getCounter, type OperationOutcome, recordOperationOutcome, @@ -941,6 +946,37 @@ export const parseAlertData = ( return { value, extraFields }; }; +/** + * Names the alert, and the tile or search it came from, on each of its + * queries. Per evaluation rather than per client, because one client is shared + * by every alert on a connection. + */ +const sourceId = (source: AlertDetails['source']): string | undefined => { + if (typeof source?.id === 'string') return source.id; + return isPopulatedRef(source) ? source._id.toString() : undefined; +}; + +export const alertQueryAttribution = ( + details: AlertDetails, +): QueryAttribution => ({ + // `?.` everywhere, even where the type says otherwise: a tag must never be + // the reason an alert fails to run. + surface: 'alert', + alert: details.alert?.id, + trace: getActiveTraceId(), + // `id` is a mongoose virtual and the tile path goes through `toObject()`, + // which leaves virtuals out, so that source only has `_id`. Try both, and + // check the type: a bare ObjectId has an `id` of its own that is a byte + // array, which would satisfy `??` and then be discarded. + source: 'source' in details ? sourceId(details.source) : undefined, + ...(details.taskType === AlertTaskType.SAVED_SEARCH + ? { search: details.savedSearch?.id } + : {}), + ...(details.taskType === AlertTaskType.TILE + ? { dashboard: details.dashboard?.id, tile: details.tile?.id } + : {}), +}); + export const processAlert = async ( now: Date, details: AlertDetails, @@ -1582,7 +1618,6 @@ export const processAlert = async ( } // We have at least one data point for this bucket - // Track the worst-case state for each group in this bucket to prevent // a subsequent OK row in the SAME bucket from overwriting an ALERT row. const bucketEvaluations = new Map< @@ -2016,13 +2051,19 @@ export default class CheckAlertTask implements HdxTask { 'processAlert', async () => { setBusinessContext({ teamId: conn.team.toString() }); - await processAlert( - alertTask.now, - alert, - clickhouseClient, - conn.id, - this.provider, - teamWebhooksById, + // Set once here, so the alert's own query, the EXPLAIN + // probes and the notification's sample rows all get tagged. + await withQueryAttribution( + alertQueryAttribution(alert), + async () => + processAlert( + alertTask.now, + alert, + clickhouseClient, + conn.id, + this.provider, + teamWebhooksById, + ), ); }, { diff --git a/packages/api/src/tasks/checkAlerts/providers/default.ts b/packages/api/src/tasks/checkAlerts/providers/default.ts index 560637a0c8..b409fafe98 100644 --- a/packages/api/src/tasks/checkAlerts/providers/default.ts +++ b/packages/api/src/tasks/checkAlerts/providers/default.ts @@ -697,6 +697,8 @@ export default class DefaultAlertProvider implements AlertProvider { password, application: `hyperdx-alerts ${config.CODE_VERSION}`, requestTimeout: requestTimeout ?? 30_000, + // Fallback for anything outside a per-alert scope. + attribution: { surface: 'alert' }, }); } } diff --git a/packages/api/src/tasks/usageStats.ts b/packages/api/src/tasks/usageStats.ts index 19704dc87c..f226f24eae 100644 --- a/packages/api/src/tasks/usageStats.ts +++ b/packages/api/src/tasks/usageStats.ts @@ -96,6 +96,8 @@ const getClickhouseTableSize = async () => { host: connection.host, username: connection.username, password: connection.password, + // This client is only used by this job. + attribution: { surface: 'api', label: 'usage-stats' }, }); const _rows = await clickhouseClient.query({ query: ` diff --git a/packages/api/src/utils/instrumentation.ts b/packages/api/src/utils/instrumentation.ts index 260acbf27a..096bb2dc5a 100644 --- a/packages/api/src/utils/instrumentation.ts +++ b/packages/api/src/utils/instrumentation.ts @@ -141,6 +141,25 @@ export function setBusinessContext(context: BusinessContext): void { opentelemetry.trace.getActiveSpan()?.setAttributes(attributes); } +/** + * Trace id of the span in flight, if any. Tagged onto ClickHouse queries so a + * row in `system.query_log` links back to the request that ran it. + */ +export function getActiveTraceId(): string | undefined { + try { + const ctx = opentelemetry.trace.getActiveSpan()?.spanContext(); + // An invalid context (e.g. the all-zero id of a non-recording span) would + // give us rows that look traced but link to nothing. + return ctx && opentelemetry.trace.isSpanContextValid(ctx) + ? ctx.traceId + : undefined; + } catch { + // Runs on every query, and the id is only a convenience. If tracing is + // off or set up oddly, go without it rather than fail the query. + return undefined; + } +} + /** * Returns the static (env / compile-time) feature flag states as span/trace * attributes. The repo has no dynamic flag service, so these are the toggles diff --git a/packages/app/src/AlertDetailPage.tsx b/packages/app/src/AlertDetailPage.tsx index ba69412ce4..787b76f41b 100644 --- a/packages/app/src/AlertDetailPage.tsx +++ b/packages/app/src/AlertDetailPage.tsx @@ -32,7 +32,7 @@ import { getAlertSourceLabel, getAlertSourceUrl } from '@/utils/alerts'; import { useBrandDisplayName } from './theme/ThemeProvider'; import api from './api'; -import { withAppNav } from './layout'; +import { withAppNavForSurface } from './layout'; import { useDefaultTimeRange, useNewTimeQuery } from './timeQuery'; import type { AlertsPageItem } from './types'; @@ -280,4 +280,4 @@ export default function AlertDetailPage() { ); } -AlertDetailPage.getLayout = withAppNav; +AlertDetailPage.getLayout = withAppNavForSurface('alert', 'alert-detail'); diff --git a/packages/app/src/ClickhousePage.tsx b/packages/app/src/ClickhousePage.tsx index cb6de4a6f0..aec44c2bdc 100644 --- a/packages/app/src/ClickhousePage.tsx +++ b/packages/app/src/ClickhousePage.tsx @@ -36,7 +36,7 @@ import { ConnectionSelectControlled } from '@/components/ConnectionSelect'; import { DBTimeChart } from '@/components/DBTimeChart'; import { PageLayout } from '@/components/PageLayout'; import { TimePicker } from '@/components/TimePicker'; -import { withAppNav } from '@/layout'; +import { withAppNavForSurface } from '@/layout'; import { ChartCard } from './components/charts/ChartCard'; import ChartContainer from './components/charts/ChartContainer'; @@ -871,6 +871,9 @@ const ClickhousePageDynamic = dynamic(async () => ClickhousePage, { }); // @ts-expect-error next/dynamic component type does not include the getLayout static -ClickhousePageDynamic.getLayout = withAppNav; +ClickhousePageDynamic.getLayout = withAppNavForSurface( + 'dashboard', + 'clickhouse', +); export default ClickhousePageDynamic; diff --git a/packages/app/src/DBChartPage.tsx b/packages/app/src/DBChartPage.tsx index 85d4a3b797..32dac8ce2d 100644 --- a/packages/app/src/DBChartPage.tsx +++ b/packages/app/src/DBChartPage.tsx @@ -37,7 +37,7 @@ import { IS_ALERT_DETAILS_ENABLED } from '@/config'; import { useChartAssistant } from '@/hooks/ai'; import { useAlertSeededChartConfig } from '@/hooks/useAlertSeededChartConfig'; import { useResolvedSourceParam } from '@/hooks/useResolvedSourceParam'; -import { withAppNav } from '@/layout'; +import { withAppNavForSurface } from '@/layout'; import { useSources } from '@/source'; import { useBrandDisplayName } from '@/theme/ThemeProvider'; import { useDefaultTimeRange, useNewTimeQuery } from '@/timeQuery'; @@ -362,6 +362,6 @@ const DBChartExplorerPageDynamic = dynamic(async () => DBChartExplorerPage, { }); // @ts-expect-error next/dynamic component type does not include the getLayout static -DBChartExplorerPageDynamic.getLayout = withAppNav; +DBChartExplorerPageDynamic.getLayout = withAppNavForSurface('chart-explorer'); export default DBChartExplorerPageDynamic; diff --git a/packages/app/src/DBDashboardPage.tsx b/packages/app/src/DBDashboardPage.tsx index da2338bfa8..7a54cbc8f8 100644 --- a/packages/app/src/DBDashboardPage.tsx +++ b/packages/app/src/DBDashboardPage.tsx @@ -165,6 +165,7 @@ import useDashboardContainers, { } from '@/hooks/useDashboardContainers'; import { useDashboardKioskMode } from '@/hooks/useDashboardKioskMode'; import { useReleaseAnnotations } from '@/hooks/useReleaseAnnotations'; +import { QueryAttributionProvider } from '@/queryAttribution'; import { calculateNextTilePosition, makeId } from '@/utils/tilePositioning'; import ChartContainer, { @@ -205,7 +206,7 @@ import { GranularityPickerControlled, } from './GranularityPicker'; import HDXMarkdownChart from './HDXMarkdownChart'; -import { withAppNav } from './layout'; +import { withAppNavForSurface } from './layout'; import { getEventBody, getFirstTimestampValueExpression, @@ -1499,7 +1500,7 @@ const Tile = ({ ); return ( - <> +
)} - + ); }; @@ -1666,36 +1667,43 @@ const EditTileModal = ({ zIndex={modalZIndex} > {chart != null && ( - - {/* Isolate chart cross-syncing to this edit modal: the preview chart + + + {/* Isolate chart cross-syncing to this edit modal: the preview chart must not drive shadow tooltips on the dashboard tiles behind it. */} - - {/* Offers the dashboard's variables as completions in every + + {/* Offers the dashboard's variables as completions in every expression input the editor renders. */} - - { - onSave({ - ...chart, - config: config, - }); - }} - onClose={handleClose} - onDirtyChange={setHasUnsavedChanges} - isDashboardForm - autoRun - /> - - - + + { + onSave({ + ...chart, + config: config, + }); + }} + onClose={handleClose} + onDirtyChange={setHasUnsavedChanges} + isDashboardForm + autoRun + /> + + + + )} ); @@ -3538,10 +3546,14 @@ function DBDashboardPageGuarded({ if (!dashboardProps || !router.isReady) return ; return ( - + + + ); } @@ -3550,6 +3562,6 @@ const DBDashboardPageDynamic = dynamic(async () => DBDashboardPageGuarded, { }); // @ts-expect-error for getLayout -DBDashboardPageDynamic.getLayout = withAppNav; +DBDashboardPageDynamic.getLayout = withAppNavForSurface('dashboard'); export default DBDashboardPageDynamic; diff --git a/packages/app/src/DBSearchPage.tsx b/packages/app/src/DBSearchPage.tsx index ee30e8c46d..f0dad9cf55 100644 --- a/packages/app/src/DBSearchPage.tsx +++ b/packages/app/src/DBSearchPage.tsx @@ -114,6 +114,7 @@ import { useExplainQuery } from '@/hooks/useExplainQuery'; import { useResolvedSourceParam } from '@/hooks/useResolvedSourceParam'; import { withAppNav } from '@/layout'; import { isNonTrivialSearch } from '@/OnboardingChecklist/onboardingTasks'; +import { QueryAttributionProvider } from '@/queryAttribution'; import { useCreateSavedSearch, useDeleteSavedSearch, @@ -1002,15 +1003,41 @@ export function useSearchTelemetry({ return { searchElapsedMs: completedSearch?.latency_ms ?? null }; } +/** + * The saved search being shown, or null for an ad-hoc one. + * + * Read from the URL directly because the Next router lags window.location, + * which races with useQueryStates. + */ +function getSavedSearchIdFromPath(): string | null { + const paths = window.location.pathname.split('/'); + return paths.length === 3 ? paths[2] : null; +} + +/** + * Tags the page's ClickHouse queries with the saved search they belong to. + * A wrapper so the page component's own JSX stays where it is. + */ export function DBSearchPage() { + return ( + + + + ); +} + +function DBSearchPageContent() { + // Read again here, not passed down: this component re-renders from its own + // query-state hooks without the wrapper, and must see the current path. + const savedSearchId = getSavedSearchIdFromPath(); const brandName = useBrandDisplayName(); const defaultTimeRange = useDefaultTimeRange('Past 15m'); - // Next router is laggy behind window.location, which causes race - // conditions with useQueryStates, so we'll parse it directly - const paths = window.location.pathname.split('/'); - const savedSearchId = paths.length === 3 ? paths[2] : null; - const [rawSearchedConfig, setSearchedConfig] = useQueryStates(queryStateMap); // `?source=` accepts a source name as well as a source ID. Resolve it to an ID diff --git a/packages/app/src/DBServiceMapPage.tsx b/packages/app/src/DBServiceMapPage.tsx index 1234420913..54791edbfe 100644 --- a/packages/app/src/DBServiceMapPage.tsx +++ b/packages/app/src/DBServiceMapPage.tsx @@ -37,7 +37,7 @@ import SearchWhereInput, { import { IS_LOCAL_MODE } from '@/config'; import { useGetKeyValues } from '@/hooks/useMetadata'; import { useResolvedSourceParam } from '@/hooks/useResolvedSourceParam'; -import { withAppNav } from '@/layout'; +import { withAppNavForSurface } from '@/layout'; import { parseAsStringEncoded } from '@/utils/queryParsers'; import OnboardingModal from './components/OnboardingModal'; @@ -410,6 +410,9 @@ const DBServiceMapPageDynamic = dynamic(async () => DBServiceMapPage, { }); // @ts-expect-error next/dynamic component type does not include the getLayout static -DBServiceMapPageDynamic.getLayout = withAppNav; +DBServiceMapPageDynamic.getLayout = withAppNavForSurface( + 'service-dashboard', + 'service-map', +); export default DBServiceMapPageDynamic; diff --git a/packages/app/src/KubernetesDashboardPage.tsx b/packages/app/src/KubernetesDashboardPage.tsx index 7897d8a548..062abb09ee 100644 --- a/packages/app/src/KubernetesDashboardPage.tsx +++ b/packages/app/src/KubernetesDashboardPage.tsx @@ -64,7 +64,7 @@ import { K8S_CPU_PERCENTAGE_NUMBER_FORMAT, K8S_MEM_NUMBER_FORMAT, } from './ChartUtils'; -import { withAppNav } from './layout'; +import { withAppNavForSurface } from './layout'; import NamespaceDetailsSidePanel from './NamespaceDetailsSidePanel'; import NodeDetailsSidePanel from './NodeDetailsSidePanel'; import PodDetailsSidePanel from './PodDetailsSidePanel'; @@ -1762,6 +1762,9 @@ const KubernetesDashboardPageDynamic = dynamic( ); // @ts-expect-error next/dynamic component type does not include the getLayout static -KubernetesDashboardPageDynamic.getLayout = withAppNav; +KubernetesDashboardPageDynamic.getLayout = withAppNavForSurface( + 'dashboard', + 'kubernetes', +); export default KubernetesDashboardPageDynamic; diff --git a/packages/app/src/ServicesDashboardPage/ServicesDashboardPage.tsx b/packages/app/src/ServicesDashboardPage/ServicesDashboardPage.tsx index 006c93ab2f..82896e1959 100644 --- a/packages/app/src/ServicesDashboardPage/ServicesDashboardPage.tsx +++ b/packages/app/src/ServicesDashboardPage/ServicesDashboardPage.tsx @@ -54,7 +54,7 @@ import { useQueriedChartConfig } from '@/hooks/useChartConfig'; import { useDashboardRefresh } from '@/hooks/useDashboardRefresh'; import usePresetDashboardFilters from '@/hooks/usePresetDashboardFilters'; import { useResolvedSourceParam } from '@/hooks/useResolvedSourceParam'; -import { withAppNav } from '@/layout'; +import { withAppNavForSurface } from '@/layout'; import { useServiceDashboardExpressions } from '@/serviceDashboard'; import { useSource, useSources } from '@/source'; import { useBrandDisplayName } from '@/theme/ThemeProvider'; @@ -473,6 +473,9 @@ const ServicesDashboardPageDynamic = dynamic( ); // @ts-expect-error Next.js layout typing -ServicesDashboardPageDynamic.getLayout = withAppNav; +ServicesDashboardPageDynamic.getLayout = withAppNavForSurface( + 'service-dashboard', + 'services', +); export default ServicesDashboardPageDynamic; diff --git a/packages/app/src/ServicesDashboardPage/__tests__/ServicesDashboardPage.test.tsx b/packages/app/src/ServicesDashboardPage/__tests__/ServicesDashboardPage.test.tsx index d049d74d45..7f4d58301b 100644 --- a/packages/app/src/ServicesDashboardPage/__tests__/ServicesDashboardPage.test.tsx +++ b/packages/app/src/ServicesDashboardPage/__tests__/ServicesDashboardPage.test.tsx @@ -159,6 +159,7 @@ jest.mock('@/hooks/useChartConfig', () => ({ jest.mock('@/layout', () => ({ __esModule: true, withAppNav: (page: unknown) => page, + withAppNavForSurface: () => (page: unknown) => page, })); jest.mock('@/theme/ThemeProvider', () => ({ diff --git a/packages/app/src/SessionsPage.tsx b/packages/app/src/SessionsPage.tsx index e613237330..b7c5492d8b 100644 --- a/packages/app/src/SessionsPage.tsx +++ b/packages/app/src/SessionsPage.tsx @@ -54,7 +54,7 @@ import SearchWhereInput, { getStoredLanguage, } from './components/SearchInput/SearchWhereInput'; import { useBrandDisplayName } from './theme/ThemeProvider'; -import { withAppNav } from './layout'; +import { withAppNavForSurface } from './layout'; import { Session, useSessions } from './sessions'; import SessionSidePanel from './SessionSidePanel'; import { useSource, useSources } from './source'; @@ -538,7 +538,7 @@ const SessionsPageDynamic = dynamic(async () => SessionsPage, { }); // @ts-expect-error for getLayout -SessionsPageDynamic.getLayout = withAppNav; +SessionsPageDynamic.getLayout = withAppNavForSurface('session-replay'); export default SessionsPageDynamic; diff --git a/packages/app/src/__tests__/queryAttribution.test.tsx b/packages/app/src/__tests__/queryAttribution.test.tsx new file mode 100644 index 0000000000..f7fd977647 --- /dev/null +++ b/packages/app/src/__tests__/queryAttribution.test.tsx @@ -0,0 +1,70 @@ +import { renderHook } from '@testing-library/react'; + +import { + QueryAttributionProvider, + useQueryAttribution, +} from '@/queryAttribution'; + +describe('QueryAttributionProvider', () => { + it('returns an empty attribution with no provider', () => { + const { result } = renderHook(() => useQueryAttribution()); + expect(result.current).toEqual({}); + }); + + it('nests, with the inner provider winning on conflict', () => { + const { result } = renderHook(() => useQueryAttribution(), { + wrapper: ({ children }) => ( + + + {children} + + + ), + }); + + expect(result.current).toEqual({ + surface: 'chart-preview', + dashboard: 'dash-1', + tile: 'tile-1', + }); + }); + + it('keeps the value referentially stable across renders of an inline attribution', () => { + const { result, rerender } = renderHook(() => useQueryAttribution(), { + wrapper: ({ children }) => ( + // A new object every render, which is how callers write it. + + {children} + + ), + }); + + const first = result.current; + rerender(); + expect(result.current).toBe(first); + }); + + it('produces a new value when the attribution changes', () => { + let tile = 'tile-1'; + const { result, rerender } = renderHook(() => useQueryAttribution(), { + wrapper: ({ children }) => ( + + {children} + + ), + }); + + const first = result.current; + tile = 'tile-2'; + rerender(); + + expect(result.current).not.toBe(first); + expect(result.current.tile).toBe('tile-2'); + }); +}); diff --git a/packages/app/src/clickhouse.ts b/packages/app/src/clickhouse.ts index 6d5def56d8..351e25e879 100644 --- a/packages/app/src/clickhouse.ts +++ b/packages/app/src/clickhouse.ts @@ -4,11 +4,12 @@ // not specific querying/functionality logic // please move app-specific functions elsewhere in the app // ================================ - import { chSql, ClickhouseClientOptions, ColumnMeta, + mergeQueryAttribution, + type QueryAttribution, ResponseJSON, } from '@hyperdx/common-utils/dist/clickhouse'; import { ClickhouseClient } from '@hyperdx/common-utils/dist/clickhouse/browser'; @@ -16,14 +17,24 @@ import { useQuery, UseQueryOptions } from '@tanstack/react-query'; import { IS_LOCAL_MODE } from '@/config'; import { getLocalConnections } from '@/connection'; +import { + useMetadataQueryAttribution, + useQueryAttribution, +} from '@/queryAttribution'; import api from './api'; import { DEFAULT_QUERY_TIMEOUT } from './defaults'; const PROXY_CLICKHOUSE_HOST = '/api/clickhouse-proxy'; +/** + * `attribution` is required, not optional, on purpose. This cannot read React + * context the way `useClickhouseClient` below does, so a caller that forgot it + * would silently produce queries with no `log_comment` and an + * `hdx-unknown-*` id. Making it required means the compiler asks instead. + */ export const getClickhouseClient = ( - options: ClickhouseClientOptions = {}, + options: ClickhouseClientOptions & { attribution: QueryAttribution }, ): ClickhouseClient => { if (IS_LOCAL_MODE) { const localConnections = getLocalConnections(); @@ -51,21 +62,27 @@ export const useClickhouseClient = ( options: ClickhouseClientOptions = {}, ): ClickhouseClient => { const { data: me } = api.useMe(); - const teamQueryTimeout = me?.team?.queryTimeout; - if (teamQueryTimeout !== undefined) { - options.queryTimeout = teamQueryTimeout; - } else { - options.queryTimeout = DEFAULT_QUERY_TIMEOUT; - } + // The team's timeout wins over anything the caller passed, as it always has. + const queryTimeout = me?.team?.queryTimeout ?? DEFAULT_QUERY_TIMEOUT; - return getClickhouseClient(options); + // Set on the client, not at each query: everything reaches ClickHouse + // through a client from here or from the factory above. + const surfaceAttribution = useQueryAttribution(); + const attribution = mergeQueryAttribution( + surfaceAttribution, + options.attribution, + ); + + return getClickhouseClient({ ...options, queryTimeout, attribution }); }; export function useDatabasesDirect( { connectionId }: { connectionId: string }, options?: Omit, 'queryKey'>, ) { - const clickhouseClient = getClickhouseClient(); + const clickhouseClient = getClickhouseClient({ + attribution: useMetadataQueryAttribution(), + }); return useQuery, Error>({ queryKey: [`direct_datasources/databases`, connectionId], queryFn: async () => { @@ -87,7 +104,9 @@ export function useTablesDirect( { database, connectionId }: { database: string; connectionId: string }, options?: Omit, 'queryKey'>, ) { - const clickhouseClient = getClickhouseClient(); + const clickhouseClient = getClickhouseClient({ + attribution: useMetadataQueryAttribution(), + }); return useQuery, Error>({ queryKey: [`direct_datasources/databases/${database}/tables`, connectionId], queryFn: async () => { diff --git a/packages/app/src/hooks/useFetchMetricAttributeValues.tsx b/packages/app/src/hooks/useFetchMetricAttributeValues.tsx index 2bbb3e01a0..bcc397b24e 100644 --- a/packages/app/src/hooks/useFetchMetricAttributeValues.tsx +++ b/packages/app/src/hooks/useFetchMetricAttributeValues.tsx @@ -7,6 +7,7 @@ import { SourceKind, TMetricSource } from '@hyperdx/common-utils/dist/types'; import { useQuery } from '@tanstack/react-query'; import { getClickhouseClient } from '@/clickhouse'; +import { useMetadataQueryAttribution } from '@/queryAttribution'; import { getMetricTableName } from '@/utils'; import { AttributeCategory } from './useFetchMetricResourceAttrs'; @@ -54,6 +55,8 @@ export const useFetchMetricAttributeValues = ({ tableSource?.kind === SourceKind.Metric, ); + const attribution = useMetadataQueryAttribution(); + return useQuery({ queryKey: [ 'metric-attribute-values', @@ -69,7 +72,7 @@ export const useFetchMetricAttributeValues = ({ return []; } - const clickhouseClient = getClickhouseClient(); + const clickhouseClient = getClickhouseClient({ attribution }); // Build optional search filter const searchFilter = searchTerm diff --git a/packages/app/src/hooks/useFetchMetricMetadata.tsx b/packages/app/src/hooks/useFetchMetricMetadata.tsx index f4150ad737..a6a54ce82c 100644 --- a/packages/app/src/hooks/useFetchMetricMetadata.tsx +++ b/packages/app/src/hooks/useFetchMetricMetadata.tsx @@ -7,6 +7,7 @@ import { SourceKind, TMetricSource } from '@hyperdx/common-utils/dist/types'; import { useQuery } from '@tanstack/react-query'; import { getClickhouseClient } from '@/clickhouse'; +import { useMetadataQueryAttribution } from '@/queryAttribution'; import { getMetricTableName } from '@/utils'; export interface MetricMetadata { @@ -45,6 +46,8 @@ export const useFetchMetricMetadata = ({ tableSource?.kind === SourceKind.Metric, ); + const attribution = useMetadataQueryAttribution(); + return useQuery({ queryKey: ['metric-metadata', databaseName, metricType, metricName], queryFn: async ({ signal }) => { @@ -52,7 +55,7 @@ export const useFetchMetricMetadata = ({ return null; } - const clickhouseClient = getClickhouseClient(); + const clickhouseClient = getClickhouseClient({ attribution }); const sql = chSql` SELECT MetricUnit, diff --git a/packages/app/src/hooks/useFetchMetricResourceAttrs.tsx b/packages/app/src/hooks/useFetchMetricResourceAttrs.tsx index e6fcb46f23..a2c35d6591 100644 --- a/packages/app/src/hooks/useFetchMetricResourceAttrs.tsx +++ b/packages/app/src/hooks/useFetchMetricResourceAttrs.tsx @@ -7,6 +7,7 @@ import { SourceKind, TMetricSource } from '@hyperdx/common-utils/dist/types'; import { useQuery } from '@tanstack/react-query'; import { getClickhouseClient } from '@/clickhouse'; +import { useMetadataQueryAttribution } from '@/queryAttribution'; import { formatAttributeClause, getMetricTableName } from '@/utils'; const METRIC_FETCH_LIMIT = 10000; @@ -159,6 +160,8 @@ export const useFetchMetricResourceAttrs = ({ tableSource?.kind === SourceKind.Metric, ); + const attribution = useMetadataQueryAttribution(); + return useQuery({ queryKey: ['metric-attributes', metricType, metricName, isSql, tableSource], queryFn: async ({ signal }) => { @@ -166,7 +169,7 @@ export const useFetchMetricResourceAttrs = ({ return []; } - const clickhouseClient = getClickhouseClient(); + const clickhouseClient = getClickhouseClient({ attribution }); const sql = chSql` SELECT DISTINCT ScopeAttributes, diff --git a/packages/app/src/hooks/useMetricCatalog.ts b/packages/app/src/hooks/useMetricCatalog.ts index 01cd71898e..94209a2c02 100644 --- a/packages/app/src/hooks/useMetricCatalog.ts +++ b/packages/app/src/hooks/useMetricCatalog.ts @@ -15,6 +15,7 @@ import { useQuery } from '@tanstack/react-query'; import { getClickhouseClient } from '@/clickhouse'; import { useMetadataWithSettings } from '@/hooks/useMetadata'; +import { useMetadataQueryAttribution } from '@/queryAttribution'; import { QUERYABLE_KINDS } from '@/utils/metricKinds'; import { mergeMetricCatalog, @@ -87,6 +88,8 @@ export function useMetricCatalog({ return tableName ? [{ kind, tableName }] : []; }); + const attribution = useMetadataQueryAttribution(); + const query = useQuery({ queryKey: [ 'useMetricCatalog', @@ -101,7 +104,7 @@ export function useMetricCatalog({ // unstable time references leaking into render. This runs inside the // queryFn, not during render, but the rule is syntactic. const [start, end] = clampCatalogDateRange(dateRange, dayjs().toDate()); - const clickhouseClient = getClickhouseClient(); + const clickhouseClient = getClickhouseClient({ attribution }); // Settled, not all: a source can legitimately have one misconfigured or // ungranted kind table, and failing the whole catalog would hide every diff --git a/packages/app/src/hooks/useOffsetPaginatedQuery.tsx b/packages/app/src/hooks/useOffsetPaginatedQuery.tsx index 5a27d39b96..cc8dda06d6 100644 --- a/packages/app/src/hooks/useOffsetPaginatedQuery.tsx +++ b/packages/app/src/hooks/useOffsetPaginatedQuery.tsx @@ -10,6 +10,7 @@ import { ChSql, ClickHouseQueryError, ColumnMetaType, + type QueryAttribution, } from '@hyperdx/common-utils/dist/clickhouse'; import { Metadata } from '@hyperdx/common-utils/dist/core/metadata'; import { renderChartConfig } from '@hyperdx/common-utils/dist/core/renderChartConfig'; @@ -38,6 +39,7 @@ import { getClickhouseClient } from '@/clickhouse'; import { MAX_TABLE_ROWS } from '@/HDXMultiSeriesTableChart'; import { useMetadataWithSettings } from '@/hooks/useMetadata'; import { useMVOptimizationExplanation } from '@/hooks/useMVOptimizationExplanation'; +import { useQueryAttribution } from '@/queryAttribution'; import { useSource } from '@/source'; import { DEFAULT_TIME_WINDOWS_SECONDS, @@ -85,6 +87,12 @@ type QueryMeta = { metadata: Metadata; optimizedConfig?: ChartConfigWithOptTimestamp; source: TSource | undefined; + /** + * Read during render and carried here because the queryFn runs outside + * React and cannot read context. `meta` rather than `queryKey`, which would + * split the cache per surface. + */ + attribution: QueryAttribution; }; // Get time window from page param @@ -179,6 +187,7 @@ const queryFn: QueryFunction = async ({ hasPreviousQueries, optimizedConfig, source, + attribution, } = meta as QueryMeta; // Only stream incrementally if this is a fresh query with no previous @@ -188,7 +197,7 @@ const queryFn: QueryFunction = async ({ !hasPreviousQueries || pageParam.offset > 0 || pageParam.windowIndex > 0; const queryTimeout = queryKey[2]; - const clickhouseClient = getClickhouseClient({ queryTimeout }); + const clickhouseClient = getClickhouseClient({ queryTimeout, attribution }); const rawConfig = queryKey[1]; const config = optimizedConfig ?? rawConfig; @@ -458,6 +467,7 @@ export default function useOffsetPaginatedQuery( const key = queryKeyFn(queryKeyPrefix, config, meData?.team?.queryTimeout); const queryClient = useQueryClient(); const metadata = useMetadataWithSettings(); + const attribution = useQueryAttribution(); const matchedQueries = queryClient.getQueriesData({ queryKey: [queryKeyPrefix, omit(config, ['dateRange'])], }); @@ -520,6 +530,7 @@ export default function useOffsetPaginatedQuery( metadata, optimizedConfig: mvOptimizationData?.optimizedConfig, source, + attribution, } satisfies QueryMeta, queryFn, gcTime: isLive ? ms('30s') : ms('5m'), // more aggressive gc for live data, since it can end up holding lots of data diff --git a/packages/app/src/layout.tsx b/packages/app/src/layout.tsx index e2c2ef5c65..a65837a810 100644 --- a/packages/app/src/layout.tsx +++ b/packages/app/src/layout.tsx @@ -1,11 +1,13 @@ import React from 'react'; import Link from 'next/link'; import { useRouter } from 'next/router'; +import type { QuerySurface } from '@hyperdx/common-utils/dist/clickhouse'; import { Button, Center, Group, Text } from '@mantine/core'; import { IconX } from '@tabler/icons-react'; import AppNav from '@/components/AppNav'; import { IS_CLICKHOUSE_BUILD } from '@/config'; +import { QueryAttributionProvider } from '@/queryAttribution'; import { HDXSpotlightProvider } from './Spotlights'; import { useLocalStorage } from './utils'; @@ -105,3 +107,24 @@ export const withAppNav = (page: React.ReactNode): React.ReactNode => { ); }; + +/** + * `withAppNav`, and tags the page's ClickHouse queries so they stand out from + * other pages' in `system.query_log`. + * + * Use it when "which page" is all you need. A page with ids of its own, like a + * dashboard tile, adds `QueryAttributionProvider` further in instead. + * + * `label` tells a page apart from something else using the same surface, such + * as the alert detail page and the alert checker. + * + * @example KubernetesPage.getLayout = withAppNavForSurface('dashboard'); + */ +export const withAppNavForSurface = + (surface: QuerySurface, label?: string) => + (page: React.ReactNode): React.ReactNode => + withAppNav( + + {page} + , + ); diff --git a/packages/app/src/llm/dashboard/LLMDashboardPage.tsx b/packages/app/src/llm/dashboard/LLMDashboardPage.tsx index 619e186ac6..1d52766293 100644 --- a/packages/app/src/llm/dashboard/LLMDashboardPage.tsx +++ b/packages/app/src/llm/dashboard/LLMDashboardPage.tsx @@ -32,7 +32,7 @@ import SearchWhereInput, { } from '@/components/SearchInput/SearchWhereInput'; import { SourceSelectControlled } from '@/components/SourceSelect'; import { TimePicker } from '@/components/TimePicker'; -import { withAppNav } from '@/layout'; +import { withAppNavForSurface } from '@/layout'; import { useLLMDashboardExpressions, useLLMLogDashboardExpressions, @@ -377,6 +377,6 @@ const LLMDashboardPageDynamic = dynamic(async () => LLMDashboardPage, { }); // @ts-expect-error Next.js layout typing -LLMDashboardPageDynamic.getLayout = withAppNav; +LLMDashboardPageDynamic.getLayout = withAppNavForSurface('dashboard', 'llm'); export default LLMDashboardPageDynamic; diff --git a/packages/app/src/metadata.ts b/packages/app/src/metadata.ts index 0f7fc0575d..459a5bbb2f 100644 --- a/packages/app/src/metadata.ts +++ b/packages/app/src/metadata.ts @@ -5,5 +5,14 @@ import { getClickhouseClient } from '@/clickhouse'; import { DEFAULT_QUERY_TIMEOUT } from './defaults'; // TODO: Get rid of this function and convert to singleton +// +// This instance is shared across pages, so it cannot know which page caused a +// lookup. Labelling it at least separates field lookups and autocomplete from +// a user's chart queries, and there are a lot of them. export const getMetadata = () => - _getMetadata(getClickhouseClient({ queryTimeout: DEFAULT_QUERY_TIMEOUT })); + _getMetadata( + getClickhouseClient({ + queryTimeout: DEFAULT_QUERY_TIMEOUT, + attribution: { surface: 'metadata' }, + }), + ); diff --git a/packages/app/src/queryAttribution.tsx b/packages/app/src/queryAttribution.tsx new file mode 100644 index 0000000000..e7c64ece61 --- /dev/null +++ b/packages/app/src/queryAttribution.tsx @@ -0,0 +1,57 @@ +import { createContext, use, useMemo } from 'react'; +import { + mergeQueryAttribution, + QueryAttribution, +} from '@hyperdx/common-utils/dist/clickhouse'; + +/** + * What the ClickHouse queries below this point belong to. + * `useClickhouseClient` reads it and tags each query. + * + * Providers nest, so a dashboard sets its id once and each tile adds its own. + * + * This stays off the chart config on purpose: the config is the React Query + * cache key, so two pages asking the same question would stop sharing an + * answer and we would run the query twice. + */ +const QueryAttributionContext = createContext({}); + +/** What is in effect here. */ +export function useQueryAttribution(): QueryAttribution { + return use(QueryAttributionContext); +} + +/** + * For schema and metric lookups. Page context first so its ids are kept, then + * `metadata` pinned over the top: these are not the page's own chart queries. + */ +export function useMetadataQueryAttribution(): QueryAttribution { + return mergeQueryAttribution(useQueryAttribution(), { surface: 'metadata' }); +} + +export function QueryAttributionProvider({ + attribution, + children, +}: { + attribution: QueryAttribution; + children: React.ReactNode; +}) { + const parent = use(QueryAttributionContext); + + // Callers build `attribution` inline, so it is a new object every render. + // Comparing the serialized value instead avoids re-rendering every chart + // below. Serializing beats listing the fields, which we would forget to + // update when a field is added. + // + // The memo reads the string back rather than closing over `attribution`, so + // its dependencies are the whole truth. + const attributionKey = JSON.stringify(attribution); + const merged = useMemo(() => { + const parsed: QueryAttribution = JSON.parse(attributionKey); + return mergeQueryAttribution(parent, parsed); + }, [parent, attributionKey]); + + return ( + {children} + ); +} diff --git a/packages/app/src/sessions.ts b/packages/app/src/sessions.ts index 19e8e84716..f7b229452e 100644 --- a/packages/app/src/sessions.ts +++ b/packages/app/src/sessions.ts @@ -16,6 +16,7 @@ import { } from '@hyperdx/common-utils/dist/types'; import { useQuery, UseQueryOptions } from '@tanstack/react-query'; +import { useQueryAttribution } from '@/queryAttribution'; import { usePrevious } from '@/utils'; import useFieldExpressionGenerator, { @@ -427,6 +428,7 @@ export function useRRWebEventStream( const lastFetchStatusRef = useRef<'fetching' | 'idle' | undefined>(undefined); const { data: source } = useSource({ id: sourceId }); + const attribution = useQueryAttribution(); const fetchResults = useCallback( async ({ @@ -471,7 +473,7 @@ export function useRRWebEventStream( const format = 'JSONEachRow'; const fetchPromise = (async () => { - const clickhouseClient = getClickhouseClient(); + const clickhouseClient = getClickhouseClient({ attribution }); const resultSet = await clickhouseClient.query({ query: query.sql, query_params: query.params, @@ -570,6 +572,7 @@ export function useRRWebEventStream( resultsKey, metadata, getSessionSourceFieldExpression, + attribution, ], ); diff --git a/packages/cli/src/api/client.ts b/packages/cli/src/api/client.ts index fa39266c94..504401e670 100644 --- a/packages/cli/src/api/client.ts +++ b/packages/cli/src/api/client.ts @@ -299,6 +299,10 @@ export class ProxyClickhouseClient extends BaseClickhouseClient { ) { super({ host: `${apiClient.getApiUrl()}/clickhouse-proxy`, + // This extends the base client, not the node one, so it never sees the + // AsyncLocalStorage scope the API uses. A client-level surface is all + // the attribution the CLI gets. + attribution: { surface: 'cli' }, ...opts, }); this.apiClient = apiClient; diff --git a/packages/common-utils/src/__tests__/queryAttribution.int.test.ts b/packages/common-utils/src/__tests__/queryAttribution.int.test.ts new file mode 100644 index 0000000000..87a8dbb9f4 --- /dev/null +++ b/packages/common-utils/src/__tests__/queryAttribution.int.test.ts @@ -0,0 +1,122 @@ +import { ClickHouseClient, createClient } from '@clickhouse/client'; + +import { + ClickhouseClient as HdxClickhouseClient, + withQueryAttribution, +} from '@/clickhouse/node'; + +/** + * Proves the attribution actually survives the round trip into ClickHouse. + * + * The unit tests assert what the client *sends*; only a real server can + * confirm that `log_comment` is accepted as a setting and lands verbatim in + * `system.query_log`, which is the whole point of the feature. + */ +describe('query attribution integration', () => { + let client: ClickHouseClient; + let hdxClient: HdxClickhouseClient; + + const host = process.env.CLICKHOUSE_HOST || 'http://localhost:8123'; + const username = process.env.CLICKHOUSE_USER || 'default'; + const password = process.env.CLICKHOUSE_PASSWORD || ''; + + beforeAll(() => { + client = createClient({ url: host, username, password }); + hdxClient = new HdxClickhouseClient({ host, username, password }); + }); + + afterAll(async () => { + await client.close(); + await hdxClient.close(); + }); + + /** + * `system.query_log` is written asynchronously, so a lookup straight after + * the query usually misses. Flush explicitly, then poll. + */ + async function findLoggedQuery(queryId: string) { + 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 query_id = {queryId:String} AND type = 'QueryFinish' + LIMIT 1`, + query_params: { queryId }, + format: 'JSON', + }); + const { data } = await result.json<{ + query_id: string; + log_comment: string; + }>(); + if (data.length > 0) return data[0]; + await new Promise(resolve => setTimeout(resolve, 250)); + } + return undefined; + } + + it('writes the attribution into system.query_log', async () => { + const dashboard = `dash-${Date.now()}`; + + const resultSet = await withQueryAttribution( + { + surface: 'dashboard', + dashboard, + tile: 'tile-42', + }, + async () => hdxClient.query({ query: 'SELECT 1 AS one', format: 'JSON' }), + ); + await resultSet.json(); + + // The client generates the id, so it is known without reading it back. + const queryId = resultSet.query_id; + expect(queryId).toMatch(/^hdx-dashboard-/); + + const logged = await findLoggedQuery(queryId); + expect(logged).toBeDefined(); + expect(JSON.parse(logged!.log_comment)).toMatchObject({ + v: 1, + surface: 'dashboard', + dashboard, + tile: 'tile-42', + }); + }); + + it('keeps the payload queryable with JSONExtractString', async () => { + const tile = `tile-${Date.now()}`; + + const resultSet = await withQueryAttribution( + { surface: 'search', search: 'saved-1', tile }, + async () => hdxClient.query({ query: 'SELECT 2 AS two', format: 'JSON' }), + ); + await resultSet.json(); + expect(await findLoggedQuery(resultSet.query_id)).toBeDefined(); + + // The access pattern an operator would actually use to attribute load. + const result = await client.query({ + query: ` + SELECT JSONExtractString(log_comment, 'surface') AS surface + FROM system.query_log + WHERE JSONExtractString(log_comment, 'tile') = {tile:String} + AND type = 'QueryFinish' + LIMIT 1`, + query_params: { tile }, + format: 'JSON', + }); + const { data } = await result.json<{ surface: string }>(); + expect(data[0]?.surface).toBe('search'); + }); + + it('leaves log_comment empty when nothing is known', async () => { + const resultSet = await hdxClient.query({ + query: 'SELECT 3 AS three', + format: 'JSON', + }); + await resultSet.json(); + + const logged = await findLoggedQuery(resultSet.query_id); + expect(logged).toBeDefined(); + expect(logged!.log_comment).toBe(''); + }); +}); diff --git a/packages/common-utils/src/clickhouse/__tests__/attribution.test.ts b/packages/common-utils/src/clickhouse/__tests__/attribution.test.ts new file mode 100644 index 0000000000..b389259b44 --- /dev/null +++ b/packages/common-utils/src/clickhouse/__tests__/attribution.test.ts @@ -0,0 +1,294 @@ +import type { BaseResultSet, DataFormat } from '@clickhouse/client-common'; + +import { BaseClickhouseClient, QueryInputs } from '@/clickhouse'; +import { + buildLogComment, + buildQueryId, + mergeQueryAttribution, + QUERY_ATTRIBUTION_VERSION, + QueryAttribution, + QuerySurface, +} from '@/clickhouse/attribution'; + +describe('mergeQueryAttribution', () => { + it('lets later layers win', () => { + expect( + mergeQueryAttribution( + { surface: 'dashboard', dashboard: 'dash-1' }, + { tile: 'tile-1' }, + ), + ).toEqual({ surface: 'dashboard', dashboard: 'dash-1', tile: 'tile-1' }); + }); + + it('does not let a later layer blank out an earlier value', () => { + expect( + mergeQueryAttribution( + { surface: 'dashboard', dashboard: 'dash-1' }, + { dashboard: undefined, tile: '' }, + ), + ).toEqual({ surface: 'dashboard', dashboard: 'dash-1' }); + }); + + it('ignores undefined layers', () => { + expect( + mergeQueryAttribution(undefined, { search: 's-1' }, undefined), + ).toEqual({ search: 's-1' }); + }); +}); + +describe('buildLogComment', () => { + it('returns undefined when there is nothing to stamp', () => { + expect(buildLogComment(undefined)).toBeUndefined(); + expect(buildLogComment({})).toBeUndefined(); + }); + + it('returns undefined when only unknown fields are present', () => { + // A surface alone is worth stamping; an empty id is not. + expect(buildLogComment({ dashboard: ' ' })).toBeUndefined(); + }); + + it('emits versioned JSON with the populated fields only', () => { + const comment = buildLogComment({ + surface: 'dashboard', + dashboard: 'dash-1', + tile: 'tile-1', + }); + + expect(JSON.parse(comment!)).toEqual({ + v: QUERY_ATTRIBUTION_VERSION, + surface: 'dashboard', + dashboard: 'dash-1', + tile: 'tile-1', + }); + }); + + it('drops a surface that is not in the closed set', () => { + const comment = buildLogComment({ + surface: 'definitely-not-a-surface' as unknown as QuerySurface, + search: 'search-1', + }); + + expect(JSON.parse(comment!)).toEqual({ + v: QUERY_ATTRIBUTION_VERSION, + search: 'search-1', + }); + }); + + /** + * From the browser this rides in the URL query string. The proxy decodes it + * before re-parsing, so an `&` surviving here would truncate the value and + * inject a bogus parameter, failing every query on the page. + */ + it('drops characters that would break the URL it travels in', () => { + const comment = buildLogComment({ + surface: 'search', + search: 'a&b?c=d#e%f+g', + }); + + expect(JSON.parse(comment!).search).toBe('abcdefg'); + }); + + it('keeps the characters real ids and routes are made of', () => { + const comment = buildLogComment({ + surface: 'api', + source: '68b2f4c1a9e3d7b5c2a10f51', + tile: '0f7c9a2e-4d1b-4a6f-9c3e-2b8d5f1a7c40', + label: '/api/v2/charts', + }); + + const parsed = JSON.parse(comment!); + expect(parsed.source).toBe('68b2f4c1a9e3d7b5c2a10f51'); + expect(parsed.tile).toBe('0f7c9a2e-4d1b-4a6f-9c3e-2b8d5f1a7c40'); + expect(parsed.label).toBe('/api/v2/charts'); + }); + + it('drops non-ASCII rather than risk splitting it at the length cap', () => { + const comment = buildLogComment({ + surface: 'search', + label: `x${String.fromCodePoint(0x1f600)}y`, + }); + + expect(JSON.parse(comment!).label).toBe('xy'); + }); + + it('strips control characters so the payload stays parseable', () => { + const comment = buildLogComment({ + surface: 'search', + label: `a${String.fromCharCode(0)}b\nc`, + }); + + expect(JSON.parse(comment!).label).toBe('abc'); + }); + + it('caps an individual field', () => { + const comment = buildLogComment({ + surface: 'search', + label: 'x'.repeat(500), + }); + + expect(JSON.parse(comment!).label.length).toBe(128); + }); + + /** + * Every field at its longest, with the longest surface name. Exact rather + * than approximate because the allowlist is ASCII, so the per-field cap + * counts bytes too. Fails if a field is added or a cap raised without + * checking the budget the browser spends on every request's query string. + * `Required` makes a new field a compile error here until it is added. + */ + it('cannot exceed the size cap even when every field is at its longest', () => { + const tooLong = 'x'.repeat(500); + const worstCase: Required = { + surface: 'service-dashboard', + dashboard: tooLong, + tile: tooLong, + search: tooLong, + alert: tooLong, + source: tooLong, + trace: tooLong, + label: tooLong, + }; + const comment = buildLogComment(worstCase); + + expect(Buffer.byteLength(comment!)).toBeLessThanOrEqual(1024); + // Nothing is shed today, so the worst case keeps every field. + expect(Object.keys(JSON.parse(comment!))).toEqual([ + 'v', + 'surface', + 'dashboard', + 'tile', + 'search', + 'alert', + 'source', + 'trace', + 'label', + ]); + }); + + it('stays within budget for a realistic payload', () => { + const comment = buildLogComment({ + surface: 'dashboard', + dashboard: '68b2f4c1a9e3d7b5c2a10f4e', + tile: '0f7c9a2e-4d1b-4a6f-9c3e-2b8d5f1a7c40', + source: '68b2f4c1a9e3d7b5c2a10f51', + trace: 'c3f1a9b7d2e4568a0b1c2d3e4f506172', + }); + + expect(Buffer.byteLength(comment!)).toBeLessThan(250); + }); + + it('produces valid JSON for values that need escaping', () => { + const comment = buildLogComment({ + surface: 'search', + label: 'quote " brace } backslash \\ ok', + }); + + // The quote, brace and backslash are not on the allowlist, so they never + // reach JSON.stringify in the first place. + expect(() => JSON.parse(comment!)).not.toThrow(); + expect(JSON.parse(comment!).label).toBe('quote brace backslash ok'); + }); +}); + +describe('buildQueryId', () => { + it('prefixes with the surface', () => { + expect(buildQueryId({ surface: 'alert' })).toMatch(/^hdx-alert-/); + }); + + it('falls back to unknown without a surface', () => { + expect(buildQueryId(undefined)).toMatch(/^hdx-unknown-/); + + expect( + buildQueryId({ surface: 'nope' as unknown as QuerySurface }), + ).toMatch(/^hdx-unknown-/); + }); + + it('is unique per call', () => { + const attribution: QueryAttribution = { surface: 'dashboard' }; + const ids = new Set( + Array.from({ length: 200 }, () => buildQueryId(attribution)), + ); + expect(ids.size).toBe(200); + }); +}); + +/** + * Captures what the base class hands to the transport, so these assertions + * cover the stamping every concrete client (browser, node, CLI) inherits. + */ +class RecordingClient extends BaseClickhouseClient { + public lastInputs?: QueryInputs<'JSON'>; + + protected async __query( + inputs: QueryInputs, + ): Promise> { + this.lastInputs = inputs as unknown as QueryInputs<'JSON'>; + // Nothing here reads the result; only the inputs are under test. + return {} as unknown as BaseResultSet; + } +} + +describe('BaseClickhouseClient attribution', () => { + const newClient = (attribution?: QueryAttribution) => + new RecordingClient({ host: 'http://localhost:8123', attribution }); + + it('stamps log_comment and a prefixed query_id', async () => { + const client = newClient({ surface: 'dashboard', dashboard: 'dash-1' }); + await client.query({ query: 'SELECT 1' }); + + const settings = client.lastInputs?.clickhouse_settings; + expect(JSON.parse(String(settings?.log_comment))).toMatchObject({ + surface: 'dashboard', + dashboard: 'dash-1', + }); + expect(client.lastInputs?.queryId).toMatch(/^hdx-dashboard-/); + }); + + it('merges per-query attribution over the client default', async () => { + const client = newClient({ surface: 'dashboard', dashboard: 'dash-1' }); + await client.query({ + query: 'SELECT 1', + attribution: { tile: 'tile-9' }, + }); + + expect( + JSON.parse(String(client.lastInputs?.clickhouse_settings?.log_comment)), + ).toMatchObject({ + surface: 'dashboard', + dashboard: 'dash-1', + tile: 'tile-9', + }); + }); + + it('preserves settings the caller already passed', async () => { + const client = newClient({ surface: 'search' }); + await client.query({ + query: 'SELECT 1', + clickhouse_settings: { max_execution_time: 7 }, + }); + + const settings = client.lastInputs?.clickhouse_settings; + expect(settings?.max_execution_time).toBe(7); + expect(settings?.log_comment).toBeDefined(); + }); + + it('never overrides a log_comment or query_id the caller set', async () => { + const client = newClient({ surface: 'search' }); + await client.query({ + query: 'SELECT 1', + clickhouse_settings: { log_comment: 'mine' }, + queryId: 'my-id', + }); + + expect(client.lastInputs?.clickhouse_settings?.log_comment).toBe('mine'); + expect(client.lastInputs?.queryId).toBe('my-id'); + }); + + it('omits log_comment entirely when nothing is known', async () => { + const client = newClient(); + await client.query({ query: 'SELECT 1' }); + + expect(client.lastInputs?.clickhouse_settings?.log_comment).toBeUndefined(); + expect(client.lastInputs?.queryId).toMatch(/^hdx-unknown-/); + }); +}); diff --git a/packages/common-utils/src/clickhouse/__tests__/nodeAttribution.test.ts b/packages/common-utils/src/clickhouse/__tests__/nodeAttribution.test.ts new file mode 100644 index 0000000000..169302837d --- /dev/null +++ b/packages/common-utils/src/clickhouse/__tests__/nodeAttribution.test.ts @@ -0,0 +1,130 @@ +import type { BaseResultSet, DataFormat } from '@clickhouse/client-common'; + +import { QueryInputs } from '@/clickhouse'; +import { + ClickhouseClient, + getCurrentQueryAttribution, + withQueryAttribution, +} from '@/clickhouse/node'; + +/** Records the inputs the base class produced instead of hitting the network. */ +class RecordingNodeClient extends ClickhouseClient { + public lastInputs?: QueryInputs<'JSON'>; + + protected async __query( + inputs: QueryInputs, + ): Promise> { + this.lastInputs = inputs as unknown as QueryInputs<'JSON'>; + return {} as unknown as BaseResultSet; + } +} + +const newClient = () => + new RecordingNodeClient({ host: 'http://localhost:8123' }); + +const logCommentOf = (client: RecordingNodeClient) => + JSON.parse(String(client.lastInputs?.clickhouse_settings?.log_comment)); + +describe('withQueryAttribution', () => { + it('exposes nothing outside a scope', () => { + expect(getCurrentQueryAttribution()).toBeUndefined(); + }); + + it('nests, with the inner scope winning on conflict', () => { + withQueryAttribution({ surface: 'alert', label: 'outer' }, () => { + withQueryAttribution({ alert: 'alert-1', label: 'inner' }, () => { + expect(getCurrentQueryAttribution()).toEqual({ + surface: 'alert', + label: 'inner', + alert: 'alert-1', + }); + }); + expect(getCurrentQueryAttribution()).toEqual({ + surface: 'alert', + label: 'outer', + }); + }); + }); + + it('survives await boundaries', async () => { + await withQueryAttribution({ surface: 'mcp', label: 'tool' }, async () => { + await Promise.resolve(); + await new Promise(resolve => setTimeout(resolve, 1)); + expect(getCurrentQueryAttribution()?.label).toBe('tool'); + }); + }); + + it('keeps concurrent scopes isolated', async () => { + const seen = await Promise.all([ + withQueryAttribution({ surface: 'alert', alert: 'a' }, async () => { + await new Promise(resolve => setTimeout(resolve, 5)); + return getCurrentQueryAttribution()?.alert; + }), + withQueryAttribution({ surface: 'alert', alert: 'b' }, async () => { + await new Promise(resolve => setTimeout(resolve, 1)); + return getCurrentQueryAttribution()?.alert; + }), + ]); + + expect(seen).toEqual(['a', 'b']); + }); +}); + +describe('node client attribution precedence', () => { + it('stamps a query issued inside a scope', async () => { + const client = newClient(); + await withQueryAttribution( + { surface: 'alert', alert: 'alert-1' }, + async () => { + await client.query({ query: 'SELECT 1' }); + }, + ); + + expect(logCommentOf(client)).toMatchObject({ + surface: 'alert', + alert: 'alert-1', + }); + expect(client.lastInputs?.queryId).toMatch(/^hdx-alert-/); + }); + + it('lets the request scope override the client default', async () => { + const client = new RecordingNodeClient({ + host: 'http://localhost:8123', + attribution: { surface: 'api', label: 'default' }, + }); + + await withQueryAttribution({ label: 'scoped' }, async () => { + await client.query({ query: 'SELECT 1' }); + }); + + expect(logCommentOf(client)).toMatchObject({ + surface: 'api', + label: 'scoped', + }); + }); + + it('lets a per-query value override the ambient scope', async () => { + const client = newClient(); + await withQueryAttribution({ surface: 'mcp', label: 'outer' }, async () => { + await client.query({ + query: 'SELECT 1', + attribution: { label: 'inner' }, + }); + }); + + expect(logCommentOf(client)).toMatchObject({ + surface: 'mcp', + label: 'inner', + }); + }); + + it('falls back to the client default outside any scope', async () => { + const client = new RecordingNodeClient({ + host: 'http://localhost:8123', + attribution: { surface: 'api' }, + }); + await client.query({ query: 'SELECT 1' }); + + expect(logCommentOf(client)).toMatchObject({ surface: 'api' }); + }); +}); diff --git a/packages/common-utils/src/clickhouse/attribution.ts b/packages/common-utils/src/clickhouse/attribution.ts new file mode 100644 index 0000000000..f3b8f7a251 --- /dev/null +++ b/packages/common-utils/src/clickhouse/attribution.ts @@ -0,0 +1,247 @@ +/** + * Tags each ClickHouse query with what asked for it, so a row in + * `system.query_log` points back to a tile, a search, or an alert. + * + * Records what was asked for, never who asked. No team, user or session id + * goes in here. + * + * Two places carry the tag. `log_comment` holds the JSON and is what you query + * later. `query_id` gets an `hdx--` prefix, which is readable in + * `system.processes` while the query is still running. + * + * None of this changes what a query returns, so it must never throw. + * + * One exception: queries that skip settings processing (the `system.settings`, + * server version and Cloud-detection probes, and the onboarding connection + * check) get a `query_id` but an empty `log_comment`. + */ + +/** + * Which part of the product sent the query. + */ +export const QUERY_SURFACES = [ + 'alert', + 'api', + 'chart-explorer', + 'chart-preview', + 'cli', + 'dashboard', + 'mcp', + 'metadata', + 'search', + 'service-dashboard', + 'session-replay', + 'unknown', +] as const; + +export type QuerySurface = (typeof QUERY_SURFACES)[number]; + +/** + * What we know about one query. Everything is optional; callers fill in + * whatever they have. + * + * These key names end up in the log, so people write queries against them. + * Renaming one breaks those queries, which is what `QUERY_ATTRIBUTION_VERSION` + * is for. + */ +export type QueryAttribution = { + surface?: QuerySurface; + dashboard?: string; + tile?: string; + search?: string; + alert?: string; + source?: string; + trace?: string; + /** Anything more specific than the surface, e.g. an MCP tool name. */ + label?: string; +}; + +/** + * Written as `v` into every log comment. Callers cannot set it per query. + */ +export const QUERY_ATTRIBUTION_VERSION = 1; + +/** + * ClickHouse allows far more, but the browser sends settings in the URL, and a + * long URL pushes the request onto a path some proxies reject. + */ +const MAX_LOG_COMMENT_BYTES = 1024; + +/** Longest any single value may be. */ +const MAX_FIELD_LENGTH = 128; + +/** + * Most useful first. If the payload runs out of room, fields at the end are + * dropped whole, so the result is still valid JSON. + * + * Today's fields fit the budget even at full length, with only a few bytes to + * 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. + */ +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], + ]; +} + +/** + * Keep an allowlist of characters and cap the length. + * + * An allowlist rather than a denylist because from the browser this value + * travels in the URL query string. An `&`, `?` or `#` reaching the proxy ends + * the value early and injects a bogus parameter, which ClickHouse then + * rejects — so one odd character in a dashboard id would fail every query on + * the page, not just lose its tag. + * + * Restricting to ASCII also means the length cap counts bytes as well as + * characters, so the payload budget is exact, and no multi-byte character can + * be cut in half into something ClickHouse cannot parse as JSON. + * + * Everything we put here is an id, a route or a name we chose, so nothing + * legitimate is lost. Values are only ever JSON strings, never SQL. + */ +function sanitizeField(value: string | undefined): string | undefined { + if (typeof value !== 'string') return undefined; + + // Checked by code point, not a regex, so no control character has to appear + // in this file - a stray one would make it read as binary to grep. + const isAllowed = (code: number) => + (code >= 0x61 && code <= 0x7a) || // a-z + (code >= 0x41 && code <= 0x5a) || // A-Z + (code >= 0x30 && code <= 0x39) || // 0-9 + code === 0x20 || // space + code === 0x2d || // - + code === 0x2e || // . + code === 0x2f || // / + code === 0x3a || // : + code === 0x5f; // _ + + let kept = ''; + for (const char of value) { + const code = char.codePointAt(0) ?? 0; + if (isAllowed(code)) kept += char; + } + + const cleaned = kept.trim(); + if (!cleaned) return undefined; + return cleaned.slice(0, MAX_FIELD_LENGTH); +} + +function isQuerySurface(value: unknown): value is QuerySurface { + return ( + typeof value === 'string' && + (QUERY_SURFACES as readonly string[]).includes(value) + ); +} + +/** + * Later layers win, but only where they have a value. So a tile can add its own + * id without erasing the dashboard id it sits inside. + */ +export function mergeQueryAttribution( + ...layers: (QueryAttribution | undefined)[] +): QueryAttribution { + const 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; + } + return merged; +} + +/** Returns undefined when there is nothing worth recording. */ +export function buildLogComment( + attribution: QueryAttribution | undefined, +): string | undefined { + if (!attribution) return undefined; + + // A Map, not an object literal: insertion order is the drop order, and + // nothing here is written through a variable key. + const payload = new Map([ + ['v', QUERY_ATTRIBUTION_VERSION], + ]); + + const surface = isQuerySurface(attribution.surface) + ? attribution.surface + : undefined; + if (surface) { + payload.set('surface', surface); + } + + let serialized = stringify(payload); + + for (const [field, raw] of idFields(attribution)) { + const cleaned = sanitizeField(raw); + if (!cleaned) continue; + + payload.set(field, cleaned); + const candidate = stringify(payload); + // sanitizeField keeps only ASCII, so length is the byte count. + if (candidate.length > MAX_LOG_COMMENT_BYTES) { + payload.delete(field); + continue; + } + serialized = candidate; + } + + // Only the version survived, so there is nothing to say. + return payload.size > 1 ? serialized : undefined; +} + +function stringify(payload: Map): string { + return JSON.stringify(Object.fromEntries(payload)); +} + +/** + * Browsers hide `crypto.randomUUID` on plain HTTP, which is how plenty of + * self-hosted HyperDX is reached, hence the fallback. These ids only need to + * be distinct, not unguessable. + */ +function randomId(): string { + const cryptoObj = + typeof globalThis !== 'undefined' ? globalThis.crypto : undefined; + if (cryptoObj?.randomUUID) { + try { + return cryptoObj.randomUUID(); + } catch { + // fall through + } + } + return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, char => { + const rand = (Math.random() * 16) | 0; + const value = char === 'x' ? rand : (rand & 0x3) | 0x8; + return value.toString(16); + }); +} + +/** + * The random part is required: ClickHouse rejects a query whose id matches one + * already running. + */ +export function buildQueryId( + attribution: QueryAttribution | undefined, +): string { + const surface = isQuerySurface(attribution?.surface) + ? attribution.surface + : 'unknown'; + return `hdx-${surface}-${randomId()}`; +} diff --git a/packages/common-utils/src/clickhouse/index.ts b/packages/common-utils/src/clickhouse/index.ts index fd13c86d16..6e360008e6 100644 --- a/packages/common-utils/src/clickhouse/index.ts +++ b/packages/common-utils/src/clickhouse/index.ts @@ -32,6 +32,13 @@ import { import { isBuilderChartConfig } from '@/guards'; import { ChartConfigWithOptDateRange, QuerySettings } from '@/types'; +import { + buildLogComment, + buildQueryId, + mergeQueryAttribution, + QueryAttribution, +} from './attribution'; + // export @clickhouse/client-common types export type { BaseResultSet, @@ -42,6 +49,17 @@ export type { Row, }; +// Re-exported so callers get these from the same place as the client. +export { + buildLogComment, + buildQueryId, + mergeQueryAttribution, + QUERY_ATTRIBUTION_VERSION, + QUERY_SURFACES, + type QueryAttribution, + type QuerySurface, +} from './attribution'; + export enum JSDataType { Array = 'array', Date = 'date', @@ -202,7 +220,6 @@ export const chSql = ( // if (typeof value === 'string') { // console.error('Unsafe string detected', value, 'in', strings, values); // } - return ( str + (value == null @@ -424,6 +441,8 @@ export interface QueryInputs { connectionId?: string; queryId?: string; shouldSkipApplySettings?: boolean; + /** Tags just this query. Added on top of the client's default. */ + attribution?: QueryAttribution; } export type ClickhouseClientOptions = { @@ -437,6 +456,8 @@ export type ClickhouseClientOptions = { requestTimeout?: number; /** Logger for per-query SQL debug output. When omitted, query logging is silent. */ customLogger?: Logger; + /** Tags every query this client issues. */ + attribution?: QueryAttribution; }; export abstract class BaseClickhouseClient { @@ -454,6 +475,7 @@ export abstract class BaseClickhouseClient { protected maxRowReadOnly: boolean; protected requestTimeout: number = 3600000; protected readonly customLogger?: Logger; + protected readonly attribution?: QueryAttribution; constructor({ host, @@ -463,6 +485,7 @@ export abstract class BaseClickhouseClient { application, requestTimeout, customLogger, + attribution, }: ClickhouseClientOptions) { this.host = host!; this.username = username; @@ -471,6 +494,7 @@ export abstract class BaseClickhouseClient { this.maxRowReadOnly = false; this.application = application; this.customLogger = customLogger; + this.attribution = attribution; if (requestTimeout != null && requestTimeout >= 0) { this.requestTimeout = requestTimeout; } @@ -598,9 +622,39 @@ export abstract class BaseClickhouseClient { }; } - async query( + /** + * Done here, not in each subclass, so the browser, node and CLI clients all + * get it, along with every query `Metadata` makes. A caller's own + * `log_comment` or `queryId` is left alone. + */ + protected applyAttribution( props: QueryInputs, + ): QueryInputs { + const attribution = mergeQueryAttribution( + this.attribution, + props.attribution, + ); + + const logComment = buildLogComment(attribution); + const clickhouse_settings = + logComment && props.clickhouse_settings?.log_comment === undefined + ? { ...props.clickhouse_settings, log_comment: logComment } + : props.clickhouse_settings; + + return { + ...props, + clickhouse_settings, + queryId: props.queryId ?? buildQueryId(attribution), + }; + } + + async query( + inputs: QueryInputs, ): Promise> { + // Once, outside the loop, so a retry keeps the same query_id. Safe + // because the only thing we retry is a rejected setting, which means + // nothing is still running under that id. + const props = this.applyAttribution(inputs); let attempts = 0; // retry query if fails while (attempts < 2) { @@ -680,6 +734,7 @@ export abstract class BaseClickhouseClient { abort_signal: opts?.abort_signal, connectionId: config.connection, clickhouse_settings: opts?.clickhouse_settings, + attribution: { source: config.source }, }); return resp.json(); } @@ -717,6 +772,9 @@ export abstract class BaseClickhouseClient { abort_signal: opts?.abort_signal, connectionId: config.connection, clickhouse_settings: opts?.clickhouse_settings, + // No label: it would overwrite the one naming who asked, and an + // EXPLAIN is recognisable from the query text anyway. + attribution: { source: config.source }, }); const jsonResult = await result.json<{ rows: string | number }>(); diff --git a/packages/common-utils/src/clickhouse/node.ts b/packages/common-utils/src/clickhouse/node.ts index 4ab917e65f..6fd48afb1d 100644 --- a/packages/common-utils/src/clickhouse/node.ts +++ b/packages/common-utils/src/clickhouse/node.ts @@ -1,3 +1,5 @@ +import { AsyncLocalStorage } from 'node:async_hooks'; + import type { BaseResultSet, ClickHouseClient as NodeClickHouseClient, @@ -6,6 +8,7 @@ import type { } from '@clickhouse/client'; import { createClient } from '@clickhouse/client'; +import { mergeQueryAttribution, QueryAttribution } from './attribution'; import { BaseClickhouseClient, ClickhouseClientOptions, @@ -15,6 +18,35 @@ import { // for api fixtures export { createClient as createNativeClient }; +/** + * On the server, the code that knows why a query is happening sits far above + * the code that runs it, and one client is shared by concurrent work. So the + * entry point sets this once and every query below it picks it up. + * + * Node only. The browser uses React context for the same job. + */ +const attributionStore = new AsyncLocalStorage(); + +/** + * Tags every ClickHouse query `fn` issues. Wrap the whole request or job, not + * each query. + */ +export function withQueryAttribution( + attribution: QueryAttribution, + fn: () => T, +): T { + const merged = mergeQueryAttribution( + attributionStore.getStore(), + attribution, + ); + return attributionStore.run(merged, fn); +} + +/** What is in effect right now, if anything. */ +export function getCurrentQueryAttribution(): QueryAttribution | undefined { + return attributionStore.getStore(); +} + export class ClickhouseClient extends BaseClickhouseClient { constructor(options: ClickhouseClientOptions) { super(options); @@ -29,6 +61,22 @@ export class ClickhouseClient extends BaseClickhouseClient { }); } + /** + * Order of precedence: the client's own default, then this request's scope, + * then anything passed with the query. + */ + protected applyAttribution( + props: QueryInputs, + ): QueryInputs { + const ambient = attributionStore.getStore(); + if (!ambient) return super.applyAttribution(props); + + return super.applyAttribution({ + ...props, + attribution: mergeQueryAttribution(ambient, props.attribution), + }); + } + // This subclass always builds a node client, so narrow the base class's // platform-agnostic client type to the node-specific one. protected getClient(): NodeClickHouseClient {