From 73dfa49d988f780bba8d646bddb4b04361568a83 Mon Sep 17 00:00:00 2001 From: Divyansh Vijayvergia Date: Mon, 1 Jun 2026 15:12:01 +0000 Subject: [PATCH] Send X-Databricks-Workspace-Id on remaining raw client.Do callers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #5368 flipped the workspace routing header from X-Databricks-Org-Id to X-Databricks-Workspace-Id across five hand-written CLI code paths. Review on that PR identified five more raw apiClient.Do call sites that were quietly missing the routing header entirely. On unified hosts with cfg.WorkspaceID set, those calls land on the gateway with no routing information and can fail or be misrouted. The newly covered call sites are all workspace-routed endpoints: - bundle/generate/downloader.go GET /api/2.0/workspace/get-status (databricks bundle generate) - cmd/pipelines/utils.go GET /api/2.0/pipelines//events (databricks pipelines events) - libs/databrickscfg/cfgpickers/warehouses.go GET /api/2.0/sql/warehouses (warehouse picker) - libs/apps/prompt/listers.go GET /api/2.0/database/instances/.../databases (Lakebase database lister) - libs/git/info.go GET /api/2.0/workspace/get-status (DBR-runtime git info probe) Rather than copying the header-building snippet a sixth time, this PR introduces a shared helper in libs/auth: auth.WorkspaceIDHeader = "X-Databricks-Workspace-Id" auth.WorkspaceIDHeaders(cfg *sdkconfig.Config) map[string]string The helper returns nil when cfg.WorkspaceID is unset or holds the CLI-only "none" sentinel, so callers can pass the result directly to client.Do without conditional checks. The five new sites use it. As cleanup, the four existing copies of the same pattern are migrated to the helper as well: - cmd/api/api.go (workspaceIDHeader const) - bundle/deploy/filer.go ((s stateFiler) workspaceIDHeaders) - libs/filer/files_client.go ((w *FilesClient) workspaceIDHeaders) - libs/telemetry/logger.go (workspaceIDHeaders package function) libs/filer/workspace_files_client.go keeps a thin method wrapper around the shared helper because it has a legitimate nil-workspaceClient case in its existing tests; the wrapper handles that nil safety and delegates the rest. Side effect of the consolidation: callers that previously sent X-Databricks-Workspace-Id: none when cfg.WorkspaceID was the CLI sentinel "none" no longer do — the shared helper treats "none" as unset, matching the existing cmd/api/api.go normalization and the auth.ResolveWorkspaceID helper introduced in PR #5379. In practice users who set workspace_id = none in .databrickscfg are on account-scoped auth and these workspace- routed calls wouldn't succeed regardless, so this is a defensive cleanup not a behavior regression. --- bundle/deploy/filer.go | 16 +------ bundle/generate/downloader.go | 3 +- cmd/api/api.go | 9 +--- cmd/pipelines/utils.go | 3 +- libs/apps/prompt/listers.go | 3 ++ libs/auth/headers.go | 31 +++++++++++++ libs/auth/headers_test.go | 49 +++++++++++++++++++++ libs/databrickscfg/cfgpickers/warehouses.go | 8 +++- libs/filer/files_client.go | 21 ++------- libs/filer/workspace_files_client.go | 19 ++++---- libs/git/info.go | 3 +- libs/telemetry/logger.go | 19 ++------ 12 files changed, 114 insertions(+), 70 deletions(-) create mode 100644 libs/auth/headers.go create mode 100644 libs/auth/headers_test.go diff --git a/bundle/deploy/filer.go b/bundle/deploy/filer.go index e773df4e0e0..683eb9fc433 100644 --- a/bundle/deploy/filer.go +++ b/bundle/deploy/filer.go @@ -11,6 +11,7 @@ import ( "strings" "github.com/databricks/cli/bundle" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/filer" "github.com/databricks/databricks-sdk-go/client" ) @@ -25,19 +26,6 @@ type stateFiler struct { root filer.WorkspaceRootPath } -// workspaceIDHeaders returns headers with X-Databricks-Workspace-Id set if a -// workspace ID is configured. SPOG hosts require this header to route requests -// to the correct workspace. -func (s stateFiler) workspaceIDHeaders() map[string]string { - wsID := s.apiClient.Config.WorkspaceID - if wsID == "" { - return nil - } - return map[string]string{ - "X-Databricks-Workspace-Id": wsID, - } -} - func (s stateFiler) Delete(ctx context.Context, path string, mode ...filer.DeleteMode) error { return s.filer.Delete(ctx, path, mode...) } @@ -63,7 +51,7 @@ func (s stateFiler) Read(ctx context.Context, path string) (io.ReadCloser, error var buf bytes.Buffer urlPath := "/api/2.0/workspace-files/" + url.PathEscape(strings.TrimLeft(absPath, "/")) - err = s.apiClient.Do(ctx, http.MethodGet, urlPath, s.workspaceIDHeaders(), nil, nil, &buf) + err = s.apiClient.Do(ctx, http.MethodGet, urlPath, auth.WorkspaceIDHeaders(s.apiClient.Config), nil, nil, &buf) if err != nil { return nil, err } diff --git a/bundle/generate/downloader.go b/bundle/generate/downloader.go index 4376dd4ac5e..acfef1ff3db 100644 --- a/bundle/generate/downloader.go +++ b/bundle/generate/downloader.go @@ -10,6 +10,7 @@ import ( "path/filepath" "strings" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/log" "github.com/databricks/cli/libs/notebook" @@ -162,7 +163,7 @@ func (n *Downloader) markNotebookForDownload(ctx context.Context, notebookPath * ctx, http.MethodGet, "/api/2.0/workspace/get-status", - nil, + auth.WorkspaceIDHeaders(n.w.Config), nil, map[string]string{ "path": *notebookPath, diff --git a/cmd/api/api.go b/cmd/api/api.go index 59528527b53..91514789cb3 100644 --- a/cmd/api/api.go +++ b/cmd/api/api.go @@ -20,13 +20,6 @@ import ( ) const ( - // workspaceIDHeader is the workspace routing identifier sent on - // workspace-scope requests against unified hosts. Generated SDK service - // methods set this per-call when cfg.WorkspaceID is populated; we mirror - // the same idiom. The gateway also accepts the legacy X-Databricks-Org-Id - // header for rollback safety. - workspaceIDHeader = "X-Databricks-Workspace-Id" - // orgIDQueryParam and workspaceIDQueryParam are the SPOG // (single-page-of-glass) URL convention used by the Databricks UI: // "?o=" or "?w=" identifies the workspace a @@ -122,7 +115,7 @@ func makeCommand(method string) *cobra.Command { headers := map[string]string{"Content-Type": "application/json"} if orgID != "" { - headers[workspaceIDHeader] = orgID + headers[auth.WorkspaceIDHeader] = orgID } var response any diff --git a/cmd/pipelines/utils.go b/cmd/pipelines/utils.go index b67d76ad2db..be2f481db05 100644 --- a/cmd/pipelines/utils.go +++ b/cmd/pipelines/utils.go @@ -12,6 +12,7 @@ import ( configresources "github.com/databricks/cli/bundle/config/resources" "github.com/databricks/cli/bundle/resources" "github.com/databricks/cli/bundle/run" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/diag" "github.com/databricks/cli/libs/logdiag" @@ -186,7 +187,7 @@ func fetchAllPipelineEvents(ctx context.Context, w *databricks.WorkspaceClient, ctx, "GET", path, - nil, + auth.WorkspaceIDHeaders(w.Config), nil, queryParams, &response, diff --git a/libs/apps/prompt/listers.go b/libs/apps/prompt/listers.go index 5052744eb52..4bc2b1e301e 100644 --- a/libs/apps/prompt/listers.go +++ b/libs/apps/prompt/listers.go @@ -4,11 +4,13 @@ import ( "context" "errors" "fmt" + "maps" "net/http" "net/url" "strconv" "strings" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go" @@ -175,6 +177,7 @@ func ListDatabases(ctx context.Context, instanceName string) ([]ListItem, error) var resp listDatabasesResponse path := fmt.Sprintf("/api/2.0/database/instances/%s/databases", url.PathEscape(instanceName)) headers := map[string]string{"Accept": "application/json"} + maps.Copy(headers, auth.WorkspaceIDHeaders(w.Config)) err = api.Do(ctx, http.MethodGet, path, headers, nil, nil, &resp) if err != nil { return nil, err diff --git a/libs/auth/headers.go b/libs/auth/headers.go new file mode 100644 index 00000000000..86ac006f4fe --- /dev/null +++ b/libs/auth/headers.go @@ -0,0 +1,31 @@ +package auth + +import ( + sdkconfig "github.com/databricks/databricks-sdk-go/config" +) + +// WorkspaceIDHeader is the request header name used to route workspace-scoped +// API calls to the correct workspace on unified ("SPOG") hosts. The platform +// gateway also accepts the legacy X-Databricks-Org-Id header for rollback +// safety. Generated SDK service methods set this header per-call when +// cfg.WorkspaceID is populated; CLI code paths that call client.Do directly +// need to set it themselves. +const WorkspaceIDHeader = "X-Databricks-Workspace-Id" + +// WorkspaceIDHeaders returns a map suitable as the headers argument to +// client.DatabricksClient.Do, populated with the workspace routing header +// when cfg.WorkspaceID is set. Returns nil when the workspace ID is unset +// or holds the CLI-only "none" sentinel, so callers can pass the result +// through without conditional checks. +func WorkspaceIDHeaders(cfg *sdkconfig.Config) map[string]string { + if cfg == nil { + return nil + } + wsID := cfg.WorkspaceID + if wsID == "" || wsID == WorkspaceIDNone { + return nil + } + return map[string]string{ + WorkspaceIDHeader: wsID, + } +} diff --git a/libs/auth/headers_test.go b/libs/auth/headers_test.go new file mode 100644 index 00000000000..fb3ab0ea3fe --- /dev/null +++ b/libs/auth/headers_test.go @@ -0,0 +1,49 @@ +package auth_test + +import ( + "testing" + + "github.com/databricks/cli/libs/auth" + sdkconfig "github.com/databricks/databricks-sdk-go/config" + "github.com/stretchr/testify/assert" +) + +func TestWorkspaceIDHeaders(t *testing.T) { + tests := []struct { + name string + cfg *sdkconfig.Config + want map[string]string + }{ + { + name: "configured numeric workspace ID", + cfg: &sdkconfig.Config{WorkspaceID: "12345"}, + want: map[string]string{"X-Databricks-Workspace-Id": "12345"}, + }, + { + name: "configured connection-id-style workspace ID", + cfg: &sdkconfig.Config{WorkspaceID: "123e4567-e89b-12d3-a456-426614174000"}, + want: map[string]string{"X-Databricks-Workspace-Id": "123e4567-e89b-12d3-a456-426614174000"}, + }, + { + name: "empty workspace ID returns nil", + cfg: &sdkconfig.Config{WorkspaceID: ""}, + want: nil, + }, + { + name: `"none" sentinel returns nil`, + cfg: &sdkconfig.Config{WorkspaceID: auth.WorkspaceIDNone}, + want: nil, + }, + { + name: "nil config returns nil", + cfg: nil, + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, auth.WorkspaceIDHeaders(tt.cfg)) + }) + } +} diff --git a/libs/databrickscfg/cfgpickers/warehouses.go b/libs/databrickscfg/cfgpickers/warehouses.go index 22c22d09022..67bdcb7f73a 100644 --- a/libs/databrickscfg/cfgpickers/warehouses.go +++ b/libs/databrickscfg/cfgpickers/warehouses.go @@ -8,6 +8,7 @@ import ( "slices" "strings" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" @@ -137,8 +138,11 @@ func listUsableWarehouses(ctx context.Context, w *databricks.WorkspaceClient) ([ apiClient := httpclient.NewApiClient(clientCfg) var response sql.ListWarehousesResponse - err = apiClient.Do(ctx, "GET", "/api/2.0/sql/warehouses?skip_cannot_use=true", - httpclient.WithResponseUnmarshal(&response)) + opts := []httpclient.DoOption{httpclient.WithResponseUnmarshal(&response)} + for name, value := range auth.WorkspaceIDHeaders(w.Config) { + opts = append(opts, httpclient.WithRequestHeader(name, value)) + } + err = apiClient.Do(ctx, "GET", "/api/2.0/sql/warehouses?skip_cannot_use=true", opts...) if err != nil { return nil, fmt.Errorf("list warehouses: %w", err) } diff --git a/libs/filer/files_client.go b/libs/filer/files_client.go index b09d10de7af..32fc24862c6 100644 --- a/libs/filer/files_client.go +++ b/libs/filer/files_client.go @@ -7,6 +7,7 @@ import ( "fmt" "io" "io/fs" + "maps" "net/http" "net/url" "path" @@ -14,6 +15,7 @@ import ( "strings" "time" + "github.com/databricks/cli/libs/auth" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/client" @@ -109,19 +111,6 @@ func NewFilesClient(w *databricks.WorkspaceClient, root string) (Filer, error) { }, nil } -// workspaceIDHeaders returns headers with X-Databricks-Workspace-Id set if a -// workspace ID is configured. SPOG hosts require this header to route requests -// to the correct workspace. -func (w *FilesClient) workspaceIDHeaders() map[string]string { - wsID := w.workspaceClient.Config.WorkspaceID - if wsID == "" { - return nil - } - return map[string]string{ - "X-Databricks-Workspace-Id": wsID, - } -} - func (w *FilesClient) urlPath(name string) (string, string, error) { absPath, err := w.root.Join(name) if err != nil { @@ -161,9 +150,7 @@ func (w *FilesClient) Write(ctx context.Context, name string, reader io.Reader, overwrite := slices.Contains(mode, OverwriteIfExists) urlPath = fmt.Sprintf("%s?overwrite=%t", urlPath, overwrite) headers := map[string]string{"Content-Type": "application/octet-stream"} - if wsID := w.workspaceClient.Config.WorkspaceID; wsID != "" { - headers["X-Databricks-Workspace-Id"] = wsID - } + maps.Copy(headers, auth.WorkspaceIDHeaders(w.workspaceClient.Config)) err = w.apiClient.Do(ctx, http.MethodPut, urlPath, headers, nil, reader, nil) // Return early on success. @@ -192,7 +179,7 @@ func (w *FilesClient) Read(ctx context.Context, name string) (io.ReadCloser, err } var reader io.ReadCloser - err = w.apiClient.Do(ctx, http.MethodGet, urlPath, w.workspaceIDHeaders(), nil, nil, &reader) + err = w.apiClient.Do(ctx, http.MethodGet, urlPath, auth.WorkspaceIDHeaders(w.workspaceClient.Config), nil, nil, &reader) // Return early on success. if err == nil { diff --git a/libs/filer/workspace_files_client.go b/libs/filer/workspace_files_client.go index d49b10b8ace..6e54f5b4d97 100644 --- a/libs/filer/workspace_files_client.go +++ b/libs/filer/workspace_files_client.go @@ -15,6 +15,7 @@ import ( "strings" "time" + "github.com/databricks/cli/libs/auth" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/client" @@ -121,20 +122,16 @@ type WorkspaceFilesClient struct { root WorkspaceRootPath } -// workspaceIDHeaders returns headers with X-Databricks-Workspace-Id set if a -// workspace ID is configured. SPOG hosts require this header to route requests -// to the correct workspace. +// workspaceIDHeaders returns the workspace routing header map for outbound +// API calls, or nil if the workspace client is unset. Wraps the shared +// auth.WorkspaceIDHeaders helper with a nil-safe workspaceClient guard +// since this filer struct can legitimately be constructed without one in +// some test setups. func (w *WorkspaceFilesClient) workspaceIDHeaders() map[string]string { - if w.workspaceClient == nil || w.workspaceClient.Config == nil { + if w.workspaceClient == nil { return nil } - wsID := w.workspaceClient.Config.WorkspaceID - if wsID == "" { - return nil - } - return map[string]string{ - "X-Databricks-Workspace-Id": wsID, - } + return auth.WorkspaceIDHeaders(w.workspaceClient.Config) } func NewWorkspaceFilesClient(w *databricks.WorkspaceClient, root string) (Filer, error) { diff --git a/libs/git/info.go b/libs/git/info.go index f7caf64da93..c59d6a82455 100644 --- a/libs/git/info.go +++ b/libs/git/info.go @@ -6,6 +6,7 @@ import ( "path" "strings" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/dbr" "github.com/databricks/cli/libs/folders" "github.com/databricks/cli/libs/log" @@ -65,7 +66,7 @@ func fetchRepositoryInfoAPI(ctx context.Context, path string, w *databricks.Work ctx, http.MethodGet, apiEndpoint, - nil, + auth.WorkspaceIDHeaders(w.Config), nil, map[string]string{ "path": path, diff --git a/libs/telemetry/logger.go b/libs/telemetry/logger.go index 7cda184bd04..cddc407c62e 100644 --- a/libs/telemetry/logger.go +++ b/libs/telemetry/logger.go @@ -9,6 +9,7 @@ import ( "strconv" "time" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/log" @@ -170,23 +171,11 @@ func Upload(ctx context.Context, ec protos.ExecutionContext) error { return errors.New("failed to upload telemetry logs after three attempts") } -// workspaceIDHeaders returns headers with X-Databricks-Workspace-Id set if a -// workspace ID is configured. SPOG hosts require this header to route requests -// to the correct workspace; without it, telemetry is recorded in a central -// shard instead of the correct workspace. -func workspaceIDHeaders(apiClient *client.DatabricksClient) map[string]string { - wsID := apiClient.Config.WorkspaceID - if wsID == "" { - return nil - } - return map[string]string{ - "X-Databricks-Workspace-Id": wsID, - } -} - func attempt(ctx context.Context, apiClient *client.DatabricksClient, protoLogs []string) (*ResponseBody, error) { + // Without the workspace routing header, telemetry on unified hosts is + // recorded in a central shard instead of the user's workspace. resp := &ResponseBody{} - err := apiClient.Do(ctx, http.MethodPost, "/telemetry-ext", workspaceIDHeaders(apiClient), nil, RequestBody{ + err := apiClient.Do(ctx, http.MethodPost, "/telemetry-ext", auth.WorkspaceIDHeaders(apiClient.Config), nil, RequestBody{ UploadTime: time.Now().UnixMilli(), // There is a bug in the `/telemetry-ext` API which requires us to // send an empty array for the `Items` field. Otherwise the API returns