[Feature] Implement U2M Authentication in the Go SDK

config/auth_default.go

@@ -13,7 +13,7 @@ var authProviders = []CredentialsStrategy{
 	PatCredentials{},
 	BasicCredentials{},
 	M2mCredentials{},
-	DatabricksCliCredentials{},
+	U2MCredentials{},
 	MetadataServiceCredentials{},
 
 	// Attempt to configure auth from most specific to most generic (the Azure CLI).

config/auth_m2m.go

@@ -2,19 +2,15 @@ package config
 
 import (
 	"context"
-	"errors"
 	"fmt"
 
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/clientcredentials"
 
 	"github.com/databricks/databricks-sdk-go/credentials"
-	"github.com/databricks/databricks-sdk-go/httpclient"
 	"github.com/databricks/databricks-sdk-go/logger"
 )
 
-var errOAuthNotSupported = errors.New("databricks OAuth is not supported for this host")
-
 type M2mCredentials struct {
 }
 
@@ -26,7 +22,7 @@ func (c M2mCredentials) Configure(ctx context.Context, cfg *Config) (credentials
 	if cfg.ClientID == "" || cfg.ClientSecret == "" {
 		return nil, nil
 	}
-	endpoints, err := oidcEndpoints(ctx, cfg)
+	endpoints, err := cfg.refreshClient.GetOidcEndpoints(ctx, cfg.Host, cfg.AccountID)
 	if err != nil {
 		return nil, fmt.Errorf("oidc: %w", err)
 	}
@@ -41,30 +37,3 @@ func (c M2mCredentials) Configure(ctx context.Context, cfg *Config) (credentials
 	visitor := refreshableVisitor(ts)
 	return credentials.NewOAuthCredentialsProvider(visitor, ts.Token), nil
 }
-
-func oidcEndpoints(ctx context.Context, cfg *Config) (*oauthAuthorizationServer, error) {
-	prefix := cfg.Host
-	if cfg.IsAccountClient() && cfg.AccountID != "" {
-		// TODO: technically, we could use the same config profile for both workspace
-		// and account, but we have to add logic for determining accounts host from
-		// workspace host.
-		prefix := fmt.Sprintf("%s/oidc/accounts/%s", cfg.Host, cfg.AccountID)
-		return &oauthAuthorizationServer{
-			AuthorizationEndpoint: fmt.Sprintf("%s/v1/authorize", prefix),
-			TokenEndpoint:         fmt.Sprintf("%s/v1/token", prefix),
-		}, nil
-	}
-	oidc := fmt.Sprintf("%s/oidc/.well-known/oauth-authorization-server", prefix)
-	var oauthEndpoints oauthAuthorizationServer
-	err := cfg.refreshClient.Do(ctx, "GET", oidc,
-		httpclient.WithResponseUnmarshal(&oauthEndpoints))
-	if err != nil {
-		return nil, errOAuthNotSupported
-	}
-	return &oauthEndpoints, nil
-}
-
-type oauthAuthorizationServer struct {
-	AuthorizationEndpoint string `json:"authorization_endpoint"` // ../v1/authorize
-	TokenEndpoint         string `json:"token_endpoint"`         // ../v1/token
-}

config/auth_m2m_test.go

@@ -4,6 +4,7 @@ import (
 	"net/url"
 	"testing"
 
+	"github.com/databricks/databricks-sdk-go/httpclient"
 	"github.com/databricks/databricks-sdk-go/httpclient/fixtures"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/oauth2"
@@ -16,7 +17,7 @@ func TestM2mHappyFlow(t *testing.T) {
 		ClientSecret: "c",
 		HTTPTransport: fixtures.MappingTransport{
 			"GET /oidc/.well-known/oauth-authorization-server": {
-				Response: oauthAuthorizationServer{
+				Response: httpclient.OAuthAuthorizationServer{
 					AuthorizationEndpoint: "https://localhost:1234/dummy/auth",
 					TokenEndpoint:         "https://localhost:1234/dummy/token",
 				},
@@ -80,5 +81,5 @@ func TestM2mNotSupported(t *testing.T) {
 			},
 		},
 	})
-	require.ErrorIs(t, err, errOAuthNotSupported)
+	require.ErrorIs(t, err, httpclient.ErrOAuthNotSupported)
 }

config/auth_u2m.go

@@ -0,0 +1,59 @@
+package config
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+
+	"github.com/databricks/databricks-sdk-go/credentials"
+	"github.com/databricks/databricks-sdk-go/credentials/oauth"
+	"github.com/databricks/databricks-sdk-go/logger"
+)
+
+type U2MCredentials struct {
+	Auth *oauth.PersistentAuth
+}
+
+// Name implements CredentialsStrategy.
+func (u U2MCredentials) Name() string {
+	return "oauth-u2m"
+}
+
+// Configure implements CredentialsStrategy.
+func (u U2MCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
+	a := u.Auth
+	if a == nil {
+		var err error
+		a, err = oauth.NewPersistentAuth(ctx)
+		if err != nil {
+			logger.Debugf(ctx, "failed to create persistent auth: %v, continuing", err)
+			return nil, nil
+		}
+	}
+	f := func(r *http.Request) error {
+		arg := oauth.BasicOAuthArgument{
+			Host:      cfg.Host,
+			AccountID: cfg.AccountID,
+		}
+		token, err := a.Load(r.Context(), arg)
+		if err != nil {
+			return fmt.Errorf("oidc: %w", err)
+		}
+		r.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token.AccessToken))
+		return nil
+	}
+
+	r, err := http.NewRequestWithContext(ctx, http.MethodGet, "", nil)
+	if err != nil {
+		return nil, fmt.Errorf("http request: %w", err)
+	}
+	// Try to load the credential from the token cache. If absent, fall back
+	// to the next credentials strategy.
+	if err := f(r); err != nil {
+		return nil, nil
+	}
+
+	return credentials.NewCredentialsProvider(f), nil
+}
+
+var _ CredentialsStrategy = U2MCredentials{}

credentials/cache/cache.go

@@ -0,0 +1,10 @@
+package cache
+
+import (
+	"golang.org/x/oauth2"
+)
+
+type TokenCache interface {
+	Store(key string, t *oauth2.Token) error
+	Lookup(key string) (*oauth2.Token, error)
+}