Skip to content

Commit c005e4d

Browse files
authored
Merge pull request #583 from atlanhq/APP-6125
APP-6125 : Migrated token refresh `boolean` flag to the `TLS` We've introduced this change to support 401 token handling in a multithreaded environment, which is used by longtail packages.
2 parents acd5676 + 6691998 commit c005e4d

File tree

1 file changed

+5
-4
lines changed

1 file changed

+5
-4
lines changed

pyatlan/client/atlan.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -158,7 +158,7 @@ class AtlanClient(BaseSettings):
158158
retry: Retry = DEFAULT_RETRY
159159
_session: requests.Session = PrivateAttr(default_factory=get_session)
160160
_request_params: dict = PrivateAttr()
161-
_has_retried_for_401: bool = PrivateAttr(default=False)
161+
_401_tls: local = local()
162162
_user_id: Optional[str] = PrivateAttr(default=None)
163163
_workflow_client: Optional[WorkflowClient] = PrivateAttr(default=None)
164164
_credential_client: Optional[CredentialClient] = PrivateAttr(default=None)
@@ -223,6 +223,7 @@ def __init__(self, **data):
223223
session.mount(HTTPS_PREFIX, adapter)
224224
session.mount(HTTP_PREFIX, adapter)
225225
AtlanClient.set_current_client(self)
226+
self._401_tls.has_retried = False
226227

227228
@property
228229
def admin(self) -> AdminClient:
@@ -462,7 +463,7 @@ def _call_api_internal(
462463
if response is None:
463464
return None
464465
if response.status_code == api.expected_status:
465-
self._has_retried_for_401 = False
466+
self._401_tls.has_retried = False
466467
try:
467468
if (
468469
response.content is None
@@ -545,7 +546,7 @@ def _call_api_internal(
545546
# on authentication failure (token may have expired)
546547
if (
547548
self._user_id
548-
and not self._has_retried_for_401
549+
and not self._401_tls.has_retried
549550
and response.status_code
550551
== ErrorCode.AUTHENTICATION_PASSTHROUGH.http_error_code
551552
):
@@ -698,7 +699,7 @@ def _handle_401_token_refresh(
698699
"""
699700
new_token = self.impersonate.user(user_id=self._user_id)
700701
self.api_key = new_token
701-
self._has_retried_for_401 = True
702+
self._401_tls.has_retried = True
702703
params["headers"]["authorization"] = f"Bearer {self.api_key}"
703704
self._request_params["headers"]["authorization"] = f"Bearer {self.api_key}"
704705
LOGGER.debug("Successfully completed 401 automatic token refresh.")

0 commit comments

Comments
 (0)