u2m

package
v0.120.0 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Mar 10, 2026 License: Apache-2.0 Imports: 24 Imported by: 2

Documentation

Overview

Package u2m supports the user-to-machine (U2M) OAuth flow for authenticating with Databricks.

Databricks uses the authorization code flow from OAuth 2.0 to authenticate users. This flow consists of four steps:

  1. Retrieve an authorization code for a user by opening a browser and directing them to the Databricks authorization URL.
  2. Exchange the authorization code for an access token.
  3. Use the access token to authenticate with Databricks.
  4. When the access token expires, use the refresh token to get a new access token.

The token and authorization endpoints for Databricks vary depending on whether the host is an account- or workspace-level host. Account-level endpoints are fixed based on the account ID and host, while workspace-level endpoints are discovered using the OIDC discovery endpoint at /oidc/.well-known/oauth-authorization-server.

For host-agnostic login through login.databricks.com, use WithDiscoveryLogin together with a DiscoveryOAuthArgument during Challenge(). This is a bootstrap flow only. Once the callback reveals the workspace host, construct the usual host-based OAuthArgument for future PersistentAuth instances.

To trigger the authorization flow, construct a PersistentAuth object and call the Challenge() method:

auth, err := oauth.NewPersistentAuth(ctx)
if err != nil {
	log.Fatalf("failed to create persistent auth: %v", err)
}
token, err := auth.Challenge(ctx, oauth.BasicAccountOAuthArgument{
	AccountHost: "https://accounts.cloud.databricks.com",
	AccountID: "xyz",
})

Because the U2M flow requires user interaction, the resulting access token and refresh token can be stored in a persistent cache to avoid prompting the user for credentials on every authentication attempt. By default, the cache is stored in ~/.databricks/token-cache.json. Retrieve the cached token by calling the Load() method:

token, err := auth.Load(ctx, oauth.BasicAccountOAuthArgument{
	AccountHost: "https://accounts.cloud.databricks.com",
	AccountID: "xyz",
})

See the cache package for more information on customizing the cache.

Index

Constants

This section is empty.

Variables

View Source
var ErrOAuthNotSupported = errors.New("databricks OAuth is not supported for this host")

Functions

func BuildDiscoveryAuthorizeURL added in v0.120.0

func BuildDiscoveryAuthorizeURL(redirectAddr, state string, pkce PKCEParams, scopes []string) string

BuildDiscoveryAuthorizeURL builds the login.databricks.com URL that initiates the discovery OAuth flow. The OIDC authorize path with all OAuth query params is URL-encoded as the destination_url parameter.

func DeriveHostFromIssuer added in v0.120.0

func DeriveHostFromIssuer(issuer string) (string, error)

DeriveHostFromIssuer extracts the workspace host (scheme + host) from an issuer URL returned in the OAuth callback iss parameter.

Examples:

"https://adb-xxx.azuredatabricks.net/oidc"           -> "https://adb-xxx.azuredatabricks.net"
"https://nike.databricks.com/oidc/accounts/xxx"      -> "https://nike.databricks.com"

func DeriveTokenEndpoint added in v0.120.0

func DeriveTokenEndpoint(issuer string) string

DeriveTokenEndpoint derives the token endpoint from an issuer URL by appending /v1/token to the issuer path.

Example:

"https://adb-xxx.net/oidc" -> "https://adb-xxx.net/oidc/v1/token"

Types

type AccountOAuthArgument

type AccountOAuthArgument interface {
	OAuthArgument

	// GetAccountHost returns the host of the account to authenticate to.
	GetAccountHost() string

	// GetAccountId returns the account ID of the account to authenticate to.
	GetAccountId() string
}

AccountOAuthArgument is an interface that provides the necessary information to authenticate using OAuth to a specific account.

type BasicAccountOAuthArgument

type BasicAccountOAuthArgument struct {
	// contains filtered or unexported fields
}

BasicAccountOAuthArgument is a basic implementation of the AccountOAuthArgument interface that links each account with exactly one OAuth token.

func NewBasicAccountOAuthArgument

func NewBasicAccountOAuthArgument(accountsHost, accountID string) (BasicAccountOAuthArgument, error)

NewBasicAccountOAuthArgument creates a new BasicAccountOAuthArgument.

func NewProfileAccountOAuthArgument added in v0.109.0

func NewProfileAccountOAuthArgument(accountsHost, accountID, profile string) (BasicAccountOAuthArgument, error)

NewProfileAccountOAuthArgument creates a new BasicAccountOAuthArgument with a profile name. When a profile is set, GetCacheKey() returns the profile name instead of the host-based key.

func (BasicAccountOAuthArgument) GetAccountHost

func (a BasicAccountOAuthArgument) GetAccountHost() string

GetAccountHost returns the host of the account to authenticate to.

func (BasicAccountOAuthArgument) GetAccountId

func (a BasicAccountOAuthArgument) GetAccountId() string

GetAccountId returns the account ID of the account to authenticate to.

func (BasicAccountOAuthArgument) GetCacheKey

func (a BasicAccountOAuthArgument) GetCacheKey() string

GetCacheKey returns a unique key for caching the OAuth token for the account. If a profile is set, the profile name is returned as the cache key. Otherwise, the key is in the format "<accountHost>/oidc/accounts/<accountID>".

func (BasicAccountOAuthArgument) GetHostCacheKey added in v0.109.0

func (a BasicAccountOAuthArgument) GetHostCacheKey() string

GetHostCacheKey returns the host-based cache key regardless of whether a profile is set. The key is in the format "<accountHost>/oidc/accounts/<accountID>".

type BasicDiscoveryOAuthArgument added in v0.120.0

type BasicDiscoveryOAuthArgument struct {
	// contains filtered or unexported fields
}

BasicDiscoveryOAuthArgument is a basic implementation of DiscoveryOAuthArgument that uses the profile name as the cache key during discovery login bootstrap.

func NewBasicDiscoveryOAuthArgument added in v0.120.0

func NewBasicDiscoveryOAuthArgument(profile string) (*BasicDiscoveryOAuthArgument, error)

NewBasicDiscoveryOAuthArgument creates a new BasicDiscoveryOAuthArgument. The profile name is required and used as the cache key during discovery login. Use it with WithDiscoveryLogin. After discovery, construct the usual host-based OAuthArgument with the discovered host.

func (*BasicDiscoveryOAuthArgument) GetCacheKey added in v0.120.0

func (a *BasicDiscoveryOAuthArgument) GetCacheKey() string

GetCacheKey returns the profile name as the cache key.

func (*BasicDiscoveryOAuthArgument) GetDiscoveredHost added in v0.120.0

func (a *BasicDiscoveryOAuthArgument) GetDiscoveredHost() string

GetDiscoveredHost returns the workspace host discovered from the callback, or empty string if not yet discovered.

func (*BasicDiscoveryOAuthArgument) SetDiscoveredHost added in v0.120.0

func (a *BasicDiscoveryOAuthArgument) SetDiscoveredHost(host string)

SetDiscoveredHost stores the workspace host derived from the iss callback parameter.

type BasicOAuthEndpointSupplier

type BasicOAuthEndpointSupplier struct {
	// Client is the ApiClient to use for making HTTP requests.
	Client *httpclient.ApiClient
}

BasicOAuthEndpointSupplier is an implementation of the OAuthEndpointSupplier interface.

func (*BasicOAuthEndpointSupplier) GetAccountOAuthEndpoints

func (c *BasicOAuthEndpointSupplier) GetAccountOAuthEndpoints(ctx context.Context, accountHost string, accountId string) (*OAuthAuthorizationServer, error)

GetAccountOAuthEndpoints returns the OAuth2 endpoints for the account. The account-level OAuth endpoints are fixed based on the account ID and host.

func (*BasicOAuthEndpointSupplier) GetEndpointsFromURL added in v0.115.0

func (c *BasicOAuthEndpointSupplier) GetEndpointsFromURL(ctx context.Context, rawURL string) (*OAuthAuthorizationServer, error)

GetEndpointsFromURL fetches OAuth2 endpoints directly from an authorization server metadata URL, bypassing host-type detection.

func (*BasicOAuthEndpointSupplier) GetUnifiedOAuthEndpoints added in v0.90.0

func (c *BasicOAuthEndpointSupplier) GetUnifiedOAuthEndpoints(ctx context.Context, host string, accountId string) (*OAuthAuthorizationServer, error)

