Skip to content

Commit

Permalink
Move credentials in config
Browse files Browse the repository at this point in the history
  • Loading branch information
renaudhartert-db committed Jan 7, 2025
1 parent adc94ca commit 1b4d60b
Show file tree
Hide file tree
Showing 20 changed files with 100 additions and 109 deletions.
4 changes: 2 additions & 2 deletions config/api_client.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import (
"time"

"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/config/credentials"
"github.com/databricks/databricks-sdk-go/httpclient"
"github.com/databricks/databricks-sdk-go/useragent"
)
Expand Down Expand Up @@ -104,7 +104,7 @@ func (noopLoader) Configure(cfg *Config) error { return nil }
type noopAuth struct{}

func (noopAuth) Name() string { return "noop" }
func (noopAuth) Configure(context.Context, *Config) (credentials.CredentialsProvider, error) {
func (noopAuth) Configure(context.Context, *Config) (CredentialsProvider, error) {
visitor := func(r *http.Request) error { return nil }
return credentials.NewCredentialsProvider(visitor), nil
}
4 changes: 2 additions & 2 deletions config/auth_azure_cli.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ import (

"golang.org/x/oauth2"

"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/config/credentials"
"github.com/databricks/databricks-sdk-go/logger"
)

Expand Down Expand Up @@ -54,7 +54,7 @@ func (c AzureCliCredentials) getVisitor(ctx context.Context, cfg *Config, inner
return azureVisitor(cfg, serviceToServiceVisitor(inner, management, xDatabricksAzureSpManagementToken)), nil
}

func (c AzureCliCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
func (c AzureCliCredentials) Configure(ctx context.Context, cfg *Config) (CredentialsProvider, error) {
if !cfg.IsAzure() {
return nil, nil
}
Expand Down
4 changes: 2 additions & 2 deletions config/auth_azure_client_secret.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import (
"golang.org/x/oauth2"
"golang.org/x/oauth2/clientcredentials"

"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/config/credentials"
"github.com/databricks/databricks-sdk-go/logger"
)

Expand All @@ -35,7 +35,7 @@ func (c AzureClientSecretCredentials) tokenSourceFor(
// as we cannot create AKV backed secret scopes when authenticated as SP.
// If we are authenticated as SP and wish to create one we want to fail early.
// Also see https://github.com/databricks/terraform-provider-databricks/issues/1490.
func (c AzureClientSecretCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
func (c AzureClientSecretCredentials) Configure(ctx context.Context, cfg *Config) (CredentialsProvider, error) {
if cfg.AzureClientID == "" || cfg.AzureClientSecret == "" || cfg.AzureTenantID == "" {
return nil, nil
}
Expand Down
4 changes: 2 additions & 2 deletions config/auth_azure_github_oidc.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import (
"fmt"
"time"

"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/config/credentials"
"github.com/databricks/databricks-sdk-go/httpclient"
"github.com/databricks/databricks-sdk-go/logger"
"golang.org/x/oauth2"
Expand All @@ -22,7 +22,7 @@ func (c AzureGithubOIDCCredentials) Name() string {
}

// Configure implements [CredentialsStrategy.Configure].
func (c AzureGithubOIDCCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
func (c AzureGithubOIDCCredentials) Configure(ctx context.Context, cfg *Config) (CredentialsProvider, error) {
// Sanity check that the config is configured for Azure Databricks.
if !cfg.IsAzure() || cfg.AzureClientID == "" || cfg.Host == "" || cfg.AzureTenantID == "" {
return nil, nil
Expand Down
4 changes: 2 additions & 2 deletions config/auth_azure_msi.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import (
"net/http"
"time"

"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/config/credentials"
"github.com/databricks/databricks-sdk-go/httpclient"
"github.com/databricks/databricks-sdk-go/logger"
"golang.org/x/oauth2"
Expand All @@ -31,7 +31,7 @@ func (c AzureMsiCredentials) Name() string {
return "azure-msi"
}

func (c AzureMsiCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
func (c AzureMsiCredentials) Configure(ctx context.Context, cfg *Config) (CredentialsProvider, error) {
if !cfg.IsAzure() || !cfg.AzureUseMSI || (cfg.AzureResourceID == "" && !cfg.IsAccountClient()) {
return nil, nil
}
Expand Down
4 changes: 2 additions & 2 deletions config/auth_basic.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import (
"fmt"
"net/http"

"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/config/credentials"
)

type BasicCredentials struct {
Expand All @@ -16,7 +16,7 @@ func (c BasicCredentials) Name() string {
return "basic"
}

func (c BasicCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
func (c BasicCredentials) Configure(ctx context.Context, cfg *Config) (CredentialsProvider, error) {
if cfg.Username == "" || cfg.Password == "" || cfg.Host == "" {
return nil, nil
}
Expand Down
4 changes: 2 additions & 2 deletions config/auth_databricks_cli.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ import (
"path/filepath"
"strings"

"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/config/credentials"
"github.com/databricks/databricks-sdk-go/logger"
"golang.org/x/oauth2"
)
Expand All @@ -22,7 +22,7 @@ func (c DatabricksCliCredentials) Name() string {
return "databricks-cli"
}

func (c DatabricksCliCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
func (c DatabricksCliCredentials) Configure(ctx context.Context, cfg *Config) (CredentialsProvider, error) {
if cfg.Host == "" {
return nil, nil
}
Expand Down
3 changes: 1 addition & 2 deletions config/auth_default.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import (
"errors"
"fmt"

"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/logger"
)

Expand Down Expand Up @@ -44,7 +43,7 @@ var errorMessage = fmt.Sprintf("cannot configure default credentials, please che
// ErrCannotConfigureAuth (experimental) is returned when no auth is configured
var ErrCannotConfigureAuth = errors.New(errorMessage)

func (c *DefaultCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
func (c *DefaultCredentials) Configure(ctx context.Context, cfg *Config) (CredentialsProvider, error) {
for _, p := range authProviders {
if cfg.AuthType != "" && p.Name() != cfg.AuthType {
// ignore other auth types if one is explicitly enforced
Expand Down
4 changes: 2 additions & 2 deletions config/auth_gcp_google_credentials.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import (
"io/ioutil"
"os"

"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/config/credentials"
"github.com/databricks/databricks-sdk-go/logger"
"golang.org/x/oauth2/google"
"google.golang.org/api/idtoken"
Expand All @@ -20,7 +20,7 @@ func (c GoogleCredentials) Name() string {
return "google-credentials"
}

func (c GoogleCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
func (c GoogleCredentials) Configure(ctx context.Context, cfg *Config) (CredentialsProvider, error) {
if cfg.GoogleCredentials == "" || !cfg.IsGcp() {
return nil, nil
}
Expand Down
4 changes: 2 additions & 2 deletions config/auth_gcp_google_id.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import (
"context"
"fmt"

"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/config/credentials"
"github.com/databricks/databricks-sdk-go/logger"
"golang.org/x/oauth2"
"google.golang.org/api/impersonate"
Expand All @@ -20,7 +20,7 @@ func (c GoogleDefaultCredentials) Name() string {
return "google-id"
}

func (c GoogleDefaultCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
func (c GoogleDefaultCredentials) Configure(ctx context.Context, cfg *Config) (CredentialsProvider, error) {
if cfg.GoogleServiceAccount == "" || !cfg.IsGcp() {
return nil, nil
}
Expand Down
4 changes: 2 additions & 2 deletions config/auth_m2m.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import (
"golang.org/x/oauth2"
"golang.org/x/oauth2/clientcredentials"

"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/config/credentials"
"github.com/databricks/databricks-sdk-go/httpclient"
"github.com/databricks/databricks-sdk-go/logger"
)
Expand All @@ -22,7 +22,7 @@ func (c M2mCredentials) Name() string {
return "oauth-m2m"
}

func (c M2mCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
func (c M2mCredentials) Configure(ctx context.Context, cfg *Config) (CredentialsProvider, error) {
if cfg.ClientID == "" || cfg.ClientSecret == "" {
return nil, nil
}
Expand Down
4 changes: 2 additions & 2 deletions config/auth_metadata_service.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import (
"net/url"
"time"

"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/config/credentials"
"github.com/databricks/databricks-sdk-go/httpclient"
"github.com/databricks/databricks-sdk-go/logger"
"golang.org/x/oauth2"
Expand Down Expand Up @@ -49,7 +49,7 @@ func (c MetadataServiceCredentials) Name() string {
return "metadata-service"
}

func (c MetadataServiceCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
func (c MetadataServiceCredentials) Configure(ctx context.Context, cfg *Config) (CredentialsProvider, error) {
if cfg.MetadataServiceURL == "" || cfg.Host == "" {
return nil, nil
}
Expand Down
4 changes: 2 additions & 2 deletions config/auth_pat.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ import (
"fmt"
"net/http"

"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/config/credentials"
)

type PatCredentials struct {
Expand All @@ -15,7 +15,7 @@ func (c PatCredentials) Name() string {
return "pat"
}

func (c PatCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
func (c PatCredentials) Configure(ctx context.Context, cfg *Config) (CredentialsProvider, error) {
if cfg.Token == "" || cfg.Host == "" {
return nil, nil
}
Expand Down
9 changes: 4 additions & 5 deletions config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ import (

"github.com/databricks/databricks-sdk-go/common"
"github.com/databricks/databricks-sdk-go/common/environment"
"github.com/databricks/databricks-sdk-go/credentials"
"github.com/databricks/databricks-sdk-go/httpclient"
"github.com/databricks/databricks-sdk-go/logger"
"golang.org/x/oauth2"
Expand All @@ -28,7 +27,7 @@ type CredentialsStrategy interface {
// 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)
Configure(context.Context, *Config) (CredentialsProvider, error)
}

type Loader interface {
Expand All @@ -50,7 +49,7 @@ type Config struct {
WarehouseID string `name:"warehouse_id" env:"DATABRICKS_WAREHOUSE_ID"`
ServerlessComputeID string `name:"serverless_compute_id" env:"DATABRICKS_SERVERLESS_COMPUTE_ID"`

// URL of the metadata service that provides authentication credentials.
// URL of the metadata service that provides authentication
MetadataServiceURL string `name:"metadata_service_url" env:"DATABRICKS_METADATA_SERVICE_URL" auth:"metadata-service,sensitive"`

// Databricks Account ID for Accounts API. This field is used in dependencies.
Expand Down Expand Up @@ -160,7 +159,7 @@ type Config struct {
mu sync.Mutex

// HTTP request interceptor, that assigns Authorization header
credentialsProvider credentials.CredentialsProvider
credentialsProvider CredentialsProvider

// Keep track of the source of each attribute
attrSource map[string]Source
Expand Down Expand Up @@ -231,7 +230,7 @@ func (c *Config) GetToken() (*oauth2.Token, error) {
if err != nil {
return nil, err
}
if h, ok := c.credentialsProvider.(credentials.OAuthCredentialsProvider); ok {
if h, ok := c.credentialsProvider.(OAuthCredentialsProvider); ok {
return h.Token()
} else {
return nil, fmt.Errorf("OAuth Token not supported for current auth type %s", c.AuthType)
Expand Down
68 changes: 68 additions & 0 deletions config/credentials.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
package config

import (
"net/http"

"golang.org/x/oauth2"
)

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

type credentialsProvider struct {
setHeaders func(r *http.Request) error
}

func (c *credentialsProvider) SetHeaders(r *http.Request) error {
return c.setHeaders(r)
}

func NewCredentialsProvider(visitor func(r *http.Request) error) CredentialsProvider {
return &credentialsProvider{
setHeaders: visitor,
}
}

// OAuthCredentialsProvider is a specialized CredentialsProvider uses and provides an OAuth token.
type OAuthCredentialsProvider interface {
CredentialsProvider
// Token returns the OAuth token generated by the provider.
Token() (*oauth2.Token, error)
}

type oauthCredentialsProvider struct {
setHeaders func(r *http.Request) error
token func() (*oauth2.Token, error)
}

func (c *oauthCredentialsProvider) SetHeaders(r *http.Request) error {
return c.setHeaders(r)
}

func (c *oauthCredentialsProvider) Token() (*oauth2.Token, error) {
return c.token()
}

func NewOAuthCredentialsProvider(visitor func(r *http.Request) error, tokenProvider func() (*oauth2.Token, error)) OAuthCredentialsProvider {
return &oauthCredentialsProvider{
setHeaders: visitor,
token: tokenProvider,
}
}

// 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"`
// 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"`
}
26 changes: 0 additions & 26 deletions credentials/credentials_provider.go

This file was deleted.

Loading

0 comments on commit 1b4d60b

Please sign in to comment.