feat(query-orchestrator): add trace_id SQL comment to data source queries - #11376
feat(query-orchestrator): add trace_id SQL comment to data source queries#11376igorlukanin wants to merge 8 commits into
Conversation
|
Claude finished @igorlukanin's task in 1m 48s —— View job Review complete — 0 blockers, 3 low-severity notes
OverallClean, well-scoped design. The three subtle points — comment must come after the SQL for Snowflake, tagging must happen after cache key derivation, and comment delimiters must be sanitized because request ids are client-controllable — are all correctly handled and explained in comments and tests. Coverage is thorough (streaming path, tuple form, hostile input, cache-key stability, refresh keys stay untagged). Nothing blocking. Low-severity notes1. Dead code in
The allowlist regex 2.
3. Doc: "span identifies the attempt within that request" is misleading
The code comment ( Nits (optional)
|
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Workflows to automatically generate PRs for you. |
|
Claude finished @igorlukanin's task in 1m 9s —— View job Re-review after 81926a1 — 0 blockers, 1 low-severity note
VerdictThe follow-up commit cleanly addresses all three prior notes:
Bonus cleanups that were on the "nits" list also landed: the 128-char truncation is now called out in Low-severity noteTruncation runs before
Nits (optional, non-blocking)
Nothing else stands out. Ready to ship pending CI. |
|
Claude finished @igorlukanin's task in 53s —— View job Re-review after efe2c4e — 0 blockers, 0 outstanding notesPrior truncation/span note addressed cleanly
The fix reorders sanitize → strip-span → cap, exactly matching the prior note. Nothing else to flag. Ready to ship pending CI. |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #11376 +/- ##
===========================================
- Coverage 83.92% 59.49% -24.44%
===========================================
Files 257 224 -33
Lines 80755 17922 -62833
Branches 0 3643 +3643
===========================================
- Hits 67777 10663 -57114
+ Misses 12978 6738 -6240
- Partials 0 521 +521
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Claude finished @igorlukanin's task in 45s —— View job Re-review after e967fa8 — 0 blockers, 0 outstanding notesTuple-form coverage added; nothing else changed
The only change since the previous review is the addition of a
The comment above the block correctly notes that |
e967fa8 to
d74b9bc
Compare
|
Claude finished @igorlukanin's task in 2m 45s —— View job Re-review — 0 blockers, 1 medium, 1 low (both new, both about which queries get tagged)Details
What changedNo code changes since Medium —
|
| path | sets primaryQuery |
|---|---|
persistent / no refresh keys → queryWithRetryAndRelease (lines 266, 281) |
✅ |
must-revalidate and default !backgroundRenew → renewQuery (line 884) |
✅ |
background fetch → cacheQueryResult (line 346) |
❌ |
The third branch is taken whenever queryCacheOptions.backgroundRenew is on or the request carries cacheMode: 'stale-while-revalidate' — and that is client-settable via the REST cache= param (gateway.ts:337). For those requests traceQuery bails on !req.primaryQuery, the SQL reaches the warehouse untagged, and the join the feature exists for silently produces nothing. The tests only cover the default path, so nothing catches it. One-line fix plus a test; if the omission is deliberate, it needs a comment and a line in query-tracing.mdx, because "silently off for a whole cacheMode" is surprising for an opt-in feature. Inline comment posted.
Low — the lambda source-table download is tagged too
QueryCache.ts:884 hardcodes primaryQuery: true inside the public renewQuery, and PreAggregationPartitionRangeLoader.downloadLambdaTable calls it (PreAggregationPartitionRangeLoader.ts:389) with the user's requestId and no renewCycle. So one pre-aggregation helper query does get tagged, contradicting "only the query the user actually asked for". Harmless in itself — it is a data source query issued for that request, and csvQuery tags it consistently — but the invariant as written is wrong, so either thread primaryQuery through or relax the wording. Inline comment posted.
What holds up well
- Sanitization. Allowlist rather than blocklist, both delimiters stripped,
+excluded with the Hive/Spark hint reason spelled out. The hostile-request-id test asserts delimiter counts, not just absence of the payload, which is the right assertion. - Cache-key deferral.
traceQueryis called at theclient.query/client.stream/downloadQueryResultsboundary, strictly afterQueryCache.queryCacheKey; the queue dedup key is likewise derived upstream. The rewritten test now actually demonstrates this end-to-end. - Append, not prepend, with the trailing-semicolon handling, and the Snowflake reason recorded in the source rather than only in the PR description.
- Cap after span-strip ordering is locked in by a targeted test.
- Streaming path has its own coverage via the
createQueuetraceFn, with the default(req) => req.querykeeping othercreateQueuecallers untouched. renewCycleand thescheduler-prefix as deliberately redundant guards, with the comment explaining why both stay.
Both findings are about coverage of the feature, not about correctness of what it emits — nothing here can corrupt SQL or break caching.
|
Claude finished @igorlukanin's task in 2m 35s —— View job Re-review after
|
detached path in cacheQueryResult |
reachable from the background-fetch branch? |
|---|---|
isSameRequest && !renewCycle continue-wait refresh (line 1085) |
yes, but only for the same request id — correct to tag |
renewalKey && isExpired → "Renewing existing key" (line 1096) |
no — that branch passes no renewalKey |
waitForRenew → return fetchNew() |
n/a, request waits on it |
So no background refresh belonging to a stale request id can pick up a tag through the new option. The fix is exactly as narrow as it needs to be. The test's setup (cacheKeyQueries: { renewalThreshold, queries: [] } + cacheMode: 'stale-while-revalidate') does reach the background-fetch branch, so it genuinely fails without the change.
Low — lambda source-table read contradicts the "only the user's query" claim: fixed by correcting the claim. query-tracing.mdx:71-78 and the traceQuery comment now both say queries issued to serve a request, name the lambda read as the concrete second case, and tell readers to match on trace_id rather than expect one row. That's the better half to change — the lambda read really is issued for that request.
Low — docs omit pre-aggregations from "no tagged query to find"
docs-mintlify/admin/connect-to-data/query-tracing.mdx:80
A request served from a pre-aggregation goes through getExternalQueue, whose handler calls client.query(q.query, …) with no traceQuery (QueryCache.ts:642). Correct behaviour — Cube Store isn't the data source — but it's the case a production deployment hits most often, and the troubleshooting sentence lists only cache hits and coalescing. Inline suggestion posted.
Nit
The new is appended on the background fetch path test asserts only that the tagged SQL is present (toContain). The renew cycle startRenewCycle kicks off concurrently on that same branch, so an added expect(...).not.toContain on an untagged copy of that SQL would also lock in that the renew-cycle query stays untagged — the second half of what the original note asked for. Non-blocking; the existing is absent for background refreshes test covers the renew path separately.
Verified unchanged and still sound
- Sanitization: allowlist, both delimiters stripped,
+excluded with the Hive/Spark reason recorded. - Cap-after-span-strip ordering in
buildTraceComment(trace-comment.ts:47-49), locked in by a targeted test. - Cache-key deferral:
traceQueryonly runs at theclient.query/client.stream/downloadQueryResultsboundary, strictly afterQueryCache.queryCacheKeyand after the queue dedup key is derived. - Append-not-prepend with trailing-semicolon handling;
renewCycleand thescheduler-prefix as deliberately redundant guards.
Not verified: I couldn't execute the test suites in this environment (sandbox denied jest), so the assessment above is from reading the code paths, not from a green run. CI covers it.
f700157 to
12d2a0a
Compare
|
Claude finished @igorlukanin's task in 2m 10s —— View job Re-review after
|
|
Claude finished @igorlukanin's task in 1m 36s —— View job Re-review after
|
| check | result |
|---|---|
| Dependency direction | @cubejs-backend/base-driver is already a runtime dependencies entry of cubejs-query-orchestrator (package.json:32), and base-driver does not depend back — no cycle. utils.ts previously imported only types from it, so this adds a real runtime edge, but a legal one. |
| Export reachability | base-driver/src/index.ts:4 already has export * from './trace-comment', so toTraceId is reachable from the package root. |
| Name collisions in the widened surface | trace-comment's four exports don't collide with ./utils (cancelCombinator, TableName, SaveCancelFn) or the other barrel members. |
| Callers | The four extractRequestUUID call sites — QueryCache.ts:410, QueryCache.ts:1078 (same-request continue-wait match) and QueryQueue.ts:522,527 (cancel-by-request matching) — all import the aliased name and see byte-identical logic. The public static QueryCache.extractRequestUUID keeps its signature. |
Nit
toTraceId is now a shared contract but has no test that exercises it under its own name — the strip is only asserted indirectly, through buildTraceComment (trace-comment.test.ts:63-79). That indirect coverage is real and includes the load-bearing cap-after-strip case, so this is genuinely a nit; a two-line direct test would just make the invariant fail loudly at the right level if someone edits it. The non-numeric suffix case (subscription-server.ts:163 builds -span-${uuidv4()}) is the one worth pinning, since it's what forces lastIndexOf over an anchored -span-\d+$.
Verified unchanged and still sound
- Sanitization: allowlist not blocklist, both comment delimiters stripped,
+excluded with the Hive/Spark hint reason recorded; the hostile-input test asserts delimiter counts rather than mere absence of the payload. - Cap after span-strip ordering (
trace-comment.ts:49-51), locked in by a targeted test. - Cache-key deferral:
traceQueryonly runs at theclient.query/client.stream/downloadQueryResultsboundary, strictly afterQueryCache.queryCacheKeyand after the queue dedup key is derived. - Append-not-prepend with trailing-semicolon handling; the block-comment form survives driver-side SQL wrapping (
CREATE TABLE … AS (query), subquery pagination) where a trailing--would not. - Tagged-query coverage:
primaryQuery: trueon the background-fetch branch,renewCycleand thescheduler-prefix as deliberately redundant guards, and docs that correctly list cache hits, pre-aggregations and coalescing as the cases with no tagged query to find.
Not verified: dependencies aren't installed in this environment, so I couldn't typecheck or run the suites — the export/collision checks above are from reading the barrels and package.json, not from a green build. CI covers it.
|
Fair on the nit: added a direct |
|
Claude finished @igorlukanin's task in 1m 11s —— View job Re-review after
|
| case | input | expected | holds |
|---|---|---|---|
| numeric suffix | …d479-span-1 |
…d479 |
✅ |
| uuid suffix | conn1-msg2-span-f47ac10b-… |
conn1-msg2 |
✅ — lastIndexOf finds the single -span-; the uuid's own hyphens don't match the marker |
| no suffix | scheduler-abc |
unchanged | ✅ lastIndexOf → -1 |
| empty | '' |
'' |
✅ |
The uuid case is the one worth having: it's the case an anchored -span-\d+$ would silently fail, and its inline comment cites subscription-server.ts:163 as the producer, so the "why not a regex" question is answered in place rather than in review history.
I also checked the file is actually collected — jest.base-ts.config.js matches <rootDir>/test/**/*.test.ts and the package's unit script points at test/unit, so it runs.
Re-verified the re-export is behaviour-preserving
Diffed the pre-fb69fd7 extractRequestUUID against today's toTraceId: identical signature ((requestId: string): string), identical lastIndexOf('-span-') + substring, only the ternary polarity flipped. No caller can observe a difference. The four call sites (QueryCache.ts:410, QueryCache.ts:1078, QueryQueue.ts:522,527) are unchanged, and utils.ts:42 keeps the old name.
Verified unchanged and still sound
- Sanitization: allowlist not blocklist, both comment delimiters stripped,
+excluded with the Hive/Spark hint reason recorded; the hostile-input test asserts delimiter counts, not just absence of the payload. - Cap after span-strip ordering (
trace-comment.ts:49-51), locked in by the targetedcaps after stripping the span, never mid-suffixcase. - Cache-key deferral:
traceQueryonly runs at theclient.query/client.stream/downloadQueryResultsboundary, strictly afterQueryCache.queryCacheKeyand after the queue dedup key is derived. - Append-not-prepend with trailing-semicolon handling; the block-comment form survives driver-side SQL wrapping (
CREATE TABLE … AS (query), subquery pagination) where a trailing--would not. - Tagged-query coverage:
primaryQuery: trueon the background-fetch branch,renewCycleand thescheduler-prefix as deliberately redundant guards, and docs that list cache hits, pre-aggregations and coalescing as the cases with no tagged query to find.
Not verified: the sandbox denied jest, so the assertions above are traced by hand against the implementation rather than from a green run. CI covers it.
Ready to ship pending CI.
Summary
trace_idSQL comment to the queries Cube sends to a data source, so a warehouse's query history can be joined back to Cube's own query history. Off by default, enabled withCUBEJS_SQL_INCLUDE_TRACE_ID=true. Works for every driver, since it operates on the SQL text rather than on a driver-specific query id.QUERY_HISTORY, so a leading comment is silently lost on exactly the engine this is most useful for. Verified live against Snowflake — the same query prepended loses the comment and appended retains it. No supported engine requires a leading comment, so appending is both simpler and more correct than per-dialect placement.primaryQuery). Refresh keys, pre-aggregation helper queries and background renewals travel the same path but answer to no API request, so they stay untagged.x-request-id/traceparentverbatim), so it is sanitized with a character allowlist. Both comment delimiters are stripped: engines disagree on whether block comments nest, so a stray closing delimiter would leak live SQL on some and a stray opening one would swallow the statement on others.trace_iddeliberately drops the-span-Nsuffix so it matches thetrace_idcolumn of the Query History export, and the span is emitted as a separate field to tell apart the several data source queries one request fans out into.Test plan
cubejs-query-orchestratorunit suite green (78 tests), covering the queued and streaming pathsQUERY_HISTORY), Postgres (server log) and DuckDB (driver boundary): the main query is tagged, streaming queries are tagged with the correct SQL text, refresh-key queries are untagged, and results are unchanged