Add SPOG bundle support: account_id, host URL parsing, workspace_id disambiguation

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

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

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

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

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 {

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

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

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