From 6299571c4751685e79f701ff97d875459fb5f133 Mon Sep 17 00:00:00 2001 From: Atila Fassina Date: Wed, 28 Oct 2026 18:14:34 +0100 Subject: [PATCH 1/5] fix(appkit): forward queryTags to statement execution The SQL warehouse connector rebuilt the request body field by field and dropped queryTags, so tags set by callers never reached the warehouse. Co-authored-by: Isaac Signed-off-by: Atila Fassina --- .../src/connectors/sql-warehouse/client.ts | 1 + .../connectors/tests/sql-warehouse.test.ts | 28 +++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index 6afca34f8..6bbaaf8db 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -313,6 +313,7 @@ export class SQLWarehouseConnector { rowLimit: input.rowLimit, onWaitTimeout: input.onWaitTimeout || executeStatementDefaults.onWaitTimeout, + queryTags: input.queryTags, }; span.addEvent("statement.submitting", { diff --git a/packages/appkit/src/connectors/tests/sql-warehouse.test.ts b/packages/appkit/src/connectors/tests/sql-warehouse.test.ts index 00c626878..ff24e334a 100644 --- a/packages/appkit/src/connectors/tests/sql-warehouse.test.ts +++ b/packages/appkit/src/connectors/tests/sql-warehouse.test.ts @@ -255,6 +255,34 @@ describe("SQLWarehouseConnector", () => { }); }); + describe("query tags", () => { + test("forwards queryTags to the statement execution request", async () => { + const connector = new SQLWarehouseConnector({ timeout: 5000 }); + const executeStatement = vi.fn().mockResolvedValue({ + statementId: "stmt-1", + status: { state: "SUCCEEDED" }, + manifest: { format: "JSON_ARRAY", schema: { columns: [] } }, + result: { dataArray: [] }, + }); + const mockWorkspaceClient = { + statementExecution: { executeStatement }, + config: { host: "https://test.databricks.com" }, + }; + const queryTags = [{ key: "team", value: "eng" }, { key: "exp" }]; + + await connector.executeStatement(mockWorkspaceClient as any, { + statement: "SELECT 1", + warehouseId: "wh-1", + queryTags, + }); + + expect(executeStatement).toHaveBeenCalledWith( + expect.objectContaining({ queryTags }), + expect.anything(), + ); + }); + }); + describe("ensureWarehouseRunning", () => { let connector: SQLWarehouseConnector; From 6c37d4c47aeb7958f1598e03e9bd030ce32a99f3 Mon Sep 17 00:00:00 2001 From: Atila Fassina Date: Thu, 29 Oct 2026 10:06:12 +0100 Subject: [PATCH 2/5] feat(appkit): tag analytics statements for warehouse attribution Statements issued by the /query and /metric routes and the query agent tool now carry query tags, recorded in system.query.history.query_tags: - automatic appkit_app, appkit_query_key, appkit_route and appkit_lane tags (opt out with autoQueryTags: false) - custom tags from the queryTags config, static or derived per request - client-sent tags from the request body (opt out with clientQueryTags: false) Automatic tags reserve the appkit_ prefix and win on collisions, then config, then client tags; the set is capped at the warehouse's 20-tag limit. Invalid tags are dropped instead of failing the query. Co-authored-by: Isaac Signed-off-by: Atila Fassina --- .../appkit/src/plugins/analytics/analytics.ts | 45 +++- .../src/plugins/analytics/manifest.json | 10 + .../src/plugins/analytics/mv/schemas.ts | 3 + .../src/plugins/analytics/query-tags.ts | 179 +++++++++++++++ .../src/plugins/analytics/result-delivery.ts | 8 +- .../tests/analytics.readonly.test.ts | 6 +- .../plugins/analytics/tests/analytics.test.ts | 147 ++++++++++++ .../plugins/analytics/tests/metric.test.ts | 36 +++ .../analytics/tests/query-tags.test.ts | 211 ++++++++++++++++++ .../appkit/src/plugins/analytics/types.ts | 33 +++ .../shared/src/workspace-client/modular.ts | 1 + 11 files changed, 667 insertions(+), 12 deletions(-) create mode 100644 packages/appkit/src/plugins/analytics/query-tags.ts create mode 100644 packages/appkit/src/plugins/analytics/tests/query-tags.test.ts diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 21d0985d9..37f7b8eca 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -16,7 +16,7 @@ import { DEFAULT_WAREHOUSE_STARTUP_TIMEOUT_MS, type WarehouseStatusUpdate, } from "../../connectors/sql-warehouse/client"; -import { getWorkspaceClient } from "../../context"; +import { getWorkspaceClient, isInUserContext } from "../../context"; import { buildToolkitEntries } from "../../core/agent/build-toolkit"; import { defineTool, @@ -29,6 +29,7 @@ import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import { defineManifest } from "../../registry"; import { getWarehouseId } from "../../resources"; +import type { QueryTag } from "../../workspace-client"; import { queryDefaults } from "./defaults"; import manifest from "./manifest.json"; import { @@ -42,6 +43,7 @@ import { validateMetricRequest, } from "./metric"; import { QueryProcessor } from "./query"; +import { resolveQueryTags, withQueryTags } from "./query-tags"; import { type ArrowCapability, deliverArrowBytes, @@ -245,8 +247,11 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { res: express.Response, ): Promise { const { query_key } = req.params; - const { parameters, format: rawFormat = "JSON_ARRAY" } = - req.body as IAnalyticsQueryRequest; + const { + parameters, + format: rawFormat = "JSON_ARRAY", + queryTags: clientTags, + } = req.body as IAnalyticsQueryRequest; if ( rawFormat !== "JSON_ARRAY" && @@ -291,6 +296,14 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { const { query, isAsUser } = queryResult; + const queryTags = resolveQueryTags(this.config, { + route: "query", + queryKey: query_key, + lane: isAsUser ? "obo" : "sp", + req, + clientTags, + }); + // ARROW_STREAM streams the raw Arrow IPC bytes back as the HTTP response // body — no SSE, no server-side stash, no second /arrow-result request. // INLINE attachments are piped straight through (the bytes are already in @@ -305,6 +318,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { query, isAsUser, parameters, + queryTags, ); return; } @@ -401,7 +415,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { // `result` message with plain rows. (ARROW_STREAM requests are // handled earlier via `_handleArrowStreamQuery`.) return await self._executeJsonArrayPath( - executor, + withQueryTags(executor, queryTags), query, processedParams, sig, @@ -539,6 +553,14 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { return; } + const queryTags = resolveQueryTags(this.config, { + route: "metric", + queryKey: key, + lane: registration.lane, + req, + clientTags: request.queryTags, + }); + // Lane dispatch. The lane comes from the registration (the entry's // `executor` in definitions.json), NOT a URL segment or `.obo.sql` // filename: an OBO-lane metric runs on-behalf-of the requesting user @@ -692,7 +714,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { // an ARROW_STREAM-inline fallback, returning plain rows in a // `result` message — byte-identical envelope to `/query`. return await self._executeJsonArrayPath( - executor, + withQueryTags(executor, queryTags), statement, processedParams, sig, @@ -754,7 +776,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { * message. External links are never used for the JSON fallback. */ private async _executeJsonArrayPath( - executor: AnalyticsPlugin, + executor: QueryExecutor, query: string, processedParams: | Record @@ -821,6 +843,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { query: string, isAsUser: boolean, parameters: IAnalyticsQueryRequest["parameters"], + queryTags: QueryTag[] | undefined, ): Promise { const executor = isAsUser ? this._asUserScoped(req) : this; const executorKey = isAsUser ? this.resolveUserId(req) : "global"; @@ -866,7 +889,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { // are cached; EXTERNAL_LINKS carry expiring pre-signed URLs and are // never cached (see `_arrowCachingExecutor`). this._arrowCachingExecutor( - executor, + withQueryTags(executor, queryTags), query_key, query, parameters, @@ -997,7 +1020,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { * INLINE→EXTERNAL_LINKS fallback still sees the structured rejection. */ private _arrowCachingExecutor( - executor: AnalyticsPlugin, + executor: QueryExecutor, query_key: string, query: string, parameters: IAnalyticsQueryRequest["parameters"], @@ -1101,7 +1124,11 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { autoInheritable: true, execute: (args, signal) => { assertReadOnlySql(args.query); - return this.query(args.query, undefined, undefined, signal); + const queryTags = resolveQueryTags(this.config, { + route: "agent_tool", + lane: isInUserContext() ? "obo" : "sp", + }); + return this.query(args.query, undefined, { queryTags }, signal); }, }), }; diff --git a/packages/appkit/src/plugins/analytics/manifest.json b/packages/appkit/src/plugins/analytics/manifest.json index d811f061d..51666b125 100644 --- a/packages/appkit/src/plugins/analytics/manifest.json +++ b/packages/appkit/src/plugins/analytics/manifest.json @@ -53,6 +53,16 @@ "type": "boolean", "default": true, "description": "When true, a STOPPED SQL warehouse is auto-started on the first analytics request. Set to false for cost-controlled deployments where billable warehouse starts must not be triggered by user requests." + }, + "autoQueryTags": { + "type": "boolean", + "default": true, + "description": "When true, statements are tagged with appkit_app, appkit_query_key, appkit_route and appkit_lane so warehouse usage can be attributed in system.query.history." + }, + "clientQueryTags": { + "type": "boolean", + "default": true, + "description": "When true, query tags sent by the client are forwarded after validation. Server tags win on key collisions." } } } diff --git a/packages/appkit/src/plugins/analytics/mv/schemas.ts b/packages/appkit/src/plugins/analytics/mv/schemas.ts index d74084a37..7bf134db0 100644 --- a/packages/appkit/src/plugins/analytics/mv/schemas.ts +++ b/packages/appkit/src/plugins/analytics/mv/schemas.ts @@ -152,6 +152,9 @@ const metricRequestSchema = z }) .optional(), format: z.enum(["JSON_ARRAY", "ARROW_STREAM", "JSON", "ARROW"]).optional(), + // Validated leniently by `resolveQueryTags`: a malformed tag is dropped, + // never a 400, matching the `/query` route. + queryTags: z.unknown().optional(), }) .strict() .superRefine((value, ctx) => { diff --git a/packages/appkit/src/plugins/analytics/query-tags.ts b/packages/appkit/src/plugins/analytics/query-tags.ts new file mode 100644 index 000000000..ad8fe23ee --- /dev/null +++ b/packages/appkit/src/plugins/analytics/query-tags.ts @@ -0,0 +1,179 @@ +import type express from "express"; + +import { createLogger } from "../../logging/logger"; +import type { QueryTag } from "../../workspace-client"; +import type { QueryExecutor } from "./result-delivery"; +import type { IAnalyticsConfig } from "./types"; + +const logger = createLogger("analytics:query-tags"); + +/** + * Query tags as a key → value map. A `null` value records a key-only tag. + * Tags land in `system.query.history.query_tags` as plain text — never put + * secrets or PII in them. + */ +export type QueryTagMap = Record; + +/** Entry point that issued the statement, recorded as `appkit_route`. */ +export type QueryTagsRoute = "query" | "metric" | "agent_tool"; + +/** Context handed to a function-form `queryTags` config. */ +export interface QueryTagsContext { + route: QueryTagsRoute; + /** Query key (`/query`) or metric key (`/metric`); absent for agent tools. */ + queryKey?: string; + /** `"obo"` when the statement runs on behalf of the user, else `"sp"`. */ + lane: "sp" | "obo"; + /** Incoming request; absent for agent tool calls. */ + req?: express.Request; + /** Sanitized tags sent by the client in the request body. */ + clientTags: QueryTagMap; +} + +type QueryTagsConfig = Pick< + IAnalyticsConfig, + "queryTags" | "autoQueryTags" | "clientQueryTags" +>; + +/** Hard cap the warehouse applies; anything past it is dropped server-side. */ +export const MAX_QUERY_TAGS = 20; +const MAX_TAG_LENGTH = 128; +/** Prefix reserved for AppKit's automatic tags. */ +const RESERVED_PREFIX = "appkit_"; +/** Characters the warehouse rejects in tag keys. */ +const INVALID_KEY_CHARS = /[,:\-/=.]/; + +function isValidKey(key: string): boolean { + return ( + key.trim().length > 0 && + key.length <= MAX_TAG_LENGTH && + !INVALID_KEY_CHARS.test(key) && + !key.startsWith("@@") + ); +} + +/** + * Keep only well-formed, non-reserved tags. Invalid entries are dropped rather + * than failing the query: a bad tag must never break a dashboard. Values are + * truncated to the warehouse's 128-character limit. + */ +export function sanitizeQueryTags( + input: unknown, + source: "client" | "config", +): QueryTagMap { + if (input === null || typeof input !== "object" || Array.isArray(input)) { + return {}; + } + const out: QueryTagMap = {}; + const dropped: string[] = []; + for (const [key, value] of Object.entries(input)) { + if ( + !isValidKey(key) || + key.startsWith(RESERVED_PREFIX) || + (value !== null && typeof value !== "string") + ) { + dropped.push(key); + continue; + } + out[key] = value === null ? null : value.slice(0, MAX_TAG_LENGTH); + } + if (dropped.length > 0) { + const message = "Dropped %d invalid %s query tag(s): %s"; + const keys = dropped.slice(0, 5).join(", ").slice(0, 200); + // Client input is untrusted and may be noisy; config mistakes are the + // developer's to fix. + if (source === "config") { + logger.warn(message, dropped.length, source, keys); + } else { + logger.debug(message, dropped.length, source, keys); + } + } + return out; +} + +function autoTags(ctx: QueryTagsContext): QueryTagMap { + const tags: QueryTagMap = {}; + const appName = process.env.DATABRICKS_APP_NAME; + if (appName) tags.appkit_app = appName.slice(0, MAX_TAG_LENGTH); + if (ctx.queryKey) { + tags.appkit_query_key = ctx.queryKey.slice(0, MAX_TAG_LENGTH); + } + tags.appkit_route = ctx.route; + tags.appkit_lane = ctx.lane; + return tags; +} + +function configTags( + config: QueryTagsConfig, + ctx: QueryTagsContext, +): QueryTagMap { + const { queryTags } = config; + if (queryTags === undefined) return {}; + if (typeof queryTags !== "function") { + return sanitizeQueryTags(queryTags, "config"); + } + try { + return sanitizeQueryTags(queryTags(ctx), "config"); + } catch (err) { + logger.warn( + "queryTags callback threw; continuing without custom tags: %s", + err instanceof Error ? err.message : String(err), + ); + return {}; + } +} + +/** + * Resolve the tags for one statement. + * + * Precedence on key collisions: automatic `appkit_*` tags (reserved, cannot + * be overridden) > server `queryTags` config > client-sent tags. Tags are + * then capped at {@link MAX_QUERY_TAGS} in the same priority order, so the + * automatic tags always fit. + */ +export function resolveQueryTags( + config: QueryTagsConfig, + ctx: Omit & { clientTags?: unknown }, +): QueryTag[] | undefined { + const clientTags = + config.clientQueryTags === false + ? {} + : sanitizeQueryTags(ctx.clientTags, "client"); + const fullCtx: QueryTagsContext = { ...ctx, clientTags }; + + const merged: QueryTagMap = { + ...(config.autoQueryTags === false ? {} : autoTags(fullCtx)), + }; + for (const tags of [configTags(config, fullCtx), clientTags]) { + for (const [key, value] of Object.entries(tags)) { + if (!Object.hasOwn(merged, key)) merged[key] = value; + } + } + + const entries = Object.entries(merged).slice(0, MAX_QUERY_TAGS); + if (entries.length === 0) return undefined; + return entries.map(([key, value]) => + value === null ? { key } : { key, value }, + ); +} + +/** + * Wrap an executor so every statement it issues carries `queryTags`. Tags are + * not part of any cache key: a cache hit never reaches the warehouse, so there + * is nothing to attribute. + */ +export function withQueryTags( + executor: QueryExecutor, + queryTags: QueryTag[] | undefined, +): QueryExecutor { + if (!queryTags) return executor; + return { + query: (query, parameters, formatParameters, signal) => + executor.query( + query, + parameters, + { ...formatParameters, queryTags }, + signal, + ), + }; +} diff --git a/packages/appkit/src/plugins/analytics/result-delivery.ts b/packages/appkit/src/plugins/analytics/result-delivery.ts index 3f40439d2..c2804c329 100644 --- a/packages/appkit/src/plugins/analytics/result-delivery.ts +++ b/packages/appkit/src/plugins/analytics/result-delivery.ts @@ -4,7 +4,7 @@ import type { SQLTypeMarker } from "shared"; import { ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; import type { RefreshChunkLink } from "../../stream/arrow-stream-processor"; -import type { ExternalLink } from "../../workspace-client"; +import type { ExternalLink, QueryTag } from "../../workspace-client"; /** * Centralized disposition/format fallback for analytics result delivery. @@ -33,7 +33,11 @@ export interface QueryExecutor { query( query: string, parameters: Record | undefined, - formatParameters: { disposition: string; format: string }, + formatParameters: { + disposition: string; + format: string; + queryTags?: QueryTag[]; + }, signal?: AbortSignal, ): Promise< | { diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.readonly.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.readonly.test.ts index 758ac62f0..d9f5178c4 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.readonly.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.readonly.test.ts @@ -88,7 +88,11 @@ describe("AnalyticsPlugin.query agent tool — runtime enforcement", () => { expect(spy).toHaveBeenCalledWith( "SELECT * FROM main.sales.orders", undefined, - undefined, + { + queryTags: expect.arrayContaining([ + { key: "appkit_route", value: "agent_tool" }, + ]), + }, undefined, ); }); diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 3104d710b..d563fff35 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -1813,6 +1813,153 @@ describe("Analytics Plugin", () => { }); }); + describe("query tags", () => { + function setup(pluginConfig: IAnalyticsConfig = config) { + const plugin = new AnalyticsPlugin(pluginConfig); + const { router, getHandler } = createMockRouter(); + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + const executeMock = vi.fn().mockResolvedValue({ + result: { data: [{ id: 1 }] }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + plugin.injectRoutes(router); + return { + plugin, + executeMock, + handler: getHandler("POST", "/query/:query_key"), + }; + } + + function sentTags(executeMock: ReturnType) { + return executeMock.mock.calls[0][1].queryTags; + } + + test("tags JSON queries with automatic, config, and client tags", async () => { + const { executeMock, handler } = setup({ + ...config, + queryTags: { team: "finance" }, + }); + + await handler( + createMockRequest({ + params: { query_key: "revenue" }, + body: { parameters: {}, queryTags: { page: "overview" } }, + }), + createMockResponse(), + ); + + expect(sentTags(executeMock)).toEqual( + expect.arrayContaining([ + { key: "appkit_query_key", value: "revenue" }, + { key: "appkit_route", value: "query" }, + { key: "appkit_lane", value: "sp" }, + { key: "team", value: "finance" }, + { key: "page", value: "overview" }, + ]), + ); + }); + + test("tags ARROW_STREAM queries", async () => { + const { executeMock, handler } = setup(); + + await handler( + createMockRequest({ + params: { query_key: "revenue" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }), + createMockResponse(), + ); + + expect(sentTags(executeMock)).toEqual( + expect.arrayContaining([ + { key: "appkit_query_key", value: "revenue" }, + { key: "appkit_route", value: "query" }, + ]), + ); + }); + + test("tags .obo.sql queries with the obo lane", async () => { + const { plugin, executeMock, handler } = setup(); + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: true, + }); + + await handler( + createMockRequest({ + params: { query_key: "mine" }, + body: { parameters: {} }, + headers: { + "x-forwarded-access-token": "user-token", + "x-forwarded-user": "user-123", + }, + }), + createMockResponse(), + ); + + expect(sentTags(executeMock)).toEqual( + expect.arrayContaining([{ key: "appkit_lane", value: "obo" }]), + ); + }); + + test("spoofed client tags for reserved keys are ignored", async () => { + const { executeMock, handler } = setup(); + + await handler( + createMockRequest({ + params: { query_key: "revenue" }, + body: { + parameters: {}, + queryTags: { appkit_query_key: "spoofed", "bad.key": "x" }, + }, + }), + createMockResponse(), + ); + + const tags = sentTags(executeMock); + expect(tags).toContainEqual({ + key: "appkit_query_key", + value: "revenue", + }); + expect(tags.map((t: { key: string }) => t.key)).not.toContain("bad.key"); + }); + + test("sends no tags when automatic tags are off and none are configured", async () => { + const { executeMock, handler } = setup({ + ...config, + autoQueryTags: false, + }); + + await handler( + createMockRequest({ + params: { query_key: "revenue" }, + body: { parameters: {} }, + }), + createMockResponse(), + ); + + expect(sentTags(executeMock)).toBeUndefined(); + }); + + test("tags agent tool queries", async () => { + const plugin = new AnalyticsPlugin(config); + const executeMock = vi.fn().mockResolvedValue({ result: { data: [] } }); + (plugin as any).SQLClient.executeStatement = executeMock; + + await (plugin as any).tools.query.execute({ query: "SELECT 1" }); + + expect(sentTags(executeMock)).toEqual( + expect.arrayContaining([ + { key: "appkit_route", value: "agent_tool" }, + { key: "appkit_lane", value: "sp" }, + ]), + ); + }); + }); + describe("toolkit()", () => { test("produces ToolkitEntry records keyed by the plugin name", () => { const plugin = new AnalyticsPlugin({ name: "analytics" }); diff --git a/packages/appkit/src/plugins/analytics/tests/metric.test.ts b/packages/appkit/src/plugins/analytics/tests/metric.test.ts index 07b83e449..b09229adc 100644 --- a/packages/appkit/src/plugins/analytics/tests/metric.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/metric.test.ts @@ -713,6 +713,42 @@ describe("analytics metric route", () => { expect(mockRes.end).toHaveBeenCalled(); }); + test("tags metric statements with the metric key and client tags", async () => { + const plugin = pluginForDir( + config, + registryDir({ + revenue: { + key: "revenue", + source: "cat.sch.revenue_metrics", + lane: "sp", + }, + }), + ); + const { router, getHandler } = createMockRouter(); + const executeMock = vi.fn().mockResolvedValue({ + result: { data: [{ arr: 1 }] }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + plugin.injectRoutes(router); + + await getHandler("POST", "/metric/:key")( + createMockRequest({ + params: { key: "revenue" }, + body: { measures: ["arr"], queryTags: { page: "overview" } }, + }), + createMockResponse(), + ); + + expect(executeMock.mock.calls[0][1].queryTags).toEqual( + expect.arrayContaining([ + { key: "appkit_query_key", value: "revenue" }, + { key: "appkit_route", value: "metric" }, + { key: "appkit_lane", value: "sp" }, + { key: "page", value: "overview" }, + ]), + ); + }); + test("emits warehouse_status before result", async () => { const plugin = pluginForDir( config, diff --git a/packages/appkit/src/plugins/analytics/tests/query-tags.test.ts b/packages/appkit/src/plugins/analytics/tests/query-tags.test.ts new file mode 100644 index 000000000..5a8f545f6 --- /dev/null +++ b/packages/appkit/src/plugins/analytics/tests/query-tags.test.ts @@ -0,0 +1,211 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +import { + MAX_QUERY_TAGS, + resolveQueryTags, + sanitizeQueryTags, + withQueryTags, +} from "../query-tags"; + +describe("sanitizeQueryTags", () => { + test("keeps valid string and key-only tags", () => { + expect( + sanitizeQueryTags( + { team: "eng", page_name: "sales", exp: null }, + "client", + ), + ).toEqual({ team: "eng", page_name: "sales", exp: null }); + }); + + test("drops keys with characters the warehouse rejects", () => { + expect( + sanitizeQueryTags( + { + "a,b": "x", + "a:b": "x", + "a-b": "x", + "a/b": "x", + "a=b": "x", + "a.b": "x", + "@@internal": "x", + " ": "x", + ok: "x", + }, + "client", + ), + ).toEqual({ ok: "x" }); + }); + + test("drops reserved appkit_ keys so attribution cannot be spoofed", () => { + expect( + sanitizeQueryTags({ appkit_app: "spoofed", team: "eng" }, "client"), + ).toEqual({ team: "eng" }); + }); + + test("drops non-string values and over-long keys, truncates long values", () => { + const out = sanitizeQueryTags( + { n: 1, o: { x: 1 }, ["k".repeat(129)]: "x", long: "v".repeat(500) }, + "client", + ); + expect(out).toEqual({ long: "v".repeat(128) }); + }); + + test("ignores non-object input", () => { + for (const input of [undefined, null, "team:eng", ["team"], 42]) { + expect(sanitizeQueryTags(input, "client")).toEqual({}); + } + }); +}); + +describe("resolveQueryTags", () => { + const prevAppName = process.env.DATABRICKS_APP_NAME; + + beforeEach(() => { + process.env.DATABRICKS_APP_NAME = "sales-dashboard"; + }); + + afterEach(() => { + if (prevAppName === undefined) delete process.env.DATABRICKS_APP_NAME; + else process.env.DATABRICKS_APP_NAME = prevAppName; + }); + + test("adds automatic attribution tags by default", () => { + expect( + resolveQueryTags({}, { route: "query", queryKey: "revenue", lane: "sp" }), + ).toEqual([ + { key: "appkit_app", value: "sales-dashboard" }, + { key: "appkit_query_key", value: "revenue" }, + { key: "appkit_route", value: "query" }, + { key: "appkit_lane", value: "sp" }, + ]); + }); + + test("omits appkit_app when the app name is unknown (local dev)", () => { + delete process.env.DATABRICKS_APP_NAME; + const tags = resolveQueryTags({}, { route: "agent_tool", lane: "obo" }); + expect(tags).toEqual([ + { key: "appkit_route", value: "agent_tool" }, + { key: "appkit_lane", value: "obo" }, + ]); + }); + + test("autoQueryTags: false disables automatic tags", () => { + expect( + resolveQueryTags( + { autoQueryTags: false }, + { route: "query", queryKey: "revenue", lane: "sp" }, + ), + ).toBeUndefined(); + }); + + test("merges static config tags and client tags; server wins on collisions", () => { + const tags = resolveQueryTags( + { autoQueryTags: false, queryTags: { team: "finance" } }, + { + route: "query", + lane: "sp", + clientTags: { team: "spoofed", page: "overview" }, + }, + ); + expect(tags).toEqual([ + { key: "team", value: "finance" }, + { key: "page", value: "overview" }, + ]); + }); + + test("clientQueryTags: false ignores client tags", () => { + const tags = resolveQueryTags( + { autoQueryTags: false, clientQueryTags: false }, + { route: "query", lane: "sp", clientTags: { page: "overview" } }, + ); + expect(tags).toBeUndefined(); + }); + + test("passes context, including sanitized client tags, to a function config", () => { + const queryTags = vi.fn().mockReturnValue({ tenant: "acme" }); + const tags = resolveQueryTags( + { autoQueryTags: false, queryTags }, + { + route: "metric", + queryKey: "orders", + lane: "obo", + clientTags: { page: "overview", "bad.key": "x" }, + }, + ); + expect(queryTags).toHaveBeenCalledWith( + expect.objectContaining({ + route: "metric", + queryKey: "orders", + lane: "obo", + clientTags: { page: "overview" }, + }), + ); + expect(tags).toEqual([ + { key: "tenant", value: "acme" }, + { key: "page", value: "overview" }, + ]); + }); + + test("a throwing function config never fails the query", () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + const tags = resolveQueryTags( + { + queryTags: () => { + throw new Error("boom"); + }, + }, + { route: "query", queryKey: "revenue", lane: "sp" }, + ); + expect(tags?.map((t) => t.key)).toEqual([ + "appkit_app", + "appkit_query_key", + "appkit_route", + "appkit_lane", + ]); + warn.mockRestore(); + }); + + test("caps at the warehouse limit, keeping automatic tags first", () => { + const many = Object.fromEntries( + Array.from({ length: 30 }, (_, i) => [`k${i}`, "v"]), + ); + const tags = resolveQueryTags( + {}, + { route: "query", queryKey: "q", lane: "sp", clientTags: many }, + ); + expect(tags).toHaveLength(MAX_QUERY_TAGS); + expect(tags?.slice(0, 4).map((t) => t.key)).toEqual([ + "appkit_app", + "appkit_query_key", + "appkit_route", + "appkit_lane", + ]); + }); +}); + +describe("withQueryTags", () => { + test("adds queryTags to every statement's format parameters", async () => { + const query = vi.fn().mockResolvedValue({ data: [] }); + const queryTags = [{ key: "team", value: "eng" }]; + const signal = new AbortController().signal; + + await withQueryTags({ query }, queryTags).query( + "SELECT 1", + undefined, + { disposition: "INLINE", format: "JSON_ARRAY" }, + signal, + ); + + expect(query).toHaveBeenCalledWith( + "SELECT 1", + undefined, + { disposition: "INLINE", format: "JSON_ARRAY", queryTags }, + signal, + ); + }); + + test("returns the executor unchanged when there are no tags", () => { + const executor = { query: vi.fn() }; + expect(withQueryTags(executor, undefined)).toBe(executor); + }); +}); diff --git a/packages/appkit/src/plugins/analytics/types.ts b/packages/appkit/src/plugins/analytics/types.ts index 924373f07..33ad2780d 100644 --- a/packages/appkit/src/plugins/analytics/types.ts +++ b/packages/appkit/src/plugins/analytics/types.ts @@ -6,6 +6,8 @@ import type { MetricViewsMetadata, } from "shared"; +import type { QueryTagMap, QueryTagsContext } from "./query-tags"; + export type { MetricFilter, MetricFilterOperatorName, @@ -15,6 +17,33 @@ export type { export interface IAnalyticsConfig extends BasePluginConfig { timeout?: number; + /** + * Custom query tags added to every statement the analytics routes and agent + * tool issue, recorded in `system.query.history.query_tags`. Pass a static + * map, or a function to derive tags per request (e.g. a tenant header). + * Keys must not contain `, : - / = .`, must not start with the reserved + * `appkit_` prefix, and keys/values are capped at 128 characters; invalid + * tags are dropped with a warning. Never put secrets or PII in tags. + */ + queryTags?: + | QueryTagMap + | ((ctx: QueryTagsContext) => QueryTagMap | undefined); + /** + * When `true` (default), statements are tagged with `appkit_app`, + * `appkit_query_key`, `appkit_route` and `appkit_lane` so warehouse usage + * can be attributed to the app and query that issued it. + * + * @default true + */ + autoQueryTags?: boolean; + /** + * When `true` (default), tags sent by the client (`queryTags` on + * `useAnalyticsQuery` / `useMetricView`) are forwarded after validation. + * Server tags win on key collisions. Set to `false` to ignore client tags. + * + * @default true + */ + clientQueryTags?: boolean; /** * Per-column display metadata. The metric route stamps the slice scoped to a * request's measures/dimensions into the SSE `result` message. @@ -125,6 +154,8 @@ export function normalizeAnalyticsFormat( export interface IAnalyticsQueryRequest { parameters?: Record; format?: AnalyticsFormat; + /** Client-supplied query tags; validated server-side before use. */ + queryTags?: QueryTagMap; } export interface AnalyticsQueryResponse { @@ -189,4 +220,6 @@ export interface IAnalyticsMetricRequest { orderBy?: MetricOrderBy[]; limit?: number; format?: AnalyticsFormat; + /** Client-supplied query tags; validated server-side before use. */ + queryTags?: QueryTagMap; } diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index acb4b051b..d302d0f6c 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -785,6 +785,7 @@ export type { ExecuteStatementRequest, ExternalLink, Format, + QueryTag, ResultData, ResultManifest, Schema, From ffeda3fc897320c5521df7b4da664d4885590bc2 Mon Sep 17 00:00:00 2001 From: Atila Fassina Date: Thu, 29 Oct 2026 10:07:45 +0100 Subject: [PATCH 3/5] feat(appkit-ui): send queryTags from useAnalyticsQuery and useMetricView A new queryTags option forwards per-request tags (e.g. the page issuing the query) to the analytics routes, so the same query can be attributed to the different pages that run it. Requests with different tags are not deduplicated together. Co-authored-by: Isaac Signed-off-by: Atila Fassina --- .../__tests__/use-analytics-query.test.ts | 43 +++++++++++++++++++ .../hooks/__tests__/use-metric-view.test.ts | 14 ++++++ packages/appkit-ui/src/react/hooks/index.ts | 1 + packages/appkit-ui/src/react/hooks/types.ts | 16 +++++++ .../src/react/hooks/use-analytics-query.ts | 7 ++- .../src/react/hooks/use-metric-view.ts | 3 ++ 6 files changed, 82 insertions(+), 2 deletions(-) diff --git a/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts b/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts index e6697eaa5..9899e2b37 100644 --- a/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts +++ b/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts @@ -320,6 +320,49 @@ describe("useAnalyticsQuery", () => { expect(mockConnectSSE).toHaveBeenCalledTimes(1); }); + describe("queryTags", () => { + test("serializes queryTags into the POST body", () => { + renderHook(() => + useAnalyticsQuery("q", null, { queryTags: { page: "overview" } }), + ); + + expect(JSON.parse(lastConnectArgs.payload)).toEqual({ + parameters: null, + format: "JSON_ARRAY", + queryTags: { page: "overview" }, + }); + }); + + test("omits queryTags from the body when undefined", () => { + renderHook(() => useAnalyticsQuery("q", null)); + + expect(JSON.parse(lastConnectArgs.payload)).not.toHaveProperty( + "queryTags", + ); + }); + + test("does not refetch when an equal queryTags object is passed on re-render", () => { + const { rerender } = renderHook( + ({ page }) => useAnalyticsQuery("q", null, { queryTags: { page } }), + { initialProps: { page: "overview" } }, + ); + rerender({ page: "overview" }); + + expect(mockConnectSSE).toHaveBeenCalledTimes(1); + }); + + test("issues separate requests for different queryTags", () => { + renderHook(() => + useAnalyticsQuery("q", null, { queryTags: { page: "a" } }), + ); + renderHook(() => + useAnalyticsQuery("q", null, { queryTags: { page: "b" } }), + ); + + expect(mockConnectSSE).toHaveBeenCalledTimes(2); + }); + }); + describe("warehouse_status", () => { test("surfaces warehouseStatus while waiting and clears loading on result", async () => { let capturedOnMessage: diff --git a/packages/appkit-ui/src/react/hooks/__tests__/use-metric-view.test.ts b/packages/appkit-ui/src/react/hooks/__tests__/use-metric-view.test.ts index 4b8de584a..754e4698d 100644 --- a/packages/appkit-ui/src/react/hooks/__tests__/use-metric-view.test.ts +++ b/packages/appkit-ui/src/react/hooks/__tests__/use-metric-view.test.ts @@ -182,6 +182,20 @@ describe("useMetricView", () => { }); }); + test("serializes queryTags into the POST body when provided", () => { + renderHook(() => + useMetricView("orders", { + measures: ["revenue"], + queryTags: { page: "overview" }, + }), + ); + + expect(JSON.parse(lastConnectArgs.payload)).toEqual({ + measures: ["revenue"], + queryTags: { page: "overview" }, + }); + }); + test("omits orderBy from the body when undefined", () => { renderHook(() => useMetricView("orders", { diff --git a/packages/appkit-ui/src/react/hooks/index.ts b/packages/appkit-ui/src/react/hooks/index.ts index 9e4d51716..40a54c5aa 100644 --- a/packages/appkit-ui/src/react/hooks/index.ts +++ b/packages/appkit-ui/src/react/hooks/index.ts @@ -7,6 +7,7 @@ export { } from "../resource-status-indicator"; export type { AnalyticsFormat, + AnalyticsQueryTags, GrainsForSelectedTimeDims, InferDimensionKeys, InferMeasureKeys, diff --git a/packages/appkit-ui/src/react/hooks/types.ts b/packages/appkit-ui/src/react/hooks/types.ts index 92ae9f45a..213539a9f 100644 --- a/packages/appkit-ui/src/react/hooks/types.ts +++ b/packages/appkit-ui/src/react/hooks/types.ts @@ -45,6 +45,17 @@ export interface TypedArrowTable< // ============================================================================ /** Options for configuring an analytics SSE query */ +/** + * Query tags attached to the statement, recorded in + * `system.query.history.query_tags`. A `null` value records a key-only tag. + * The server validates them (keys must not contain `, : - / = .` or start with + * `appkit_`; keys/values are capped at 128 characters), drops invalid ones, and + * lets server-configured tags win on collisions. Never put secrets or PII here. + * + * @example { page: "sales_overview" } + */ +export type AnalyticsQueryTags = Record; + export interface UseAnalyticsQueryOptions< F extends AnalyticsFormat = "JSON_ARRAY", > { @@ -56,6 +67,9 @@ export interface UseAnalyticsQueryOptions< /** Whether to automatically start the query when the hook is mounted. Default is true. */ autoStart?: boolean; + + /** Query tags for this request, e.g. the page issuing it. See {@link AnalyticsQueryTags}. */ + queryTags?: AnalyticsQueryTags; } /** @@ -456,6 +470,8 @@ export type UseMetricViewOptions< limit?: number; /** Whether to automatically start the metric query. Default is true. */ autoStart?: boolean; + /** Query tags for this request, e.g. the page issuing it. See {@link AnalyticsQueryTags}. */ + queryTags?: AnalyticsQueryTags; } & MetricViewTimeOptions; export interface UseMetricViewResult[]> { diff --git a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts index e5f9e5e4b..3303af10b 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -125,12 +125,14 @@ export function useAnalyticsQuery< } const stableParameters = useStableParams(parameters); + const stableQueryTags = useStableParams(options?.queryTags); const payload = useMemo(() => { try { const serialized = JSON.stringify({ parameters: stableParameters, format, + queryTags: stableQueryTags, }); const sizeInBytes = new Blob([serialized]).size; if (sizeInBytes > maxParametersSize) { @@ -144,10 +146,11 @@ export function useAnalyticsQuery< console.error("useAnalyticsQuery: Failed to serialize parameters", error); return null; } - }, [stableParameters, format, maxParametersSize]); + }, [stableParameters, format, stableQueryTags, maxParametersSize]); // Cache key shared across hook instances. `payload` already serializes - // `{ parameters, format }`, so identical requests collapse to one key. + // `{ parameters, format, queryTags }`, so identical requests collapse to one + // key. // On a serialization failure (`payload === null`) the key stays unused: no // request is retained and the store reports the stable idle snapshot. const cacheKey = `${urlSuffix}::${payload}`; diff --git a/packages/appkit-ui/src/react/hooks/use-metric-view.ts b/packages/appkit-ui/src/react/hooks/use-metric-view.ts index 38d934c7c..93df1a851 100644 --- a/packages/appkit-ui/src/react/hooks/use-metric-view.ts +++ b/packages/appkit-ui/src/react/hooks/use-metric-view.ts @@ -119,6 +119,7 @@ export function useMetricView< timeDimension?: unknown; orderBy?: ReadonlyArray; limit?: number; + queryTags?: Record; } = { measures: options.measures }; if (options.dimensions !== undefined) body.dimensions = options.dimensions; if (options.filter !== undefined) body.filter = options.filter; @@ -127,6 +128,7 @@ export function useMetricView< body.timeDimension = options.timeDimension; if (options.orderBy !== undefined) body.orderBy = options.orderBy; if (options.limit !== undefined) body.limit = options.limit; + if (options.queryTags !== undefined) body.queryTags = options.queryTags; return JSON.stringify(body); }, [ options.measures, @@ -136,6 +138,7 @@ export function useMetricView< options.timeDimension, options.orderBy, options.limit, + options.queryTags, ]); // Detect shape changes: clear rows when metric/measures/dimensions change, From 90fb57671d171c9854759300f29925a183a66aa3 Mon Sep 17 00:00:00 2001 From: Atila Fassina Date: Thu, 29 Oct 2026 10:10:15 +0100 Subject: [PATCH 4/5] docs: document analytics query tags Co-authored-by: Isaac Signed-off-by: Atila Fassina --- docs/docs/plugins/analytics.md | 74 ++++++++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) diff --git a/docs/docs/plugins/analytics.md b/docs/docs/plugins/analytics.md index b144a5c08..85d63bc16 100644 --- a/docs/docs/plugins/analytics.md +++ b/docs/docs/plugins/analytics.md @@ -282,6 +282,78 @@ On failure it emits an `error` event instead. Editing `definitions.json` is picked up on the next request — no server restart is needed. A previously malformed file that you fix likewise starts working on the next request. +## Query tags + +Every statement issued by the `/query` and `/metric` routes and by the `query` agent tool carries [query tags](https://docs.databricks.com/aws/en/sql/user/queries/query-tags). Tags are recorded in the `query_tags` column of `system.query.history`, so warehouse owners can attribute load and latency to the app and query that caused it. + +AppKit adds these tags automatically: + +| Tag | Value | +| ------------------ | ---------------------------------------------------------------------- | +| `appkit_app` | `DATABRICKS_APP_NAME` (omitted when unset, e.g. local development) | +| `appkit_query_key` | The query key (`/query`) or metric key (`/metric`) | +| `appkit_route` | `query`, `metric`, or `agent_tool` | +| `appkit_lane` | `sp` (service principal) or `obo` (on behalf of the user) | + +The user's identity is never tagged. On-behalf-of statements already run as that user, so `system.query.history` attributes them. + +### Custom tags + +Add your own tags with the `queryTags` config, as a static map or a function evaluated per request: + +```ts +analytics({ + // static + queryTags: { team: "finance" }, +}); + +analytics({ + // per request: ctx has route, queryKey, lane, req, and the client's tags + queryTags: ({ req }) => { + const tenant = req?.header("x-tenant"); + return tenant ? { tenant } : undefined; + }, +}); +``` + +The same query can run on several pages, so the client can also tag each request. A typical use is naming the page that issued it: + +```tsx +useAnalyticsQuery("revenue", params, { queryTags: { page: "sales_overview" } }); +useMetricView("orders", { measures: ["revenue"], queryTags: { page: "orders" } }); +``` + +Client tags are untrusted input. Set `clientQueryTags: false` to ignore them, or filter them in a `queryTags` function, which receives the validated client tags as `ctx.clientTags`. + +**Rules:** + +- Precedence on collisions: automatic `appkit_*` tags, then the `queryTags` config, then client tags. The `appkit_` prefix is reserved and can't be set by config or the client. +- At most 20 tags are sent per statement, with automatic tags first. +- Keys must not contain `,` `:` `-` `/` `=` `.` or start with `@@`. Keys and values are limited to 128 characters; longer values are truncated. A `null` value records a key-only tag. +- Invalid tags are dropped, never failing the query. Invalid config tags log a warning. +- Tags are stored as plain text. Never put secrets or PII in them. +- Cached results never reach the warehouse, so they're not tagged. With a per-request `queryTags` function on a shared (service principal) cache, only the request that missed the cache is recorded. + +| Option | Type | Default | Description | +| ----------------- | ------------------------------ | ------- | ------------------------------------------------------------------ | +| `queryTags` | `QueryTagMap \| (ctx) => QueryTagMap` | — | Custom tags for every statement | +| `autoQueryTags` | `boolean` | `true` | Add the automatic `appkit_*` tags | +| `clientQueryTags` | `boolean` | `true` | Forward validated tags sent by `useAnalyticsQuery`/`useMetricView` | + +To find slow pages, for example: + +```sql +SELECT + query_tags['appkit_query_key'] AS query_key, + query_tags['page'] AS page, + count(*) AS runs, + percentile(total_duration_ms, 0.9) AS p90_ms +FROM system.query.history +WHERE query_tags['appkit_app'] = 'my-app' +GROUP BY ALL +ORDER BY p90_ms DESC; +``` + ## Frontend usage ### useAnalyticsQuery @@ -316,6 +388,7 @@ const { data, loading, error } = useAnalyticsQuery( | `format` | `"JSON" \| "ARROW"` | `"JSON"` | Response format | | `maxParametersSize` | `number` | `102400` | Max serialized parameters size in bytes | | `autoStart` | `boolean` | `true` | Start query on mount | +| `queryTags` | `AnalyticsQueryTags` | — | Query tags for this request (see [Query tags](#query-tags)) | ### Warehouse readiness @@ -562,6 +635,7 @@ Time-series queries should explicitly order their selected time dimension ascend | `orderBy` | `{field, direction?}[]` | no | Sort keys. `field` is narrowed to the measures/dimensions this call selected, so ordering by an unselected column is a type error. See [Deterministic results with `limit`](#deterministic-results-with-limit). | | `limit` | `number` | no | Positive integer row cap. | | `autoStart` | `boolean` | no | Start the metric query automatically. Defaults to `true`; set to `false` to defer it until the option becomes `true`. | +| `queryTags` | `AnalyticsQueryTags` | no | Query tags for this request (see [Query tags](#query-tags)). | **Return type:** From bc073fa318fda980fc67b5403bffe5afe89b9ab1 Mon Sep 17 00:00:00 2001 From: Atila Fassina Date: Thu, 29 Oct 2026 13:56:43 +0100 Subject: [PATCH 5/5] fix(appkit): harden query tag resolution from PR review - A function queryTags config now owns the custom tag set: client tags are only sent if it returns them (ctx.clientTags). Previously client tags were always merged in, so a conditional function (e.g. tenant from a header) let any client forge the key it left out. - Merge tags in an insertion-ordered Map. A plain object enumerated integer-like client keys first, so 20 of them pushed every appkit_* tag past the 20-tag cap; it also swallowed __proto__ tags. - Reject control characters, padded keys and the appkit_ prefix in any case; truncate values by code point. - Examine at most 64 client entries and log only dropped counts, never untrusted key names. - Sanitize a static queryTags map once instead of warning per statement. - Export QueryTagMap and QueryTagsContext; document that concurrent identical requests share the first caller's tags. Co-authored-by: Isaac Signed-off-by: Atila Fassina --- docs/docs/plugins/analytics.md | 25 ++- packages/appkit-ui/src/react/hooks/types.ts | 5 +- .../src/plugins/analytics/query-tags.ts | 168 +++++++++++------- .../plugins/analytics/tests/analytics.test.ts | 14 ++ .../analytics/tests/query-tags.test.ts | 168 +++++++++++++++--- .../appkit/src/plugins/analytics/types.ts | 11 +- 6 files changed, 293 insertions(+), 98 deletions(-) diff --git a/docs/docs/plugins/analytics.md b/docs/docs/plugins/analytics.md index 85d63bc16..33fdf911c 100644 --- a/docs/docs/plugins/analytics.md +++ b/docs/docs/plugins/analytics.md @@ -323,16 +323,33 @@ useAnalyticsQuery("revenue", params, { queryTags: { page: "sales_overview" } }); useMetricView("orders", { measures: ["revenue"], queryTags: { page: "orders" } }); ``` -Client tags are untrusted input. Set `clientQueryTags: false` to ignore them, or filter them in a `queryTags` function, which receives the validated client tags as `ctx.clientTags`. +Client tags are untrusted input, so any user of the app can set them. How they're handled depends on the `queryTags` config: + +- **No config, or a static map:** validated client tags are sent. On a key collision, the config wins. +- **A function:** its return value is the whole custom set. The function receives the validated client tags as `ctx.clientTags`, and they're only sent if it returns them. To keep a client's `page` tag but nothing else: + +```ts +analytics({ + queryTags: ({ clientTags, req }) => { + const tags: Record = {}; + if (clientTags.page) tags.page = clientTags.page; + const tenant = req?.header("x-tenant"); + if (tenant) tags.tenant = tenant; + return tags; + }, +}); +``` + +Use a function whenever a tag feeds chargeback or governance, so a client can't fill in a key you left out. Set `clientQueryTags: false` to ignore client tags entirely. **Rules:** -- Precedence on collisions: automatic `appkit_*` tags, then the `queryTags` config, then client tags. The `appkit_` prefix is reserved and can't be set by config or the client. +- Precedence on collisions: automatic `appkit_*` tags, then the `queryTags` config, then client tags. The `appkit_` prefix is reserved (in any letter case) and can't be set by config or the client. - At most 20 tags are sent per statement, with automatic tags first. -- Keys must not contain `,` `:` `-` `/` `=` `.` or start with `@@`. Keys and values are limited to 128 characters; longer values are truncated. A `null` value records a key-only tag. +- Keys must not contain `,` `:` `-` `/` `=` `.` or control characters, start with `@@`, or have leading or trailing spaces. Keys and values are limited to 128 characters; longer values are truncated. A `null` value records a key-only tag. - Invalid tags are dropped, never failing the query. Invalid config tags log a warning. - Tags are stored as plain text. Never put secrets or PII in them. -- Cached results never reach the warehouse, so they're not tagged. With a per-request `queryTags` function on a shared (service principal) cache, only the request that missed the cache is recorded. +- Tags aren't part of the cache key. Cached results never reach the warehouse, so they're not tagged. Concurrent identical requests share one in-flight statement, which carries the tags of the request that arrived first. For per-page or per-request tags, that means only the request that actually ran the statement is recorded, and other requests served by that run go uncounted. | Option | Type | Default | Description | | ----------------- | ------------------------------ | ------- | ------------------------------------------------------------------ | diff --git a/packages/appkit-ui/src/react/hooks/types.ts b/packages/appkit-ui/src/react/hooks/types.ts index 213539a9f..62f6f6bc5 100644 --- a/packages/appkit-ui/src/react/hooks/types.ts +++ b/packages/appkit-ui/src/react/hooks/types.ts @@ -44,18 +44,19 @@ export interface TypedArrowTable< // Query Options & Result Types // ============================================================================ -/** Options for configuring an analytics SSE query */ /** * Query tags attached to the statement, recorded in * `system.query.history.query_tags`. A `null` value records a key-only tag. * The server validates them (keys must not contain `, : - / = .` or start with * `appkit_`; keys/values are capped at 128 characters), drops invalid ones, and - * lets server-configured tags win on collisions. Never put secrets or PII here. + * decides which to forward (server-configured tags win on collisions). Never + * put secrets or PII here. * * @example { page: "sales_overview" } */ export type AnalyticsQueryTags = Record; +/** Options for configuring an analytics SSE query */ export interface UseAnalyticsQueryOptions< F extends AnalyticsFormat = "JSON_ARRAY", > { diff --git a/packages/appkit/src/plugins/analytics/query-tags.ts b/packages/appkit/src/plugins/analytics/query-tags.ts index ad8fe23ee..88f78682f 100644 --- a/packages/appkit/src/plugins/analytics/query-tags.ts +++ b/packages/appkit/src/plugins/analytics/query-tags.ts @@ -26,7 +26,10 @@ export interface QueryTagsContext { lane: "sp" | "obo"; /** Incoming request; absent for agent tool calls. */ req?: express.Request; - /** Sanitized tags sent by the client in the request body. */ + /** + * Validated tags sent by the client in the request body. They are only + * forwarded if the function returns them. + */ clientTags: QueryTagMap; } @@ -38,20 +41,41 @@ type QueryTagsConfig = Pick< /** Hard cap the warehouse applies; anything past it is dropped server-side. */ export const MAX_QUERY_TAGS = 20; const MAX_TAG_LENGTH = 128; -/** Prefix reserved for AppKit's automatic tags. */ +/** + * Upper bound on client entries examined per request, so a huge `queryTags` + * body can't buy unbounded work when at most 20 tags are ever sent. + */ +const MAX_CLIENT_CANDIDATES = 64; +/** Prefix reserved for AppKit's automatic tags (matched case-insensitively). */ const RESERVED_PREFIX = "appkit_"; -/** Characters the warehouse rejects in tag keys. */ -const INVALID_KEY_CHARS = /[,:\-/=.]/; +/** Characters the warehouse rejects in tag keys, plus control characters. */ +const INVALID_KEY_CHARS = /[,:\-/=.\p{Cc}]/u; + +/** + * Tags in priority order. A `Map` keeps insertion order for every key — + * a plain object would enumerate integer-like keys (`"0"`, `"1"`) first and + * treat `__proto__` as the prototype setter. + */ +type TagEntries = Map; function isValidKey(key: string): boolean { return ( - key.trim().length > 0 && + key.length > 0 && + key === key.trim() && key.length <= MAX_TAG_LENGTH && !INVALID_KEY_CHARS.test(key) && - !key.startsWith("@@") + !key.startsWith("@@") && + !key.toLowerCase().startsWith(RESERVED_PREFIX) ); } +/** Truncate by code point so a surrogate pair is never split. */ +function truncate(value: string): string { + return value.length <= MAX_TAG_LENGTH + ? value + : Array.from(value).slice(0, MAX_TAG_LENGTH).join(""); +} + /** * Keep only well-formed, non-reserved tags. Invalid entries are dropped rather * than failing the query: a bad tag must never break a dashboard. Values are @@ -60,58 +84,66 @@ function isValidKey(key: string): boolean { export function sanitizeQueryTags( input: unknown, source: "client" | "config", -): QueryTagMap { +): TagEntries { + const out: TagEntries = new Map(); if (input === null || typeof input !== "object" || Array.isArray(input)) { - return {}; + return out; } - const out: QueryTagMap = {}; - const dropped: string[] = []; - for (const [key, value] of Object.entries(input)) { - if ( - !isValidKey(key) || - key.startsWith(RESERVED_PREFIX) || - (value !== null && typeof value !== "string") - ) { - dropped.push(key); + let examined = 0; + let dropped = 0; + for (const key in input) { + if (!Object.hasOwn(input, key)) continue; + if (source === "client" && ++examined > MAX_CLIENT_CANDIDATES) { + dropped++; + break; + } + const value = (input as Record)[key]; + if (!isValidKey(key) || (value !== null && typeof value !== "string")) { + dropped++; continue; } - out[key] = value === null ? null : value.slice(0, MAX_TAG_LENGTH); + out.set(key, value === null ? null : truncate(value)); } - if (dropped.length > 0) { - const message = "Dropped %d invalid %s query tag(s): %s"; - const keys = dropped.slice(0, 5).join(", ").slice(0, 200); - // Client input is untrusted and may be noisy; config mistakes are the - // developer's to fix. + if (dropped > 0) { + // Only the count is logged: key names are untrusted and could forge log + // lines. Client input may be noisy; config mistakes are the developer's. if (source === "config") { - logger.warn(message, dropped.length, source, keys); + logger.warn("Dropped %d invalid config query tag(s)", dropped); } else { - logger.debug(message, dropped.length, source, keys); + logger.debug("Dropped %d invalid client query tag(s)", dropped); } } return out; } -function autoTags(ctx: QueryTagsContext): QueryTagMap { - const tags: QueryTagMap = {}; +function autoTags( + ctx: Pick, +): TagEntries { + const tags: TagEntries = new Map(); const appName = process.env.DATABRICKS_APP_NAME; - if (appName) tags.appkit_app = appName.slice(0, MAX_TAG_LENGTH); - if (ctx.queryKey) { - tags.appkit_query_key = ctx.queryKey.slice(0, MAX_TAG_LENGTH); + if (appName) tags.set("appkit_app", truncate(appName)); + if (ctx.queryKey) tags.set("appkit_query_key", truncate(ctx.queryKey)); + tags.set("appkit_route", ctx.route); + tags.set("appkit_lane", ctx.lane); + return tags; +} + +/** Static maps are sanitized (and warned about) once, not per statement. */ +const staticTagsCache = new WeakMap(); + +function staticConfigTags(queryTags: QueryTagMap): TagEntries { + let tags = staticTagsCache.get(queryTags); + if (!tags) { + tags = sanitizeQueryTags(queryTags, "config"); + staticTagsCache.set(queryTags, tags); } - tags.appkit_route = ctx.route; - tags.appkit_lane = ctx.lane; return tags; } -function configTags( - config: QueryTagsConfig, +function functionConfigTags( + queryTags: (ctx: QueryTagsContext) => QueryTagMap | undefined, ctx: QueryTagsContext, -): QueryTagMap { - const { queryTags } = config; - if (queryTags === undefined) return {}; - if (typeof queryTags !== "function") { - return sanitizeQueryTags(queryTags, "config"); - } +): TagEntries { try { return sanitizeQueryTags(queryTags(ctx), "config"); } catch (err) { @@ -119,48 +151,62 @@ function configTags( "queryTags callback threw; continuing without custom tags: %s", err instanceof Error ? err.message : String(err), ); - return {}; + return new Map(); } } /** * Resolve the tags for one statement. * - * Precedence on key collisions: automatic `appkit_*` tags (reserved, cannot - * be overridden) > server `queryTags` config > client-sent tags. Tags are - * then capped at {@link MAX_QUERY_TAGS} in the same priority order, so the - * automatic tags always fit. + * Automatic `appkit_*` tags come first and can't be overridden. Then: + * - static `queryTags` map: config tags, then client tags (config wins on + * collisions); + * - function `queryTags`: its return value is the whole custom set. Client + * tags are only sent if it returns them, e.g. `{ ...ctx.clientTags }`. + * + * The result is capped at {@link MAX_QUERY_TAGS} in that priority order, so + * the automatic tags always fit. */ export function resolveQueryTags( config: QueryTagsConfig, ctx: Omit & { clientTags?: unknown }, ): QueryTag[] | undefined { - const clientTags = + const clientTags: TagEntries = config.clientQueryTags === false - ? {} + ? new Map() : sanitizeQueryTags(ctx.clientTags, "client"); - const fullCtx: QueryTagsContext = { ...ctx, clientTags }; - const merged: QueryTagMap = { - ...(config.autoQueryTags === false ? {} : autoTags(fullCtx)), - }; - for (const tags of [configTags(config, fullCtx), clientTags]) { - for (const [key, value] of Object.entries(tags)) { - if (!Object.hasOwn(merged, key)) merged[key] = value; + const merged: TagEntries = + config.autoQueryTags === false ? new Map() : autoTags(ctx); + const { queryTags } = config; + const sources: TagEntries[] = + typeof queryTags === "function" + ? [ + functionConfigTags(queryTags, { + ...ctx, + clientTags: Object.fromEntries(clientTags), + }), + ] + : [queryTags ? staticConfigTags(queryTags) : new Map(), clientTags]; + for (const tags of sources) { + for (const [key, value] of tags) { + if (!merged.has(key)) merged.set(key, value); } } - const entries = Object.entries(merged).slice(0, MAX_QUERY_TAGS); - if (entries.length === 0) return undefined; - return entries.map(([key, value]) => - value === null ? { key } : { key, value }, - ); + const result: QueryTag[] = []; + for (const [key, value] of merged) { + if (result.length === MAX_QUERY_TAGS) break; + result.push(value === null ? { key } : { key, value }); + } + return result.length > 0 ? result : undefined; } /** * Wrap an executor so every statement it issues carries `queryTags`. Tags are - * not part of any cache key: a cache hit never reaches the warehouse, so there - * is nothing to attribute. + * not part of any cache key: a cache hit never reaches the warehouse, and + * concurrent identical requests share one in-flight statement tagged by the + * first caller. */ export function withQueryTags( executor: QueryExecutor, diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index d563fff35..00e2a6594 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -785,6 +785,13 @@ describe("Analytics Plugin", () => { disposition: "EXTERNAL_LINKS", format: "ARROW_STREAM", }); + // Both attempts carry the query tags, not just the first. + for (const call of executeMock.mock.calls) { + expect(call[1].queryTags).toContainEqual({ + key: "appkit_query_key", + value: "test_query", + }); + } // Streams the pre-signed links from the execute response (no re-fetch). expect(streamExternalLinksMock).toHaveBeenCalledTimes(1); expect(streamExternalLinksMock.mock.calls[0][0]).toBe(links); @@ -1336,6 +1343,13 @@ describe("Analytics Plugin", () => { disposition: "INLINE", format: "ARROW_STREAM", }); + // Both attempts carry the query tags, not just the first. + for (const call of executeMock.mock.calls) { + expect(call[1].queryTags).toContainEqual({ + key: "appkit_query_key", + value: "test_query", + }); + } // The SSE wire payload must look like a JSON_ARRAY result, not an // arrow message — the caller asked for JSON_ARRAY and the server has diff --git a/packages/appkit/src/plugins/analytics/tests/query-tags.test.ts b/packages/appkit/src/plugins/analytics/tests/query-tags.test.ts index 5a8f545f6..6cbc9eed8 100644 --- a/packages/appkit/src/plugins/analytics/tests/query-tags.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/query-tags.test.ts @@ -7,54 +7,90 @@ import { withQueryTags, } from "../query-tags"; +function sanitized(input: unknown, source: "client" | "config" = "client") { + return Object.fromEntries(sanitizeQueryTags(input, source)); +} + describe("sanitizeQueryTags", () => { test("keeps valid string and key-only tags", () => { - expect( - sanitizeQueryTags( - { team: "eng", page_name: "sales", exp: null }, - "client", - ), - ).toEqual({ team: "eng", page_name: "sales", exp: null }); + expect(sanitized({ team: "eng", page_name: "sales", exp: null })).toEqual({ + team: "eng", + page_name: "sales", + exp: null, + }); }); test("drops keys with characters the warehouse rejects", () => { expect( - sanitizeQueryTags( - { - "a,b": "x", - "a:b": "x", - "a-b": "x", - "a/b": "x", - "a=b": "x", - "a.b": "x", - "@@internal": "x", - " ": "x", - ok: "x", - }, - "client", - ), + sanitized({ + "a,b": "x", + "a:b": "x", + "a-b": "x", + "a/b": "x", + "a=b": "x", + "a.b": "x", + "@@internal": "x", + " ": "x", + " lead": "x", + "trail ": "x", + "new\nline": "x", + "tab\there": "x", + ok: "x", + }), ).toEqual({ ok: "x" }); }); test("drops reserved appkit_ keys so attribution cannot be spoofed", () => { expect( - sanitizeQueryTags({ appkit_app: "spoofed", team: "eng" }, "client"), + sanitized({ appkit_app: "spoofed", APPKIT_route: "x", team: "eng" }), ).toEqual({ team: "eng" }); }); test("drops non-string values and over-long keys, truncates long values", () => { - const out = sanitizeQueryTags( - { n: 1, o: { x: 1 }, ["k".repeat(129)]: "x", long: "v".repeat(500) }, - "client", - ); + const out = sanitized({ + n: 1, + o: { x: 1 }, + ["k".repeat(129)]: "x", + long: "v".repeat(500), + }); expect(out).toEqual({ long: "v".repeat(128) }); }); test("ignores non-object input", () => { for (const input of [undefined, null, "team:eng", ["team"], 42]) { - expect(sanitizeQueryTags(input, "client")).toEqual({}); + expect(sanitized(input)).toEqual({}); } }); + + test("truncates by code point without splitting surrogate pairs", () => { + const value = sanitized({ emoji: "😀".repeat(200) }).emoji as string; + expect(Array.from(value)).toHaveLength(128); + expect(value).toBe("😀".repeat(128)); + }); + + test("keeps a JSON-parsed __proto__ key as a plain tag", () => { + expect( + sanitizeQueryTags(JSON.parse('{"__proto__":"exp"}'), "client").get( + "__proto__", + ), + ).toBe("exp"); + }); + + test("examines a bounded number of client entries", () => { + const huge = Object.fromEntries( + Array.from({ length: 10_000 }, (_, i) => [`k${i}`, "v"]), + ); + expect(sanitizeQueryTags(huge, "client").size).toBe(64); + }); + + test("logs only the dropped count, never client key names", () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + sanitizeQueryTags({ "bad:\nFORGED": "x" }, "config"); + const logged = warn.mock.calls.flat().join(" "); + expect(logged).toContain("Dropped 1 invalid config query tag(s)"); + expect(logged).not.toContain("FORGED"); + warn.mockRestore(); + }); }); describe("resolveQueryTags", () => { @@ -121,7 +157,7 @@ describe("resolveQueryTags", () => { expect(tags).toBeUndefined(); }); - test("passes context, including sanitized client tags, to a function config", () => { + test("passes context, including validated client tags, to a function config", () => { const queryTags = vi.fn().mockReturnValue({ tenant: "acme" }); const tags = resolveQueryTags( { autoQueryTags: false, queryTags }, @@ -140,12 +176,71 @@ describe("resolveQueryTags", () => { clientTags: { page: "overview" }, }), ); + // The function's return value is the whole custom set: client tags it + // didn't return are not sent. + expect(tags).toEqual([{ key: "tenant", value: "acme" }]); + }); + + test("a function config that returns nothing drops client tags", () => { + // The docs' conditional-tenant example: no header means no tenant, and a + // client must not be able to fill the gap. + const tags = resolveQueryTags( + { + queryTags: ({ req }) => { + const tenant = req?.header("x-tenant"); + return tenant ? { tenant } : undefined; + }, + }, + { + route: "query", + queryKey: "revenue", + lane: "sp", + clientTags: { tenant: "acme", page: "overview" }, + }, + ); + expect(tags?.map((t) => t.key)).toEqual([ + "appkit_app", + "appkit_query_key", + "appkit_route", + "appkit_lane", + ]); + }); + + test("a function config can pass client tags through", () => { + const tags = resolveQueryTags( + { + autoQueryTags: false, + queryTags: ({ clientTags }) => ({ ...clientTags, team: "finance" }), + }, + { route: "query", lane: "sp", clientTags: { page: "overview" } }, + ); expect(tags).toEqual([ - { key: "tenant", value: "acme" }, { key: "page", value: "overview" }, + { key: "team", value: "finance" }, ]); }); + test("a function config returning a non-object adds no tags", () => { + const tags = resolveQueryTags( + { autoQueryTags: false, queryTags: () => "team:eng" as never }, + { route: "query", lane: "sp" }, + ); + expect(tags).toBeUndefined(); + }); + + test("an invalid static config warns once, not per statement", () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + const config = { queryTags: { "cost-center": "x", team: "eng" } }; + for (let i = 0; i < 3; i++) { + resolveQueryTags(config, { route: "query", lane: "sp" }); + } + const drops = warn.mock.calls.filter((call) => + call.join(" ").includes("Dropped"), + ); + expect(drops).toHaveLength(1); + warn.mockRestore(); + }); + test("a throwing function config never fails the query", () => { const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); const tags = resolveQueryTags( @@ -165,6 +260,23 @@ describe("resolveQueryTags", () => { warn.mockRestore(); }); + test("integer-like client keys cannot push out automatic tags", () => { + const numeric = Object.fromEntries( + Array.from({ length: 20 }, (_, i) => [String(i), "v"]), + ); + const tags = resolveQueryTags( + {}, + { route: "query", queryKey: "q", lane: "sp", clientTags: numeric }, + ); + expect(tags).toHaveLength(MAX_QUERY_TAGS); + expect(tags?.slice(0, 4).map((t) => t.key)).toEqual([ + "appkit_app", + "appkit_query_key", + "appkit_route", + "appkit_lane", + ]); + }); + test("caps at the warehouse limit, keeping automatic tags first", () => { const many = Object.fromEntries( Array.from({ length: 30 }, (_, i) => [`k${i}`, "v"]), diff --git a/packages/appkit/src/plugins/analytics/types.ts b/packages/appkit/src/plugins/analytics/types.ts index 33ad2780d..139560291 100644 --- a/packages/appkit/src/plugins/analytics/types.ts +++ b/packages/appkit/src/plugins/analytics/types.ts @@ -8,6 +8,8 @@ import type { import type { QueryTagMap, QueryTagsContext } from "./query-tags"; +export type { QueryTagMap, QueryTagsContext } from "./query-tags"; + export type { MetricFilter, MetricFilterOperatorName, @@ -20,7 +22,9 @@ export interface IAnalyticsConfig extends BasePluginConfig { /** * Custom query tags added to every statement the analytics routes and agent * tool issue, recorded in `system.query.history.query_tags`. Pass a static - * map, or a function to derive tags per request (e.g. a tenant header). + * map (client tags are merged in, config wins on collisions), or a function + * to derive tags per request. A function's return value is the whole custom + * set: client tags are only sent if it returns them (`ctx.clientTags`). * Keys must not contain `, : - / = .`, must not start with the reserved * `appkit_` prefix, and keys/values are capped at 128 characters; invalid * tags are dropped with a warning. Never put secrets or PII in tags. @@ -38,8 +42,9 @@ export interface IAnalyticsConfig extends BasePluginConfig { autoQueryTags?: boolean; /** * When `true` (default), tags sent by the client (`queryTags` on - * `useAnalyticsQuery` / `useMetricView`) are forwarded after validation. - * Server tags win on key collisions. Set to `false` to ignore client tags. + * `useAnalyticsQuery` / `useMetricView`) are validated and forwarded — merged + * under a static `queryTags` map, or exposed as `ctx.clientTags` to a + * function. Set to `false` to ignore client tags entirely. * * @default true */