From 8d2337276dcf1836c6d3ea849ed32c3690f7a811 Mon Sep 17 00:00:00 2001 From: eric wang Date: Thu, 14 Nov 2024 16:07:55 -0800 Subject: [PATCH 1/2] fix token cache for oauth --- databricks/sdk/credentials_provider.py | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index e91e37af..ce5f1759 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -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, @@ -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) From 61f0ddaa186d303f1489433696fc657f6ee41822 Mon Sep 17 00:00:00 2001 From: Eric Wang Date: Tue, 10 Dec 2024 09:15:40 -0800 Subject: [PATCH 2/2] address comments --- databricks/sdk/credentials_provider.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/databricks/sdk/credentials_provider.py b/databricks/sdk/credentials_provider.py index ce5f1759..02f83cb0 100644 --- a/databricks/sdk/credentials_provider.py +++ b/databricks/sdk/credentials_provider.py @@ -215,8 +215,9 @@ def external_browser(cfg: 'Config') -> Optional[CredentialsProvider]: try: credentials.token() return credentials(cfg) + # TODO: we should ideally use more specific exceptions. except Exception as e: - logger.warning(f'Failed to refresh cached token: {e}, will init new OAuth login flow') + logger.warning(f'Failed to refresh cached token: {e}. Initiating new OAuth login flow') oauth_client = OAuthClient(oidc_endpoints=oidc_endpoints, client_id=client_id,