Skip to content
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

[Feature] Implement U2M Authentication in the Go SDK #1108

Open
wants to merge 31 commits into
base: main
Choose a base branch
from

Conversation

mgyucht
Copy link
Contributor

@mgyucht mgyucht commented Jan 3, 2025

What changes are proposed in this pull request?

This PR moves logic about U2M OAuth login from the CLI to the Go SDK. This eliminates a cyclic dependency between the SDK and CLI for interacting with the OAuth token cache and enables U2M support directly in the Go SDK without need for the CLI to be installed.

Most of this code is carried over from the CLI, but I have made specific refactors to generalize it where needed.

Currently, the token cache key follows a specific structure:

  • https://<accounts-host>/oidc/accounts/<account-id> for account-based sessions
  • https://<workspace host> for workspace-based sessions
    This can be generalized to allow callers to cache tokens in other manners. For example, users may want to cache tokens per principal or per set of OAuth scopes for their own OAuth applications.

Additionally, products should be able to use an isolated token cache or a token cache backed by a different storage medium, like a database. This is simple to do by introducing a token cache interface. Implementers can define their own credential strategy and reuse the PersistentAuth type to handle the negotiation.

How is this tested?

Carried over tests from the CLI.

More tests are forthcoming, once a decision is made on this approach.

@mgyucht mgyucht temporarily deployed to test-trigger-is January 3, 2025 10:51 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 3, 2025 10:52 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 3, 2025 12:18 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 3, 2025 12:18 — with GitHub Actions Inactive
Copy link
Contributor

@renaudhartert-db renaudhartert-db left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR — always good to get the SDK more self-contained.

tokenCacheVersion = 1
)

var ErrNotConfigured = errors.New("databricks OAuth is not configured for this host")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Move in cache.go?

)

const (
// where the token cache is stored
Copy link
Contributor

@renaudhartert-db renaudhartert-db Jan 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I realize you didn't write these comments. Though, let's remain consistent by writing them as full sentences (with majuscule and period). Ref: https://go.dev/wiki/CodeReviewComments#comment-sentences

Comment on lines 24 to 25
// format versioning leaves some room for format improvement
tokenCacheVersion = 1
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe explain a little more what problem this is solving? It's not super clear from looking at this definition.

// https://datatracker.ietf.org/doc/html/rfc6749
type OAuthToken struct {
// The access token issued by the authorization server. This is the token that will be used to authenticate requests.
AccessToken string `json:"access_token" auth:",sensitive"`
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is auth:",sensitive"? I don't remember seeing this annotation before.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In some cases (like failed authentication), debug logs will include configuration parameters. Config attributes marked as sensitive are redacted in these logs. I don't think that also applies to this struct though. I'll take a closer look here.

@@ -22,6 +21,19 @@ type GetOAuthTokenRequest struct {
Assertion string `url:"assertion"`
}

// OAuthToken represents an OAuth token as defined by the OAuth 2.0 Authorization Framework.
// https://datatracker.ietf.org/doc/html/rfc6749
type OAuthToken struct {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use oauth2.Token instead?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oauth2.Token has an "Expiry" field, which is computed based on the "expires_in" field. This mimics the unexported tokenJSON type from that library.

Comment on lines 85 to 88
loc, err := c.location()
if err != nil {
return err
}
Copy link
Contributor

@renaudhartert-db renaudhartert-db Jan 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nit] What location was doing wasn't obvious to me reading this line. We could make that clearer with a better name. For example tokenCacheFilepath(). Though, I think it might just be simpler to inline the code since it results in the same number of line minus the location method.

Suggested change
loc, err := c.location()
if err != nil {
return err
}
home, err := os.UserHomeDir()
if err != nil {
return fmt.Errorf("could not access home dir: %w")
}
c.fileLocation = filepath.Join(home, tokenCacheFileName)

Comment on lines 30 to 32
// this implementation requires the calling code to do a machine-wide lock,
// otherwise the file might get corrupt.
type FileTokenCache struct {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we briefly explain what this cache does? Something like:

// FileTokenCache caches tokens in "~/.databricks/token-cache.json". FileTokenCache
// implements the TokenCache interface.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great, will use this description. Sorry, most of this is copied directly from the CLI without carefully reviewing for style or comments, but thanks for pointing these out, and I'm happy to correct them here.

Comment on lines 41 to 49
if errors.Is(err, fs.ErrNotExist) {
dir := filepath.Dir(c.fileLocation)
err = os.MkdirAll(dir, ownerExecReadWrite)
if err != nil {
return fmt.Errorf("mkdir: %w", err)
}
} else if err != nil {
return fmt.Errorf("load: %w", err)
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[optional] This is slightly more idiomatic to test for nil error first.

Suggested change
if errors.Is(err, fs.ErrNotExist) {
dir := filepath.Dir(c.fileLocation)
err = os.MkdirAll(dir, ownerExecReadWrite)
if err != nil {
return fmt.Errorf("mkdir: %w", err)
}
} else if err != nil {
return fmt.Errorf("load: %w", err)
}
if err != nil {
if !errors.Is(err, fs.ErrNotExist) {
return fmt.Errorf("load: %w", err)
}
dir := filepath.Dir(c.fileLocation)
if err := os.MkdirAll(dir, ownerExecReadWrite); err != nil {
return fmt.Errorf("mkdir: %w", err)
}
}

var ErrNotConfigured = errors.New("databricks OAuth is not configured for this host")

// this implementation requires the calling code to do a machine-wide lock,
// otherwise the file might get corrupt.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Naively, I would have expected the lock to be managed by this cache implementation as an internal detail. Having the lock managed by the client sounds a little bug prone — especially given that the client does not have a programmatic way to know in what file the cache is storing the tokens. Am I missing something?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good question. It needs to be at least locked in PersistentAuth: the Load() method will read from the cache and then, if the access token is expired, refresh it and update the cache, which replaces it entirely. That requires some form of user-wide mutual exclusion that lasts longer than an individual method call, thus the lock file. I don't see any issues with adding an in-memory mutex within FileTokenCache as well in case an application is directly managing and updating tokens, considering what that would be protecting against.

// Cache is the token cache to store and lookup tokens.
cache cache.TokenCache
// Locker is the lock to synchronize token cache access.
locker sync.Locker
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like locker is never locked. Am I missing something?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're not missing anything, I forgot to add it before, but I have it locally and will include it in my next update.

@mgyucht mgyucht temporarily deployed to test-trigger-is January 7, 2025 14:43 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 7, 2025 14:56 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 7, 2025 14:58 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 7, 2025 16:21 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 7, 2025 16:22 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 8, 2025 11:01 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 8, 2025 11:03 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 8, 2025 11:04 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 8, 2025 11:04 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 8, 2025 11:12 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 8, 2025 11:12 — with GitHub Actions Inactive
@mgyucht mgyucht temporarily deployed to test-trigger-is January 20, 2025 12:39 — with GitHub Actions Inactive
Copy link

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-go

Inputs:

  • PR number: 1108
  • Commit SHA: 356a7c9aefe20a753a62b4f4d7fa6654529d4feb

Checks will be approved automatically on success.

@mgyucht mgyucht temporarily deployed to test-trigger-is January 20, 2025 12:40 — with GitHub Actions Inactive
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants