diff --git a/.codegen/workspaces.go.tmpl b/.codegen/workspaces.go.tmpl index f42308cd4..d4db23cfc 100644 --- a/.codegen/workspaces.go.tmpl +++ b/.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 diff --git a/config/auth_azure_client_secret.go b/config/auth_azure_client_secret.go index fc6725846..beb91898d 100644 --- a/config/auth_azure_client_secret.go +++ b/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 } diff --git a/config/auth_gcp_google_id.go b/config/auth_gcp_google_id.go index 767efc996..4dd291c17 100644 --- a/config/auth_gcp_google_id.go +++ b/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, diff --git a/config/config.go b/config/config.go index 6cbdd7b2e..a00ab8679 100644 --- a/config/config.go +++ b/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) } } diff --git a/credentials/credentials_provider.go b/credentials/credentials_provider.go index e1bb202c5..4e82ab3b6 100644 --- a/credentials/credentials_provider.go +++ b/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 } diff --git a/credentials/oauth_credentials_provider.go b/credentials/oauth_credentials_provider.go index 4fea7c2ee..fe6d2dd06 100644 --- a/credentials/oauth_credentials_provider.go +++ b/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) } diff --git a/credentials/oauth_token.go b/credentials/oauth_token.go index 88d9a218f..a1f6c131e 100644 --- a/credentials/oauth_token.go +++ b/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"` } diff --git a/httpclient/oauth_token.go b/httpclient/oauth_token.go index 5b33793a3..691877a53 100644 --- a/httpclient/oauth_token.go +++ b/httpclient/oauth_token.go @@ -8,10 +8,17 @@ 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. @@ -19,23 +26,19 @@ type GetOAuthTokenRequest struct { // // **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 } diff --git a/httpclient/request.go b/httpclient/request.go index 856d23cc2..4da4e0039 100644 --- a/httpclient/request.go +++ b/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 diff --git a/httpclient/request_test.go b/httpclient/request_test.go index 2a28794d7..64b38d0cc 100644 --- a/httpclient/request_test.go +++ b/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) diff --git a/service/iam/model.go b/service/iam/model.go index c0e595616..76321c4b8 100755 --- a/service/iam/model.go +++ b/service/iam/model.go @@ -329,7 +329,7 @@ type Group struct { Groups []ComplexValue `json:"groups,omitempty"` // Databricks group ID - Id string `json:"id,omitempty" url:"-"` + Id string `json:"id,omitempty"` Members []ComplexValue `json:"members,omitempty"` // Container for the group identifier. Workspace local versus account. diff --git a/workspace_client.go b/workspace_client.go index a86402657..195c63a07 100755 --- a/workspace_client.go +++ b/workspace_client.go @@ -3,6 +3,7 @@ package databricks import ( + "context" "errors" "github.com/databricks/databricks-sdk-go/client" @@ -976,12 +977,12 @@ type WorkspaceClient struct { // // **NOTE:** Experimental: This API may change or be removed in a future release // without warning. -func (a *WorkspaceClient) GetOAuthToken(authorizationDetails string) (*credentials.OAuthToken, error) { +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(authorizationDetails, originalToken) + return a.apiClient.GetOAuthToken(ctx, authorizationDetails, originalToken) } var ErrNotWorkspaceClient = errors.New("invalid Databricks Workspace configuration")