Send X-Databricks-Workspace-Id on remaining raw client.Do callers

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
 	}

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,

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=<workspace-id>" or "?w=<workspace-id>" 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

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,

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

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,
+	}
+}

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))
+		})
+	}
+}

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)
 	}

libs/filer/files_client.go

@@ -7,13 +7,15 @@ import (
 	"fmt"
 	"io"
 	"io/fs"
+	"maps"
 	"net/http"
 	"net/url"
 	"path"
 	"slices"
 	"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 {

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) {

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,

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