diff --git a/acceptance/auth/credentials/unified-host/out.requests.txt b/acceptance/auth/credentials/unified-host/out.requests.txt new file mode 100644 index 0000000000..6e204f5dba --- /dev/null +++ b/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" +} diff --git a/acceptance/auth/credentials/unified-host/out.test.toml b/acceptance/auth/credentials/unified-host/out.test.toml new file mode 100644 index 0000000000..d560f1de04 --- /dev/null +++ b/acceptance/auth/credentials/unified-host/out.test.toml @@ -0,0 +1,5 @@ +Local = true +Cloud = false + +[EnvMatrix] + DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/auth/credentials/unified-host/output.txt b/acceptance/auth/credentials/unified-host/output.txt new file mode 100644 index 0000000000..89aa4c891c --- /dev/null +++ b/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 diff --git a/acceptance/auth/credentials/unified-host/script b/acceptance/auth/credentials/unified-host/script new file mode 100644 index 0000000000..f785987219 --- /dev/null +++ b/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 diff --git a/acceptance/auth/credentials/unified-host/test.toml b/acceptance/auth/credentials/unified-host/test.toml new file mode 100644 index 0000000000..fd0cd96421 --- /dev/null +++ b/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"] diff --git a/bundle/config/workspace.go b/bundle/config/workspace.go index 7969177f5f..dcdebf9781 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 4a9c477f8a..39abfa16a7 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": |- + 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": |- diff --git a/bundle/schema/jsonschema.json b/bundle/schema/jsonschema.json index c564dce172..b027fbe5e1 100644 --- a/bundle/schema/jsonschema.json +++ b/bundle/schema/jsonschema.json @@ -2748,6 +2748,10 @@ "description": "The client ID for the workspace", "$ref": "#/$defs/string" }, + "experimental_is_unified_host": { + "description": "Experimental feature 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" @@ -2775,6 +2779,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..fc2c45abb8 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,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() +} diff --git a/cmd/auth/login.go b/cmd/auth/login.go index a76cf103b8..a59ec21ee9 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 @@ -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 } diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 4a12dedbdd..eb8635f133 100644 --- a/cmd/auth/login_test.go +++ b/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 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..702103d66c 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..6d1bf84434 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..15b50c2a9f 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: "123456789", + WorkspaceId: "123456789", + 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..a9ab53a8e8 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..0b198fba51 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..d68accaddd 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 {