[change] Removed bool contextvar for managing 401 token refresh

- Also removed `PyAtlanThreadPoolExecutor` as it was not used anymore.

Author: Aryamanz29

pyatlan/client/atlan.py

@@ -12,20 +12,8 @@
 from contextvars import ContextVar
 from http import HTTPStatus
 from importlib.resources import read_text
-from threading import local
 from types import SimpleNamespace
-from typing import (
-    Any,
-    ClassVar,
-    Dict,
-    Generator,
-    List,
-    Literal,
-    Optional,
-    Set,
-    Type,
-    Union,
-)
+from typing import Any, Dict, Generator, List, Literal, Optional, Set, Type, Union
 from urllib.parse import urljoin
 from warnings import warn
 
@@ -150,17 +138,14 @@ def get_session():
 
 
 class AtlanClient(BaseSettings):
-    _401_has_retried_ctx: ClassVar[ContextVar] = ContextVar(
-        "_401_has_retried_ctx", default=False
-    )
     base_url: Union[Literal["INTERNAL"], HttpUrl]
     api_key: str
     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)
     _session: requests.Session = PrivateAttr(default_factory=get_session)
     _request_params: dict = PrivateAttr()
-    _401_tls: local = local()
     _user_id: Optional[str] = PrivateAttr(default=None)
     _workflow_client: Optional[WorkflowClient] = PrivateAttr(default=None)
     _credential_client: Optional[CredentialClient] = PrivateAttr(default=None)
@@ -203,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_ctx.set(False)
+        self._401_has_retried = False
 
     @property
     def admin(self) -> AdminClient:
@@ -426,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_ctx.get()
+                self._401_has_retried
                 and response.status_code
                 != ErrorCode.AUTHENTICATION_PASSTHROUGH.http_error_code
             ):
-                self._401_has_retried_ctx.set(False)
+                self._401_has_retried = False
 
             if response.status_code == api.expected_status:
                 try:
@@ -521,7 +506,7 @@ def _call_api_internal(
                     # on authentication failure (token may have expired)
                     if (
                         self._user_id
-                        and not self._401_has_retried_ctx.get()
+                        and not self._401_has_retried
                         and response.status_code
                         == ErrorCode.AUTHENTICATION_PASSTHROUGH.http_error_code
                     ):
@@ -680,7 +665,7 @@ def _handle_401_token_refresh(
             )
             raise
         self.api_key = new_token
-        self._401_has_retried_ctx.set(True)
+        self._401_has_retried = 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.")

pyatlan/utils.py

@@ -8,8 +8,7 @@
 import random
 import re
 import time
-from concurrent.futures import ThreadPoolExecutor
-from contextvars import ContextVar, copy_context
+from contextvars import ContextVar
 from datetime import datetime
 from enum import Enum, EnumMeta
 from functools import reduce, wraps
@@ -474,53 +473,6 @@ def validate_single_required_field(field_names: List[str], values: List[Any]):
         )
 
 
-class PyAtlanThreadPoolExecutor(ThreadPoolExecutor):
-    """
-    A ThreadPoolExecutor that preserves context variables (e.g: `AtlanClient`)
-    across threads—useful when running SDK methods in multithreading and async environments.
-
-    For example:
-
-    ```
-    import asyncio
-    from functools import partial
-    from pyatlan.model.assets import Table
-    from pyatlan.client.atlan import AtlanClient
-    from pyatlan.utils import PyAtlanThreadPoolExecutor
-
-    client = AtlanClient()
-
-    async def fetch_asset():
-        loop = asyncio.get_running_loop()
-        sdk_func = partial(
-            client.asset.get_by_guid, "ef1ffe2c-8fc9-433a-8cf8-b4583f2d2375", Table
-        )
-        return await loop.run_in_executor(
-            executor=PyAtlanThreadPoolExecutor(), func=sdk_func
-        )
-
-    result = asyncio.run(fetch_asset())
-    ```
-    """
-
-    _SDK_CONTEXT_VAR_NAMES = [
-        "_current_client_ctx",
-        "_401_has_retried_ctx",
-    ]
-
-    def submit(self, fn, /, *args, **kwargs):
-        ctx_vars = copy_context().items()
-
-        def _fn():
-            for var, value in ctx_vars:
-                # Only set the context variables that are used by the SDK
-                if var.name in self._SDK_CONTEXT_VAR_NAMES:
-                    var.set(value)
-            return fn(*args, **kwargs)
-
-        return super().submit(_fn)
-
-
 class ExtendableEnumMeta(EnumMeta):
     def __init__(cls, name, bases, namespace):
         super().__init__(name, bases, namespace)

tests/unit/test_client.py

@@ -2152,62 +2152,6 @@ def test_atlan_call_api_server_error_messages_with_causes(
             client.asset.save(glossary)
 
 
-# @pytest.mark.parametrize("thread_count", [3])  # Run with three threads
-# def test_atlan_client_tls(thread_count):
-#     """Tests that AtlanClient instances remain isolated across multiple threads."""
-#     validation_results = {}
-#     results_lock = threading.Lock()
-
-#     def _test_atlan_client_isolation(name, api_key1, api_key2, api_key3):
-#         """Creates three AtlanClient instances within the same thread and verifies isolation."""
-#         # Instantiate three separate AtlanClient instances
-#         client1 = AtlanClient(base_url="https://test.atlan.com", api_key=api_key1)
-#         time.sleep(0.2)
-#         observed1 = client1.get_current_client().api_key  # Should match api_key1
-
-#         client2 = AtlanClient(base_url="https://test.atlan.com", api_key=api_key2)
-#         time.sleep(0.2)
-#         observed2 = client2.get_current_client().api_key  # Should match api_key2
-
-#         client3 = AtlanClient(base_url="https://test.atlan.com", api_key=api_key3)
-#         time.sleep(0.2)
-#         observed3 = client3.get_current_client().api_key  # Should match api_key3
-
-#         # Store results in a thread-safe way
-#         with results_lock:
-#             validation_results[name] = (observed1, observed2, observed3)
-
-#     # Define unique API keys for each thread
-#     api_keys = [
-#         ("API_KEY_1A", "API_KEY_1B", "API_KEY_1C"),
-#         ("API_KEY_2A", "API_KEY_2B", "API_KEY_2C"),
-#         ("API_KEY_3A", "API_KEY_3B", "API_KEY_3C"),
-#     ]
-
-#     threads = []
-#     for i in range(thread_count):
-#         thread = threading.Thread(
-#             target=_test_atlan_client_isolation,
-#             args=(f"thread{i + 1}", *api_keys[i]),
-#         )
-#         threads.append(thread)
-#         thread.start()
-
-#     # Wait for all threads to finish
-#     for thread in threads:
-#         thread.join()
-
-#     # Validate that each thread's clients retained their assigned API keys
-#     for i in range(thread_count):
-#         thread_name = f"thread{i + 1}"
-#         expected_keys = api_keys[i]
-
-#         assert validation_results[thread_name] == expected_keys, (
-#             f"Clients were overwritten across threads! "
-#             f"{thread_name} saw {validation_results[thread_name]} instead of {expected_keys}"
-#         )
-
-
 class TestBatch:
     def test_init(self, mock_atlan_client):
         sut = Batch(client=mock_atlan_client, max_size=10)