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

mgyucht · 2025-01-03T10:51:45Z

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.

renaudhartert-db

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

renaudhartert-db · 2025-01-03T13:18:17Z

credentials/cache/file.go

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


Move in cache.go?

renaudhartert-db · 2025-01-03T13:19:41Z

credentials/cache/file.go

+)
+
+const (
+	// where the token cache is stored


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

renaudhartert-db · 2025-01-03T13:22:31Z

credentials/cache/file.go

+	// format versioning leaves some room for format improvement
+	tokenCacheVersion = 1


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

renaudhartert-db · 2025-01-03T13:28:26Z

httpclient/oauth_token.go

+// 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"`


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

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.

Good observation: this is only used for attributes of the Config object, to determine whether they are omitted when the config is printed out. I will remove this here.

renaudhartert-db · 2025-01-03T13:28:38Z

httpclient/oauth_token.go

@@ -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 {


Use oauth2.Token instead?

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

renaudhartert-db · 2025-01-03T14:50:55Z

credentials/cache/file.go

+	loc, err := c.location()
+	if err != nil {
+		return err
+	}


[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)

renaudhartert-db · 2025-01-03T14:54:58Z

credentials/cache/file.go

+// this implementation requires the calling code to do a machine-wide lock,
+// otherwise the file might get corrupt.
+type FileTokenCache struct {


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

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

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.

renaudhartert-db · 2025-01-03T14:59:42Z

credentials/cache/file.go

+	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)
+	}


[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)

}

}

renaudhartert-db · 2025-01-03T15:58:48Z

credentials/cache/file.go

+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.


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?

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.

renaudhartert-db · 2025-01-03T16:05:16Z

credentials/oauth/persistent_auth.go

+	// 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


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

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

renaudhartert-db · 2025-01-29T14:17:06Z

Added a few comments, specifically around the cache package.

My remaining concerns are about the new oauth package and I think it would take time to review it properly, taking the whole "auth" library in consideration.

If there's an easy way to provide a first implementation of u2m that does not rely on the oauth package, then I would recommend doing this first, and consider introducing the oauth package as a follow-up.

Discussed offline:

I misunderstood the scope of this package as being more general than u2m.
We agreed to make it clear that the package is specifically designed for u2m.

renaudhartert-db · 2025-01-30T17:20:33Z

credentials/u2m/cache/file.go

+	fileLocation string
+
+	// locker protects the token cache file from concurrent reads and writes.
+	locker sync.Locker


No need for an interface, let's use a mutex directly. The zero value of a mutex is a valid mutex so we do not even need to think about initialization anymore.

Suggested change

locker sync.Locker

locker sync.Mutex

renaudhartert-db · 2025-01-30T19:36:11Z

credentials/u2m/cache/file.go

