WIP?

Author: jh-db

account_client.go

@@ -10,12 +10,12 @@ import (
 	"golang.org/x/oauth2"
 )
 
-func New(cfg *config.Config, isWorkspaceClient bool) (*DatabricksClient, error) {
+func New(cfg *config.Config) (*DatabricksClient, error) {
 	err := cfg.EnsureResolved()
 	if err != nil {
 		return nil, err
 	}
-	clientCfg, err := config.HTTPClientConfigFromConfig(cfg, isWorkspaceClient)
+	clientCfg, err := config.HTTPClientConfigFromConfig(cfg)
 	if err != nil {
 		return nil, err
 	}

client/client.go

@@ -16,6 +16,7 @@ import (
 
 // workspaceOrgIdVisitor creates a visitor that adds X-Databricks-Org-Id header
 // with the workspace ID to all requests made by workspace clients
+// It relies on the workspace ID being set if and only if a workspace client is being used.
 func workspaceOrgIdVisitor(cfg *Config) func(r *http.Request) error {
 	return func(r *http.Request) error {
 		if cfg.WorkspaceId != "" {
@@ -25,14 +26,7 @@ func workspaceOrgIdVisitor(cfg *Config) func(r *http.Request) error {
 	}
 }
 
-// noopVisitor creates a visitor that does nothing
-func noopVisitor() func(r *http.Request) error {
-	return func(r *http.Request) error {
-		return nil
-	}
-}
-
-func HTTPClientConfigFromConfig(cfg *Config, isWorkspaceClient bool) (httpclient.ClientConfig, error) {
+func HTTPClientConfigFromConfig(cfg *Config) (httpclient.ClientConfig, error) {
 	if skippable, ok := cfg.HTTPTransport.(interface {
 		SkipRetryOnIO() bool
 	}); ok && skippable.SkipRetryOnIO() {
@@ -92,12 +86,7 @@ func HTTPClientConfigFromConfig(cfg *Config, isWorkspaceClient bool) (httpclient
 				*r = *r.WithContext(ctx) // replace request
 				return nil
 			},
-			func() httpclient.RequestVisitor {
-				if isWorkspaceClient {
-					return workspaceOrgIdVisitor(cfg)
-				}
-				return noopVisitor()
-			}(),
+			workspaceOrgIdVisitor(cfg),
 		},
 		TransientErrors: []string{
 			// This is temporary workaround for SCIM API returning 500.
@@ -131,8 +120,8 @@ func (noopAuth) Configure(context.Context, *Config) (credentials.CredentialsProv
 }
 
 // Deprecated: use [HTTPClientConfigFromConfig] with [httpclient.NewApiClient].
-func (c *Config) NewApiClient(isWorkspaceClient bool) (*httpclient.ApiClient, error) {
-	cfg, err := HTTPClientConfigFromConfig(c, isWorkspaceClient)
+func (c *Config) NewApiClient() (*httpclient.ApiClient, error) {
+	cfg, err := HTTPClientConfigFromConfig(c)
 	if err != nil {
 		return nil, err
 	}

config/api_client.go

@@ -41,9 +41,11 @@ type Loader interface {
 
 type HostTypeEnum string
 
-const WorkspaceHost HostTypeEnum = `WORKSPACE_HOST`
-const AccountHost HostTypeEnum = `ACCOUNT_HOST`
-const UnifiedHost HostTypeEnum = `UNIFIED_HOST`
+const (
+	WorkspaceHost HostTypeEnum = `WORKSPACE_HOST`
+	AccountHost   HostTypeEnum = `ACCOUNT_HOST`
+	UnifiedHost   HostTypeEnum = `UNIFIED_HOST`
+)
 
 // Config represents configuration for Databricks Connectivity
 type Config struct {
@@ -193,8 +195,8 @@ type Config struct {
 	// Keep track of the source of each attribute
 	attrSource map[string]Source
 
-	// Marker for unified hosts. Will be unnecessary once we've settled on a way to determine if a host is unified.
-	isUnifiedHost bool `name:"is_unified_host" env:"DATABRICKS_IS_UNIFIED_HOST" auth:"-"`
+	// Marker for unified hosts. Will be redundant once we can recognize unified hosts by their hostname.
+	Experimental_IsUnifiedHost bool `name:"experimental_is_unified_host" env:"DATABRICKS_EXPERIMENTAL_IS_UNIFIED_HOST" auth:"-"`
 }
 
 // NewWithWorkspaceHost returns a new instance of the Config with the host set to
@@ -207,7 +209,7 @@ func (c *Config) NewWithWorkspaceHost(host string) (*Config, error) {
 		return nil, err
 	}
 
-	var fieldsToSkip = map[string]struct{}{
+	fieldsToSkip := map[string]struct{}{
 		"Host":            {},
 		"AzureResourceID": {},
 		"AccountID":       {},
@@ -303,7 +305,7 @@ func (c *Config) IsAws() bool {
 
 // GetHostType returns one of WORKSPACE_HOST, ACCOUNT_HOST, or UNIFIED HOST
 func (c *Config) GetHostType() HostTypeEnum {
-	if c.isUnifiedHost {
+	if c.Experimental_IsUnifiedHost {
 		return UnifiedHost
 	}
 
@@ -344,7 +346,6 @@ func (c *Config) EnsureResolved() error {
 		logger.Tracef(ctx, "Loading config via %s", loader.Name())
 		err := loader.Configure(c)
 		if err != nil {
-
 			return c.wrapDebug(fmt.Errorf("resolve: %w", err))
 		}
 	}
@@ -501,6 +502,10 @@ func (c *Config) getOidcEndpoints(ctx context.Context) (*u2m.OAuthAuthorizationS
 }
 
 func (c *Config) getOAuthArgument() (u2m.OAuthArgument, error) {
+	err := c.EnsureResolved()
+	if err != nil {
+		return nil, err
+	}
 	host := c.CanonicalHostName()
 	if c.GetHostType() == AccountHost {
 		return u2m.NewBasicAccountOAuthArgument(host, c.AccountID)

config/config.go

@@ -365,8 +365,9 @@ func (a *PersistentAuth) validateArg() error {
 	}
 	_, isWorkspaceArg := a.oAuthArgument.(WorkspaceOAuthArgument)
 	_, isAccountArg := a.oAuthArgument.(AccountOAuthArgument)
-	if !isWorkspaceArg && !isAccountArg {
-		return fmt.Errorf("unsupported OAuthArgument type: %T, must implement either WorkspaceOAuthArgument or AccountOAuthArgument interface", a.oAuthArgument)
+	_, isUnifiedArg := a.oAuthArgument.(UnifiedOAuthArgument)
+	if !isWorkspaceArg && !isAccountArg && !isUnifiedArg {
+		return fmt.Errorf("unsupported OAuthArgument type: %T, must implement either WorkspaceOAuthArgument, AccountOAuthArgument or UnifiedOAuthArgument interface", a.oAuthArgument)
 	}
 	return nil
 }
@@ -385,8 +386,10 @@ func (a *PersistentAuth) oauth2Config() (*oauth2.Config, error) {
 	case AccountOAuthArgument:
 		endpoints, err = a.endpointSupplier.GetAccountOAuthEndpoints(
 			a.ctx, argg.GetAccountHost(), argg.GetAccountId())
+	case UnifiedOAuthArgument:
+		endpoints, err = a.endpointSupplier.GetUnifiedOAuthEndpoints(a.ctx, argg.GetHost(), argg.GetAccountId())
 	default:
-		return nil, fmt.Errorf("unsupported OAuthArgument type: %T, must implement either WorkspaceOAuthArgument or AccountOAuthArgument interface", a.oAuthArgument)
+		return nil, fmt.Errorf("unsupported OAuthArgument type: %T, must implement either WorkspaceOAuthArgument, AccountOAuthArgument or UnifiedOAuthArgument interface", a.oAuthArgument)
 	}
 	if err != nil {
 		return nil, fmt.Errorf("fetching OAuth endpoints: %w", err)