From 3cfbac2cf30b21cd4c2a003248c64bf39e376fe6 Mon Sep 17 00:00:00 2001 From: simon Date: Tue, 24 Mar 2026 11:01:02 +0100 Subject: [PATCH 1/2] Add SPOG bundle support: account_id field, host URL parsing, workspace_id profile disambiguation Bundles need to support SPOG (Single Pane of Glass) hosts where multiple workspaces share the same host URL. This adds three capabilities: 1. account_id field in bundle workspace config, passed through to the SDK. 2. Host URL query parameter extraction (?o= for workspace_id, ?a= for account_id). Runs in Workspace.Client() before profile resolution so the extracted fields are available for disambiguation. A mutator in the initialize phase provides secondary cleanup. 3. Profile resolution by workspace_id when multiple profiles share the same host. Host matching runs first (preserving existing behavior), then workspace_id disambiguates if available. Also adds workspace_id and account_id to the auth interpolation warning since they now participate in profile resolution at auth time. Co-authored-by: Isaac --- bundle/config/mutator/normalize_host_url.go | 30 ++++++ .../config/mutator/normalize_host_url_test.go | 94 +++++++++++++++++++ .../validate/interpolation_in_auth_config.go | 4 + bundle/config/workspace.go | 50 ++++++++++ bundle/config/workspace_test.go | 61 ++++++++++++ bundle/internal/schema/annotations.yml | 3 + bundle/phases/initialize.go | 4 + bundle/schema/jsonschema.json | 4 + libs/databrickscfg/loader.go | 43 ++++++++- libs/databrickscfg/loader_test.go | 79 ++++++++++++++++ libs/databrickscfg/profile/file_test.go | 2 +- .../profile/testdata/databrickscfg | 22 +++++ 12 files changed, 394 insertions(+), 2 deletions(-) create mode 100644 bundle/config/mutator/normalize_host_url.go create mode 100644 bundle/config/mutator/normalize_host_url_test.go diff --git a/bundle/config/mutator/normalize_host_url.go b/bundle/config/mutator/normalize_host_url.go new file mode 100644 index 0000000000..2ca64a3849 --- /dev/null +++ b/bundle/config/mutator/normalize_host_url.go @@ -0,0 +1,30 @@ +package mutator + +import ( + "context" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/libs/diag" +) + +type normalizeHostURL struct{} + +// NormalizeHostURL extracts query parameters from the workspace host URL +// and strips them from the host. This allows users to paste SPOG URLs +// (e.g. https://host.databricks.com/?o=12345) directly into their bundle config. +// +// The primary extraction happens in [config.Workspace.Client] before the SDK +// config is built. This mutator serves as a secondary normalization pass to +// ensure the bundle config is clean for any later code that reads it directly. +func NormalizeHostURL() bundle.Mutator { + return &normalizeHostURL{} +} + +func (m *normalizeHostURL) Name() string { + return "NormalizeHostURL" +} + +func (m *normalizeHostURL) Apply(_ context.Context, b *bundle.Bundle) diag.Diagnostics { + b.Config.Workspace.NormalizeHostURL() + return nil +} diff --git a/bundle/config/mutator/normalize_host_url_test.go b/bundle/config/mutator/normalize_host_url_test.go new file mode 100644 index 0000000000..65788ce576 --- /dev/null +++ b/bundle/config/mutator/normalize_host_url_test.go @@ -0,0 +1,94 @@ +package mutator_test + +import ( + "testing" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/config/mutator" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNormalizeHostURL(t *testing.T) { + tests := []struct { + name string + host string + workspaceID string + accountID string + wantHost string + wantWorkspaceID string + wantAccountID string + }{ + { + name: "empty host", + host: "", + wantHost: "", + }, + { + name: "host without query params", + host: "https://api.databricks.com", + wantHost: "https://api.databricks.com", + }, + { + name: "host with ?o= extracts workspace_id", + host: "https://dogfood.staging.databricks.com/?o=6051921418418893", + wantHost: "https://dogfood.staging.databricks.com/", + wantWorkspaceID: "6051921418418893", + }, + { + name: "host with ?o= and ?a= extracts both", + host: "https://dogfood.staging.databricks.com/?o=605&a=abc123", + wantHost: "https://dogfood.staging.databricks.com/", + wantWorkspaceID: "605", + wantAccountID: "abc123", + }, + { + name: "host with ?account_id= extracts account_id", + host: "https://accounts.cloud.databricks.com/?account_id=968367da-1234", + wantHost: "https://accounts.cloud.databricks.com/", + wantAccountID: "968367da-1234", + }, + { + name: "host with ?workspace_id= extracts workspace_id", + host: "https://dogfood.staging.databricks.com/?workspace_id=12345", + wantHost: "https://dogfood.staging.databricks.com/", + wantWorkspaceID: "12345", + }, + { + name: "explicit workspace_id takes precedence over ?o=", + host: "https://dogfood.staging.databricks.com/?o=999", + workspaceID: "explicit-id", + wantHost: "https://dogfood.staging.databricks.com/", + wantWorkspaceID: "explicit-id", + }, + { + name: "explicit account_id takes precedence over ?a=", + host: "https://dogfood.staging.databricks.com/?a=from-url", + accountID: "explicit-account", + wantHost: "https://dogfood.staging.databricks.com/", + wantAccountID: "explicit-account", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + b := &bundle.Bundle{ + Config: config.Root{ + Workspace: config.Workspace{ + Host: tt.host, + WorkspaceID: tt.workspaceID, + AccountID: tt.accountID, + }, + }, + } + + diags := bundle.Apply(t.Context(), b, mutator.NormalizeHostURL()) + require.NoError(t, diags.Error()) + + assert.Equal(t, tt.wantHost, b.Config.Workspace.Host) + assert.Equal(t, tt.wantWorkspaceID, b.Config.Workspace.WorkspaceID) + assert.Equal(t, tt.wantAccountID, b.Config.Workspace.AccountID) + }) + } +} diff --git a/bundle/config/validate/interpolation_in_auth_config.go b/bundle/config/validate/interpolation_in_auth_config.go index 1a5b64a268..1598ef49d8 100644 --- a/bundle/config/validate/interpolation_in_auth_config.go +++ b/bundle/config/validate/interpolation_in_auth_config.go @@ -42,6 +42,10 @@ func (f *noInterpolationInAuthConfig) Apply(ctx context.Context, b *bundle.Bundl "azure_tenant_id", "azure_environment", "azure_login_app_id", + + // Unified host specific attributes. + "account_id", + "workspace_id", } diags := diag.Diagnostics{} diff --git a/bundle/config/workspace.go b/bundle/config/workspace.go index 32e2fdd38a..137c669bf9 100644 --- a/bundle/config/workspace.go +++ b/bundle/config/workspace.go @@ -1,6 +1,7 @@ package config import ( + "net/url" "os" "path/filepath" @@ -43,6 +44,7 @@ type Workspace struct { // Unified host specific attributes. ExperimentalIsUnifiedHost bool `json:"experimental_is_unified_host,omitempty"` + AccountID string `json:"account_id,omitempty"` WorkspaceID string `json:"workspace_id,omitempty"` // CurrentUser holds the current user. @@ -124,6 +126,7 @@ func (w *Workspace) Config() *config.Config { // Unified host Experimental_IsUnifiedHost: w.ExperimentalIsUnifiedHost, + AccountID: w.AccountID, WorkspaceID: w.WorkspaceID, } @@ -137,7 +140,54 @@ func (w *Workspace) Config() *config.Config { return cfg } +// NormalizeHostURL extracts query parameters from the host URL and populates +// the corresponding fields if not already set. This allows users to paste SPOG +// URLs (e.g. https://host.databricks.com/?o=12345) directly into their bundle +// config. Must be called before Config() so the extracted fields are included +// in the SDK config used for profile resolution and authentication. +func (w *Workspace) NormalizeHostURL() { + if w.Host == "" { + return + } + u, err := url.Parse(w.Host) + if err != nil { + return + } + q := u.Query() + if len(q) == 0 { + return + } + + // Extract workspace_id from ?o= or ?workspace_id= if not already set. + if w.WorkspaceID == "" { + if v := q.Get("o"); v != "" { + w.WorkspaceID = v + } else if v := q.Get("workspace_id"); v != "" { + w.WorkspaceID = v + } + } + + // Extract account_id from ?a= or ?account_id= if not already set. + if w.AccountID == "" { + if v := q.Get("a"); v != "" { + w.AccountID = v + } else if v := q.Get("account_id"); v != "" { + w.AccountID = v + } + } + + // Strip query parameters from the host. + u.RawQuery = "" + u.Fragment = "" + w.Host = u.String() +} + func (w *Workspace) Client() (*databricks.WorkspaceClient, error) { + // Extract query parameters (?o=, ?a=) from the host URL before building + // the SDK config. This ensures workspace_id and account_id are available + // for profile resolution during EnsureResolved(). + w.NormalizeHostURL() + cfg := w.Config() // If only the host is configured, we try and unambiguously match it to diff --git a/bundle/config/workspace_test.go b/bundle/config/workspace_test.go index 14947e2894..c1bb35641a 100644 --- a/bundle/config/workspace_test.go +++ b/bundle/config/workspace_test.go @@ -73,6 +73,67 @@ func TestWorkspaceResolveProfileFromHost(t *testing.T) { }) } +func TestWorkspaceNormalizeHostURL(t *testing.T) { + t.Run("extracts workspace_id from query param", func(t *testing.T) { + w := Workspace{ + Host: "https://spog.databricks.com/?o=12345", + } + w.NormalizeHostURL() + assert.Equal(t, "https://spog.databricks.com/", w.Host) + assert.Equal(t, "12345", w.WorkspaceID) + }) + + t.Run("explicit workspace_id takes precedence", func(t *testing.T) { + w := Workspace{ + Host: "https://spog.databricks.com/?o=999", + WorkspaceID: "explicit", + } + w.NormalizeHostURL() + assert.Equal(t, "https://spog.databricks.com/", w.Host) + assert.Equal(t, "explicit", w.WorkspaceID) + }) + + t.Run("no-op for host without query params", func(t *testing.T) { + w := Workspace{ + Host: "https://normal.databricks.com", + } + w.NormalizeHostURL() + assert.Equal(t, "https://normal.databricks.com", w.Host) + assert.Empty(t, w.WorkspaceID) + }) +} + +func TestWorkspaceClientNormalizesHostBeforeProfileResolution(t *testing.T) { + // Regression test: Client() must normalize the host URL (strip ?o= and + // populate WorkspaceID) before building the SDK config and resolving + // profiles. This ensures workspace_id is available for disambiguation. + setupWorkspaceTest(t) + + err := databrickscfg.SaveToProfile(t.Context(), &config.Config{ + Profile: "ws1", + Host: "https://spog.databricks.com", + Token: "token1", + WorkspaceID: "111", + }) + require.NoError(t, err) + + err = databrickscfg.SaveToProfile(t.Context(), &config.Config{ + Profile: "ws2", + Host: "https://spog.databricks.com", + Token: "token2", + WorkspaceID: "222", + }) + require.NoError(t, err) + + // Host with ?o= should be normalized and workspace_id used to disambiguate. + w := Workspace{ + Host: "https://spog.databricks.com/?o=222", + } + client, err := w.Client() + require.NoError(t, err) + assert.Equal(t, "ws2", client.Config.Profile) +} + func TestWorkspaceVerifyProfileForHost(t *testing.T) { // If both a workspace host and a profile are specified, // verify that the host configured in the profile matches diff --git a/bundle/internal/schema/annotations.yml b/bundle/internal/schema/annotations.yml index 8dcf37f7e8..f30e054ed8 100644 --- a/bundle/internal/schema/annotations.yml +++ b/bundle/internal/schema/annotations.yml @@ -415,6 +415,9 @@ github.com/databricks/cli/bundle/config.Target: "description": |- The Databricks workspace for the target. github.com/databricks/cli/bundle/config.Workspace: + "account_id": + "description": |- + The Databricks account ID. "artifact_path": "description": |- The artifact path to use within the workspace for both deployments and workflow runs diff --git a/bundle/phases/initialize.go b/bundle/phases/initialize.go index 64eb18943a..8e237ce728 100644 --- a/bundle/phases/initialize.go +++ b/bundle/phases/initialize.go @@ -34,6 +34,10 @@ func Initialize(ctx context.Context, b *bundle.Bundle) { validate.ValidateEngine(), validate.Scripts(), + // Extract query parameters (?o=, ?a=) from workspace host URL and strip them. + // Must run before any mutator that uses the host for authentication or API calls. + mutator.NormalizeHostURL(), + // Updates (dynamic): sync.{paths,include,exclude} (makes them relative to bundle root rather than to definition file) // Rewrites sync paths to be relative to the bundle root instead of the file they were defined in. mutator.RewriteSyncPaths(), diff --git a/bundle/schema/jsonschema.json b/bundle/schema/jsonschema.json index abfdb55233..da37fcd786 100644 --- a/bundle/schema/jsonschema.json +++ b/bundle/schema/jsonschema.json @@ -2627,6 +2627,10 @@ { "type": "object", "properties": { + "account_id": { + "description": "The Databricks account ID.", + "$ref": "#/$defs/string" + }, "artifact_path": { "description": "The artifact path to use within the workspace for both deployments and workflow runs", "$ref": "#/$defs/string" diff --git a/libs/databrickscfg/loader.go b/libs/databrickscfg/loader.go index 4113c80f1e..8f4f2a38db 100644 --- a/libs/databrickscfg/loader.go +++ b/libs/databrickscfg/loader.go @@ -98,7 +98,21 @@ func (l profileFromHostLoader) Configure(cfg *config.Config) error { if err == errNoMatchingProfiles { return nil } - if err, ok := err.(errMultipleProfiles); ok { + + // If multiple profiles match the same host and we have a workspace_id, + // try to disambiguate by matching workspace_id. + if names, ok := AsMultipleProfiles(err); ok && cfg.WorkspaceID != "" { + originalErr := err + match, err = l.disambiguateByWorkspaceID(ctx, configFile, host, cfg.WorkspaceID, names) + if err == errNoMatchingProfiles { + // workspace_id didn't match any of the host-matching profiles. + // Fall back to the original ambiguity error. + log.Debugf(ctx, "workspace_id=%s did not match any profiles for host %s: %v", cfg.WorkspaceID, host, names) + err = originalErr + } + } + + if _, ok := AsMultipleProfiles(err); ok { return fmt.Errorf( "%s: %w: please set DATABRICKS_CONFIG_PROFILE or provide --profile flag to specify one", host, err) @@ -120,6 +134,33 @@ func (l profileFromHostLoader) Configure(cfg *config.Config) error { return nil } +// disambiguateByWorkspaceID filters the profiles that matched a host by workspace_id. +func (l profileFromHostLoader) disambiguateByWorkspaceID( + ctx context.Context, + configFile *config.File, + host string, + workspaceID string, + profileNames []string, +) (*ini.Section, error) { + log.Debugf(ctx, "Multiple profiles matched host %s, disambiguating by workspace_id=%s", host, workspaceID) + + nameSet := make(map[string]bool, len(profileNames)) + for _, name := range profileNames { + nameSet[name] = true + } + + return findMatchingProfile(configFile, func(s *ini.Section) bool { + if !nameSet[s.Name()] { + return false + } + key, err := s.GetKey("workspace_id") + if err != nil { + return false + } + return key.Value() == workspaceID + }) +} + func (l profileFromHostLoader) isAnyAuthConfigured(cfg *config.Config) bool { // If any of the auth-specific attributes are set, we can skip profile resolution. for _, a := range config.ConfigAttributes { diff --git a/libs/databrickscfg/loader_test.go b/libs/databrickscfg/loader_test.go index 425620abb5..c53351ae42 100644 --- a/libs/databrickscfg/loader_test.go +++ b/libs/databrickscfg/loader_test.go @@ -164,3 +164,82 @@ func TestAsMultipleProfilesReturnsFalseForNil(t *testing.T) { assert.False(t, ok) assert.Nil(t, names) } + +func TestLoaderDisambiguatesByWorkspaceID(t *testing.T) { + cfg := config.Config{ + Loaders: []config.Loader{ + ResolveProfileFromHost, + }, + ConfigFile: "profile/testdata/databrickscfg", + Host: "https://spog.databricks.com", + WorkspaceID: "111", + } + + err := cfg.EnsureResolved() + require.NoError(t, err) + assert.Equal(t, "spog-ws1", cfg.Profile) + assert.Equal(t, "spog-ws1", cfg.Token) +} + +func TestLoaderDisambiguatesByWorkspaceIDSecondProfile(t *testing.T) { + cfg := config.Config{ + Loaders: []config.Loader{ + ResolveProfileFromHost, + }, + ConfigFile: "profile/testdata/databrickscfg", + Host: "https://spog.databricks.com", + WorkspaceID: "222", + } + + err := cfg.EnsureResolved() + require.NoError(t, err) + assert.Equal(t, "spog-ws2", cfg.Profile) + assert.Equal(t, "spog-ws2", cfg.Token) +} + +func TestLoaderErrorsOnMultipleMatchesWithSameWorkspaceID(t *testing.T) { + cfg := config.Config{ + Loaders: []config.Loader{ + ResolveProfileFromHost, + }, + ConfigFile: "profile/testdata/databrickscfg", + Host: "https://spog-dup.databricks.com", + WorkspaceID: "333", + } + + err := cfg.EnsureResolved() + require.Error(t, err) + assert.ErrorContains(t, err, "multiple profiles matched: spog-dup1, spog-dup2") +} + +func TestLoaderErrorsOnMultipleMatchesWithoutWorkspaceID(t *testing.T) { + // Without workspace_id, multiple host matches still error as before. + cfg := config.Config{ + Loaders: []config.Loader{ + ResolveProfileFromHost, + }, + ConfigFile: "profile/testdata/databrickscfg", + Host: "https://spog.databricks.com", + } + + err := cfg.EnsureResolved() + require.Error(t, err) + assert.ErrorContains(t, err, "multiple profiles matched: spog-ws1, spog-ws2") +} + +func TestLoaderNoWorkspaceIDMatchFallsThrough(t *testing.T) { + // workspace_id doesn't match any of the host-matching profiles. + // Falls back to the original host ambiguity error. + cfg := config.Config{ + Loaders: []config.Loader{ + ResolveProfileFromHost, + }, + ConfigFile: "profile/testdata/databrickscfg", + Host: "https://spog.databricks.com", + WorkspaceID: "999", + } + + err := cfg.EnsureResolved() + require.Error(t, err) + assert.ErrorContains(t, err, "multiple profiles matched: spog-ws1, spog-ws2") +} diff --git a/libs/databrickscfg/profile/file_test.go b/libs/databrickscfg/profile/file_test.go index e207877cf4..8f6c5ad790 100644 --- a/libs/databrickscfg/profile/file_test.go +++ b/libs/databrickscfg/profile/file_test.go @@ -70,7 +70,7 @@ func TestLoadProfilesMatchWorkspace(t *testing.T) { profiler := FileProfilerImpl{} profiles, err := profiler.LoadProfiles(ctx, MatchWorkspaceProfiles) require.NoError(t, err) - assert.Equal(t, []string{"DEFAULT", "query", "foo1", "foo2"}, profiles.Names()) + assert.Equal(t, []string{"DEFAULT", "query", "foo1", "foo2", "spog-ws1", "spog-ws2", "spog-dup1", "spog-dup2"}, profiles.Names()) } func TestLoadProfilesMatchAccount(t *testing.T) { diff --git a/libs/databrickscfg/profile/testdata/databrickscfg b/libs/databrickscfg/profile/testdata/databrickscfg index ba045c6c28..c88350a752 100644 --- a/libs/databrickscfg/profile/testdata/databrickscfg +++ b/libs/databrickscfg/profile/testdata/databrickscfg @@ -22,3 +22,25 @@ account_id = abc [foo2] host = https://foo token = foo2 + +# SPOG profiles sharing the same host but with different workspace_ids +[spog-ws1] +host = https://spog.databricks.com +workspace_id = 111 +token = spog-ws1 + +[spog-ws2] +host = https://spog.databricks.com +workspace_id = 222 +token = spog-ws2 + +# SPOG profiles with same host and same workspace_id (ambiguous) +[spog-dup1] +host = https://spog-dup.databricks.com +workspace_id = 333 +token = spog-dup1 + +[spog-dup2] +host = https://spog-dup.databricks.com +workspace_id = 333 +token = spog-dup2 From 66289fd23d9ce5f2f4c42b2b91cd10e7bf4e57d7 Mon Sep 17 00:00:00 2001 From: simon Date: Tue, 24 Mar 2026 15:12:15 +0100 Subject: [PATCH 2/2] Reuse shared ExtractHostQueryParams, remove redundant mutator Refactor Workspace.NormalizeHostURL() to use auth.ExtractHostQueryParams() instead of duplicating URL query parameter parsing. Remove the NormalizeHostURL mutator since Client() already calls NormalizeHostURL() before building the SDK config, making the Initialize-phase mutator redundant. Addresses review feedback from andrewnester on #4825. Co-authored-by: Isaac --- bundle/config/mutator/normalize_host_url.go | 30 ------ .../config/mutator/normalize_host_url_test.go | 94 ------------------- bundle/config/workspace.go | 36 +------ bundle/config/workspace_test.go | 24 ++++- bundle/phases/initialize.go | 4 - 5 files changed, 27 insertions(+), 161 deletions(-) delete mode 100644 bundle/config/mutator/normalize_host_url.go delete mode 100644 bundle/config/mutator/normalize_host_url_test.go diff --git a/bundle/config/mutator/normalize_host_url.go b/bundle/config/mutator/normalize_host_url.go deleted file mode 100644 index 2ca64a3849..0000000000 --- a/bundle/config/mutator/normalize_host_url.go +++ /dev/null @@ -1,30 +0,0 @@ -package mutator - -import ( - "context" - - "github.com/databricks/cli/bundle" - "github.com/databricks/cli/libs/diag" -) - -type normalizeHostURL struct{} - -// NormalizeHostURL extracts query parameters from the workspace host URL -// and strips them from the host. This allows users to paste SPOG URLs -// (e.g. https://host.databricks.com/?o=12345) directly into their bundle config. -// -// The primary extraction happens in [config.Workspace.Client] before the SDK -// config is built. This mutator serves as a secondary normalization pass to -// ensure the bundle config is clean for any later code that reads it directly. -func NormalizeHostURL() bundle.Mutator { - return &normalizeHostURL{} -} - -func (m *normalizeHostURL) Name() string { - return "NormalizeHostURL" -} - -func (m *normalizeHostURL) Apply(_ context.Context, b *bundle.Bundle) diag.Diagnostics { - b.Config.Workspace.NormalizeHostURL() - return nil -} diff --git a/bundle/config/mutator/normalize_host_url_test.go b/bundle/config/mutator/normalize_host_url_test.go deleted file mode 100644 index 65788ce576..0000000000 --- a/bundle/config/mutator/normalize_host_url_test.go +++ /dev/null @@ -1,94 +0,0 @@ -package mutator_test - -import ( - "testing" - - "github.com/databricks/cli/bundle" - "github.com/databricks/cli/bundle/config" - "github.com/databricks/cli/bundle/config/mutator" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestNormalizeHostURL(t *testing.T) { - tests := []struct { - name string - host string - workspaceID string - accountID string - wantHost string - wantWorkspaceID string - wantAccountID string - }{ - { - name: "empty host", - host: "", - wantHost: "", - }, - { - name: "host without query params", - host: "https://api.databricks.com", - wantHost: "https://api.databricks.com", - }, - { - name: "host with ?o= extracts workspace_id", - host: "https://dogfood.staging.databricks.com/?o=6051921418418893", - wantHost: "https://dogfood.staging.databricks.com/", - wantWorkspaceID: "6051921418418893", - }, - { - name: "host with ?o= and ?a= extracts both", - host: "https://dogfood.staging.databricks.com/?o=605&a=abc123", - wantHost: "https://dogfood.staging.databricks.com/", - wantWorkspaceID: "605", - wantAccountID: "abc123", - }, - { - name: "host with ?account_id= extracts account_id", - host: "https://accounts.cloud.databricks.com/?account_id=968367da-1234", - wantHost: "https://accounts.cloud.databricks.com/", - wantAccountID: "968367da-1234", - }, - { - name: "host with ?workspace_id= extracts workspace_id", - host: "https://dogfood.staging.databricks.com/?workspace_id=12345", - wantHost: "https://dogfood.staging.databricks.com/", - wantWorkspaceID: "12345", - }, - { - name: "explicit workspace_id takes precedence over ?o=", - host: "https://dogfood.staging.databricks.com/?o=999", - workspaceID: "explicit-id", - wantHost: "https://dogfood.staging.databricks.com/", - wantWorkspaceID: "explicit-id", - }, - { - name: "explicit account_id takes precedence over ?a=", - host: "https://dogfood.staging.databricks.com/?a=from-url", - accountID: "explicit-account", - wantHost: "https://dogfood.staging.databricks.com/", - wantAccountID: "explicit-account", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - b := &bundle.Bundle{ - Config: config.Root{ - Workspace: config.Workspace{ - Host: tt.host, - WorkspaceID: tt.workspaceID, - AccountID: tt.accountID, - }, - }, - } - - diags := bundle.Apply(t.Context(), b, mutator.NormalizeHostURL()) - require.NoError(t, diags.Error()) - - assert.Equal(t, tt.wantHost, b.Config.Workspace.Host) - assert.Equal(t, tt.wantWorkspaceID, b.Config.Workspace.WorkspaceID) - assert.Equal(t, tt.wantAccountID, b.Config.Workspace.AccountID) - }) - } -} diff --git a/bundle/config/workspace.go b/bundle/config/workspace.go index 137c669bf9..c699dc070b 100644 --- a/bundle/config/workspace.go +++ b/bundle/config/workspace.go @@ -1,10 +1,10 @@ package config import ( - "net/url" "os" "path/filepath" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" @@ -146,40 +146,14 @@ func (w *Workspace) Config() *config.Config { // config. Must be called before Config() so the extracted fields are included // in the SDK config used for profile resolution and authentication. func (w *Workspace) NormalizeHostURL() { - if w.Host == "" { - return - } - u, err := url.Parse(w.Host) - if err != nil { - return - } - q := u.Query() - if len(q) == 0 { - return - } - - // Extract workspace_id from ?o= or ?workspace_id= if not already set. + params := auth.ExtractHostQueryParams(w.Host) + w.Host = params.Host if w.WorkspaceID == "" { - if v := q.Get("o"); v != "" { - w.WorkspaceID = v - } else if v := q.Get("workspace_id"); v != "" { - w.WorkspaceID = v - } + w.WorkspaceID = params.WorkspaceID } - - // Extract account_id from ?a= or ?account_id= if not already set. if w.AccountID == "" { - if v := q.Get("a"); v != "" { - w.AccountID = v - } else if v := q.Get("account_id"); v != "" { - w.AccountID = v - } + w.AccountID = params.AccountID } - - // Strip query parameters from the host. - u.RawQuery = "" - u.Fragment = "" - w.Host = u.String() } func (w *Workspace) Client() (*databricks.WorkspaceClient, error) { diff --git a/bundle/config/workspace_test.go b/bundle/config/workspace_test.go index c1bb35641a..4181d17170 100644 --- a/bundle/config/workspace_test.go +++ b/bundle/config/workspace_test.go @@ -79,20 +79,40 @@ func TestWorkspaceNormalizeHostURL(t *testing.T) { Host: "https://spog.databricks.com/?o=12345", } w.NormalizeHostURL() - assert.Equal(t, "https://spog.databricks.com/", w.Host) + assert.Equal(t, "https://spog.databricks.com", w.Host) assert.Equal(t, "12345", w.WorkspaceID) }) + t.Run("extracts both workspace_id and account_id", func(t *testing.T) { + w := Workspace{ + Host: "https://spog.databricks.com/?o=605&a=abc123", + } + w.NormalizeHostURL() + assert.Equal(t, "https://spog.databricks.com", w.Host) + assert.Equal(t, "605", w.WorkspaceID) + assert.Equal(t, "abc123", w.AccountID) + }) + t.Run("explicit workspace_id takes precedence", func(t *testing.T) { w := Workspace{ Host: "https://spog.databricks.com/?o=999", WorkspaceID: "explicit", } w.NormalizeHostURL() - assert.Equal(t, "https://spog.databricks.com/", w.Host) + assert.Equal(t, "https://spog.databricks.com", w.Host) assert.Equal(t, "explicit", w.WorkspaceID) }) + t.Run("explicit account_id takes precedence", func(t *testing.T) { + w := Workspace{ + Host: "https://spog.databricks.com/?a=from-url", + AccountID: "explicit-account", + } + w.NormalizeHostURL() + assert.Equal(t, "https://spog.databricks.com", w.Host) + assert.Equal(t, "explicit-account", w.AccountID) + }) + t.Run("no-op for host without query params", func(t *testing.T) { w := Workspace{ Host: "https://normal.databricks.com", diff --git a/bundle/phases/initialize.go b/bundle/phases/initialize.go index 8e237ce728..64eb18943a 100644 --- a/bundle/phases/initialize.go +++ b/bundle/phases/initialize.go @@ -34,10 +34,6 @@ func Initialize(ctx context.Context, b *bundle.Bundle) { validate.ValidateEngine(), validate.Scripts(), - // Extract query parameters (?o=, ?a=) from workspace host URL and strip them. - // Must run before any mutator that uses the host for authentication or API calls. - mutator.NormalizeHostURL(), - // Updates (dynamic): sync.{paths,include,exclude} (makes them relative to bundle root rather than to definition file) // Rewrites sync paths to be relative to the bundle root instead of the file they were defined in. mutator.RewriteSyncPaths(),