-
Notifications
You must be signed in to change notification settings - Fork 50
Implement Databricks OIDC as Token Source #1204
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from 22 commits
Commits
Show all changes
36 commits
Select commit
Hold shift + click to select a range
0b11b82
WIP
hectorcast-db ba4d103
Merge branch 'main' into databricks-auth
hectorcast-db 200f16e
Test
hectorcast-db 452e56f
CLI test
hectorcast-db 780ddc7
Other stuff
hectorcast-db 9494d40
refresh support
hectorcast-db e94811a
Test
hectorcast-db bac4c39
PR Comments
hectorcast-db bbc8f0c
Merge branch 'main' into databricks-auth
hectorcast-db 83619c2
Default values
hectorcast-db 66d8057
Merge branch 'main' into databricks-auth2
hectorcast-db 9b7f636
WIP
hectorcast-db 878c349
WIP
hectorcast-db 8738b46
Remove conf + test
hectorcast-db 506a146
More tests
hectorcast-db b8d28dd
almost
hectorcast-db b2c2ce4
last
hectorcast-db 0bf2d2f
fixes
hectorcast-db 3d39805
Rename files
hectorcast-db 35df156
Enabe
hectorcast-db c5652f3
Remove config from tests
hectorcast-db 635752b
comments
hectorcast-db 78db765
PR comments
hectorcast-db 47f7a2d
Undo endpoint supply
hectorcast-db 10e5615
Cleanup
hectorcast-db d5cb4cc
Merge branch 'main' into databricks-auth2
hectorcast-db 2431d67
PR comments
hectorcast-db a6e1769
Name
hectorcast-db 72d715e
Revert mistake
hectorcast-db 7f32a67
Ensure credentials order
hectorcast-db b09f4e8
More comments
hectorcast-db 3e040c9
Rename
hectorcast-db 4f59d7d
Remove duplicate code
hectorcast-db fa9fd25
Fix comment
hectorcast-db a7c2ab0
Merge branch 'main' into databricks-auth2
hectorcast-db 41c2e7b
Merge branch 'main' into databricks-auth2
hectorcast-db File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,113 @@ | ||
| package config | ||
|
|
||
| import ( | ||
| "context" | ||
| "errors" | ||
| "net/url" | ||
|
|
||
| "github.com/databricks/databricks-sdk-go/credentials/u2m" | ||
| "github.com/databricks/databricks-sdk-go/logger" | ||
| "golang.org/x/oauth2" | ||
| "golang.org/x/oauth2/clientcredentials" | ||
| ) | ||
|
|
||
| // Constructs all Workload Identity Federation Credentials Strategies | ||
| func OidcTokenCredentialStrategies(cfg *Config) []CredentialsStrategy { | ||
| providers := map[string]TokenProvider{ | ||
| "github-oidc": &GithubProvider{ | ||
| actionsIDTokenRequestURL: cfg.ActionsIDTokenRequestURL, | ||
| actionsIDTokenRequestToken: cfg.ActionsIDTokenRequestToken, | ||
| refreshClient: cfg.refreshClient, | ||
| }, | ||
| // Add new providers at the end of the list | ||
| } | ||
| strategies := []CredentialsStrategy{} | ||
| for name, provider := range providers { | ||
| strategies = append(strategies, newOidcTokenStrategy(name, cfg, provider)) | ||
| } | ||
| return strategies | ||
| } | ||
|
|
||
| func newOidcTokenStrategy( | ||
| name string, | ||
| cfg *Config, | ||
| tokenProvider TokenProvider, | ||
| ) CredentialsStrategy { | ||
| accountId := "" | ||
| // NOTE: cfg.AccountID can be present even if the client is not an account client. | ||
| if cfg.IsAccountClient() { | ||
| accountId = cfg.AccountID | ||
| } | ||
| oidcTokenExchange := &oidcTokenExchange{ | ||
| clientID: cfg.ClientID, | ||
| accountID: accountId, | ||
| host: cfg.Host, | ||
| tokenEndpointProvider: cfg.getOidcEndpoints, | ||
| audience: cfg.TokenAudience, | ||
| tokenProvider: tokenProvider, | ||
| } | ||
| return NewTokenSourceStrategy(name, oidcTokenExchange) | ||
| } | ||
|
|
||
| // oidcTokenExchange is a auth.TokenSource which exchanges a token using | ||
| // Workload Identity Federation. | ||
| type oidcTokenExchange struct { | ||
| clientID string | ||
| accountID string | ||
| host string | ||
| tokenEndpointProvider func(ctx context.Context) (*u2m.OAuthAuthorizationServer, error) | ||
| audience string | ||
| tokenProvider TokenProvider | ||
| } | ||
|
|
||
| // Token implements [TokenSource.Token] | ||
| func (w *oidcTokenExchange) Token(ctx context.Context) (*oauth2.Token, error) { | ||
| if w.clientID == "" { | ||
| logger.Debugf(ctx, "Missing ClientID") | ||
| return nil, errors.New("missing ClientID") | ||
| } | ||
| if w.host == "" { | ||
| logger.Debugf(ctx, "Missing Host") | ||
| return nil, errors.New("missing Host") | ||
| } | ||
| audience, err := w.determineAudience(ctx) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| idToken, err := w.tokenProvider.IDToken(ctx, audience) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| endpoints, err := w.tokenEndpointProvider(ctx) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| c := &clientcredentials.Config{ | ||
| ClientID: w.clientID, | ||
| AuthStyle: oauth2.AuthStyleInParams, | ||
| TokenURL: endpoints.TokenEndpoint, | ||
| Scopes: []string{"all-apis"}, | ||
| EndpointParams: url.Values{ | ||
| "subject_token_type": {"urn:ietf:params:oauth:token-type:jwt"}, | ||
| "subject_token": {idToken.Value}, | ||
| "grant_type": {"urn:ietf:params:oauth:grant-type:token-exchange"}, | ||
| }, | ||
| } | ||
| return c.Token(ctx) | ||
| } | ||
|
|
||
| func (w *oidcTokenExchange) determineAudience(ctx context.Context) (string, error) { | ||
| if w.audience != "" { | ||
| return w.audience, nil | ||
| } | ||
| // For Databricks Accounts, the account id is the default audience. | ||
| if w.accountID != "" { | ||
| return w.accountID, nil | ||
| } | ||
| endpoints, err := w.tokenEndpointProvider(ctx) | ||
| if err != nil { | ||
| return "", err | ||
| } | ||
| // For Databricks Workspaces, the auth endpoint is the default audience. | ||
| return endpoints.TokenEndpoint, nil | ||
| } |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.