Merge pull request #670 from atlanhq/APP-7350

APP-7350: Fixed auto `401` token refresh (lets use `ContextVar` for bool flag to avoid race condition)

Author: Aryamanz29

pyatlan/client/atlan.py

@@ -143,7 +143,7 @@ class AtlanClient(BaseSettings):
     connect_timeout: float = 30.0  # 30 secs
     read_timeout: float = 900.0  # 15 mins
     retry: Retry = DEFAULT_RETRY
-    _401_has_retried: bool = PrivateAttr(default=False)
+    _401_has_retried: bool = ContextVar("_401_has_retried", default=False)
     _session: requests.Session = PrivateAttr(default_factory=get_session)
     _request_params: dict = PrivateAttr()
     _user_id: Optional[str] = PrivateAttr(default=None)
@@ -188,7 +188,7 @@ def __init__(self, **data):
         adapter = HTTPAdapter(max_retries=self.retry)
         session.mount(HTTPS_PREFIX, adapter)
         session.mount(HTTP_PREFIX, adapter)
-        self._401_has_retried = False
+        self._401_has_retried.set(False)
 
     @property
     def admin(self) -> AdminClient:
@@ -411,11 +411,11 @@ def _call_api_internal(
             # - But if the next response is != 401 (e.g. 403), and `has_retried = True`,
             # then we should reset `has_retried = False` so that future 401s can trigger a new token refresh.
             if (
-                self._401_has_retried
+                self._401_has_retried.get()
                 and response.status_code
                 != ErrorCode.AUTHENTICATION_PASSTHROUGH.http_error_code
             ):
-                self._401_has_retried = False
+                self._401_has_retried.set(False)
 
             if response.status_code == api.expected_status:
                 try:
@@ -511,7 +511,7 @@ def _call_api_internal(
                     # on authentication failure (token may have expired)
                     if (
                         self._user_id
-                        and not self._401_has_retried
+                        and not self._401_has_retried.get()
                         and response.status_code
                         == ErrorCode.AUTHENTICATION_PASSTHROUGH.http_error_code
                     ):
@@ -673,7 +673,7 @@ def _handle_401_token_refresh(
             )
             raise
         self.api_key = new_token
-        self._401_has_retried = True
+        self._401_has_retried.set(True)
         params["headers"]["authorization"] = f"Bearer {self.api_key}"
         self._request_params["headers"]["authorization"] = f"Bearer {self.api_key}"
         LOGGER.debug("Successfully completed 401 automatic token refresh.")