GetUnifiedOAuthEndpoints returns the OAuth2 endpoints for the unified host

func (*BasicOAuthEndpointSupplier) GetWorkspaceOAuthEndpoints

func (c *BasicOAuthEndpointSupplier) GetWorkspaceOAuthEndpoints(ctx context.Context, workspaceHost string) (*OAuthAuthorizationServer, error)

GetWorkspaceOAuthEndpoints returns the OAuth endpoints for the given workspace.

type BasicUnifiedOAuthArgument added in v0.90.0

type BasicUnifiedOAuthArgument struct {
	// contains filtered or unexported fields
}

BasicUnifiedOAuthArgument is a basic implementation of the UnifiedOAuthArgument interface that links each account with exactly one OAuth token.

func NewBasicUnifiedOAuthArgument added in v0.90.0

func NewBasicUnifiedOAuthArgument(accountsHost, accountID string) (BasicUnifiedOAuthArgument, error)

NewBasicUnifiedOAuthArgument creates a new BasicUnifiedOAuthArgument.

func NewProfileUnifiedOAuthArgument added in v0.109.0

func NewProfileUnifiedOAuthArgument(accountsHost, accountID, profile string) (BasicUnifiedOAuthArgument, error)

NewProfileUnifiedOAuthArgument creates a new BasicUnifiedOAuthArgument with a profile name. When a profile is set, GetCacheKey() returns the profile name instead of the host-based key.

func (BasicUnifiedOAuthArgument) GetAccountId added in v0.90.0

func (a BasicUnifiedOAuthArgument) GetAccountId() string

GetAccountId returns the account ID of the account to authenticate to.

func (BasicUnifiedOAuthArgument) GetCacheKey added in v0.90.0

func (a BasicUnifiedOAuthArgument) GetCacheKey() string

GetCacheKey returns a unique key for caching the OAuth token for the account. If a profile is set, the profile name is returned as the cache key. Otherwise, the key is in the format "<hostName>/oidc/accounts/<accountID>".

func (BasicUnifiedOAuthArgument) GetHost added in v0.90.0

func (a BasicUnifiedOAuthArgument) GetHost() string

GetHost returns the host to authenticate to.

func (BasicUnifiedOAuthArgument) GetHostCacheKey added in v0.109.0

func (a BasicUnifiedOAuthArgument) GetHostCacheKey() string

GetHostCacheKey returns the host-based cache key regardless of whether a profile is set. The key is in the format "<hostName>/oidc/accounts/<accountID>".

type BasicWorkspaceOAuthArgument

type BasicWorkspaceOAuthArgument struct {
	// contains filtered or unexported fields
}

BasicWorkspaceOAuthArgument is a basic implementation of the WorkspaceOAuthArgument interface that links each host with exactly one OAuth token.

func NewBasicWorkspaceOAuthArgument

func NewBasicWorkspaceOAuthArgument(host string) (BasicWorkspaceOAuthArgument, error)

NewBasicWorkspaceOAuthArgument creates a new BasicWorkspaceOAuthArgument.

func NewProfileWorkspaceOAuthArgument added in v0.109.0

func NewProfileWorkspaceOAuthArgument(host, profile string) (BasicWorkspaceOAuthArgument, error)

NewProfileWorkspaceOAuthArgument creates a new BasicWorkspaceOAuthArgument with a profile name. When a profile is set, GetCacheKey() returns the profile name instead of the host-based key.

func (BasicWorkspaceOAuthArgument) GetCacheKey

func (a BasicWorkspaceOAuthArgument) GetCacheKey() string

GetCacheKey returns a unique key for caching the OAuth token for the workspace. If a profile is set, the profile name is returned as the cache key. Otherwise, the key is in the format "<host>".

func (BasicWorkspaceOAuthArgument) GetHostCacheKey added in v0.109.0

func (a BasicWorkspaceOAuthArgument) GetHostCacheKey() string

GetHostCacheKey returns the host-based cache key regardless of whether a profile is set. The key is in the format "<host>".

func (BasicWorkspaceOAuthArgument) GetWorkspaceHost

func (a BasicWorkspaceOAuthArgument) GetWorkspaceHost() string

GetWorkspaceHost returns the host of the workspace to authenticate to.

