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

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.

github-actions · 2025-01-20T12:40:27Z

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.

renaudhartert-db

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.

renaudhartert-db · 2025-01-27T08:57:08Z

credentials/cache/file.go

+	Tokens  map[string]*oauth2.Token `json:"tokens"`
+}
+
+type FileTokenCacheOpt func(*FileTokenCache)


Suggested change

type FileTokenCacheOpt func(*FileTokenCache)

type FileTokenCacheOption func(*FileTokenCache)

renaudhartert-db · 2025-01-27T08:58:01Z

credentials/cache/file.go

+func WithLocker(locker sync.Locker) FileTokenCacheOpt {
+	return func(c *FileTokenCache) {
+		c.locker = locker
+	}
+}


I would prefer to keep the locker internal. What use case do you have in mind?

renaudhartert-db · 2025-01-27T09:04:16Z

credentials/cache/file.go

+// 0600 and the directory is created with owner permissions 0700. If the cache
+// file is corrupt or if its version does not match tokenCacheVersion, an error
+// is returned.
+func NewFileTokenCache(opts ...FileTokenCacheOpt) (*FileTokenCache, error) {


I think we should return a TokenCache here rather than a FileTokenCache. The reasons are:

we do not want users to create a FileTokenCache without using this constructor,

the content of the FileTokenCache is private anyway, and

FileTokenCache does not expose new methods compared to the interface it implements.

Suggested change

func NewFileTokenCache(opts ...FileTokenCacheOpt) (*FileTokenCache, error) {

func NewFileTokenCache(opts ...FileTokenCacheOpt) (TokenCache, error) {

renaudhartert-db · 2025-01-27T09:06:21Z

credentials/cache/file.go

+	if err := c.init(); err != nil {
+		return nil, err
+	}
+	// verify the cache is working


[nit] Let's try to follow the same rules for all comments: https://go.dev/wiki/CodeReviewComments#comment-sentences

renaudhartert-db · 2025-01-27T09:18:27Z

credentials/cache/file.go

+	return f, nil
+}
+
+var _ TokenCache = (*FileTokenCache)(nil)


Not necessary if the constructor returns a TokenCache.

renaudhartert-db · 2025-01-27T09:30:10Z

credentials/cache/file.go

+	// create the directory if it doesn't already exist
+	if _, err := os.Stat(filepath.Dir(c.fileLocation)); err != nil {
+		if !os.IsNotExist(err) {
+			return fmt.Errorf("stat directory: %w", err)
+		}
+		// create the directory
+		if err := os.MkdirAll(filepath.Dir(c.fileLocation), ownerExecReadWrite); err != nil {
+			return fmt.Errorf("mkdir: %w", err)
+		}
+	}
+	// create the file if it doesn't already exist
+	if _, err := os.Stat(c.fileLocation); err != nil {
+		if !os.IsNotExist(err) {
+			return fmt.Errorf("stat file: %w", err)
+		}
+		f := &tokenCacheFile{
+			Version: tokenCacheVersion,
+			Tokens:  map[string]*oauth2.Token{},
+		}
+		raw, err := json.MarshalIndent(f, "", "  ")
+		if err != nil {
+			return fmt.Errorf("marshal: %w", err)
+		}
+		if err := os.WriteFile(c.fileLocation, raw, ownerReadWrite); err != nil {
+			return fmt.Errorf("write: %w", err)
+		}
+	}


I think we can simplify this slightly.

Suggested change

// create the directory if it doesn't already exist

if _, err := os.Stat(filepath.Dir(c.fileLocation)); err != nil {

if !os.IsNotExist(err) {

return fmt.Errorf("stat directory: %w", err)

}

// create the directory

if err := os.MkdirAll(filepath.Dir(c.fileLocation), ownerExecReadWrite); err != nil {

return fmt.Errorf("mkdir: %w", err)

}

}

// create the file if it doesn't already exist

if _, err := os.Stat(c.fileLocation); err != nil {

if !os.IsNotExist(err) {

return fmt.Errorf("stat file: %w", err)

}

f := &tokenCacheFile{

Version: tokenCacheVersion,

Tokens: map[string]*oauth2.Token{},

}

raw, err := json.MarshalIndent(f, "", " ")

if err != nil {

return fmt.Errorf("marshal: %w", err)

}

if err := os.WriteFile(c.fileLocation, raw, ownerReadWrite); err != nil {

return fmt.Errorf("write: %w", err)

}

}

// create the file if it doesn't already exist

if _, err := os.Stat(c.fileLocation); err != nil {

if !os.IsNotExist(err) {

return fmt.Errorf("stat file: %w", err)

}

// Create the parent directories if needed.

if err := os.MkdirAll(filepath.Dir(c.fileLocation), ownerExecReadWrite); err != nil {

return fmt.Errorf("mkdir: %w", err)

}

// Create an empty file cache.

f := &tokenCacheFile{

Version: tokenCacheVersion,

Tokens: map[string]*oauth2.Token{},

}

raw, err := json.MarshalIndent(f, "", " ")

if err != nil {

return fmt.Errorf("marshal: %w", err)

}

if err := os.WriteFile(c.fileLocation, raw, ownerReadWrite); err != nil {

return fmt.Errorf("write: %w", err)

}

}

renaudhartert-db · 2025-01-27T09:31:06Z

credentials/cache/file.go

+	if c.locker == nil {
+		home, err := os.UserHomeDir()
+		if err != nil {
+			return fmt.Errorf("home: %w", err)
+		}
+
+		c.locker, err = newLocker(filepath.Join(home, lockFilePath))
+		if err != nil {
+			return fmt.Errorf("locker: %w", err)
+		}
+	}


See comment above, I would prefer to keep the locker internal.

renaudhartert-db · 2025-01-27T09:32:14Z

credentials/cache/file.go

+	err = json.Unmarshal(raw, &f)
+	if err != nil {


[nit] This is more idiomatic and slightly more robust as we do not expect this error to "exist" outside the if statement.

Suggested change

err = json.Unmarshal(raw, &f)

if err != nil {

if err := json.Unmarshal(raw, &f); err != nil {

renaudhartert-db · 2025-01-27T09:44:30Z

config/auth_u2m.go

+	a := u.auth
+	if a == nil {
+		var err error
+		a, err = oauth.NewPersistentAuth(ctx)


Should we update u.auth?

pietern · 2025-01-27T11:15:41Z

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)


Why is this configurable?

pietern · 2025-01-27T11:17:56Z

config/auth_u2m.go

+	r, err := http.NewRequestWithContext(ctx, http.MethodGet, "", nil)
+	if err != nil {
+		return nil, fmt.Errorf("http request: %w", err)
+	}


I was puzzled why this seemingly invalid request was around, but now makes sense. Could you move it below and include a comment to say this is a dummy request to force loading a token?

pietern · 2025-01-27T11:21:05Z

credentials/cache/file.go

+
+	// lockFilePath is the path of the lock file used to prevent concurrent
+	// reads and writes to the token cache file.
+	lockFilePath = ".databricks/token-cache.lock"


The diff from the version in the CLI repo is that this uses file-based locking, correct?

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 temporarily deployed to test-trigger-is January 7, 2025 14:56 — with GitHub Actions Inactive

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

work

6c2ed96

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

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

better error message

e9f3732

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

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

fix test

19a34ce

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

fix

c7b5155

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

mgyucht added 2 commits January 8, 2025 13:53

tweak

37b44b1

work

36ea3dc

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

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

fix test

fc87393

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

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

undo move

945151f

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

remove extra context

356a7c9

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

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

renaudhartert-db requested changes Jan 27, 2025

View reviewed changes

pietern reviewed Jan 27, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

mgyucht commented Jan 3, 2025

renaudhartert-db left a comment

renaudhartert-db Jan 3, 2025

renaudhartert-db Jan 3, 2025 •

edited

Loading

renaudhartert-db Jan 3, 2025

renaudhartert-db Jan 3, 2025

mgyucht Jan 3, 2025

renaudhartert-db Jan 3, 2025

mgyucht Jan 3, 2025

renaudhartert-db Jan 3, 2025 •

edited

Loading

renaudhartert-db Jan 3, 2025

mgyucht Jan 3, 2025

renaudhartert-db Jan 3, 2025

renaudhartert-db Jan 3, 2025

mgyucht Jan 3, 2025

renaudhartert-db Jan 3, 2025

mgyucht Jan 3, 2025

github-actions bot commented Jan 20, 2025

renaudhartert-db left a comment •

edited

Loading

renaudhartert-db Jan 27, 2025

renaudhartert-db Jan 27, 2025

renaudhartert-db Jan 27, 2025

renaudhartert-db Jan 27, 2025

renaudhartert-db Jan 27, 2025

renaudhartert-db Jan 27, 2025

renaudhartert-db Jan 27, 2025

renaudhartert-db Jan 27, 2025

renaudhartert-db Jan 27, 2025

pietern Jan 27, 2025

pietern Jan 27, 2025

pietern Jan 27, 2025

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

	type FileTokenCacheOpt func(*FileTokenCache)
	type FileTokenCacheOption func(*FileTokenCache)

	func NewFileTokenCache(opts ...FileTokenCacheOpt) (*FileTokenCache, error) {
	func NewFileTokenCache(opts ...FileTokenCacheOpt) (TokenCache, error) {

	err = json.Unmarshal(raw, &f)
	if err != nil {
	if err := json.Unmarshal(raw, &f); err != nil {

[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

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

github-actions bot commented Jan 20, 2025

renaudhartert-db left a comment • 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

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

renaudhartert-db left a comment •

edited

Loading