Skip to content

Commit 86ccb6e

Browse files
Add X-Databricks-Org-Id header to telemetry API calls for SPOG hosts (#5099)
## Why SPOG hosts serve multiple workspaces from a single URL and require the `X-Databricks-Org-Id` header to route requests to the correct workspace. Without it, telemetry recorded by `/telemetry-ext` lands in a central shard rather than the user's workspace. The CLI's telemetry code makes a direct `apiClient.Do()` call to `/telemetry-ext` that bypasses the SDK's generated service methods (which set this header from `cfg.WorkspaceID`) and was passing `nil` for headers. This mirrors the fix already applied for filer in #4985. SDK PR databricks/databricks-sdk-go#1634 only patched `Workspace.Download()` / `Workspace.Upload()`; it is not a transport-level visitor, so direct `apiClient.Do()` callers in the CLI still need to set the header explicitly. ## Changes - `libs/telemetry/logger.go`: add `orgIDHeaders(apiClient)` helper that returns `{"X-Databricks-Org-Id": cfg.WorkspaceID}` when set, `nil` otherwise; pass it to `apiClient.Do()` in `attempt()`. - `acceptance/telemetry/test.toml`: append `X-Databricks-Org-Id` to `IncludeRequestHeaders` so the existing telemetry acceptance tests capture the header on `POST /telemetry-ext`. The header value is resolved from `.well-known/databricks-config` via the SDK and now flows into the recorded request, giving end-to-end coverage. ## Test plan - [x] `go test ./acceptance -run TestAccept/telemetry` passes; `out.requests.txt` for success/partial-success/failure/timeout now includes `X-Databricks-Org-Id`. - [x] `go test ./libs/telemetry/...` passes. This pull request was AI-assisted by Isaac.
1 parent 80f7a54 commit 86ccb6e

6 files changed

Lines changed: 40 additions & 2 deletions

File tree

acceptance/telemetry/failure/out.requests.txt

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@
66
"headers": {
77
"Authorization": [
88
"Bearer [DATABRICKS_TOKEN]"
9+
],
10+
"X-Databricks-Org-Id": [
11+
"[NUMID]"
912
]
1013
},
1114
"method": "POST",
@@ -23,6 +26,9 @@
2326
"headers": {
2427
"Authorization": [
2528
"Bearer [DATABRICKS_TOKEN]"
29+
],
30+
"X-Databricks-Org-Id": [
31+
"[NUMID]"
2632
]
2733
},
2834
"method": "POST",
@@ -40,6 +46,9 @@
4046
"headers": {
4147
"Authorization": [
4248
"Bearer [DATABRICKS_TOKEN]"
49+
],
50+
"X-Databricks-Org-Id": [
51+
"[NUMID]"
4352
]
4453
},
4554
"method": "POST",

acceptance/telemetry/partial-success/out.requests.txt

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@
66
"headers": {
77
"Authorization": [
88
"Bearer [DATABRICKS_TOKEN]"
9+
],
10+
"X-Databricks-Org-Id": [
11+
"[NUMID]"
912
]
1013
},
1114
"method": "POST",
@@ -23,6 +26,9 @@
2326
"headers": {
2427
"Authorization": [
2528
"Bearer [DATABRICKS_TOKEN]"
29+
],
30+
"X-Databricks-Org-Id": [
31+
"[NUMID]"
2632
]
2733
},
2834
"method": "POST",
@@ -40,6 +46,9 @@
4046
"headers": {
4147
"Authorization": [
4248
"Bearer [DATABRICKS_TOKEN]"
49+
],
50+
"X-Databricks-Org-Id": [
51+
"[NUMID]"
4352
]
4453
},
4554
"method": "POST",

acceptance/telemetry/success/out.requests.txt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,9 @@
1414
],
1515
"User-Agent": [
1616
"cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS] cmd/selftest_send-telemetry cmd-exec-id/[CMD-EXEC-ID] interactive/none auth/pat"
17+
],
18+
"X-Databricks-Org-Id": [
19+
"[NUMID]"
1720
]
1821
},
1922
"method": "POST",

acceptance/telemetry/test.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
IncludeRequestHeaders = ["Authorization"]
1+
IncludeRequestHeaders = ["Authorization", "X-Databricks-Org-Id"]
22
RecordRequests = true
33

44
Local = true

acceptance/telemetry/timeout/out.requests.txt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@
66
"headers": {
77
"Authorization": [
88
"Bearer [DATABRICKS_TOKEN]"
9+
],
10+
"X-Databricks-Org-Id": [
11+
"[NUMID]"
912
]
1013
},
1114
"method": "POST",

libs/telemetry/logger.go

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -171,9 +171,23 @@ func Upload(ctx context.Context, ec protos.ExecutionContext) error {
171171
return errors.New("failed to upload telemetry logs after three attempts")
172172
}
173173

174+
// orgIDHeaders returns headers with X-Databricks-Org-Id set if a workspace ID
175+
// is configured. SPOG hosts require this header to route requests to the
176+
// correct workspace; without it, telemetry is recorded in a central shard
177+
// instead of the correct workspace.
178+
func orgIDHeaders(apiClient *client.DatabricksClient) map[string]string {
179+
wsID := apiClient.Config.WorkspaceID
180+
if wsID == "" {
181+
return nil
182+
}
183+
return map[string]string{
184+
"X-Databricks-Org-Id": wsID,
185+
}
186+
}
187+
174188
func attempt(ctx context.Context, apiClient *client.DatabricksClient, protoLogs []string) (*ResponseBody, error) {
175189
resp := &ResponseBody{}
176-
err := apiClient.Do(ctx, http.MethodPost, "/telemetry-ext", nil, nil, RequestBody{
190+
err := apiClient.Do(ctx, http.MethodPost, "/telemetry-ext", orgIDHeaders(apiClient), nil, RequestBody{
177191
UploadTime: time.Now().UnixMilli(),
178192
// There is a bug in the `/telemetry-ext` API which requires us to
179193
// send an empty array for the `Items` field. Otherwise the API returns

0 commit comments

Comments
 (0)