From 9c323cb2f8617b7db527499006d94a34a9e8ce71 Mon Sep 17 00:00:00 2001 From: tejas-kochar Date: Mon, 22 Dec 2025 11:12:09 +0000 Subject: [PATCH 1/7] custom scopes support in wif --- config/auth_default.go | 1 + config/auth_default_test.go | 102 ++++++++++++++ config/experimental/auth/oidc/tokensource.go | 5 +- .../auth/oidc/tokensource_test.go | 132 ++++++++++++++++++ 4 files changed, 239 insertions(+), 1 deletion(-) diff --git a/config/auth_default.go b/config/auth_default.go index 1b34a16fb..6a7c38369 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -145,6 +145,7 @@ func oidcStrategy(cfg *Config, name string, ts oidc.IDTokenSource) CredentialsSt TokenEndpointProvider: cfg.getOidcEndpoints, Audience: cfg.TokenAudience, IDTokenSource: ts, + Scopes: cfg.GetScopes(), } if cfg.HostType() != WorkspaceHost { oidcConfig.AccountID = cfg.AccountID diff --git a/config/auth_default_test.go b/config/auth_default_test.go index fbcb67116..adbf81d8e 100644 --- a/config/auth_default_test.go +++ b/config/auth_default_test.go @@ -2,8 +2,13 @@ package config import ( "context" + "encoding/json" + "net/http" + "net/http/httptest" "strings" "testing" + + "github.com/databricks/databricks-sdk-go/credentials/u2m" ) func TestDefaultCredentials_Configure(t *testing.T) { @@ -47,3 +52,100 @@ func TestDefaultCredentials_Configure(t *testing.T) { }) } } + +func TestGithubOIDC_Scopes(t *testing.T) { + tests := []struct { + name string + scopes []string + expectedScope string + }{ + { + name: "default scopes", + scopes: nil, + expectedScope: "all-apis", + }, + { + name: "custom scopes", + scopes: []string{"clusters", "jobs"}, + expectedScope: "clusters jobs", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + githubTokenCalled := false + tokenExchangeCalled := false + + // Simulates the GitHub Actions OIDC token endpoint. + githubServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + githubTokenCalled = true + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]string{"value": "github-id-token"}) + })) + defer githubServer.Close() + + // Simulates a Databricks workspace. + // Asserts whether the right scopes are passed to the token exchange endpoint. + var databricksServer *httptest.Server + databricksServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/oidc/.well-known/oauth-authorization-server": + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(u2m.OAuthAuthorizationServer{ + AuthorizationEndpoint: "https://host.com/oidc/v1/authorize", + TokenEndpoint: databricksServer.URL + "/oidc/v1/token", + }) + + case "/oidc/v1/token": + tokenExchangeCalled = true + if err := r.ParseForm(); err != nil { + t.Fatalf("Failed to parse form: %v", err) + } + // Verify scope is passed correctly to token exchange. + if got := r.Form.Get("scope"); got != tt.expectedScope { + t.Errorf("scope: got %q, want %q", got, tt.expectedScope) + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]interface{}{ + "token_type": "Bearer", + "access_token": "databricks-access-token", + "expires_in": 3600, + }) + + default: + t.Errorf("Unexpected request: %s %s", r.Method, r.URL.Path) + http.Error(w, "Not found", http.StatusNotFound) + } + })) + defer databricksServer.Close() + + cfg := &Config{ + Host: databricksServer.URL, + ClientID: "test-client-id", + ActionsIDTokenRequestURL: githubServer.URL + "/github-token?version=1", + ActionsIDTokenRequestToken: "github-request-token", + TokenAudience: "databricks-test-audience", + AuthType: "github-oidc", + } + if tt.scopes != nil { + cfg.Scopes = tt.scopes + } + + req, _ := http.NewRequest("GET", databricksServer.URL+"/api/test", nil) + err := cfg.Authenticate(req) + if err != nil { + t.Fatalf("Authenticate(): got error %v, want none", err) + } + + if got := req.Header.Get("Authorization"); got != "Bearer databricks-access-token" { + t.Errorf("Authorization header: got %q, want %q", got, "Bearer databricks-access-token") + } + if !githubTokenCalled { + t.Error("GitHub token endpoint was not called") + } + if !tokenExchangeCalled { + t.Error("Token exchange endpoint was not called") + } + }) + } +} diff --git a/config/experimental/auth/oidc/tokensource.go b/config/experimental/auth/oidc/tokensource.go index 25eb1b6f4..a13c4b784 100644 --- a/config/experimental/auth/oidc/tokensource.go +++ b/config/experimental/auth/oidc/tokensource.go @@ -39,6 +39,9 @@ type DatabricksOIDCTokenSourceConfig struct { // IDTokenSource returns the IDToken to be used for the token exchange. IDTokenSource IDTokenSource + + // Scopes is the list of OAuth scopes to request. + Scopes []string } // NewDatabricksOIDCTokenSource returns a new Databricks OIDC TokenSource. @@ -81,7 +84,7 @@ func (w *databricksOIDCTokenSource) Token(ctx context.Context) (*oauth2.Token, e ClientID: w.cfg.ClientID, AuthStyle: oauth2.AuthStyleInParams, TokenURL: endpoints.TokenEndpoint, - Scopes: []string{"all-apis"}, + Scopes: w.cfg.Scopes, EndpointParams: url.Values{ "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, "subject_token": {idToken.Value}, diff --git a/config/experimental/auth/oidc/tokensource_test.go b/config/experimental/auth/oidc/tokensource_test.go index 3410978c6..78b9e1896 100644 --- a/config/experimental/auth/oidc/tokensource_test.go +++ b/config/experimental/auth/oidc/tokensource_test.go @@ -275,6 +275,7 @@ func TestDatabricksOidcTokenSource(t *testing.T) { ClientID: tc.clientID, AccountID: tc.accountID, Host: tc.host, + Scopes: []string{"all-apis"}, TokenEndpointProvider: tc.oidcEndpointProvider, Audience: tc.tokenAudience, IDTokenSource: IDTokenSourceFn(func(ctx context.Context, aud string) (*IDToken, error) { @@ -319,3 +320,134 @@ func TestDatabricksOidcTokenSource(t *testing.T) { }) } } + +func TestWIF_Scopes(t *testing.T) { + tests := []struct { + name string + clientID string + accountID string + host string + audience string + scopes []string + tokenEndpoint string + expectedClientID string + expectedScope string + expectedAccessToken string + }{ + { + name: "single scope", + clientID: "client-id", + host: "http://host.com", + audience: "token-audience", + scopes: []string{"dashboards"}, + tokenEndpoint: "https://host.com/oidc/v1/token", + expectedClientID: "client-id", + expectedScope: "dashboards", + expectedAccessToken: "test-token", + }, + { + name: "multiple scopes sorted", + clientID: "client-id", + host: "http://host.com", + audience: "token-audience", + scopes: []string{"files", "jobs", "mlflow"}, + tokenEndpoint: "https://host.com/oidc/v1/token", + expectedClientID: "client-id", + expectedScope: "files jobs mlflow", + expectedAccessToken: "test-token", + }, + { + name: "workspace-level WIF", + clientID: "client-id", + host: "https://my-workspace.cloud.databricks.com", + audience: "workspace-audience", + scopes: []string{"genie"}, + tokenEndpoint: "https://my-workspace.cloud.databricks.com/oidc/v1/token", + expectedClientID: "client-id", + expectedScope: "genie", + expectedAccessToken: "workspace-token", + }, + { + name: "account-level WIF", + clientID: "client-id", + accountID: "my-account", + host: "https://accounts.cloud.databricks.com", + audience: "account-audience", + scopes: []string{"files", "iam"}, + tokenEndpoint: "https://accounts.cloud.databricks.com/oidc/accounts/my-account/v1/token", + expectedClientID: "client-id", + expectedScope: "files iam", + expectedAccessToken: "account-token", + }, + { + name: "account-wide token federation (no ClientID)", + clientID: "", + host: "http://host.com", + audience: "token-audience", + scopes: []string{"workspaces"}, + tokenEndpoint: "https://host.com/oidc/v1/token", + expectedClientID: "", + expectedScope: "workspaces", + expectedAccessToken: "account-wide-token", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := DatabricksOIDCTokenSourceConfig{ + ClientID: tt.clientID, + AccountID: tt.accountID, + Host: tt.host, + TokenEndpointProvider: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { + return &u2m.OAuthAuthorizationServer{ + TokenEndpoint: tt.tokenEndpoint, + }, nil + }, + Audience: tt.audience, + IDTokenSource: IDTokenSourceFn(func(ctx context.Context, aud string) (*IDToken, error) { + return &IDToken{Value: "id-token"}, nil + }), + Scopes: tt.scopes, + } + + ts := NewDatabricksOIDCTokenSource(cfg) + + expectedRequest := url.Values{ + "scope": {tt.expectedScope}, + "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, + "subject_token": {"id-token"}, + "grant_type": {"urn:ietf:params:oauth:grant-type:token-exchange"}, + } + if tt.expectedClientID != "" { + expectedRequest["client_id"] = []string{tt.expectedClientID} + } + + endpointURL, _ := url.Parse(tt.tokenEndpoint) + endpointPath := "POST " + endpointURL.Path + + ctx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{ + Transport: fixtures.MappingTransport{ + endpointPath: { + Status: http.StatusOK, + ExpectedHeaders: map[string]string{ + "Content-Type": "application/x-www-form-urlencoded", + }, + ExpectedRequest: expectedRequest, + Response: map[string]string{ + "token_type": "Bearer", + "access_token": tt.expectedAccessToken, + }, + }, + }, + }) + + token, err := ts.Token(ctx) + if err != nil { + t.Fatalf("Token(ctx): got error %q, want none", err) + } + if token.AccessToken != tt.expectedAccessToken { + t.Errorf("Token(ctx): got access token %q, want %q", token.AccessToken, tt.expectedAccessToken) + } + }) + } +} From 7f8f1b5635902b9803ec1dd1905e1daf13463413 Mon Sep 17 00:00:00 2001 From: tejas-kochar Date: Sun, 4 Jan 2026 13:49:38 +0000 Subject: [PATCH 2/7] improve comments --- config/auth_default_test.go | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/config/auth_default_test.go b/config/auth_default_test.go index adbf81d8e..15215ef4a 100644 --- a/config/auth_default_test.go +++ b/config/auth_default_test.go @@ -76,7 +76,7 @@ func TestGithubOIDC_Scopes(t *testing.T) { githubTokenCalled := false tokenExchangeCalled := false - // Simulates the GitHub Actions OIDC token endpoint. + // Mock GitHub server to verify the SDK requests an OIDC token during auth flow. githubServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { githubTokenCalled = true w.Header().Set("Content-Type", "application/json") @@ -84,8 +84,7 @@ func TestGithubOIDC_Scopes(t *testing.T) { })) defer githubServer.Close() - // Simulates a Databricks workspace. - // Asserts whether the right scopes are passed to the token exchange endpoint. + // Mock Databricks server to verify the SDK passes the correct scopes during token exchange. var databricksServer *httptest.Server databricksServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { @@ -101,7 +100,6 @@ func TestGithubOIDC_Scopes(t *testing.T) { if err := r.ParseForm(); err != nil { t.Fatalf("Failed to parse form: %v", err) } - // Verify scope is passed correctly to token exchange. if got := r.Form.Get("scope"); got != tt.expectedScope { t.Errorf("scope: got %q, want %q", got, tt.expectedScope) } From 40d201d766182fe8d108eea6210e77696dac1dde Mon Sep 17 00:00:00 2001 From: tejas-kochar Date: Mon, 5 Jan 2026 08:26:00 +0000 Subject: [PATCH 3/7] simplify tests --- config/auth_default_test.go | 59 ++++++----- .../auth/oidc/tokensource_test.go | 98 ++++--------------- 2 files changed, 48 insertions(+), 109 deletions(-) diff --git a/config/auth_default_test.go b/config/auth_default_test.go index 15215ef4a..d4ef67e89 100644 --- a/config/auth_default_test.go +++ b/config/auth_default_test.go @@ -55,36 +55,42 @@ func TestDefaultCredentials_Configure(t *testing.T) { func TestGithubOIDC_Scopes(t *testing.T) { tests := []struct { - name string - scopes []string - expectedScope string + name string + scopes []string + want string }{ { - name: "default scopes", - scopes: nil, - expectedScope: "all-apis", + name: "nil scopes uses default", + scopes: nil, + want: "all-apis", }, { - name: "custom scopes", - scopes: []string{"clusters", "jobs"}, - expectedScope: "clusters jobs", + name: "empty scopes uses default", + scopes: []string{}, + want: "all-apis", + }, + { + name: "single scope", + scopes: []string{"clusters"}, + want: "clusters", + }, + { + name: "multiple scopes are sorted", + scopes: []string{"jobs", "clusters", "files:read"}, + want: "clusters files:read jobs", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - githubTokenCalled := false - tokenExchangeCalled := false - - // Mock GitHub server to verify the SDK requests an OIDC token during auth flow. + // Mock GitHub server for OIDC token requests. githubServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - githubTokenCalled = true w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]string{"value": "github-id-token"}) })) defer githubServer.Close() - // Mock Databricks server to verify the SDK passes the correct scopes during token exchange. + // Mock Databricks server to verify the SDK passes the correct scopes. var databricksServer *httptest.Server databricksServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { @@ -96,12 +102,12 @@ func TestGithubOIDC_Scopes(t *testing.T) { }) case "/oidc/v1/token": - tokenExchangeCalled = true if err := r.ParseForm(); err != nil { t.Fatalf("Failed to parse form: %v", err) } - if got := r.Form.Get("scope"); got != tt.expectedScope { - t.Errorf("scope: got %q, want %q", got, tt.expectedScope) + // The scope assertion: verifies the SDK sends the correct scope parameter. + if got := r.Form.Get("scope"); got != tt.want { + t.Errorf("scope: got %q, want %q", got, tt.want) } w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(map[string]interface{}{ @@ -124,25 +130,16 @@ func TestGithubOIDC_Scopes(t *testing.T) { ActionsIDTokenRequestToken: "github-request-token", TokenAudience: "databricks-test-audience", AuthType: "github-oidc", - } - if tt.scopes != nil { - cfg.Scopes = tt.scopes + Scopes: tt.scopes, } req, _ := http.NewRequest("GET", databricksServer.URL+"/api/test", nil) err := cfg.Authenticate(req) if err != nil { - t.Fatalf("Authenticate(): got error %v, want none", err) - } - - if got := req.Header.Get("Authorization"); got != "Bearer databricks-access-token" { - t.Errorf("Authorization header: got %q, want %q", got, "Bearer databricks-access-token") - } - if !githubTokenCalled { - t.Error("GitHub token endpoint was not called") + t.Fatalf("Authenticate(): unexpected error: %v", err) } - if !tokenExchangeCalled { - t.Error("Token exchange endpoint was not called") + if !strings.HasPrefix(req.Header.Get("Authorization"), "Bearer ") { + t.Errorf("Authorization header missing Bearer prefix: got %q", req.Header.Get("Authorization")) } }) } diff --git a/config/experimental/auth/oidc/tokensource_test.go b/config/experimental/auth/oidc/tokensource_test.go index 78b9e1896..6df6f7a62 100644 --- a/config/experimental/auth/oidc/tokensource_test.go +++ b/config/experimental/auth/oidc/tokensource_test.go @@ -323,87 +323,33 @@ func TestDatabricksOidcTokenSource(t *testing.T) { func TestWIF_Scopes(t *testing.T) { tests := []struct { - name string - clientID string - accountID string - host string - audience string - scopes []string - tokenEndpoint string - expectedClientID string - expectedScope string - expectedAccessToken string + name string + scopes []string + want string }{ { - name: "single scope", - clientID: "client-id", - host: "http://host.com", - audience: "token-audience", - scopes: []string{"dashboards"}, - tokenEndpoint: "https://host.com/oidc/v1/token", - expectedClientID: "client-id", - expectedScope: "dashboards", - expectedAccessToken: "test-token", + name: "single scope", + scopes: []string{"dashboards"}, + want: "dashboards", }, { - name: "multiple scopes sorted", - clientID: "client-id", - host: "http://host.com", - audience: "token-audience", - scopes: []string{"files", "jobs", "mlflow"}, - tokenEndpoint: "https://host.com/oidc/v1/token", - expectedClientID: "client-id", - expectedScope: "files jobs mlflow", - expectedAccessToken: "test-token", - }, - { - name: "workspace-level WIF", - clientID: "client-id", - host: "https://my-workspace.cloud.databricks.com", - audience: "workspace-audience", - scopes: []string{"genie"}, - tokenEndpoint: "https://my-workspace.cloud.databricks.com/oidc/v1/token", - expectedClientID: "client-id", - expectedScope: "genie", - expectedAccessToken: "workspace-token", - }, - { - name: "account-level WIF", - clientID: "client-id", - accountID: "my-account", - host: "https://accounts.cloud.databricks.com", - audience: "account-audience", - scopes: []string{"files", "iam"}, - tokenEndpoint: "https://accounts.cloud.databricks.com/oidc/accounts/my-account/v1/token", - expectedClientID: "client-id", - expectedScope: "files iam", - expectedAccessToken: "account-token", - }, - { - name: "account-wide token federation (no ClientID)", - clientID: "", - host: "http://host.com", - audience: "token-audience", - scopes: []string{"workspaces"}, - tokenEndpoint: "https://host.com/oidc/v1/token", - expectedClientID: "", - expectedScope: "workspaces", - expectedAccessToken: "account-wide-token", + name: "multiple scopes", + scopes: []string{"jobs", "files:read", "mlflow"}, + want: "jobs files:read mlflow", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { cfg := DatabricksOIDCTokenSourceConfig{ - ClientID: tt.clientID, - AccountID: tt.accountID, - Host: tt.host, + ClientID: "client-id", + Host: "http://host.com", TokenEndpointProvider: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { return &u2m.OAuthAuthorizationServer{ - TokenEndpoint: tt.tokenEndpoint, + TokenEndpoint: "https://host.com/oidc/v1/token", }, nil }, - Audience: tt.audience, + Audience: "token-audience", IDTokenSource: IDTokenSourceFn(func(ctx context.Context, aud string) (*IDToken, error) { return &IDToken{Value: "id-token"}, nil }), @@ -412,22 +358,18 @@ func TestWIF_Scopes(t *testing.T) { ts := NewDatabricksOIDCTokenSource(cfg) + // The scope assertion: verifies the token source sends the correct scope parameter. expectedRequest := url.Values{ - "scope": {tt.expectedScope}, + "client_id": {"client-id"}, + "scope": {tt.want}, "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, "subject_token": {"id-token"}, "grant_type": {"urn:ietf:params:oauth:grant-type:token-exchange"}, } - if tt.expectedClientID != "" { - expectedRequest["client_id"] = []string{tt.expectedClientID} - } - - endpointURL, _ := url.Parse(tt.tokenEndpoint) - endpointPath := "POST " + endpointURL.Path ctx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{ Transport: fixtures.MappingTransport{ - endpointPath: { + "POST /oidc/v1/token": { Status: http.StatusOK, ExpectedHeaders: map[string]string{ "Content-Type": "application/x-www-form-urlencoded", @@ -435,7 +377,7 @@ func TestWIF_Scopes(t *testing.T) { ExpectedRequest: expectedRequest, Response: map[string]string{ "token_type": "Bearer", - "access_token": tt.expectedAccessToken, + "access_token": "test-token", }, }, }, @@ -445,8 +387,8 @@ func TestWIF_Scopes(t *testing.T) { if err != nil { t.Fatalf("Token(ctx): got error %q, want none", err) } - if token.AccessToken != tt.expectedAccessToken { - t.Errorf("Token(ctx): got access token %q, want %q", token.AccessToken, tt.expectedAccessToken) + if !strings.HasPrefix(token.AccessToken, "test-") { + t.Errorf("Token(ctx): got unexpected access token %q", token.AccessToken) } }) } From bf6ebead885aa1207c1d58463bbf452867f8c68d Mon Sep 17 00:00:00 2001 From: tejas-kochar Date: Wed, 7 Jan 2026 08:10:51 +0000 Subject: [PATCH 4/7] address comments --- config/auth_default_test.go | 12 ++++--- config/experimental/auth/oidc/tokensource.go | 8 ++++- .../auth/oidc/tokensource_test.go | 34 ++++++++++++++----- 3 files changed, 40 insertions(+), 14 deletions(-) diff --git a/config/auth_default_test.go b/config/auth_default_test.go index d4ef67e89..f5e714c20 100644 --- a/config/auth_default_test.go +++ b/config/auth_default_test.go @@ -133,13 +133,17 @@ func TestGithubOIDC_Scopes(t *testing.T) { Scopes: tt.scopes, } - req, _ := http.NewRequest("GET", databricksServer.URL+"/api/test", nil) - err := cfg.Authenticate(req) + req, err := http.NewRequest("GET", databricksServer.URL+"/api/test", nil) + if err != nil { + t.Fatalf("http.NewRequest(): unexpected error: %v", err) + } + err = cfg.Authenticate(req) if err != nil { t.Fatalf("Authenticate(): unexpected error: %v", err) } - if !strings.HasPrefix(req.Header.Get("Authorization"), "Bearer ") { - t.Errorf("Authorization header missing Bearer prefix: got %q", req.Header.Get("Authorization")) + wantAuthHeader := "Bearer databricks-access-token" + if got := req.Header.Get("Authorization"); got != wantAuthHeader { + t.Errorf("Authorization header: got %q, want %q", got, wantAuthHeader) } }) } diff --git a/config/experimental/auth/oidc/tokensource.go b/config/experimental/auth/oidc/tokensource.go index a13c4b784..d0e5621d2 100644 --- a/config/experimental/auth/oidc/tokensource.go +++ b/config/experimental/auth/oidc/tokensource.go @@ -80,11 +80,17 @@ func (w *databricksOIDCTokenSource) Token(ctx context.Context) (*oauth2.Token, e return nil, err } + // This nil check is to ensure backwards compatibility for users implementing their own + // OIDC token source. + scopes := w.cfg.Scopes + if len(scopes) == 0 { + scopes = []string{"all-apis"} + } c := &clientcredentials.Config{ ClientID: w.cfg.ClientID, AuthStyle: oauth2.AuthStyleInParams, TokenURL: endpoints.TokenEndpoint, - Scopes: w.cfg.Scopes, + Scopes: scopes, EndpointParams: url.Values{ "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, "subject_token": {idToken.Value}, diff --git a/config/experimental/auth/oidc/tokensource_test.go b/config/experimental/auth/oidc/tokensource_test.go index 6df6f7a62..d7c77b0d9 100644 --- a/config/experimental/auth/oidc/tokensource_test.go +++ b/config/experimental/auth/oidc/tokensource_test.go @@ -275,7 +275,6 @@ func TestDatabricksOidcTokenSource(t *testing.T) { ClientID: tc.clientID, AccountID: tc.accountID, Host: tc.host, - Scopes: []string{"all-apis"}, TokenEndpointProvider: tc.oidcEndpointProvider, Audience: tc.tokenAudience, IDTokenSource: IDTokenSourceFn(func(ctx context.Context, aud string) (*IDToken, error) { @@ -322,11 +321,28 @@ func TestDatabricksOidcTokenSource(t *testing.T) { } func TestWIF_Scopes(t *testing.T) { + const ( + testClientID = "test-client-id" + testIDToken = "test-id-token" + testAccessToken = "test-access-token" + testTokenURL = "https://host.com/oidc/v1/token" + ) + tests := []struct { name string scopes []string want string }{ + { + name: "nil scopes uses default", + scopes: nil, + want: "all-apis", + }, + { + name: "empty scopes uses default", + scopes: []string{}, + want: "all-apis", + }, { name: "single scope", scopes: []string{"dashboards"}, @@ -342,16 +358,16 @@ func TestWIF_Scopes(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { cfg := DatabricksOIDCTokenSourceConfig{ - ClientID: "client-id", + ClientID: testClientID, Host: "http://host.com", TokenEndpointProvider: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { return &u2m.OAuthAuthorizationServer{ - TokenEndpoint: "https://host.com/oidc/v1/token", + TokenEndpoint: testTokenURL, }, nil }, Audience: "token-audience", IDTokenSource: IDTokenSourceFn(func(ctx context.Context, aud string) (*IDToken, error) { - return &IDToken{Value: "id-token"}, nil + return &IDToken{Value: testIDToken}, nil }), Scopes: tt.scopes, } @@ -360,10 +376,10 @@ func TestWIF_Scopes(t *testing.T) { // The scope assertion: verifies the token source sends the correct scope parameter. expectedRequest := url.Values{ - "client_id": {"client-id"}, + "client_id": {testClientID}, "scope": {tt.want}, "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, - "subject_token": {"id-token"}, + "subject_token": {testIDToken}, "grant_type": {"urn:ietf:params:oauth:grant-type:token-exchange"}, } @@ -377,7 +393,7 @@ func TestWIF_Scopes(t *testing.T) { ExpectedRequest: expectedRequest, Response: map[string]string{ "token_type": "Bearer", - "access_token": "test-token", + "access_token": testAccessToken, }, }, }, @@ -387,8 +403,8 @@ func TestWIF_Scopes(t *testing.T) { if err != nil { t.Fatalf("Token(ctx): got error %q, want none", err) } - if !strings.HasPrefix(token.AccessToken, "test-") { - t.Errorf("Token(ctx): got unexpected access token %q", token.AccessToken) + if token.AccessToken != testAccessToken { + t.Errorf("Token(ctx): got access token %q, want %q", token.AccessToken, testAccessToken) } }) } From a77d3050ab507a47198de8f3cb523c7efd4c6265 Mon Sep 17 00:00:00 2001 From: tejas-kochar Date: Wed, 7 Jan 2026 08:22:26 +0000 Subject: [PATCH 5/7] address comments --- config/auth_default_test.go | 6 ++++-- config/experimental/auth/oidc/tokensource_test.go | 9 +++++---- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/config/auth_default_test.go b/config/auth_default_test.go index f5e714c20..ede5dfb3a 100644 --- a/config/auth_default_test.go +++ b/config/auth_default_test.go @@ -54,6 +54,8 @@ func TestDefaultCredentials_Configure(t *testing.T) { } func TestGithubOIDC_Scopes(t *testing.T) { + const oidcTokenPath = "/oidc/v1/token" + tests := []struct { name string scopes []string @@ -98,10 +100,10 @@ func TestGithubOIDC_Scopes(t *testing.T) { w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(u2m.OAuthAuthorizationServer{ AuthorizationEndpoint: "https://host.com/oidc/v1/authorize", - TokenEndpoint: databricksServer.URL + "/oidc/v1/token", + TokenEndpoint: databricksServer.URL + oidcTokenPath, }) - case "/oidc/v1/token": + case oidcTokenPath: if err := r.ParseForm(); err != nil { t.Fatalf("Failed to parse form: %v", err) } diff --git a/config/experimental/auth/oidc/tokensource_test.go b/config/experimental/auth/oidc/tokensource_test.go index d7c77b0d9..03fa3dcbc 100644 --- a/config/experimental/auth/oidc/tokensource_test.go +++ b/config/experimental/auth/oidc/tokensource_test.go @@ -325,7 +325,8 @@ func TestWIF_Scopes(t *testing.T) { testClientID = "test-client-id" testIDToken = "test-id-token" testAccessToken = "test-access-token" - testTokenURL = "https://host.com/oidc/v1/token" + testTokenPath = "/oidc/v1/token" + testHost = "https://host.com" ) tests := []struct { @@ -359,10 +360,10 @@ func TestWIF_Scopes(t *testing.T) { t.Run(tt.name, func(t *testing.T) { cfg := DatabricksOIDCTokenSourceConfig{ ClientID: testClientID, - Host: "http://host.com", + Host: testHost, TokenEndpointProvider: func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) { return &u2m.OAuthAuthorizationServer{ - TokenEndpoint: testTokenURL, + TokenEndpoint: testHost + testTokenPath, }, nil }, Audience: "token-audience", @@ -385,7 +386,7 @@ func TestWIF_Scopes(t *testing.T) { ctx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{ Transport: fixtures.MappingTransport{ - "POST /oidc/v1/token": { + "POST " + testTokenPath: { Status: http.StatusOK, ExpectedHeaders: map[string]string{ "Content-Type": "application/x-www-form-urlencoded", From c26f841d909923da86fbc2a227a4d6a5e6bb53af Mon Sep 17 00:00:00 2001 From: tejas-kochar Date: Thu, 8 Jan 2026 15:10:44 +0000 Subject: [PATCH 6/7] Use getter for scopes that sets default --- config/auth_default.go | 2 +- config/experimental/auth/oidc/tokensource.go | 20 +++++++++++++------ .../auth/oidc/tokensource_test.go | 2 +- 3 files changed, 16 insertions(+), 8 deletions(-) diff --git a/config/auth_default.go b/config/auth_default.go index 6a7c38369..8ec3ce119 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -145,11 +145,11 @@ func oidcStrategy(cfg *Config, name string, ts oidc.IDTokenSource) CredentialsSt TokenEndpointProvider: cfg.getOidcEndpoints, Audience: cfg.TokenAudience, IDTokenSource: ts, - Scopes: cfg.GetScopes(), } if cfg.HostType() != WorkspaceHost { oidcConfig.AccountID = cfg.AccountID } + oidcConfig.SetScopes(cfg.GetScopes()) tokenSource := oidc.NewDatabricksOIDCTokenSource(oidcConfig) return NewTokenSourceStrategy(name, tokenSource) } diff --git a/config/experimental/auth/oidc/tokensource.go b/config/experimental/auth/oidc/tokensource.go index d0e5621d2..8e0d19385 100644 --- a/config/experimental/auth/oidc/tokensource.go +++ b/config/experimental/auth/oidc/tokensource.go @@ -40,8 +40,19 @@ type DatabricksOIDCTokenSourceConfig struct { // IDTokenSource returns the IDToken to be used for the token exchange. IDTokenSource IDTokenSource - // Scopes is the list of OAuth scopes to request. - Scopes []string + // scopes is the list of OAuth scopes to request. + scopes []string +} + +func (c *DatabricksOIDCTokenSourceConfig) GetScopes() []string { + if len(c.scopes) == 0 { + return []string{"all-apis"} + } + return c.scopes +} + +func (c *DatabricksOIDCTokenSourceConfig) SetScopes(scopes []string) { + c.scopes = scopes } // NewDatabricksOIDCTokenSource returns a new Databricks OIDC TokenSource. @@ -82,10 +93,7 @@ func (w *databricksOIDCTokenSource) Token(ctx context.Context) (*oauth2.Token, e // This nil check is to ensure backwards compatibility for users implementing their own // OIDC token source. - scopes := w.cfg.Scopes - if len(scopes) == 0 { - scopes = []string{"all-apis"} - } + scopes := w.cfg.GetScopes() c := &clientcredentials.Config{ ClientID: w.cfg.ClientID, AuthStyle: oauth2.AuthStyleInParams, diff --git a/config/experimental/auth/oidc/tokensource_test.go b/config/experimental/auth/oidc/tokensource_test.go index 03fa3dcbc..604ea8ca9 100644 --- a/config/experimental/auth/oidc/tokensource_test.go +++ b/config/experimental/auth/oidc/tokensource_test.go @@ -370,7 +370,7 @@ func TestWIF_Scopes(t *testing.T) { IDTokenSource: IDTokenSourceFn(func(ctx context.Context, aud string) (*IDToken, error) { return &IDToken{Value: testIDToken}, nil }), - Scopes: tt.scopes, + scopes: tt.scopes, } ts := NewDatabricksOIDCTokenSource(cfg) From 74189d788ccd918f0142cdb5e156028b943a0621 Mon Sep 17 00:00:00 2001 From: tejas-kochar Date: Mon, 12 Jan 2026 07:14:54 +0000 Subject: [PATCH 7/7] add doccoments --- config/experimental/auth/oidc/tokensource.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/config/experimental/auth/oidc/tokensource.go b/config/experimental/auth/oidc/tokensource.go index 8e0d19385..5d5e1a114 100644 --- a/config/experimental/auth/oidc/tokensource.go +++ b/config/experimental/auth/oidc/tokensource.go @@ -44,6 +44,8 @@ type DatabricksOIDCTokenSourceConfig struct { scopes []string } +// GetScopes returns the OAuth scopes to request. If no scopes have been set, +// it returns the default scope "all-apis". func (c *DatabricksOIDCTokenSourceConfig) GetScopes() []string { if len(c.scopes) == 0 { return []string{"all-apis"} @@ -51,6 +53,7 @@ func (c *DatabricksOIDCTokenSourceConfig) GetScopes() []string { return c.scopes } +// SetScopes sets the OAuth scopes to request. func (c *DatabricksOIDCTokenSourceConfig) SetScopes(scopes []string) { c.scopes = scopes }