[WIP] support unified hosts

jh-db · jh-db · commit b8242caf847a · 2025-10-24T13:48:45.000+02:00
diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go
@@ -25,6 +25,7 @@ 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().BoolVarP(&authArguments.IsUnifiedHost, "experimental-is-unified-host", "u", false, "Whether the host is a unified host")
 
 	cmd.AddCommand(newEnvCommand())
 	cmd.AddCommand(newLoginCommand(&authArguments))
@@ -55,3 +56,19 @@ func promptForAccountID(ctx context.Context) (string, error) {
 	prompt.AllowEdit = true
 	return prompt.Run()
 }
+
+func promptForWorkspaceId(ctx context.Context) (string, error) {
+	if !cmdio.IsInTTY(ctx) {
+		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
+	result, err := prompt.Run()
+	if err != nil {
+		return "", err
+	}
+	return result, nil
+}
diff --git a/cmd/auth/login.go b/cmd/auth/login.go
@@ -132,7 +132,7 @@ depends on the existing profiles you have set in your configuration file
 		if err != nil {
 			return err
 		}
-		err = setHostAndAccountId(ctx, existingProfile, authArguments, args)
+		err = setHostAndAccountId(ctx, cmd, existingProfile, authArguments, args)
 		if err != nil {
 			return err
 		}
@@ -170,6 +170,12 @@ depends on the existing profiles you have set in your configuration file
 			return err
 		}
 
