Skip to content

Commit

Permalink
fix token cache for oauth
Browse files Browse the repository at this point in the history
  • Loading branch information
eric-wang-1990 committed Nov 15, 2024
1 parent ee6e70a commit 8d23372
Showing 1 changed file with 16 additions and 10 deletions.
26 changes: 16 additions & 10 deletions databricks/sdk/credentials_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,7 @@ def oauth_service_principal(cfg: 'Config') -> Optional[CredentialsProvider]:
oidc = cfg.oidc_endpoints
if oidc is None:
return None

token_source = ClientCredentials(client_id=cfg.client_id,
client_secret=cfg.client_secret,
token_url=oidc.token_endpoint,
Expand Down Expand Up @@ -210,16 +211,21 @@ def external_browser(cfg: 'Config') -> Optional[CredentialsProvider]:
credentials = token_cache.load()
if credentials:
# Force a refresh in case the loaded credentials are expired.
credentials.token()
else:
oauth_client = OAuthClient(oidc_endpoints=oidc_endpoints,
client_id=client_id,
redirect_url=redirect_url,
client_secret=client_secret)
consent = oauth_client.initiate_consent()
if not consent:
return None
credentials = consent.launch_external_browser()
# If the refresh fails, rather than throw exception we will initiate a new OAuth login flow.
try:
credentials.token()
return credentials(cfg)
except Exception as e:
logger.warning(f'Failed to refresh cached token: {e}, will init new OAuth login flow')

oauth_client = OAuthClient(oidc_endpoints=oidc_endpoints,
client_id=client_id,
redirect_url=redirect_url,
client_secret=client_secret)
consent = oauth_client.initiate_consent()
if not consent:
return None
credentials = consent.launch_external_browser()
token_cache.save(credentials)
return credentials(cfg)

Expand Down

0 comments on commit 8d23372

Please sign in to comment.