remove clientType mutation from config

Author: jh-db

account_client.go

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

client/client.go

@@ -18,14 +18,21 @@ import (
 // with the workspace ID to all requests made by workspace clients
 func workspaceOrgIdVisitor(cfg *Config) func(r *http.Request) error {
 	return func(r *http.Request) error {
-		if cfg.GetClientType() == WorkspaceClient && cfg.WorkspaceId != "" {
+		if cfg.WorkspaceId != "" {
 			r.Header.Set("X-Databricks-Org-Id", cfg.WorkspaceId)
 		}
 		return nil
 	}
 }
 
-func HTTPClientConfigFromConfig(cfg *Config) (httpclient.ClientConfig, 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) {
 	if skippable, ok := cfg.HTTPTransport.(interface {
 		SkipRetryOnIO() bool
 	}); ok && skippable.SkipRetryOnIO() {
@@ -85,7 +92,12 @@ func HTTPClientConfigFromConfig(cfg *Config) (httpclient.ClientConfig, error) {
 				*r = *r.WithContext(ctx) // replace request
 				return nil
 			},
-			workspaceOrgIdVisitor(cfg),
+			func() httpclient.RequestVisitor {
+				if isWorkspaceClient {
+					return workspaceOrgIdVisitor(cfg)
+				}
+				return noopVisitor()
+			}(),
 		},
 		TransientErrors: []string{
 			// This is temporary workaround for SCIM API returning 500.
@@ -119,8 +131,8 @@ func (noopAuth) Configure(context.Context, *Config) (credentials.CredentialsProv
 }
 
 // Deprecated: use [HTTPClientConfigFromConfig] with [httpclient.NewApiClient].
-func (c *Config) NewApiClient() (*httpclient.ApiClient, error) {
-	cfg, err := HTTPClientConfigFromConfig(c)
+func (c *Config) NewApiClient(isWorkspaceClient bool) (*httpclient.ApiClient, error) {
+	cfg, err := HTTPClientConfigFromConfig(c, isWorkspaceClient)
 	if err != nil {
 		return nil, err
 	}

config/api_client.go

@@ -36,11 +36,9 @@ func (c AzureMsiCredentials) Configure(ctx context.Context, cfg *Config) (creden
 		return nil, nil
 	}
 	env := cfg.Environment()
-	if cfg.GetClientType() == WorkspaceClient {
-		err := cfg.azureEnsureWorkspaceUrl(ctx, c)
-		if err != nil {
-			return nil, fmt.Errorf("resolve host: %w", err)
-		}
+	err := cfg.azureEnsureWorkspaceUrl(ctx, c)
+	if err != nil {
+		return nil, fmt.Errorf("resolve host: %w", err)
 	}
 	logger.Debugf(ctx, "Generating AAD token via Azure MSI")
 	inner := azureReuseTokenSource(nil, c.tokenSourceFor(ctx, cfg, "", env.AzureApplicationID))

config/auth_azure_msi.go

@@ -28,11 +28,6 @@ func (c GoogleDefaultCredentials) Configure(ctx context.Context, cfg *Config) (c
 	if err != nil {
 		return nil, err
 	}
-	if cfg.GetClientType() == WorkspaceClient {
-		logger.Infof(ctx, "Using Google Default Application Credentials for Workspace")
-		visitor := refreshableVisitor(inner)
-		return credentials.CredentialsProviderFn(visitor), nil
-	}
 	// source for generateAccessToken
 	platform, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{
 		TargetPrincipal: cfg.GoogleServiceAccount,
@@ -44,7 +39,7 @@ func (c GoogleDefaultCredentials) Configure(ctx context.Context, cfg *Config) (c
 	if err != nil {
 		return nil, err
 	}
-	logger.Infof(ctx, "Using Google Default Application Credentials for Accounts API")
+	logger.Infof(ctx, "Using Google Default Application Credentials")
 	visitor := serviceToServiceVisitor(inner, platform, "X-Databricks-GCP-SA-Access-Token")
 	return credentials.NewOAuthCredentialsProvider(visitor, inner.Token), nil
 }

config/auth_gcp_google_id.go

@@ -39,11 +39,6 @@ type Loader interface {
 	Configure(*Config) error
 }
 
-type ClientTypeEnum string
-
-const WorkspaceClient ClientTypeEnum = `WORKSPACE_CLIENT`
-const AccountClient ClientTypeEnum = `ACCOUNT_CLIENT`
-
 type HostTypeEnum string
 
 const WorkspaceHost HostTypeEnum = `WORKSPACE_HOST`
@@ -309,18 +304,6 @@ func (c *Config) IsAws() bool {
 	return c.Host != "" && !c.IsAzure() && !c.IsGcp()
 }
 
-// GetClientType returns one of WORKSPACE_CLIENT or ACCOUNT_CLIENT
-func (c *Config) GetClientType() ClientTypeEnum {
-	// TODO: Implement this so it relies on a flag in Config that we set when
-	// we create the client
-	return c.clientType
-}
-
-// SetClientType sets the client type to one of WORKSPACE_CLIENT or ACCOUNT_CLIENT
-func (c *Config) SetClientType(clientType ClientTypeEnum) {
-	c.clientType = clientType
-}
-
 // GetHostType returns one of WORKSPACE_HOST, ACCOUNT_HOST, or UNIFIED HOST
 func (c *Config) GetHostType() HostTypeEnum {
 	if c.isUnifiedHost {

config/config.go

@@ -12,7 +12,7 @@ import (
 )
 
 // serviceToServiceVisitor returns a visitor that sets the Authorization header
-// to the token from the auth token sourcevand the provided secondary header to
+// to the token from the auth token source and the provided secondary header to
 // the token from the secondary token source.
 func serviceToServiceVisitor(primary, secondary oauth2.TokenSource, secondaryHeader string) func(r *http.Request) error {
 	refreshableAuth := auth.NewCachedTokenSource(authconv.AuthTokenSource(primary))