From 1bf8921bc98fe50034a0f5190b92ce5b12fb10e0 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Mon, 12 Jan 2026 14:16:27 +0100 Subject: [PATCH 1/4] [Feature] Add support for Unified Host --- NEXT_CHANGELOG.md | 1 + bundle/config/workspace.go | 8 +++ bundle/internal/schema/annotations.yml | 6 ++ bundle/schema/jsonschema.json | 6 ++ cmd/auth/auth.go | 14 +++++ cmd/auth/login.go | 81 +++++++++++++++++++------- cmd/auth/profiles.go | 9 ++- cmd/auth/testdata/.databrickscfg | 10 ++++ cmd/auth/token.go | 10 ++++ cmd/labs/project/entrypoint.go | 3 +- cmd/labs/project/installer.go | 4 +- cmd/labs/project/login.go | 3 +- libs/auth/arguments.go | 26 ++++++--- libs/auth/error.go | 7 ++- libs/databrickscfg/profile/file.go | 2 + libs/databrickscfg/profile/profile.go | 2 + libs/databrickscfg/profile/profiler.go | 10 +++- 17 files changed, 163 insertions(+), 39 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 86e288dfd9..fad4cd8012 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -7,6 +7,7 @@ To disable this, set the environment variable DATABRICKS_CACHE_ENABLED to false. ### CLI +* Add support for unified host with experimental flag ### Bundles * Enable caching user identity by default ([#4202](https://github.com/databricks/cli/pull/4202)) diff --git a/bundle/config/workspace.go b/bundle/config/workspace.go index 7969177f5f..82f1cc07ca 100644 --- a/bundle/config/workspace.go +++ b/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 { diff --git a/bundle/internal/schema/annotations.yml b/bundle/internal/schema/annotations.yml index 329b1a8925..2f4b9f1b90 100644 --- a/bundle/internal/schema/annotations.yml +++ b/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": |- + PLACEHOLDER "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": |- + PLACEHOLDER github.com/databricks/cli/bundle/config/resources.Alert: "create_time": "description": |- diff --git a/bundle/schema/jsonschema.json b/bundle/schema/jsonschema.json index b746626664..33590e787a 100644 --- a/bundle/schema/jsonschema.json +++ b/bundle/schema/jsonschema.json @@ -2699,6 +2699,9 @@ "description": "The client ID for the workspace", "$ref": "#/$defs/string" }, + "experimental_is_unified_host": { + "$ref": "#/$defs/bool" + }, "file_path": { "description": "The file path to use within the workspace for both deployments and workflow runs", "$ref": "#/$defs/string" @@ -2726,6 +2729,9 @@ "state_path": { "description": "The workspace state path", "$ref": "#/$defs/string" + }, + "workspace_id": { + "$ref": "#/$defs/string" } }, "additionalProperties": false diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index e00a1934a4..b1e4708675 100644 --- a/cmd/auth/auth.go +++ b/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, "Whether the host is a unified host that supports both workspace and account APIs") + cmd.PersistentFlags().StringVar(&authArguments.WorkspaceID, "workspace-id", "", "Databricks Workspace ID (for unified hosts)") cmd.AddCommand(newEnvCommand()) cmd.AddCommand(newLoginCommand(&authArguments)) @@ -55,3 +57,15 @@ func promptForAccountID(ctx context.Context) (string, error) { prompt.AllowEdit = true return prompt.Run() } + +func promptForWorkspaceID(ctx context.Context) (string, error) { + if !cmdio.IsPromptSupported(ctx) { + return "", errors.New("the command is being run in a non-interactive environment, please specify a workspace ID using --workspace-id") + } + + prompt := cmdio.Prompt(ctx) + prompt.Label = "Databricks workspace ID" + prompt.Default = "" + prompt.AllowEdit = true + return prompt.Run() +} diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 733c404e33..0193f269c7 100644 --- a/cmd/auth/login.go +++ b/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,54 @@ 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) + // 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 - prompt for either workspace ID or account ID + // User must specify exactly one of them + hasWorkspaceID := authArguments.WorkspaceID != "" + hasAccountID := authArguments.AccountID != "" + + if hasWorkspaceID && hasAccountID { + return errors.New("unified host cannot have both workspace_id and account_id; please specify only one") + } + + if !hasWorkspaceID && !hasAccountID { + // Neither provided - prompt user to choose + // For now, prompt for workspace ID by default (most common case) + // TODO: Could add a prompt to ask which one the user wants + workspaceId, err := promptForWorkspaceID(ctx) if err != nil { return err } - authArguments.AccountID = accountId + authArguments.WorkspaceID = workspaceId } + case config.WorkspaceHost: + // Workspace host - no additional prompts needed + default: + return fmt.Errorf("unknown host type: %v", cfg.HostType()) } + return nil } diff --git a/cmd/auth/profiles.go b/cmd/auth/profiles.go index 7181f2b1c6..14281e9e56 100644 --- a/cmd/auth/profiles.go +++ b/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 } } diff --git a/cmd/auth/testdata/.databrickscfg b/cmd/auth/testdata/.databrickscfg index 192839b9be..093d588528 100644 --- a/cmd/auth/testdata/.databrickscfg +++ b/cmd/auth/testdata/.databrickscfg @@ -15,3 +15,13 @@ 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 +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 diff --git a/cmd/auth/token.go b/cmd/auth/token.go index cf3873cb89..4987915e21 100644 --- a/cmd/auth/token.go +++ b/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 diff --git a/cmd/labs/project/entrypoint.go b/cmd/labs/project/entrypoint.go index b6150bd6fd..a0a26065f9 100644 --- a/cmd/labs/project/entrypoint.go +++ b/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 diff --git a/cmd/labs/project/installer.go b/cmd/labs/project/installer.go index 5025ae28dd..ebfc433b0f 100644 --- a/cmd/labs/project/installer.go +++ b/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} diff --git a/cmd/labs/project/login.go b/cmd/labs/project/login.go index b8245e1ac9..efa7a8ee2e 100644 --- a/cmd/labs/project/login.go +++ b/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) diff --git a/libs/auth/arguments.go b/libs/auth/arguments.go index d0242992a0..b123287cd7 100644 --- a/libs/auth/arguments.go +++ b/libs/auth/arguments.go @@ -1,6 +1,8 @@ package auth import ( + "fmt" + "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/credentials/u2m" ) @@ -8,20 +10,30 @@ import ( // AuthArguments is a struct that contains the common arguments passed to // `databricks auth` commands. type AuthArguments struct { - Host string - AccountID string + Host string + AccountID string + WorkspaceID string + IsUnifiedHost bool } // ToOAuthArgument converts the AuthArguments to an OAuthArgument from the Go SDK. func (a AuthArguments) ToOAuthArgument() (u2m.OAuthArgument, error) { cfg := &config.Config{ - Host: a.Host, - AccountID: a.AccountID, + Host: a.Host, + AccountID: a.AccountID, + WorkspaceId: a.WorkspaceID, + Experimental_IsUnifiedHost: a.IsUnifiedHost, } host := cfg.CanonicalHostName() - //nolint:staticcheck // SA1019: IsAccountClient is deprecated but is still used here to avoid breaking changes - if cfg.IsAccountClient() { + + switch cfg.HostType() { + case config.AccountHost: return u2m.NewBasicAccountOAuthArgument(host, cfg.AccountID) + case config.WorkspaceHost: + return u2m.NewBasicWorkspaceOAuthArgument(host) + case config.UnifiedHost: + return u2m.NewBasicUnifiedOAuthArgument(host, cfg.AccountID) + default: + return nil, fmt.Errorf("unknown host type: %v", cfg.HostType()) } - return u2m.NewBasicWorkspaceOAuthArgument(host) } diff --git a/libs/auth/error.go b/libs/auth/error.go index 5311bea965..ca08a2d41e 100644 --- a/libs/auth/error.go +++ b/libs/auth/error.go @@ -13,7 +13,10 @@ import ( func RewriteAuthError(ctx context.Context, host, accountId, profile string, err error) (bool, error) { target := &u2m.InvalidRefreshTokenError{} if errors.As(err, &target) { - oauthArgument, err := AuthArguments{host, accountId}.ToOAuthArgument() + oauthArgument, err := AuthArguments{ + Host: host, + AccountID: accountId, + }.ToOAuthArgument() if err != nil { return false, err } @@ -35,6 +38,8 @@ func BuildLoginCommand(ctx context.Context, profile string, arg u2m.OAuthArgumen cmd = append(cmd, "--profile", profile) } else { switch arg := arg.(type) { + case u2m.UnifiedOAuthArgument: + cmd = append(cmd, "--host", arg.GetHost(), "--account-id", arg.GetAccountId(), "--experimental-is-unified-host") case u2m.AccountOAuthArgument: cmd = append(cmd, "--host", arg.GetAccountHost(), "--account-id", arg.GetAccountId()) case u2m.WorkspaceOAuthArgument: diff --git a/libs/databrickscfg/profile/file.go b/libs/databrickscfg/profile/file.go index e9a5aa3a2a..a9d45d0825 100644 --- a/libs/databrickscfg/profile/file.go +++ b/libs/databrickscfg/profile/file.go @@ -82,6 +82,8 @@ func (f FileProfilerImpl) LoadProfiles(ctx context.Context, fn ProfileMatchFunct Name: v.Name(), Host: host, AccountID: all["account_id"], + WorkspaceID: all["workspace_id"], + IsUnifiedHost: all["experimental_is_unified_host"] == "true", ClusterID: all["cluster_id"], ServerlessComputeID: all["serverless_compute_id"], } diff --git a/libs/databrickscfg/profile/profile.go b/libs/databrickscfg/profile/profile.go index d2c3a88d5a..0358b8f7ec 100644 --- a/libs/databrickscfg/profile/profile.go +++ b/libs/databrickscfg/profile/profile.go @@ -13,6 +13,8 @@ type Profile struct { Name string Host string AccountID string + WorkspaceID string + IsUnifiedHost bool ClusterID string ServerlessComputeID string } diff --git a/libs/databrickscfg/profile/profiler.go b/libs/databrickscfg/profile/profiler.go index c0a5492561..5d1ea0e72f 100644 --- a/libs/databrickscfg/profile/profiler.go +++ b/libs/databrickscfg/profile/profiler.go @@ -7,11 +7,17 @@ import ( type ProfileMatchFunction func(Profile) bool func MatchWorkspaceProfiles(p Profile) bool { - return p.AccountID == "" + // Match workspace profiles: regular workspace profiles (no account ID) + // or unified hosts with workspace ID + return (p.AccountID == "" && !p.IsUnifiedHost) || + (p.IsUnifiedHost && p.WorkspaceID != "") } func MatchAccountProfiles(p Profile) bool { - return p.Host != "" && p.AccountID != "" + // Match account profiles: regular account profiles (with account ID) + // or unified hosts with account ID but no workspace ID + return (p.Host != "" && p.AccountID != "" && !p.IsUnifiedHost) || + (p.IsUnifiedHost && p.AccountID != "" && p.WorkspaceID == "") } func MatchAllProfiles(p Profile) bool { From 88eb8541a8ab3327f6bae6d5658ea66a25a435de Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Wed, 14 Jan 2026 14:35:55 +0100 Subject: [PATCH 2/4] - --- cmd/auth/login.go | 41 +++++++++++++++++++++++-------------- libs/auth/arguments.go | 2 ++ libs/auth/arguments_test.go | 30 ++++++++++++++++++++++++--- 3 files changed, 55 insertions(+), 18 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 0193f269c7..3435a472c0 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -294,24 +294,35 @@ func setHostAndAccountId(ctx context.Context, existingProfile *profile.Profile, } } case config.UnifiedHost: - // Unified host - prompt for either workspace ID or account ID - // User must specify exactly one of them - hasWorkspaceID := authArguments.WorkspaceID != "" - hasAccountID := authArguments.AccountID != "" - - if hasWorkspaceID && hasAccountID { - return errors.New("unified host cannot have both workspace_id and account_id; please specify only one") + // 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 + } } - if !hasWorkspaceID && !hasAccountID { - // Neither provided - prompt user to choose - // For now, prompt for workspace ID by default (most common case) - // TODO: Could add a prompt to ask which one the user wants - workspaceId, err := promptForWorkspaceID(ctx) - if err != nil { - return err + // 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.WorkspaceID = workspaceId } case config.WorkspaceHost: // Workspace host - no additional prompts needed diff --git a/libs/auth/arguments.go b/libs/auth/arguments.go index b123287cd7..c1112cf265 100644 --- a/libs/auth/arguments.go +++ b/libs/auth/arguments.go @@ -32,6 +32,8 @@ func (a AuthArguments) ToOAuthArgument() (u2m.OAuthArgument, error) { case config.WorkspaceHost: return u2m.NewBasicWorkspaceOAuthArgument(host) case config.UnifiedHost: + // For unified hosts, always use the unified OAuth argument with account ID. + // The workspace ID is stored in the config for API routing, not OAuth. return u2m.NewBasicUnifiedOAuthArgument(host, cfg.AccountID) default: return nil, fmt.Errorf("unknown host type: %v", cfg.HostType()) diff --git a/libs/auth/arguments_test.go b/libs/auth/arguments_test.go index d75827a771..a8a51381b1 100644 --- a/libs/auth/arguments_test.go +++ b/libs/auth/arguments_test.go @@ -58,6 +58,25 @@ func TestToOAuthArgument(t *testing.T) { }, wantHost: "https://my-workspace.cloud.databricks.com", }, + { + name: "unified host with account ID only", + args: AuthArguments{ + Host: "https://unified.cloud.databricks.com", + AccountID: "123456789", + IsUnifiedHost: true, + }, + wantHost: "https://unified.cloud.databricks.com", + }, + { + name: "unified host with both account ID and workspace ID", + args: AuthArguments{ + Host: "https://unified.cloud.databricks.com", + AccountID: "968367da-7edd-44f7-9dea-3e0b20b0ec97", + WorkspaceID: "470576644108500", + IsUnifiedHost: true, + }, + wantHost: "https://unified.cloud.databricks.com", + }, } for _, tt := range tests { @@ -70,13 +89,18 @@ func TestToOAuthArgument(t *testing.T) { assert.NoError(t, err) // Check if we got the right type of argument and verify the hostname - if tt.args.AccountID != "" { + if tt.args.IsUnifiedHost { + // Unified hosts always return UnifiedOAuthArgument which implements AccountOAuthArgument + arg, ok := got.(u2m.AccountOAuthArgument) + assert.True(t, ok, "expected AccountOAuthArgument (UnifiedOAuthArgument) for unified host") + assert.Equal(t, tt.wantHost, arg.GetAccountHost()) + } else if tt.args.AccountID != "" { arg, ok := got.(u2m.AccountOAuthArgument) - assert.True(t, ok, "expected AccountOAuthArgument for account ID") + assert.True(t, ok, "expected AccountOAuthArgument for account host") assert.Equal(t, tt.wantHost, arg.GetAccountHost()) } else { arg, ok := got.(u2m.WorkspaceOAuthArgument) - assert.True(t, ok, "expected WorkspaceOAuthArgument for workspace") + assert.True(t, ok, "expected WorkspaceOAuthArgument for workspace host") assert.Equal(t, tt.wantHost, arg.GetWorkspaceHost()) } }) From 6927440aeaf67cda5d011c069058b4dc848d0523 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Wed, 14 Jan 2026 15:13:52 +0100 Subject: [PATCH 3/4] - --- libs/auth/arguments_test.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/libs/auth/arguments_test.go b/libs/auth/arguments_test.go index a8a51381b1..2c24b6c9d9 100644 --- a/libs/auth/arguments_test.go +++ b/libs/auth/arguments_test.go @@ -90,10 +90,10 @@ func TestToOAuthArgument(t *testing.T) { // Check if we got the right type of argument and verify the hostname if tt.args.IsUnifiedHost { - // Unified hosts always return UnifiedOAuthArgument which implements AccountOAuthArgument - arg, ok := got.(u2m.AccountOAuthArgument) - assert.True(t, ok, "expected AccountOAuthArgument (UnifiedOAuthArgument) for unified host") - assert.Equal(t, tt.wantHost, arg.GetAccountHost()) + // Unified hosts return UnifiedOAuthArgument (distinct from Account/Workspace) + arg, ok := got.(u2m.UnifiedOAuthArgument) + assert.True(t, ok, "expected UnifiedOAuthArgument for unified host") + assert.Equal(t, tt.wantHost, arg.GetHost()) } else if tt.args.AccountID != "" { arg, ok := got.(u2m.AccountOAuthArgument) assert.True(t, ok, "expected AccountOAuthArgument for account host") From e9130cbb595cb4772163d7cc9401cb8b942ec56a Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Wed, 14 Jan 2026 15:22:12 +0100 Subject: [PATCH 4/4] - --- bundle/internal/schema/annotations.yml | 4 ++-- bundle/schema/jsonschema.json | 2 ++ cmd/auth/auth.go | 4 ++-- cmd/auth/testdata/.databrickscfg | 1 + 4 files changed, 7 insertions(+), 4 deletions(-) diff --git a/bundle/internal/schema/annotations.yml b/bundle/internal/schema/annotations.yml index cf9ba96d5e..91250f9306 100644 --- a/bundle/internal/schema/annotations.yml +++ b/bundle/internal/schema/annotations.yml @@ -426,7 +426,7 @@ github.com/databricks/cli/bundle/config.Workspace: The client ID for the workspace "experimental_is_unified_host": "description": |- - PLACEHOLDER + 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 @@ -450,7 +450,7 @@ github.com/databricks/cli/bundle/config.Workspace: The workspace state path "workspace_id": "description": |- - PLACEHOLDER + The Databricks workspace ID github.com/databricks/cli/bundle/config/resources.Alert: "create_time": "description": |- diff --git a/bundle/schema/jsonschema.json b/bundle/schema/jsonschema.json index 9bda7681f0..b267ca790c 100644 --- a/bundle/schema/jsonschema.json +++ b/bundle/schema/jsonschema.json @@ -2747,6 +2747,7 @@ "$ref": "#/$defs/string" }, "experimental_is_unified_host": { + "description": "Flag to indicate if the host is a unified host", "$ref": "#/$defs/bool" }, "file_path": { @@ -2778,6 +2779,7 @@ "$ref": "#/$defs/string" }, "workspace_id": { + "description": "The Databricks workspace ID", "$ref": "#/$defs/string" } }, diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index b1e4708675..387175e58d 100644 --- a/cmd/auth/auth.go +++ b/cmd/auth/auth.go @@ -25,8 +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, "Whether the host is a unified host that supports both workspace and account APIs") - cmd.PersistentFlags().StringVar(&authArguments.WorkspaceID, "workspace-id", "", "Databricks Workspace ID (for unified hosts)") + 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)) diff --git a/cmd/auth/testdata/.databrickscfg b/cmd/auth/testdata/.databrickscfg index 093d588528..ca1a063076 100644 --- a/cmd/auth/testdata/.databrickscfg +++ b/cmd/auth/testdata/.databrickscfg @@ -18,6 +18,7 @@ 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