From 7a7d723bfffb29ac232ae74e6b0424cf4b330c56 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Wed, 20 Apr 2022 16:00:13 -0700 Subject: [PATCH] Always try to refresh token, start flow if fail (#965) Signed-off-by: Yee Hing Tong --- flytekit/clients/raw.py | 6 +++--- flytekit/clis/auth/auth.py | 6 +----- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/flytekit/clients/raw.py b/flytekit/clients/raw.py index d37a5a87ea..6bee91613a 100644 --- a/flytekit/clients/raw.py +++ b/flytekit/clients/raw.py @@ -54,7 +54,6 @@ def handler(*args, **kwargs): elif e.code() == grpc.StatusCode.NOT_FOUND: raise _user_exceptions.FlyteEntityNotExistException(e) else: - print(e) # No more retries if retry=False or max_retries reached. if (retry is False) or i == (max_retries - 1): raise @@ -195,9 +194,10 @@ def _refresh_credentials_standard(self): # metadata field yet. Therefore, if there's a mismatch, copy it over. self.set_access_token(client.credentials.access_token, authorization_header_key) return - elif client.can_refresh_token: + + try: client.refresh_access_token() - else: + except ValueError: client.start_authorization_flow() self.set_access_token(client.credentials.access_token, authorization_header_key) diff --git a/flytekit/clis/auth/auth.py b/flytekit/clis/auth/auth.py index 241f5e0a3a..f54379485a 100644 --- a/flytekit/clis/auth/auth.py +++ b/flytekit/clis/auth/auth.py @@ -190,10 +190,6 @@ def set_tokens_from_store(self): def has_valid_credentials(self) -> bool: return self._credentials is not None - @property - def can_refresh_token(self) -> bool: - return self._refresh_token is not None - def start_authorization_flow(self): # In the absence of globally-set token values, initiate the token request flow ctx = _mp_get_context("fork") @@ -292,7 +288,7 @@ def refresh_access_token(self): _keyring.delete_password(_keyring_service_name, _keyring_access_token_storage_key) _keyring.delete_password(_keyring_service_name, _keyring_refresh_token_storage_key) - return + raise ValueError(f"Non-200 returned from refresh token endpoint {resp.status_code}") self._initialize_credentials(resp) @property