[Feature] Add support for Unified Host with experimental flag

NEXT_CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Notable Changes
 
 ### CLI
+* Add support for unified host with experimental flag ([#4260](https://github.com/databricks/cli/pull/4260))
 
 ### Bundles
 * engine/direct: Support bind & unbind. ([#4279](https://github.com/databricks/cli/pull/4279))

bundle/config/workspace.go

@@ -41,6 +41,10 @@ type Workspace struct {
 	AzureEnvironment string `json:"azure_environment,omitempty"`
 	AzureLoginAppID  string `json:"azure_login_app_id,omitempty"`
 
+	// Unified host specific attributes.
+	ExperimentalIsUnifiedHost bool   `json:"experimental_is_unified_host,omitempty"`
+	WorkspaceID               string `json:"workspace_id,omitempty"`
+
 	// CurrentUser holds the current user.
 	// This is set after configuration initialization.
 	CurrentUser *User `json:"current_user,omitempty" bundle:"readonly"`
@@ -117,6 +121,10 @@ func (w *Workspace) Config() *config.Config {
 		AzureTenantID:    w.AzureTenantID,
 		AzureEnvironment: w.AzureEnvironment,
 		AzureLoginAppID:  w.AzureLoginAppID,
+
+		// Unified host
+		Experimental_IsUnifiedHost: w.ExperimentalIsUnifiedHost,
+		WorkspaceId:                w.WorkspaceID,
 	}
 
 	for k := range config.ConfigAttributes {

bundle/internal/schema/annotations.yml

@@ -424,6 +424,9 @@ github.com/databricks/cli/bundle/config.Workspace:
   "client_id":
     "description": |-
       The client ID for the workspace
+  "experimental_is_unified_host":
+    "description": |-
+      Flag to indicate if the host is a unified host
   "file_path":
     "description": |-
       The file path to use within the workspace for both deployments and workflow runs
@@ -445,6 +448,9 @@ github.com/databricks/cli/bundle/config.Workspace:
   "state_path":
     "description": |-
       The workspace state path
+  "workspace_id":
+    "description": |-
+      The Databricks workspace ID
 github.com/databricks/cli/bundle/config/resources.Alert:
   "create_time":
     "description": |-

cmd/auth/auth.go

@@ -25,6 +25,8 @@ GCP: https://docs.gcp.databricks.com/dev-tools/auth/index.html`,
 	var authArguments auth.AuthArguments
 	cmd.PersistentFlags().StringVar(&authArguments.Host, "host", "", "Databricks Host")
 	cmd.PersistentFlags().StringVar(&authArguments.AccountID, "account-id", "", "Databricks Account ID")
+	cmd.PersistentFlags().BoolVar(&authArguments.IsUnifiedHost, "experimental-is-unified-host", false, "Flag to indicate if the host is a unified host")
+	cmd.PersistentFlags().StringVar(&authArguments.WorkspaceID, "workspace-id", "", "Databricks Workspace ID")
 
 	cmd.AddCommand(newEnvCommand())
 	cmd.AddCommand(newLoginCommand(&authArguments))
@@ -55,3 +57,16 @@ func promptForAccountID(ctx context.Context) (string, error) {
 	prompt.AllowEdit = true
 	return prompt.Run()
 }
+
+func promptForWorkspaceID(ctx context.Context) (string, error) {
+	if !cmdio.IsPromptSupported(ctx) {
+		// Workspace ID is optional for unified hosts, so return empty string in non-interactive mode
+		return "", nil
+	}
+
+	prompt := cmdio.Prompt(ctx)
+	prompt.Label = "Databricks workspace ID (optional - provide only if using this profile for workspace operations, leave empty for account operations)"
+	prompt.Default = ""
+	prompt.AllowEdit = true
+	return prompt.Run()
+}

cmd/auth/login.go

@@ -133,6 +133,15 @@ depends on the existing profiles you have set in your configuration file
 		if err != nil {
 			return err
 		}
+
+		// Load unified host flags from the profile if not explicitly set via CLI flag
+		if !cmd.Flag("experimental-is-unified-host").Changed && existingProfile != nil {
+			authArguments.IsUnifiedHost = existingProfile.IsUnifiedHost
+		}
+		if !cmd.Flag("workspace-id").Changed && existingProfile != nil {
+			authArguments.WorkspaceID = existingProfile.WorkspaceID
+		}
+
 		err = setHostAndAccountId(ctx, existingProfile, authArguments, args)
 		if err != nil {
 			return err
@@ -202,13 +211,15 @@ depends on the existing profiles you have set in your configuration file
 
 		if profileName != "" {
 			err = databrickscfg.SaveToProfile(ctx, &config.Config{
-				Profile:             profileName,
-				Host:                cfg.Host,
-				AuthType:            cfg.AuthType,
-				AccountID:           cfg.AccountID,
-				ClusterID:           cfg.ClusterID,
-				ConfigFile:          cfg.ConfigFile,
-				ServerlessComputeID: cfg.ServerlessComputeID,
+				Profile:                    profileName,
+				Host:                       cfg.Host,
+				AuthType:                   cfg.AuthType,
+				AccountID:                  cfg.AccountID,
+				WorkspaceId:                authArguments.WorkspaceID,
+				Experimental_IsUnifiedHost: authArguments.IsUnifiedHost,
+				ClusterID:                  cfg.ClusterID,
+				ConfigFile:                 cfg.ConfigFile,
+				ServerlessComputeID:        cfg.ServerlessComputeID,
 			})
 			if err != nil {
 				return err
@@ -260,24 +271,65 @@ func setHostAndAccountId(ctx context.Context, existingProfile *profile.Profile,
 		}
 	}
 
-	// If the account-id was not provided as a cmd line flag, try to read it from
-	// the specified profile.
-	//nolint:staticcheck // SA1019: IsAccountClient is deprecated but is still used here to avoid breaking changes
-	isAccountClient := (&config.Config{Host: authArguments.Host}).IsAccountClient()
-	accountID := authArguments.AccountID
-	if isAccountClient && accountID == "" {
-		if existingProfile != nil && existingProfile.AccountID != "" {
-			authArguments.AccountID = existingProfile.AccountID
-		} else {
-			// Prompt user for the account-id if it we could not get it from a
-			// profile.
-			accountId, err := promptForAccountID(ctx)
-			if err != nil {
-				return err
+	// Determine the host type and handle account ID / workspace ID accordingly
+	cfg := &config.Config{
+		Host:                       authArguments.Host,
+		AccountID:                  authArguments.AccountID,
+		WorkspaceId:                authArguments.WorkspaceID,
+		Experimental_IsUnifiedHost: authArguments.IsUnifiedHost,
+	}
+
+	switch cfg.HostType() {
+	case config.AccountHost:
+		// Account host - prompt for account ID if not provided
+		if authArguments.AccountID == "" {
+			if existingProfile != nil && existingProfile.AccountID != "" {
+				authArguments.AccountID = existingProfile.AccountID
+			} else {
+				accountId, err := promptForAccountID(ctx)
+				if err != nil {
+					return err
+				}
+				authArguments.AccountID = accountId
+			}
+		}
+	case config.UnifiedHost:
+		// Unified host requires an account ID for OAuth URL construction
+		if authArguments.AccountID == "" {
+			if existingProfile != nil && existingProfile.AccountID != "" {
+				authArguments.AccountID = existingProfile.AccountID
+			} else {
+				accountId, err := promptForAccountID(ctx)
+				if err != nil {
+					return err
+				}
+				authArguments.AccountID = accountId
+			}
+		}
+
+		// Workspace ID is optional and determines API access level:
+		// - With workspace ID: workspace-level APIs
+		// - Without workspace ID: account-level APIs
+		// If neither is provided via flags, prompt for workspace ID (most common case)
+		hasWorkspaceID := authArguments.WorkspaceID != ""
+		if !hasWorkspaceID {
+			if existingProfile != nil && existingProfile.WorkspaceID != "" {
+				authArguments.WorkspaceID = existingProfile.WorkspaceID
+			} else {
+				// Prompt for workspace ID for workspace-level access
+				workspaceId, err := promptForWorkspaceID(ctx)
+				if err != nil {
+					return err
+				}
+				authArguments.WorkspaceID = workspaceId
 			}
-			authArguments.AccountID = accountId
 		}
+	case config.WorkspaceHost:
+		// Workspace host - no additional prompts needed
+	default:
+		return fmt.Errorf("unknown host type: %v", cfg.HostType())
 	}
+
 	return nil
 }
 

cmd/auth/login_test.go

@@ -101,6 +101,74 @@ func TestSetAccountId(t *testing.T) {
 	assert.EqualError(t, err, "the command is being run in a non-interactive environment, please specify an account ID using --account-id")
 }
 
+func TestSetWorkspaceIdForUnifiedHost(t *testing.T) {
+	var authArguments auth.AuthArguments
+	t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg")
+	ctx, _ := cmdio.SetupTest(context.Background(), cmdio.TestOptions{})
+
+	unifiedWorkspaceProfile := loadTestProfile(t, ctx, "unified-workspace")
+	unifiedAccountProfile := loadTestProfile(t, ctx, "unified-account")
+
+	// Test setting workspace-id from flag for unified host
+	authArguments = auth.AuthArguments{
+		Host:          "https://unified.databricks.com",
+		AccountID:     "test-unified-account",
+		WorkspaceID:   "val from --workspace-id",
+		IsUnifiedHost: true,
+	}
+	err := setHostAndAccountId(ctx, unifiedWorkspaceProfile, &authArguments, []string{})
+	assert.NoError(t, err)
+	assert.Equal(t, "https://unified.databricks.com", authArguments.Host)
+	assert.Equal(t, "test-unified-account", authArguments.AccountID)
+	assert.Equal(t, "val from --workspace-id", authArguments.WorkspaceID)
+
+	// Test setting workspace_id from profile for unified host
+	authArguments = auth.AuthArguments{
+		Host:          "https://unified.databricks.com",
+		AccountID:     "test-unified-account",
+		IsUnifiedHost: true,
+	}
+	err = setHostAndAccountId(ctx, unifiedWorkspaceProfile, &authArguments, []string{})
+	assert.NoError(t, err)
+	assert.Equal(t, "https://unified.databricks.com", authArguments.Host)
+	assert.Equal(t, "test-unified-account", authArguments.AccountID)
+	assert.Equal(t, "123456789", authArguments.WorkspaceID)
+
+	// Test workspace_id is optional - should default to empty in non-interactive mode
+	authArguments = auth.AuthArguments{
+		Host:          "https://unified.databricks.com",
+		AccountID:     "test-unified-account",
+		IsUnifiedHost: true,
+	}
+	err = setHostAndAccountId(ctx, unifiedAccountProfile, &authArguments, []string{})
+	assert.NoError(t, err)
+	assert.Equal(t, "https://unified.databricks.com", authArguments.Host)
+	assert.Equal(t, "test-unified-account", authArguments.AccountID)
+	assert.Equal(t, "", authArguments.WorkspaceID) // Empty is valid for account-level access
+
+	// Test workspace_id is optional - should default to empty when no profile exists
+	authArguments = auth.AuthArguments{
+		Host:          "https://unified.databricks.com",
+		AccountID:     "test-unified-account",
+		IsUnifiedHost: true,
+	}
+	err = setHostAndAccountId(ctx, nil, &authArguments, []string{})
+	assert.NoError(t, err)
+	assert.Equal(t, "https://unified.databricks.com", authArguments.Host)
+	assert.Equal(t, "test-unified-account", authArguments.AccountID)
+	assert.Equal(t, "", authArguments.WorkspaceID) // Empty is valid for account-level access
+}
+
+func TestPromptForWorkspaceIDInNonInteractiveMode(t *testing.T) {
+	// Setup non-interactive context
+	ctx, _ := cmdio.SetupTest(context.Background(), cmdio.TestOptions{})
+
+	// Test that promptForWorkspaceID returns empty string (no error) in non-interactive mode
+	workspaceID, err := promptForWorkspaceID(ctx)
+	assert.NoError(t, err)
+	assert.Equal(t, "", workspaceID)
+}
+
 func TestLoadProfileByNameAndClusterID(t *testing.T) {
 	testCases := []struct {
 		name              string

cmd/auth/profiles.go

@@ -51,8 +51,8 @@ func (c *profileMetadata) Load(ctx context.Context, configFilePath string, skipV
 		return
 	}
 
-	//nolint:staticcheck // SA1019: IsAccountClient is deprecated but is still used here to avoid breaking changes
-	if cfg.IsAccountClient() {
+	switch cfg.ConfigType() {
+	case config.AccountConfig:
 		a, err := databricks.NewAccountClient((*databricks.Config)(cfg))
 		if err != nil {
 			return
@@ -64,7 +64,7 @@ func (c *profileMetadata) Load(ctx context.Context, configFilePath string, skipV
 			return
 		}
 		c.Valid = true
-	} else {
+	case config.WorkspaceConfig:
 		w, err := databricks.NewWorkspaceClient((*databricks.Config)(cfg))
 		if err != nil {
 			return
@@ -76,6 +76,9 @@ func (c *profileMetadata) Load(ctx context.Context, configFilePath string, skipV
 			return
 		}
 		c.Valid = true
+	case config.InvalidConfig:
+		// Invalid configuration, skip validation
+		return
 	}
 }
 

cmd/auth/testdata/.databrickscfg

@@ -15,3 +15,14 @@ cluster_id = cluster-from-config
 [invalid-profile]
 # This profile is missing the required 'host' field
 cluster_id = some-cluster-id
+
+[unified-workspace]
+host = https://unified.databricks.com
+account_id = test-unified-account
+workspace_id = 123456789
+experimental_is_unified_host = true
+
+[unified-account]
+host = https://unified.databricks.com
+account_id = test-unified-account
+experimental_is_unified_host = true

cmd/auth/token.go

@@ -98,6 +98,16 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
 		return nil, err
 	}
 
+	// Load unified host flags from the profile if available
+	if existingProfile != nil {
+		if !args.authArguments.IsUnifiedHost && existingProfile.IsUnifiedHost {
+			args.authArguments.IsUnifiedHost = existingProfile.IsUnifiedHost
+		}
+		if args.authArguments.WorkspaceID == "" && existingProfile.WorkspaceID != "" {
+			args.authArguments.WorkspaceID = existingProfile.WorkspaceID
+		}
+	}
+
 	err = setHostAndAccountId(ctx, existingProfile, args.authArguments, args.args)
 	if err != nil {
 		return nil, err

cmd/labs/project/entrypoint.go

@@ -248,8 +248,7 @@ func (e *Entrypoint) validLogin(cmd *cobra.Command) (*config.Config, error) {
 	// an account profile during installation (anymore) and just prompt for it, when context
 	// does require it. This also means that we always prompt for account-level commands, unless
 	// users specify a `--profile` flag.
-	//nolint:staticcheck // SA1019: IsAccountClient is deprecated but is still used here to avoid breaking changes
-	isACC := cfg.IsAccountClient()
+	isACC := cfg.ConfigType() == config.AccountConfig
 	if e.IsAccountLevel && cfg.Profile == "" {
 		if !cmdio.IsPromptSupported(ctx) {
 			return nil, config.ErrCannotConfigureDefault

cmd/labs/project/installer.go

@@ -16,6 +16,7 @@ import (
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/cli/libs/process"
 	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 	"github.com/fatih/color"
@@ -177,8 +178,7 @@ func (i *installer) login(ctx context.Context) (*databricks.WorkspaceClient, err
 	} else if err != nil {
 		return nil, fmt.Errorf("valid: %w", err)
 	}
-	//nolint:staticcheck // SA1019: IsAccountClient is deprecated but is still used here to avoid breaking changes
-	if !i.HasAccountLevelCommands() && cfg.IsAccountClient() {
+	if !i.HasAccountLevelCommands() && cfg.ConfigType() == config.AccountConfig {
 		return nil, errors.New("got account-level client, but no account-level commands")
 	}
 	lc := &loginConfig{Entrypoint: i.Installer.Entrypoint}

cmd/labs/project/login.go

@@ -23,8 +23,7 @@ type loginConfig struct {
 }
 
 func (lc *loginConfig) askWorkspace(ctx context.Context, cfg *config.Config) (*databricks.WorkspaceClient, error) {
-	//nolint:staticcheck // SA1019: IsAccountClient is deprecated but is still used here to avoid breaking changes
-	if cfg.IsAccountClient() {
+	if cfg.ConfigType() == config.AccountConfig {
 		return nil, nil
 	}
 	err := lc.askWorkspaceProfile(ctx, cfg)