+		// Prompt for workspace_id if the host is unified
+		workspaceId, err := promptForWorkspaceIdIfUnified(ctx, authArguments, existingProfile)
+		if err != nil {
+			return err
+		}
+
 		switch {
 		case configureCluster:
 			w, err := databricks.NewWorkspaceClient((*databricks.Config)(&cfg))
@@ -201,13 +207,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:                workspaceId,
+				ClusterID:                  cfg.ClusterID,
+				ConfigFile:                 cfg.ConfigFile,
+				ServerlessComputeID:        cfg.ServerlessComputeID,
+				Experimental_IsUnifiedHost: authArguments.IsUnifiedHost,
 			})
 			if err != nil {
 				return err
@@ -233,7 +241,7 @@ depends on the existing profiles you have set in your configuration file
 // 1. --account-id flag.
 // 2. account-id from the specified profile, if available.
 // 3. Prompt the user for the account-id.
-func setHostAndAccountId(ctx context.Context, existingProfile *profile.Profile, authArguments *auth.AuthArguments, args []string) error {
+func setHostAndAccountId(ctx context.Context, cmd *cobra.Command, existingProfile *profile.Profile, authArguments *auth.AuthArguments, args []string) error {
 	// If both [HOST] and --host are provided, return an error.
 	host := authArguments.Host
 	if len(args) > 0 && host != "" {
@@ -259,11 +267,21 @@ func setHostAndAccountId(ctx context.Context, existingProfile *profile.Profile,
 		}
 	}
 
+	// Determine if the host is a unified host in the following order of precedence:
+	// 1. --experimental-is-unified-host flag (if explicitly set)
+	// 2. experimental_is_unified_host from the specified profile, if available
+	// 3. default to false if neither is provided.
+	if !cmd.Flag("experimental-is-unified-host").Changed {
+		if existingProfile != nil {
+			authArguments.IsUnifiedHost = existingProfile.Experimental_IsUnifiedHost
+		}
+	}
+
 	// If the account-id was not provided as a cmd line flag, try to read it from
 	// the specified profile.
-	isAccountClient := (&config.Config{Host: authArguments.Host}).IsAccountClient()
+	isAccountHost := (&config.Config{Host: authArguments.Host, Experimental_IsUnifiedHost: authArguments.IsUnifiedHost}).GetHostType() != config.WorkspaceHost
 	accountID := authArguments.AccountID
-	if isAccountClient && accountID == "" {
+	if isAccountHost && accountID == "" {
 		if existingProfile != nil && existingProfile.AccountID != "" {
 			authArguments.AccountID = existingProfile.AccountID
 		} else {
@@ -279,6 +297,28 @@ func setHostAndAccountId(ctx context.Context, existingProfile *profile.Profile,
 	return nil
 }
 
+// promptForWorkspaceIdIfUnified prompts for workspace ID if the host is unified.
+// Returns the workspace ID from the profile if available, otherwise prompts the user.
+func promptForWorkspaceIdIfUnified(ctx context.Context, authArguments *auth.AuthArguments, existingProfile *profile.Profile) (string, error) {
+	// Only prompt for workspace_id if the host is a unified host
+	if !authArguments.IsUnifiedHost {
+		return "", nil
+	}
+
+	// If the existing profile has a workspace_id, use it
+	if existingProfile != nil && existingProfile.WorkspaceId != "" {
+		return existingProfile.WorkspaceId, nil
+	}
+
+	// Prompt the user for workspace_id
+	workspaceId, err := promptForWorkspaceId(ctx)
+	if err != nil {
+		return "", err
+	}
+
+	return workspaceId, nil
+}
+
 // getProfileName returns the default profile name for a given host/account ID.
 // If the account ID is provided, the profile name is "ACCOUNT-<account-id>".
 // Otherwise, the profile name is the first part of the host URL.
diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go
@@ -26,7 +26,8 @@ func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) {
 	existingProfile, err := loadProfileByName(ctx, "foo", profile.DefaultProfiler)
 	assert.NoError(t, err)
 
-	err = setHostAndAccountId(ctx, existingProfile, &auth.AuthArguments{Host: "test"}, []string{})
+	cmd := newLoginCommand(&auth.AuthArguments{})
+	err = setHostAndAccountId(ctx, cmd, existingProfile, &auth.AuthArguments{Host: "test"}, []string{})
 	assert.NoError(t, err)
 }
 
@@ -38,38 +39,40 @@ func TestSetHost(t *testing.T) {
 	profile1 := loadTestProfile(t, ctx, "profile-1")
 	profile2 := loadTestProfile(t, ctx, "profile-2")
 
+	cmd := newLoginCommand(&auth.AuthArguments{})
+
 	// Test error when both flag and argument are provided
 	authArguments.Host = "val from --host"
-	err := setHostAndAccountId(ctx, profile1, &authArguments, []string{"val from [HOST]"})
+	err := setHostAndAccountId(ctx, cmd, profile1, &authArguments, []string{"val from [HOST]"})
 	assert.EqualError(t, err, "please only provide a host as an argument or a flag, not both")
 
 	// Test setting host from flag
 	authArguments.Host = "val from --host"
-	err = setHostAndAccountId(ctx, profile1, &authArguments, []string{})
+	err = setHostAndAccountId(ctx, cmd, profile1, &authArguments, []string{})
 	assert.NoError(t, err)
 	assert.Equal(t, "val from --host", authArguments.Host)
 
 	// Test setting host from argument
 	authArguments.Host = ""
-	err = setHostAndAccountId(ctx, profile1, &authArguments, []string{"val from [HOST]"})
+	err = setHostAndAccountId(ctx, cmd, profile1, &authArguments, []string{"val from [HOST]"})
 	assert.NoError(t, err)
 	assert.Equal(t, "val from [HOST]", authArguments.Host)
 
 	// Test setting host from profile
 	authArguments.Host = ""
-	err = setHostAndAccountId(ctx, profile1, &authArguments, []string{})
+	err = setHostAndAccountId(ctx, cmd, profile1, &authArguments, []string{})
 	assert.NoError(t, err)
 	assert.Equal(t, "https://www.host1.com", authArguments.Host)
 
 	// Test setting host from profile
 	authArguments.Host = ""
-	err = setHostAndAccountId(ctx, profile2, &authArguments, []string{})
+	err = setHostAndAccountId(ctx, cmd, profile2, &authArguments, []string{})
 	assert.NoError(t, err)
 	assert.Equal(t, "https://www.host2.com", authArguments.Host)
 
 	// Test host is not set. Should prompt.
 	authArguments.Host = ""
-	err = setHostAndAccountId(ctx, nil, &authArguments, []string{})
+	err = setHostAndAccountId(ctx, cmd, nil, &authArguments, []string{})
 	assert.EqualError(t, err, "the command is being run in a non-interactive environment, please specify a host using --host")
 }
 
@@ -80,24 +83,26 @@ func TestSetAccountId(t *testing.T) {
 
 	accountProfile := loadTestProfile(t, ctx, "account-profile")
 
+	cmd := newLoginCommand(&auth.AuthArguments{})
+
 	// Test setting account-id from flag
 	authArguments.AccountID = "val from --account-id"
-	err := setHostAndAccountId(ctx, accountProfile, &authArguments, []string{})
+	err := setHostAndAccountId(ctx, cmd, accountProfile, &authArguments, []string{})
 	assert.NoError(t, err)
 	assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host)
 	assert.Equal(t, "val from --account-id", authArguments.AccountID)
 
 	// Test setting account_id from profile
 	authArguments.AccountID = ""
-	err = setHostAndAccountId(ctx, accountProfile, &authArguments, []string{})
+	err = setHostAndAccountId(ctx, cmd, accountProfile, &authArguments, []string{})
 	require.NoError(t, err)
 	assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host)
 	assert.Equal(t, "id-from-profile", authArguments.AccountID)
 
 	// Neither flag nor profile account-id is set, should prompt
 	authArguments.AccountID = ""
 	authArguments.Host = "https://accounts.cloud.databricks.com"
-	err = setHostAndAccountId(ctx, nil, &authArguments, []string{})
+	err = setHostAndAccountId(ctx, cmd, nil, &authArguments, []string{})
 	assert.EqualError(t, err, "the command is being run in a non-interactive environment, please specify an account ID using --account-id")
 }
 
diff --git a/cmd/auth/profiles.go b/cmd/auth/profiles.go
@@ -51,7 +51,7 @@ func (c *profileMetadata) Load(ctx context.Context, configFilePath string, skipV
 		return
 	}
 
-	if cfg.IsAccountClient() {
+	if cfg.GetHostType() != config.WorkspaceHost {
 		a, err := databricks.NewAccountClient((*databricks.Config)(cfg))
 		if err != nil {
 			return
diff --git a/cmd/auth/token.go b/cmd/auth/token.go
@@ -43,6 +43,7 @@ using a client ID and secret is not supported.`,
 		}
 
 		t, err := loadToken(ctx, loadTokenArgs{
+			cmd:                cmd,
 			authArguments:      authArguments,
 			profileName:        profileName,
 			args:               args,
@@ -65,6 +66,9 @@ using a client ID and secret is not supported.`,
 }
 
 type loadTokenArgs struct {
+	// cmd is the cobra command, used to access flags like experimental-is-unified-host.
+	cmd *cobra.Command
+
 	// authArguments is the parsed auth arguments, including the host and optionally the account ID.
 	authArguments *auth.AuthArguments
 
@@ -98,7 +102,7 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
 		return nil, err
 	}
 
-	err = setHostAndAccountId(ctx, existingProfile, args.authArguments, args.args)
+	err = setHostAndAccountId(ctx, args.cmd, existingProfile, args.authArguments, args.args)
 	if err != nil {
 		return nil, err
 	}
diff --git a/go.mod b/go.mod
@@ -75,3 +75,5 @@ require (
 	google.golang.org/grpc v1.75.1 // indirect
 	google.golang.org/protobuf v1.36.9 // indirect
 )
+
+replace github.com/databricks/databricks-sdk-go => ../databricks-sdk-go
diff --git a/go.sum b/go.sum
@@ -29,8 +29,6 @@ github.com/cloudflare/circl v1.6.1/go.mod h1:uddAzsPgqdMAYatqJ0lsjX1oECcQLIlRpzZ
 github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
 github.com/cyphar/filepath-securejoin v0.4.1 h1:JyxxyPEaktOD+GAnqIqTf9A8tHyAG22rowi7HkoSU1s=
 github.com/cyphar/filepath-securejoin v0.4.1/go.mod h1:Sdj7gXlvMcPZsbhwhQ33GguGLDGQL7h7bg04C/+u9jI=
-github.com/databricks/databricks-sdk-go v0.86.0 h1:Di1+NBQlfzBMUhY6w6gS2mtmNXIWycowoCsLCGFQPyU=
-github.com/databricks/databricks-sdk-go v0.86.0/go.mod h1:hWoHnHbNLjPKiTm5K/7bcIv3J3Pkgo5x9pPzh8K3RVE=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
diff --git a/libs/auth/arguments.go b/libs/auth/arguments.go
@@ -8,19 +8,23 @@ 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
+	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,
+		Experimental_IsUnifiedHost: a.IsUnifiedHost,
 	}
 	host := cfg.CanonicalHostName()
-	if cfg.IsAccountClient() {
+	if cfg.GetHostType() == config.AccountHost {
 		return u2m.NewBasicAccountOAuthArgument(host, cfg.AccountID)
+	} else if cfg.GetHostType() == config.UnifiedHost {
+		return u2m.NewBasicUnifiedOAuthArgument(host, cfg.AccountID)
 	}
 	return u2m.NewBasicWorkspaceOAuthArgument(host)
 }
diff --git a/libs/auth/error.go b/libs/auth/error.go
@@ -13,7 +13,7 @@ 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, accountId, false}.ToOAuthArgument() // TODO: pass the IsUnifiedHost flag
 		if err != nil {
 			return false, err
 		}
diff --git a/libs/databrickscfg/profile/file.go b/libs/databrickscfg/profile/file.go
@@ -79,11 +79,13 @@ func (f FileProfilerImpl) LoadProfiles(ctx context.Context, fn ProfileMatchFunct
 			continue
 		}
 		profile := Profile{
-			Name:                v.Name(),
-			Host:                host,
-			AccountID:           all["account_id"],
-			ClusterID:           all["cluster_id"],
-			ServerlessComputeID: all["serverless_compute_id"],
+			Name:                       v.Name(),
+			Host:                       host,
+			AccountID:                  all["account_id"],
+			WorkspaceId:                all["workspace_id"],
+			ClusterID:                  all["cluster_id"],
+			ServerlessComputeID:        all["serverless_compute_id"],
+			Experimental_IsUnifiedHost: all["experimental_is_unified_host"] == "true",
 		}
 		if fn(profile) {
 			profiles = append(profiles, profile)
diff --git a/libs/databrickscfg/profile/file_test.go b/libs/databrickscfg/profile/file_test.go
@@ -80,5 +80,32 @@ func TestLoadProfilesMatchAccount(t *testing.T) {
 	profiler := FileProfilerImpl{}
 	profiles, err := profiler.LoadProfiles(ctx, MatchAccountProfiles)
 	require.NoError(t, err)
-	assert.Equal(t, []string{"acc"}, profiles.Names())
+	assert.Equal(t, []string{"acc", "unified", "unified-with-workspace"}, profiles.Names())
+}
+
+func TestLoadProfilesWithExperimentalIsUnifiedHost(t *testing.T) {
+	ctx := context.Background()
+	ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
+	profiler := FileProfilerImpl{}
+	profiles, err := profiler.LoadProfiles(ctx, WithName("unified"))
+	require.NoError(t, err)
+	require.Len(t, profiles, 1)
+	assert.Equal(t, "unified", profiles[0].Name)
+	assert.Equal(t, "https://unified.databricks.com", profiles[0].Host)
+	assert.Equal(t, "def", profiles[0].AccountID)
+	assert.True(t, profiles[0].Experimental_IsUnifiedHost)
+}
+
+func TestLoadProfilesWithWorkspaceId(t *testing.T) {
+	ctx := context.Background()
+	ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./testdata/databrickscfg")
+	profiler := FileProfilerImpl{}
+	profiles, err := profiler.LoadProfiles(ctx, WithName("unified-with-workspace"))
+	require.NoError(t, err)
+	require.Len(t, profiles, 1)
+	assert.Equal(t, "unified-with-workspace", profiles[0].Name)
+	assert.Equal(t, "https://unified.databricks.com", profiles[0].Host)
+	assert.Equal(t, "def", profiles[0].AccountID)
+	assert.Equal(t, "123456789", profiles[0].WorkspaceId)
+	assert.True(t, profiles[0].Experimental_IsUnifiedHost)
 }
diff --git a/libs/databrickscfg/profile/profile.go b/libs/databrickscfg/profile/profile.go
@@ -10,11 +10,13 @@ import (
 // It should only be used for prompting and filtering.
 // Use its name to construct a config.Config.
 type Profile struct {
-	Name                string
-	Host                string
-	AccountID           string
-	ClusterID           string
-	ServerlessComputeID string
+	Name                       string
+	Host                       string
+	AccountID                  string
+	WorkspaceId                string
+	ClusterID                  string
+	ServerlessComputeID        string
+	Experimental_IsUnifiedHost bool
 }
 
 func (p Profile) Cloud() string {
diff --git a/libs/databrickscfg/profile/testdata/databrickscfg b/libs/databrickscfg/profile/testdata/databrickscfg

Original file line number	Diff line number	Diff line change
`@@ -51,7 +51,7 @@ func (c *profileMetadata) Load(ctx context.Context, configFilePath string, skipV`
`51`	`51`	`return`
`52`	`52`	`}`
`53`	`53`
`54`		`- if cfg.IsAccountClient() {`
	`54`	`+ if cfg.GetHostType() != config.WorkspaceHost {`
`55`	`55`	`a, err := databricks.NewAccountClient((*databricks.Config)(cfg))`
`56`	`56`	`if err != nil {`
`57`	`57`	`return`
Original file line number	Diff line number	Diff line change
@@ -43,6 +43,7 @@ using a client ID and secret is not supported.`,
`43`	`43`	`}`
`44`	`44`
`45`	`45`	`t, err := loadToken(ctx, loadTokenArgs{`
	`46`	`+ cmd: cmd,`
`46`	`47`	`authArguments: authArguments,`
`47`	`48`	`profileName: profileName,`
`48`	`49`	`args: args,`
@@ -65,6 +66,9 @@ using a client ID and secret is not supported.`,
`65`	`66`	`}`
`66`	`67`
`67`	`68`	`type loadTokenArgs struct {`
	`69`	`+ // cmd is the cobra command, used to access flags like experimental-is-unified-host.`
	`70`	`+ cmd *cobra.Command`
	`71`	`+`
`68`	`72`	`// authArguments is the parsed auth arguments, including the host and optionally the account ID.`
`69`	`73`	`authArguments *auth.AuthArguments`
`70`	`74`
`@@ -98,7 +102,7 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {`
`98`	`102`	`return nil, err`
`99`	`103`	`}`
`100`	`104`
`101`		`- err = setHostAndAccountId(ctx, existingProfile, args.authArguments, args.args)`
	`105`	`+ err = setHostAndAccountId(ctx, args.cmd, existingProfile, args.authArguments, args.args)`
`102`	`106`	`if err != nil {`
`103`	`107`	`return nil, err`
`104`	`108`	`}`
Original file line number	Diff line number	Diff line change
`@@ -75,3 +75,5 @@ require (`
`75`	`75`	`google.golang.org/grpc v1.75.1 // indirect`
`76`	`76`	`google.golang.org/protobuf v1.36.9 // indirect`
`77`	`77`	`)`
	`78`	`+`
	`79`	`+replace github.com/databricks/databricks-sdk-go => ../databricks-sdk-go`
Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,7 @@ import (`
`13`	`13`	`func RewriteAuthError(ctx context.Context, host, accountId, profile string, err error) (bool, error) {`
`14`	`14`	`target := &u2m.InvalidRefreshTokenError{}`
`15`	`15`	`if errors.As(err, &target) {`
`16`		`- oauthArgument, err := AuthArguments{host, accountId}.ToOAuthArgument()`
	`16`	`+ oauthArgument, err := AuthArguments{host, accountId, false}.ToOAuthArgument() // TODO: pass the IsUnifiedHost flag`
`17`	`17`	`if err != nil {`
`18`	`18`	`return false, err`
`19`	`19`	`}`