type DiscoveryOAuthArgument added in v0.120.0

type DiscoveryOAuthArgument interface {
	OAuthArgument

	// SetDiscoveredHost stores the workspace host derived from the iss
	// callback parameter.
	SetDiscoveredHost(host string)

	// GetDiscoveredHost returns the workspace host discovered from the
	// callback, or empty string if not yet discovered.
	GetDiscoveredHost() string
}

DiscoveryOAuthArgument is an OAuthArgument for bootstrapping the login.databricks.com discovery flow. Unlike other OAuthArgument types, it has no host at construction time. The host is discovered from the iss parameter in the OAuth callback after the user authenticates and selects a workspace.

DiscoveryOAuthArgument is only intended for use with WithDiscoveryLogin during PersistentAuth.Challenge. Once the host has been discovered, callers should construct the usual host-based OAuthArgument (for example WorkspaceOAuthArgument) for future PersistentAuth instances.

type HostCacheKeyProvider added in v0.109.0

type HostCacheKeyProvider interface {
	GetHostCacheKey() string
}

HostCacheKeyProvider is an interface for OAuthArgument implementations that can return a host-based cache key regardless of whether a profile is set.

This is used in two places:

  1. dualWrite: writing tokens to both the profile key and the legacy host key so that older Python/Java SDKs (which only know host keys) can still find the token.
  2. Token() fallback: when a profile key miss occurs, the host key is tried as a read-only fallback for tokens written before profile-based keys existed.

This interface (and the host-key write path) can be removed once all SDKs sharing the token cache have migrated to profile-based keys.

type InvalidRefreshTokenError

type InvalidRefreshTokenError struct {
	// contains filtered or unexported fields
}

InvalidRefreshTokenError is returned from PersistentAuth's Load() method if the access token has expired and the refresh token in the token cache is invalid.

type OAuthArgument

type OAuthArgument interface {
	// GetCacheKey returns a unique key for the OAuthArgument. This key is used
	// to store and retrieve the token from the token cache.
	GetCacheKey() string
}

OAuthArgument is an interface that provides the necessary information to authenticate with PersistentAuth. Implementations of this interface must implement either the WorkspaceOAuthArgument or AccountOAuthArgument interface.

type OAuthAuthorizationServer

type OAuthAuthorizationServer struct {
	// AuthorizationEndpoint is the URL to redirect users to for authorization.
	// It typically ends with /v1/authroize.
	AuthorizationEndpoint string `json:"authorization_endpoint"`

	// TokenEndpoint is the URL to exchange an authorization code for an access token.
	// It typically ends with /v1/token.
	TokenEndpoint string `json:"token_endpoint"`
}

OAuthAuthorizationServer contains the OAuth endpoints for a Databricks account or workspace.

type OAuthEndpointSupplier

type OAuthEndpointSupplier interface {
	// GetWorkspaceOAuthEndpoints returns the OAuth2 endpoints for the workspace.
	GetWorkspaceOAuthEndpoints(ctx context.Context, workspaceHost string) (*OAuthAuthorizationServer, error)

	// GetAccountOAuthEndpoints returns the OAuth2 endpoints for the account.
	GetAccountOAuthEndpoints(ctx context.Context, accountHost string, accountId string) (*OAuthAuthorizationServer, error)

	// GetUnifiedOAuthEndpoints returns the OAuth2 endpoints for the unified host.
	GetUnifiedOAuthEndpoints(ctx context.Context, host string, accountId string) (*OAuthAuthorizationServer, error)

	// GetEndpointsFromURL fetches OAuth2 endpoints directly from an authorization
	// server metadata URL, bypassing host-type detection.
	GetEndpointsFromURL(ctx context.Context, rawURL string) (*OAuthAuthorizationServer, error)
}

OAuthEndpointSupplier provides the http functionality needed for interacting with the Databricks OAuth APIs.

type PKCEParams added in v0.120.0

type PKCEParams struct {
	Challenge       string
	ChallengeMethod string
	Verifier        string
}

PKCEParams holds the PKCE challenge parameters used to build the discovery authorize URL. This mirrors authhandler.PKCEParams but is used directly so the caller does not need to import authhandler.

type PersistentAuth

type PersistentAuth struct {
	// contains filtered or unexported fields
}

