[Feature] Add support for Unified Host with experimental flag (#4260)

## Changes
<!-- Brief summary of your changes that is easy to understand -->
- Add support for unified host with experimental flag. 
- Prompt for workspace id as it's not part of the host
- Write the IDs and flag in the created profile. 
- Depends on databricks/databricks-sdk-go#1307 

## Why
<!-- Why are these changes needed? Provide the context that the reviewer
might be missing. -->
This support is required for enabling unified hosts which Databricks
free edition uses.

## Tests
<!-- How have you tested the changes? -->
- Unit tests
- Manual E2E tests -- CUJ for CLI U2M `databricks auth login --host
<spog-host> --experimental-is-unified-host` prompts for account and
workspace id followed by opening web browser. The IDs and flag is stores
in the config. The subsequent cli operations eg: databricks clusters
list --profile "above-profile" works.

Author: tanmay-db

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/out.test.toml

@@ -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/output.txt

@@ -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/script

@@ -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"]

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

@@ -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/config/workspace.go

@@ -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": |-

bundle/internal/schema/annotations.yml

@@ -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()
+}

bundle/schema/jsonschema.json

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