From 08ef98912abfbb82b5c949014ae7f773686b5fbd Mon Sep 17 00:00:00 2001 From: Thijs Nijhuis Date: Mon, 10 Jun 2024 14:40:59 +0200 Subject: [PATCH] Added retry logic to token refresh to deal with both spa and non-spa azure app registrations --- databricks/sdk/oauth.py | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/databricks/sdk/oauth.py b/databricks/sdk/oauth.py index e9a3afb90..ed05dfc87 100644 --- a/databricks/sdk/oauth.py +++ b/databricks/sdk/oauth.py @@ -252,16 +252,23 @@ def refresh(self) -> Token: raise ValueError('oauth2: token expired and refresh token is not set') params = {'grant_type': 'refresh_token', 'refresh_token': refresh_token} headers = {} - if 'microsoft' in self._client.token_url: - # Tokens issued for the 'Single-Page Application' client-type may - # only be redeemed via cross-origin requests - headers = {'Origin': self._client.redirect_url} - return retrieve_token(client_id=self._client.client_id, - client_secret=self._client.client_secret, - token_url=self._client.token_url, - params=params, - use_params=True, - headers=headers) + while True: + try: + return retrieve_token(client_id=self._client.client_id, + client_secret=self._client.client_secret, + token_url=self._client.token_url, + params=params, + use_params=True, + headers=headers) + except ValueError as e: + if NO_ORIGIN_FOR_SPA_CLIENT_ERROR in str(e): + # Retry in cases of 'Single-Page Application' client-type with + # 'Origin' header equal to client's redirect URL. + headers = {'Origin': self._client.redirect_url} + msg = f'Retrying OAuth token exchange with {self._client.redirect_url} origin' + logger.debug(msg) + continue + raise e class Consent: