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..c699dc070b 100644 --- a/bundle/config/workspace.go +++ b/bundle/config/workspace.go @@ -4,6 +4,7 @@ import ( "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" @@ -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,28 @@ 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() { + params := auth.ExtractHostQueryParams(w.Host) + w.Host = params.Host + if w.WorkspaceID == "" { + w.WorkspaceID = params.WorkspaceID + } + if w.AccountID == "" { + w.AccountID = params.AccountID + } +} + 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..4181d17170 100644 --- a/bundle/config/workspace_test.go +++ b/bundle/config/workspace_test.go @@ -73,6 +73,87 @@ 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("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, "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", + } + 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/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