Comments

.codegen/workspaces.go.tmpl

@@ -3,6 +3,8 @@
 package databricks
 
 import (
+	"context"
+
 	"github.com/databricks/databricks-sdk-go/client"
 	{{range .Packages}}
 	"github.com/databricks/databricks-sdk-go/service/{{.Name}}"{{end}}
@@ -18,6 +20,19 @@ type WorkspaceClient struct {
 	{{end}}{{end}}
 }
 
+// Returns a new OAuth scoped to the authorization details provided.
+// It will return an error if the CredentialStrategy does not support OAuth tokens.
+//
+// **NOTE:** Experimental: This API may change or be removed in a future release
+// without warning.
+func (a *WorkspaceClient) GetOAuthToken(ctx context.Context, authorizationDetails string) (*credentials.OAuthToken, error) {
+	originalToken, err := a.Config.GetToken()
+	if err != nil {
+		return nil, err
+	}
+	return a.apiClient.GetOAuthToken(ctx, authorizationDetails, originalToken)
+}
+
 var ErrNotWorkspaceClient = errors.New("invalid Databricks Workspace configuration")
 
 // NewWorkspaceClient creates new Databricks SDK client for Workspaces or

config/auth_azure_client_secret.go

@@ -53,5 +53,5 @@ func (c AzureClientSecretCredentials) Configure(ctx context.Context, cfg *Config
 	inner := azureReuseTokenSource(nil, c.tokenSourceFor(ctx, cfg, aadEndpoint, env.AzureApplicationID))
 	management := azureReuseTokenSource(nil, c.tokenSourceFor(ctx, cfg, aadEndpoint, managementEndpoint))
 	visitor := azureVisitor(cfg, serviceToServiceVisitor(inner, management, xDatabricksAzureSpManagementToken))
-	return credentials.NewCredentialsProvider(visitor), nil
+	return credentials.NewOAuthCredentialsProvider(visitor, inner.Token), nil
 }

config/auth_gcp_google_id.go

@@ -46,7 +46,7 @@ func (c GoogleDefaultCredentials) Configure(ctx context.Context, cfg *Config) (c
 	}
 	logger.Infof(ctx, "Using Google Default Application Credentials for Accounts API")
 	visitor := serviceToServiceVisitor(inner, platform, "X-Databricks-GCP-SA-Access-Token")
-	return credentials.NewCredentialsProvider(visitor), nil
+	return credentials.NewOAuthCredentialsProvider(visitor, inner.Token), nil
 }
 
 func (c GoogleDefaultCredentials) idTokenSource(ctx context.Context, host, serviceAccount string,

config/config.go

@@ -20,11 +20,11 @@ import (
 // CredentialsStrategy responsible for configuring static or refreshable
 // authentication credentials for Databricks REST APIs
 type CredentialsStrategy interface {
-	// Name returns human-addressable name of this credentials provider name
+	// Name returns human-addressable name of this credentials provider strategy
 	Name() string
 
-	// Configure creates CredentialsProvider or returns nil if a given credetials
-	// strategy ßare not configured. It returns an error if credentials are misconfigured.
+	// Configure creates CredentialsProvider or returns nil if a given credentials
+	// strategy are not configured. It returns an error if credentials are misconfigured.
 	// Takes a context and a pointer to a Config instance, that holds auth mutex.
 	Configure(context.Context, *Config) (credentials.CredentialsProvider, error)
 }
@@ -223,7 +223,7 @@ func (c *Config) GetToken() (*oauth2.Token, error) {
 	if h, ok := c.credentialsProvider.(credentials.OAuthCredentialsProvider); ok {
 		return h.Token()
 	} else {
-		return nil, fmt.Errorf("OAuth Token can only be retrieved for OauthCredentialsProvider")
+		return nil, fmt.Errorf("OAuth Token not supported for current auth type %s", c.AuthType)
 	}
 }
 

credentials/credentials_provider.go

@@ -4,7 +4,10 @@ import (
 	"net/http"
 )
 
+// CredentialsProvider is an interface for providing credentials to the client.
+// Implementations of this interface should set the necessary headers on the request.
 type CredentialsProvider interface {
+	// SetHeaders sets the necessary headers on the request.
 	SetHeaders(r *http.Request) error
 }
 

credentials/oauth_credentials_provider.go

@@ -6,8 +6,10 @@ import (
 	"golang.org/x/oauth2"
 )
 
+// OAuthCredentialsProvider is a specialized CredentialsProvider uses and provides an OAuth token.
 type OAuthCredentialsProvider interface {
-	SetHeaders(r *http.Request) error
+	CredentialsProvider
+	// Token returns the OAuth token generated by the provider.
 	Token() (*oauth2.Token, error)
 }
 

credentials/oauth_token.go

@@ -1,8 +1,14 @@
 package credentials
 
+// OAuthToken represents an OAuth token as defined by the OAuth 2.0 Authorization Framework.
+// 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"`
-	ExpiresIn   int    `json:"expires_in"`
-	Scope       string `json:"scope"`
-	TokenType   string `json:"token_type"`
+	// Time in seconds until the token expires.
+	ExpiresIn int `json:"expires_in"`
+	// The scope of the token. This is a space-separated list of strings that represent the permissions granted by the token.
+	Scope string `json:"scope"`
+	// The type of token that was issued.
+	TokenType string `json:"token_type"`
 }

httpclient/oauth_token.go

@@ -8,34 +8,37 @@ import (
 	"golang.org/x/oauth2"
 )
 
+const JWTGrantType = "urn:ietf:params:oauth:grant-type:jwt-bearer"
+
+// GetOAuthTokenRequest is the request to get an OAuth token. It follows the OAuth 2.0 Rich Authorization Requests specification.
+// https://datatracker.ietf.org/doc/html/rfc9396
 type GetOAuthTokenRequest struct {
-	GrantType            string `url:"grant_type"`
+	// Defines the method used to get the token.
+	GrantType string `url:"grant_type"`
+	// An array of authorization details that the token should be scoped to. This needs to be passed in string format.
 	AuthorizationDetails string `url:"authorization_details"`
-	Assertion            string `url:"assertion"`
+	// The token that will be exchanged for an OAuth token.
+	Assertion string `url:"assertion"`
 }
 
 // Returns a new OAuth token using the provided token. The token must be a JWT token.
 // The resulting token is scoped to the authorization details provided.
 //
 // **NOTE:** Experimental: This API may change or be removed in a future release
 // without warning.
-func (c *ApiClient) GetOAuthToken(authDetails string, token *oauth2.Token) (*credentials.OAuthToken, error) {
+func (c *ApiClient) GetOAuthToken(ctx context.Context, authDetails string, token *oauth2.Token) (*credentials.OAuthToken, error) {
 	path := "/oidc/v1/token"
-	headers := map[string]string{
-		"Content-Type": "application/x-www-form-urlencoded",
-	}
 	data := GetOAuthTokenRequest{
-		GrantType:            "urn:ietf:params:oauth:grant-type:jwt-bearer",
+		GrantType:            JWTGrantType,
 		AuthorizationDetails: authDetails,
 		Assertion:            token.AccessToken,
 	}
 	var response credentials.OAuthToken
 	opts := []DoOption{
-		WithRequestHeaders(headers),
 		WithUrlEncodedData(data),
 		WithResponseUnmarshal(&response),
 	}
-	err := c.Do(c.InContextForOAuth2(context.Background()), http.MethodPost, path, opts...)
+	err := c.Do(ctx, http.MethodPost, path, opts...)
 	if err != nil {
 		return nil, err
 	}

httpclient/request.go

@@ -13,6 +13,8 @@ import (
 	"golang.org/x/oauth2"
 )
 
+const UrlEncodedContentType = "application/x-www-form-urlencoded"
+
 // WithRequestHeader adds a request visitor, that sets a header on a request
 func WithRequestHeader(k, v string) DoOption {
 	return WithRequestVisitor(func(r *http.Request) error {
@@ -73,11 +75,11 @@ func WithRequestData(body any) DoOption {
 func WithUrlEncodedData(body any) DoOption {
 	return DoOption{
 		in: func(r *http.Request) error {
-			r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+			r.Header.Set("Content-Type", UrlEncodedContentType)
 			return nil
 		},
 		body:        body,
-		contentType: "application/x-www-form-urlencoded",
+		contentType: UrlEncodedContentType,
 	}
 }
 
@@ -151,7 +153,7 @@ func makeRequestBody(method string, requestURL *string, data interface{}, conten
 		*requestURL += "?" + qs
 		return common.NewRequestBody([]byte{})
 	}
-	if contentType == "application/x-www-form-urlencoded" {
+	if contentType == UrlEncodedContentType {
 		qs, err := makeQueryString(data)
 		if err != nil {
 			return common.RequestBody{}, err

httpclient/request_test.go

@@ -61,7 +61,7 @@ func TestUrlEncoding(t *testing.T) {
 		GrantType:            "grant",
 	}
 	requestURL := "/a/b/c"
-	body, err := makeRequestBody("POST", &requestURL, data, "application/x-www-form-urlencoded")
+	body, err := makeRequestBody("POST", &requestURL, data, UrlEncodedContentType)
 	require.NoError(t, err)
 	bodyBytes, err := io.ReadAll(body.Reader)
 	require.NoError(t, err)