+}
+
+// Store implements the TokenCache interface.
+func (c *fileTokenCache) Store(key string, t *oauth2.Token) error {


What's the intended behavior when storing a nil token? Should we delete it from the cache?

delete(key, f.Tokens)

Good question. I think that's a reasonable implementation.

renaudhartert-db · 2025-01-30T20:09:07Z

credentials/u2m/error.go

+// InvalidRefreshTokenError is returned from PersistentAuth's Load() method
+// if the access token has expired and the refresh token in the token cache
+// is invalid.
+type InvalidRefreshTokenError struct {
+	err error
+}
+
+func (e *InvalidRefreshTokenError) Error() string {
+	return e.err.Error()
+}
+
+func (e *InvalidRefreshTokenError) Unwrap() error {
+	return e.err
+}
+
+var _ error = &InvalidRefreshTokenError{}


I think this could just be this:

Suggested change

// InvalidRefreshTokenError is returned from PersistentAuth's Load() method

// if the access token has expired and the refresh token in the token cache

// is invalid.

type InvalidRefreshTokenError struct {

err error

}

func (e *InvalidRefreshTokenError) Error() string {

return e.err.Error()

}

func (e *InvalidRefreshTokenError) Unwrap() error {

return e.err

}

var _ error = &InvalidRefreshTokenError{}

// InvalidRefreshTokenError is returned from PersistentAuth's Load() method

// if the access token has expired and the refresh token in the token cache

// is invalid.

type InvalidRefreshTokenError struct {

error

}

renaudhartert-db · 2025-01-30T20:37:05Z

config/auth_u2m.go

+// a special error message that instructs the user how to reauthenticate.
+type u2mCredentials struct {
+	// auth is the persistent auth object that manages the token cache.
+	auth *u2m.PersistentAuth


Should this just be a token source?

Suggested change

auth *u2m.PersistentAuth

ts oauth2.TokenSource

renaudhartert-db · 2025-01-30T20:43:02Z

credentials/u2m/persistent_auth.go

+// callback server listens for the redirect from the identity provider and
+// exchanges the authorization code for an access token. It returns the OAuth2
+// token on success.
+func (a *PersistentAuth) Challenge() (*oauth2.Token, error) {


This method does not seem to be used. Also, its documentation does not look up to date. Is it necessary to implement some interface?

This is the core routine that handles the U2M OAuth login. There isn't a specific interface that it implements. It might be reasonable to consider this to be a token source as well

renaudhartert-db · 2025-01-30T20:43:34Z

credentials/u2m/persistent_auth.go

+	return nil
+}
+
+func (a *PersistentAuth) Close() error {


This method does not seem to be used publicly, can we make it private?

It will be used by the CLI and by any clients who want to use U2M but want to manage the token cache independently from the CLI.

renaudhartert-db · 2025-01-30T20:45:01Z

credentials/u2m/persistent_auth.go

+}
+
+// NewPersistentAuth creates a new PersistentAuth with the provided options.
+func NewPersistentAuth(ctx context.Context, opts ...PersistentAuthOption) (*PersistentAuth, error) {


Should this return a TokenSource so that the implementation is hidden? Ultimately, it does not look like the Challenge and Close methods are used outside of this package.

renaudhartert-db · 2025-01-30T20:50:25Z

config/auth_u2m_test.go

+			cfg: &Config{
+				Host: "https://myworkspace.cloud.databricks.com",
+			},
+			auth: must(


I would have naturally leaned toward mocking PersistentAuth rather than the OAuthClient — especially if PersitentAuth is treated as a TokenSource (see other comments in persistent_auth.go).

renaudhartert-db · 2025-01-30T20:51:44Z

credentials/u2m/client.go

+	// GetHttpClient returns an HTTP client for OAuth2 requests.
+	GetHttpClient(context.Context) *http.Client


This method bothers me. I'm wondering if we would not be better off passing the HTTPClient directly to PersistentAuth. I'm also wondering if this interface needs to be public (see comments about mocking PersistentAuth in auth_u2m_test.go).

We can split the OAuth endpoints and the HTTP client into separate parameters. However, because oauth2 depends on an http.Client, and to preserve the existing behavior, we do need to pass in the existing ApiClient somehow.

In an ideal world, we would have a pipelined API client specifically for OAuth that would be independent of the ApiClient package. We could actually just duplicate the code in order to break that dependency.

github-actions · 2025-02-14T15:59:14Z

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: 2f0ebbb9153a2c180fe9f5c034e61639f398ce72

Checks will be approved automatically on success.

mgyucht · 2025-01-14T09:49:51Z

config/auth_u2m.go

+	// Auth is the persistent auth object to use. If not specified, a new one will
+	// be created, using the default cache and locker.
+	Auth *oauth.PersistentAuth
+
+	// GetOAuthArg is a function that returns the OAuth argument to use for
+	// loading the OAuth session token. If not specified, the OAuth argument is
+	// determined by the account host and account ID or workspace host in the
+	// Config.
+	GetOAuthArg func(context.Context, *Config) (oauth.OAuthArgument, error)
+
+	// ErrorHandler controls the behavior of Configure() when loading the OAuth
+	// token fails. If not specified, any error will cause Configure() to return
+	// said error.
+	ErrorHandler func(context.Context, *Config, oauth.OAuthArgument, error) error


Let's make these unexported (ditto for other types introduced in this PR).

mgyucht · 2025-01-14T10:08:16Z

credentials/cache/file.go

+	// mu protects the token cache file from concurrent reads and writes.
+	mu *sync.Mutex


This should be the filemutex

mgyucht · 2025-01-14T10:08:31Z

credentials/oauth/persistent_auth.go

+	// 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


We remove this.

mgyucht · 2025-01-14T10:17:38Z

httpclient/oauth_token.go

@@ -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 {


Unexport, explicit reason why this exists

mgyucht · 2025-01-14T10:19:50Z

credentials/oauth/persistent_auth.go

+
+func (a *PersistentAuth) randomString(size int) string {
+	raw := make([]byte, size)
+	_, _ = rand.Read(raw)


mgyucht · 2025-01-28T16:17:15Z

config/auth_u2m.go

+	// loading the OAuth session token. If not specified, the OAuth argument is
+	// determined by the account host and account ID or workspace host in the
+	// Config.
+	getOAuthArg func(context.Context, *Config) (oauth.OAuthArgument, error)


This is how we will support scopes. We'll need a way to determine which key in the token cache the user is intending to use based on the config. It's intentionally unexported now to indicate that we don't really want to support it yet, but it would be the narrow place to make the change later.

mgyucht · 2025-02-14T13:00:18Z

credentials/u2m/cache/file.go

+}
+
+// Store implements the TokenCache interface.
+func (c *fileTokenCache) Store(key string, t *oauth2.Token) error {


Good question. I think that's a reasonable implementation.

mgyucht · 2025-02-14T13:02:57Z

credentials/u2m/persistent_auth.go

+// callback server listens for the redirect from the identity provider and
+// exchanges the authorization code for an access token. It returns the OAuth2
+// token on success.
+func (a *PersistentAuth) Challenge() (*oauth2.Token, error) {


This is the core routine that handles the U2M OAuth login. There isn't a specific interface that it implements. It might be reasonable to consider this to be a token source as well

mgyucht · 2025-02-14T16:00:48Z

credentials/u2m/persistent_auth.go

+	return nil
+}
+
+func (a *PersistentAuth) Close() error {


It will be used by the CLI and by any clients who want to use U2M but want to manage the token cache independently from the CLI.

mgyucht · 2025-02-14T16:05:07Z

credentials/u2m/client.go

+	// GetHttpClient returns an HTTP client for OAuth2 requests.
+	GetHttpClient(context.Context) *http.Client


We can split the OAuth endpoints and the HTTP client into separate parameters. However, because oauth2 depends on an http.Client, and to preserve the existing behavior, we do need to pass in the existing ApiClient somehow.

In an ideal world, we would have a pipelined API client specifically for OAuth that would be independent of the ApiClient package. We could actually just duplicate the code in order to break that dependency.

work

4a538e3

mgyucht temporarily deployed to test-trigger-is January 3, 2025 10:51 — with GitHub Actions Inactive

mgyucht temporarily deployed to test-trigger-is January 3, 2025 10:52 — with GitHub Actions Inactive

mgyucht added 3 commits January 3, 2025 11:54

Merge branch 'main' into move-u2m-to-sdks

70ef496

fix tests

3550329

Merge branch 'main' into move-u2m-to-sdks

05bd576

mgyucht temporarily deployed to test-trigger-is January 3, 2025 12:18 — with GitHub Actions Inactive

fix test

0688c3b

mgyucht temporarily deployed to test-trigger-is January 3, 2025 13:37 — with GitHub Actions Inactive

mgyucht temporarily deployed to test-trigger-is January 3, 2025 13:38 — with GitHub Actions Inactive

mgyucht mentioned this pull request Jan 3, 2025

[WIP] Migrate U2M Code to SDKs databricks/cli#2076

Open

renaudhartert-db requested changes Jan 3, 2025

View reviewed changes

mgyucht added 2 commits January 6, 2025 13:20

respond to some changes

35b8365

more work

47f2fb2

mgyucht temporarily deployed to test-trigger-is January 6, 2025 12:57 — with GitHub Actions Inactive

mgyucht had a problem deploying to test-trigger-is January 6, 2025 12:57 — with GitHub Actions Failure

work

e883616

mgyucht temporarily deployed to test-trigger-is January 6, 2025 13:34 — with GitHub Actions Inactive

mgyucht had a problem deploying to test-trigger-is January 6, 2025 13:35 — with GitHub Actions Failure

basic tests

181f88b

mgyucht temporarily deployed to test-trigger-is January 6, 2025 14:00 — with GitHub Actions Inactive

mgyucht had a problem deploying to test-trigger-is January 6, 2025 14:00 — with GitHub Actions Failure

some tests

74592c3

mgyucht temporarily deployed to test-trigger-is January 6, 2025 15:58 — with GitHub Actions Inactive

mgyucht had a problem deploying to test-trigger-is January 6, 2025 15:58 — with GitHub Actions Failure

one more test case

9b5913c

mgyucht temporarily deployed to test-trigger-is January 6, 2025 16:01 — with GitHub Actions Inactive

mgyucht had a problem deploying to test-trigger-is January 6, 2025 16:02 — with GitHub Actions Failure

more tests

73c73af

mgyucht added 2 commits January 30, 2025 10:30

address comment

3550bba

fix

fcb031f

mgyucht temporarily deployed to test-trigger-is January 30, 2025 09:40 — with GitHub Actions Inactive

more tweaks

83e4141

mgyucht temporarily deployed to test-trigger-is January 30, 2025 09:57 — with GitHub Actions Inactive

work

51a5b08

mgyucht temporarily deployed to test-trigger-is January 30, 2025 10:02 — with GitHub Actions Inactive

simpler

62acd68

mgyucht temporarily deployed to test-trigger-is January 30, 2025 11:07 — with GitHub Actions Inactive

tweaks

f506510

mgyucht temporarily deployed to test-trigger-is January 30, 2025 11:11 — with GitHub Actions Inactive

fix test

d70dac4

mgyucht temporarily deployed to test-trigger-is January 30, 2025 11:14 — with GitHub Actions Inactive

mgyucht requested a review from renaudhartert-db January 30, 2025 11:14

mgyucht temporarily deployed to test-trigger-is January 30, 2025 11:14 — with GitHub Actions Inactive

renaudhartert-db reviewed Jan 30, 2025

View reviewed changes

Merge branch 'main' into move-u2m-to-sdks

373daf3

mgyucht temporarily deployed to test-trigger-is February 14, 2025 12:59 — with GitHub Actions Inactive

work

2f0ebbb

mgyucht temporarily deployed to test-trigger-is February 14, 2025 15:59 — with GitHub Actions Inactive

mgyucht deployed to test-trigger-is February 14, 2025 15:59 — with GitHub Actions Active

mgyucht commented Feb 14, 2025

View reviewed changes

		// format versioning leaves some room for format improvement
		tokenCacheVersion = 1

		// GetHttpClient returns an HTTP client for OAuth2 requests.
		GetHttpClient(context.Context) *http.Client

		// mu protects the token cache file from concurrent reads and writes.
		mu *sync.Mutex

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

Are you sure you want to change the base?

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

Conversation

mgyucht commented Jan 3, 2025

What changes are proposed in this pull request?

How is this tested?

renaudhartert-db left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

renaudhartert-db Jan 3, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

renaudhartert-db Jan 3, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

renaudhartert-db commented Jan 29, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented Feb 14, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

renaudhartert-db Jan 3, 2025 •

edited

Loading

renaudhartert-db Jan 3, 2025 •

edited

Loading