Skip to content

Commit 01d8ab7

Browse files
committed
Resolve TokenAudience from host metadata for account hosts
Port of Python SDK PR #1321. When host metadata indicates an account host (no workspace_id) and account_id is present, automatically set TokenAudience to the account_id if not already configured by the user. Co-authored-by: Isaac
1 parent 6e27bb2 commit 01d8ab7

2 files changed

Lines changed: 67 additions & 0 deletions

File tree

config/config.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -699,6 +699,10 @@ func (c *Config) resolveHostMetadata(ctx context.Context) {
699699
c.Cloud = c.Environment().Cloud
700700
logger.Debugf(ctx, "Resolved cloud from hostname: %q", c.Cloud)
701701
}
702+
if c.TokenAudience == "" && meta.WorkspaceID == "" && c.AccountID != "" {
703+
logger.Debugf(ctx, "Setting token_audience to account_id for account host: %q", c.AccountID)
704+
c.TokenAudience = c.AccountID
705+
}
702706
if c.DiscoveryURL == "" {
703707
if meta.OIDCEndpoint == "" {
704708
logger.Warnf(ctx, "Host metadata missing oidc_endpoint; skipping discovery URL resolution")

config/config_test.go

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -724,6 +724,69 @@ func TestEnsureResolved_HostMetadata_MissingAccountIdWithPlaceholder_Warns(t *te
724724
assert.Empty(t, cfg.DiscoveryURL)
725725
}
726726

727+
func TestApplyHostMetadata_SetsTokenAudienceForAccountHost(t *testing.T) {
728+
noopLoader := mockLoader(func(cfg *Config) error { return nil })
729+
cfg := &Config{
730+
Host: testHMHost,
731+
Experimental_IsUnifiedHost: true,
732+
Loaders: []Loader{noopLoader},
733+
HTTPTransport: fixtures.SliceTransport{
734+
{
735+
Method: "GET",
736+
Resource: "/.well-known/databricks-config",
737+
Status: 200,
738+
Response: `{"oidc_endpoint": "` + testHMHost + `/oidc", "account_id": "` + testHMAccountID + `", "cloud": "AWS"}`,
739+
},
740+
},
741+
}
742+
err := cfg.EnsureResolved()
743+
require.NoError(t, err)
744+
// No workspace_id and has account_id → token audience should be set
745+
assert.Equal(t, testHMAccountID, cfg.TokenAudience)
746+
}
747+
748+
func TestApplyHostMetadata_NoTokenAudienceForWorkspaceHost(t *testing.T) {
749+
noopLoader := mockLoader(func(cfg *Config) error { return nil })
750+
cfg := &Config{
751+
Host: testHMHost,
752+
Experimental_IsUnifiedHost: true,
753+
Loaders: []Loader{noopLoader},
754+
HTTPTransport: fixtures.SliceTransport{
755+
{
756+
Method: "GET",
757+
Resource: "/.well-known/databricks-config",
758+
Status: 200,
759+
Response: `{"oidc_endpoint": "` + testHMHost + `/oidc", "account_id": "` + testHMAccountID + `", "workspace_id": "` + testHMWorkspaceID + `", "cloud": "AWS"}`,
760+
},
761+
},
762+
}
763+
err := cfg.EnsureResolved()
764+
require.NoError(t, err)
765+
// Has workspace_id → token audience should NOT be set
766+
assert.Empty(t, cfg.TokenAudience)
767+
}
768+
769+
func TestApplyHostMetadata_DoesNotOverrideExistingTokenAudience(t *testing.T) {
770+
noopLoader := mockLoader(func(cfg *Config) error { return nil })
771+
cfg := &Config{
772+
Host: testHMHost,
773+
TokenAudience: "custom-audience",
774+
Experimental_IsUnifiedHost: true,
775+
Loaders: []Loader{noopLoader},
776+
HTTPTransport: fixtures.SliceTransport{
777+
{
778+
Method: "GET",
779+
Resource: "/.well-known/databricks-config",
780+
Status: 200,
781+
Response: `{"oidc_endpoint": "` + testHMHost + `/oidc", "account_id": "` + testHMAccountID + `", "cloud": "AWS"}`,
782+
},
783+
},
784+
}
785+
err := cfg.EnsureResolved()
786+
require.NoError(t, err)
787+
assert.Equal(t, "custom-audience", cfg.TokenAudience)
788+
}
789+
727790
func TestConfig_ResolveHostMetadata_Clouds(t *testing.T) {
728791
tests := []struct {
729792
name string

0 commit comments

Comments
 (0)