PersistentAuth is an OAuth manager that handles the U2M OAuth flow. Tokens are stored in and looked up from the provided cache. Tokens include the refresh token. On load, if the access token is expired, it is refreshed using the refresh token.

The PersistentAuth is safe for concurrent use. The token cache is locked during token retrieval, refresh and storage.

func NewPersistentAuth

func NewPersistentAuth(ctx context.Context, opts ...PersistentAuthOption) (*PersistentAuth, error)

NewPersistentAuth creates a new PersistentAuth with the provided options.

func (*PersistentAuth) Challenge

func (a *PersistentAuth) Challenge() error

Challenge initiates the OAuth2 login flow for the given OAuthArgument. The OAuth2 flow is started by opening the browser to the OAuth2 authorization URL. The user is redirected to the callback server on appRedirectAddr. The callback server listens for the redirect from the identity provider and exchanges the authorization code for an access token.

func (*PersistentAuth) Close

func (a *PersistentAuth) Close() error

func (*PersistentAuth) Token

func (a *PersistentAuth) Token() (t *oauth2.Token, err error)

Token loads the OAuth2 token for the given OAuthArgument from the cache. If the token is expired, it is refreshed using the refresh token.

When a profile is set, lookup is by profile key first. If the profile key is not found and the OAuthArgument implements HostCacheKeyProvider, a fallback lookup by host key is attempted. If found, the token is returned without being migrated to the profile key (see inline comment for rationale).

type PersistentAuthOption

type PersistentAuthOption func(*PersistentAuth)

func WithBrowser

func WithBrowser(b func(url string) error) PersistentAuthOption

WithBrowser sets the browser function for the PersistentAuth.

func WithDisableOfflineAccess added in v0.97.0

func WithDisableOfflineAccess(disable bool) PersistentAuthOption

WithDisableOfflineAccess controls whether offline_access scope is requested.

func WithDiscoveryLogin added in v0.120.0

func WithDiscoveryLogin() PersistentAuthOption

WithDiscoveryLogin enables the login.databricks.com discovery flow. When enabled, Challenge() routes through login.databricks.com instead of directly to a workspace OIDC endpoint.

This option is only valid with DiscoveryOAuthArgument, which is a bootstrap-only argument type for discovery login. Once the workspace host has been discovered, callers should construct the usual host-based OAuthArgument for future PersistentAuth instances.

func WithHttpClient

func WithHttpClient(c *http.Client) PersistentAuthOption

WithHttpClient sets the HTTP client for the PersistentAuth.

func WithOAuthArgument

func WithOAuthArgument(arg OAuthArgument) PersistentAuthOption

WithOAuthArgument sets the OAuthArgument for the PersistentAuth.

func WithOAuthEndpointSupplier

func WithOAuthEndpointSupplier(c OAuthEndpointSupplier) PersistentAuthOption

WithOAuthEndpointSupplier sets the OAuth endpoint supplier for the PersistentAuth.

func WithPort added in v0.81.0

func WithPort(port int) PersistentAuthOption

WithPort sets the port for the PersistentAuth.

func WithScopes added in v0.97.0

func WithScopes(scopes []string) PersistentAuthOption

WithScopes sets the OAuth scopes for the PersistentAuth.

func WithTokenCache

func WithTokenCache(c cache.TokenCache) PersistentAuthOption

WithTokenCache sets the token cache for the PersistentAuth.

type UnifiedOAuthArgument added in v0.90.0

type UnifiedOAuthArgument interface {
	OAuthArgument

	// GetHost returns the host to authenticate to.
	GetHost() string

	// GetAccountId returns the account ID of the account to authenticate to.
	GetAccountId() string
}

UnifiedOAuthArgument is an interface that provides the necessary information to authenticate using OAuth to a host that supports both account and workspace APIs.

type WorkspaceOAuthArgument

type WorkspaceOAuthArgument interface {
	OAuthArgument

	// GetWorkspaceHost returns the host of the workspace to authenticate to.
	GetWorkspaceHost() string
}

WorkspaceOAuthArgument is an interface that provides the necessary information to authenticate using OAuth to a specific workspace.

Directories

Path Synopsis
Package cache provides an interface for storing and looking up OAuth tokens.
Package cache provides an interface for storing and looking up OAuth tokens.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL