[Feature] Add support for Unified Host with experimental flag

acceptance/auth/credentials/unified-host/out.requests.txt

@@ -0,0 +1,15 @@
+{
+  "headers": {
+    "Authorization": [
+      "Bearer dapi-unified-token"
+    ],
+    "User-Agent": [
+      "cli/[DEV_VERSION] databricks-sdk-go/[SDK_VERSION] go/[GO_VERSION] os/[OS] cmd/current-user_me cmd-exec-id/[UUID] auth/pat"
+    ],
+    "X-Databricks-Org-Id": [
+      "[NUMID]"
+    ]
+  },
+  "method": "GET",
+  "path": "/api/2.0/preview/scim/v2/Me"
+}

acceptance/auth/credentials/unified-host/output.txt

@@ -0,0 +1,11 @@
+
+=== With workspace_id
+{
+  "id":"[USERID]",
+  "userName":"[USERNAME]"
+}
+
+=== Without workspace_id (should error)
+Error: WorkspaceId must be set when using WorkspaceClient with unified host
+
+Exit code: 1

acceptance/auth/credentials/unified-host/script

@@ -0,0 +1,12 @@
+# Test unified host authentication with PAT token
+export DATABRICKS_TOKEN=dapi-unified-token
+export DATABRICKS_ACCOUNT_ID=test-account-123
+export DATABRICKS_WORKSPACE_ID=1234567890
+export DATABRICKS_EXPERIMENTAL_IS_UNIFIED_HOST=true
+
+title "With workspace_id\n"
+$CLI current-user me
+
+title "Without workspace_id (should error)\n"
+unset DATABRICKS_WORKSPACE_ID
+errcode $CLI current-user me

acceptance/auth/credentials/unified-host/test.toml

@@ -0,0 +1,3 @@
+# Test unified host authentication with PAT tokens
+# Include X-Databricks-Org-Id header to verify workspace_id is sent
+IncludeRequestHeaders = ["Authorization", "User-Agent", "X-Databricks-Org-Id"]

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": |-
+      Experimental feature 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
@@ -155,9 +164,11 @@ depends on the existing profiles you have set in your configuration file
 		// We need the config without the profile before it's used to initialise new workspace client below.
 		// Otherwise it will complain about non existing profile because it was not yet saved.
 		cfg := config.Config{
-			Host:      authArguments.Host,
-			AccountID: authArguments.AccountID,
-			AuthType:  "databricks-cli",
+			Host:                       authArguments.Host,
+			AccountID:                  authArguments.AccountID,
+			WorkspaceId:                authArguments.WorkspaceId,
+			Experimental_IsUnifiedHost: authArguments.IsUnifiedHost,
+			AuthType:                   "databricks-cli",
 		}
 		databricksCfgFile := os.Getenv("DATABRICKS_CONFIG_FILE")
 		if databricksCfgFile != "" {
@@ -202,13 +213,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 +273,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