Skip to content

Commit 2e61308

Browse files
committed
WIP?
1 parent f91efff commit 2e61308

6 files changed

Lines changed: 34 additions & 33 deletions

File tree

account_client.go

Lines changed: 4 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

client/client.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,12 @@ import (
1010
"golang.org/x/oauth2"
1111
)
1212

13-
func New(cfg *config.Config, isWorkspaceClient bool) (*DatabricksClient, error) {
13+
func New(cfg *config.Config) (*DatabricksClient, error) {
1414
err := cfg.EnsureResolved()
1515
if err != nil {
1616
return nil, err
1717
}
18-
clientCfg, err := config.HTTPClientConfigFromConfig(cfg, isWorkspaceClient)
18+
clientCfg, err := config.HTTPClientConfigFromConfig(cfg)
1919
if err != nil {
2020
return nil, err
2121
}

config/api_client.go

Lines changed: 5 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616

1717
// workspaceOrgIdVisitor creates a visitor that adds X-Databricks-Org-Id header
1818
// with the workspace ID to all requests made by workspace clients
19+
// It relies on the workspace ID being set if and only if a workspace client is being used.
1920
func workspaceOrgIdVisitor(cfg *Config) func(r *http.Request) error {
2021
return func(r *http.Request) error {
2122
if cfg.WorkspaceId != "" {
@@ -25,14 +26,7 @@ func workspaceOrgIdVisitor(cfg *Config) func(r *http.Request) error {
2526
}
2627
}
2728

28-
// noopVisitor creates a visitor that does nothing
29-
func noopVisitor() func(r *http.Request) error {
30-
return func(r *http.Request) error {
31-
return nil
32-
}
33-
}
34-
35-
func HTTPClientConfigFromConfig(cfg *Config, isWorkspaceClient bool) (httpclient.ClientConfig, error) {
29+
func HTTPClientConfigFromConfig(cfg *Config) (httpclient.ClientConfig, error) {
3630
if skippable, ok := cfg.HTTPTransport.(interface {
3731
SkipRetryOnIO() bool
3832
}); ok && skippable.SkipRetryOnIO() {
@@ -92,12 +86,7 @@ func HTTPClientConfigFromConfig(cfg *Config, isWorkspaceClient bool) (httpclient
9286
*r = *r.WithContext(ctx) // replace request
9387
return nil
9488
},
95-
func() httpclient.RequestVisitor {
96-
if isWorkspaceClient {
97-
return workspaceOrgIdVisitor(cfg)
98-
}
99-
return noopVisitor()
100-
}(),
89+
workspaceOrgIdVisitor(cfg),
10190
},
10291
TransientErrors: []string{
10392
// This is temporary workaround for SCIM API returning 500.
@@ -131,8 +120,8 @@ func (noopAuth) Configure(context.Context, *Config) (credentials.CredentialsProv
131120
}
132121

133122
// Deprecated: use [HTTPClientConfigFromConfig] with [httpclient.NewApiClient].
134-
func (c *Config) NewApiClient(isWorkspaceClient bool) (*httpclient.ApiClient, error) {
135-
cfg, err := HTTPClientConfigFromConfig(c, isWorkspaceClient)
123+
func (c *Config) NewApiClient() (*httpclient.ApiClient, error) {
124+
cfg, err := HTTPClientConfigFromConfig(c)
136125
if err != nil {
137126
return nil, err
138127
}

config/config.go

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,11 @@ type Loader interface {
4141

4242
type HostTypeEnum string
4343

44-
const WorkspaceHost HostTypeEnum = `WORKSPACE_HOST`
45-
const AccountHost HostTypeEnum = `ACCOUNT_HOST`
46-
const UnifiedHost HostTypeEnum = `UNIFIED_HOST`
44+
const (
45+
WorkspaceHost HostTypeEnum = `WORKSPACE_HOST`
46+
AccountHost HostTypeEnum = `ACCOUNT_HOST`
47+
UnifiedHost HostTypeEnum = `UNIFIED_HOST`
48+
)
4749

4850
// Config represents configuration for Databricks Connectivity
4951
type Config struct {
@@ -193,8 +195,8 @@ type Config struct {
193195
// Keep track of the source of each attribute
194196
attrSource map[string]Source
195197

196-
// Marker for unified hosts. Will be unnecessary once we've settled on a way to determine if a host is unified.
197-
isUnifiedHost bool `name:"is_unified_host" env:"DATABRICKS_IS_UNIFIED_HOST" auth:"-"`
198+
// Marker for unified hosts. Will be redundant once we can recognize unified hosts by their hostname.
199+
Experimental_IsUnifiedHost bool `name:"experimental_is_unified_host" env:"DATABRICKS_EXPERIMENTAL_IS_UNIFIED_HOST" auth:"-"`
198200
}
199201

200202
// 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) {
207209
return nil, err
208210
}
209211

210-
var fieldsToSkip = map[string]struct{}{
212+
fieldsToSkip := map[string]struct{}{
211213
"Host": {},
212214
"AzureResourceID": {},
213215
"AccountID": {},
@@ -303,7 +305,7 @@ func (c *Config) IsAws() bool {
303305

304306
// GetHostType returns one of WORKSPACE_HOST, ACCOUNT_HOST, or UNIFIED HOST
305307
func (c *Config) GetHostType() HostTypeEnum {
306-
if c.isUnifiedHost {
308+
if c.Experimental_IsUnifiedHost {
307309
return UnifiedHost
308310
}
309311

@@ -344,7 +346,6 @@ func (c *Config) EnsureResolved() error {
344346
logger.Tracef(ctx, "Loading config via %s", loader.Name())
345347
err := loader.Configure(c)
346348
if err != nil {
347-
348349
return c.wrapDebug(fmt.Errorf("resolve: %w", err))
349350
}
350351
}
@@ -501,6 +502,10 @@ func (c *Config) getOidcEndpoints(ctx context.Context) (*u2m.OAuthAuthorizationS
501502
}
502503

503504
func (c *Config) getOAuthArgument() (u2m.OAuthArgument, error) {
505+
err := c.EnsureResolved()
506+
if err != nil {
507+
return nil, err
508+
}
504509
host := c.CanonicalHostName()
505510
if c.GetHostType() == AccountHost {
506511
return u2m.NewBasicAccountOAuthArgument(host, c.AccountID)

credentials/u2m/persistent_auth.go

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -365,8 +365,9 @@ func (a *PersistentAuth) validateArg() error {
365365
}
366366
_, isWorkspaceArg := a.oAuthArgument.(WorkspaceOAuthArgument)
367367
_, isAccountArg := a.oAuthArgument.(AccountOAuthArgument)
368-
if !isWorkspaceArg && !isAccountArg {
369-
return fmt.Errorf("unsupported OAuthArgument type: %T, must implement either WorkspaceOAuthArgument or AccountOAuthArgument interface", a.oAuthArgument)
368+
_, isUnifiedArg := a.oAuthArgument.(UnifiedOAuthArgument)
369+
if !isWorkspaceArg && !isAccountArg && !isUnifiedArg {
370+
return fmt.Errorf("unsupported OAuthArgument type: %T, must implement either WorkspaceOAuthArgument, AccountOAuthArgument or UnifiedOAuthArgument interface", a.oAuthArgument)
370371
}
371372
return nil
372373
}
@@ -385,8 +386,10 @@ func (a *PersistentAuth) oauth2Config() (*oauth2.Config, error) {
385386
case AccountOAuthArgument:
386387
endpoints, err = a.endpointSupplier.GetAccountOAuthEndpoints(
387388
a.ctx, argg.GetAccountHost(), argg.GetAccountId())
389+
case UnifiedOAuthArgument:
390+
endpoints, err = a.endpointSupplier.GetUnifiedOAuthEndpoints(a.ctx, argg.GetHost(), argg.GetAccountId())
388391
default:
389-
return nil, fmt.Errorf("unsupported OAuthArgument type: %T, must implement either WorkspaceOAuthArgument or AccountOAuthArgument interface", a.oAuthArgument)
392+
return nil, fmt.Errorf("unsupported OAuthArgument type: %T, must implement either WorkspaceOAuthArgument, AccountOAuthArgument or UnifiedOAuthArgument interface", a.oAuthArgument)
390393
}
391394
if err != nil {
392395
return nil, fmt.Errorf("fetching OAuth endpoints: %w", err)

workspace_client.go

Lines changed: 4 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)