Merge pull request #624 from atlanhq/APP-6579

APP-6156: Migrated from `TLS` to `ContextVars` to support both multithreading and async environments

Author: Aryamanz29

pyatlan/client/atlan.py

@@ -150,7 +150,12 @@ def get_session():
 
 
 class AtlanClient(BaseSettings):
-    _current_client_tls: ClassVar[local] = local()  # Thread-local storage (TLS)
+    _current_client_ctx: ClassVar[ContextVar] = ContextVar(
+        "_current_client_ctx", default=None
+    )
+    _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
@@ -190,37 +195,24 @@ class AtlanClient(BaseSettings):
     class Config:
         env_prefix = "atlan_"
 
-    @classmethod
-    def init_for_multithreading(cls, client: AtlanClient):
-        """
-        Prepares the given client for use in multi-threaded environments.
-
-        This sets the thread-local context and resets internal retry flags
-        to ensure correct behavior when using the client across multiple threads.
-        """
-        AtlanClient.set_current_client(client)
-        client._401_tls.has_retried = False
-
     @classmethod
     def set_current_client(cls, client: AtlanClient):
         """
-        Sets the current client to thread-local storage (TLS)
+        Sets the current client context
         """
         if not isinstance(client, AtlanClient):
             raise ErrorCode.MISSING_ATLAN_CLIENT.exception_with_parameters()
-        cls._current_client_tls.client = client
+        cls._current_client_ctx.set(client)
 
     @classmethod
     def get_current_client(cls) -> AtlanClient:
         """
-        Retrieves the current client
+        Retrieves the current client context
         """
-        if (
-            not hasattr(cls._current_client_tls, "client")
-            or not cls._current_client_tls.client
-        ):
+        client = cls._current_client_ctx.get()
+        if not client:
             raise ErrorCode.NO_ATLAN_CLIENT_AVAILABLE.exception_with_parameters()
-        return cls._current_client_tls.client
+        return client
 
     def __init__(self, **data):
         super().__init__(**data)
@@ -233,7 +225,8 @@ def __init__(self, **data):
         adapter = HTTPAdapter(max_retries=self.retry)
         session.mount(HTTPS_PREFIX, adapter)
         session.mount(HTTP_PREFIX, adapter)
-        AtlanClient.init_for_multithreading(self)
+        AtlanClient.set_current_client(self)
+        self._401_has_retried_ctx.set(False)
 
     @property
     def admin(self) -> AdminClient:
@@ -482,11 +475,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_tls.has_retried
+                self._401_has_retried_ctx.get()
                 and response.status_code
                 != ErrorCode.AUTHENTICATION_PASSTHROUGH.http_error_code
             ):
-                self._401_tls.has_retried = False
+                self._401_has_retried_ctx.set(False)
 
             if response.status_code == api.expected_status:
                 try:
@@ -571,7 +564,7 @@ def _call_api_internal(
                     # on authentication failure (token may have expired)
                     if (
                         self._user_id
-                        and not self._401_tls.has_retried
+                        and not self._401_has_retried_ctx.get()
                         and response.status_code
                         == ErrorCode.AUTHENTICATION_PASSTHROUGH.http_error_code
                     ):
@@ -732,7 +725,7 @@ def _handle_401_token_refresh(
             )
             raise
         self.api_key = new_token
-        self._401_tls.has_retried = True
+        self._401_has_retried_ctx.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.")

pyatlan/utils.py

@@ -8,7 +8,8 @@
 import random
 import re
 import time
-from contextvars import ContextVar
+from concurrent.futures import ThreadPoolExecutor
+from contextvars import ContextVar, copy_context
 from datetime import datetime
 from enum import Enum, EnumMeta
 from functools import reduce, wraps
@@ -473,6 +474,51 @@ 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_event_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)
+
+
 class ExtendableEnumMeta(EnumMeta):
     def __init__(cls, name, bases, namespace):
         super().__init__(name, bases, namespace)