diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index aef3b4b5bd..ef2e9795be 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,6 +8,7 @@ To disable this, set the environment variable DATABRICKS_CACHE_ENABLED to false. ### CLI * Add commands to pipelines command group ([#4275](https://github.com/databricks/cli/pull/4275)) +* Add support for unified host with experimental flag ([#4260](https://github.com/databricks/cli/pull/4260)) ### Bundles * Add support for configuring app.yaml options for apps via bundle config ([#4271](https://github.com/databricks/cli/pull/4271)) 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 c324d24157..91250f9306 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": |- + 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": |- diff --git a/bundle/schema/jsonschema.json b/bundle/schema/jsonschema.json index 5150ab9c9f..b267ca790c 100644 --- a/bundle/schema/jsonschema.json +++ b/bundle/schema/jsonschema.json @@ -2746,6 +2746,10 @@ "description": "The client ID for the workspace", "$ref": "#/$defs/string" }, + "experimental_is_unified_host": { + "description": "Flag to indicate if the host is a 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" @@ -2773,6 +2777,10 @@ "state_path": { "description": "The workspace state path", "$ref": "#/$defs/string" + }, + "workspace_id": { + "description": "The Databricks workspace ID", + "$ref": "#/$defs/string" } }, "additionalProperties": false diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index e00a1934a4..387175e58d 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, "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,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..3435a472c0 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,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 } 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..ca1a063076 100644 --- a/cmd/auth/testdata/.databrickscfg +++ b/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 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..c1112cf265 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,32 @@ 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: + // 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()) } - return u2m.NewBasicWorkspaceOAuthArgument(host) } diff --git a/libs/auth/arguments_test.go b/libs/auth/arguments_test.go index d75827a771..2c24b6c9d9 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 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 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()) } }) 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 {