[change] Removed unused _current_client_context var

Author: Aryamanz29

pyatlan/client/atlan.py

@@ -150,9 +150,6 @@ def get_session():
 
 
 class AtlanClient(BaseSettings):
-    _current_client_ctx: ClassVar[ContextVar] = ContextVar(
-        "_current_client_ctx", default=None
-    )
     _401_has_retried_ctx: ClassVar[ContextVar] = ContextVar(
         "_401_has_retried_ctx", default=False
     )
@@ -195,25 +192,6 @@ class AtlanClient(BaseSettings):
     class Config:
         env_prefix = "atlan_"
 
-    @classmethod
-    def set_current_client(cls, client: AtlanClient):
-        """
-        Sets the current client context
-        """
-        if not isinstance(client, AtlanClient):
-            raise ErrorCode.MISSING_ATLAN_CLIENT.exception_with_parameters()
-        cls._current_client_ctx.set(client)
-
-    @classmethod
-    def get_current_client(cls) -> AtlanClient:
-        """
-        Retrieves the current client context
-        """
-        client = cls._current_client_ctx.get()
-        if not client:
-            raise ErrorCode.NO_ATLAN_CLIENT_AVAILABLE.exception_with_parameters()
-        return client
-
     def __init__(self, **data):
         super().__init__(**data)
         self._request_params = {
@@ -225,188 +203,161 @@ def __init__(self, **data):
         adapter = HTTPAdapter(max_retries=self.retry)
         session.mount(HTTPS_PREFIX, adapter)
         session.mount(HTTP_PREFIX, adapter)
-        AtlanClient.set_current_client(self)
         self._401_has_retried_ctx.set(False)
 
     @property
     def admin(self) -> AdminClient:
         if self._admin_client is None:
-            AtlanClient.set_current_client(self)
             self._admin_client = AdminClient(client=self)
         return self._admin_client
 
     @property
     def audit(self) -> AuditClient:
         if self._audit_client is None:
-            AtlanClient.set_current_client(self)
             self._audit_client = AuditClient(client=self)
         return self._audit_client
 
     @property
     def search_log(self) -> SearchLogClient:
         if self._search_log_client is None:
-            AtlanClient.set_current_client(self)
             self._search_log_client = SearchLogClient(client=self)
         return self._search_log_client
 
     @property
     def workflow(self) -> WorkflowClient:
         if self._workflow_client is None:
-            AtlanClient.set_current_client(self)
             self._workflow_client = WorkflowClient(client=self)
         return self._workflow_client
 
     @property
     def credentials(self) -> CredentialClient:
         if self._credential_client is None:
-            AtlanClient.set_current_client(self)
             self._credential_client = CredentialClient(client=self)
         return self._credential_client
 
     @property
     def group(self) -> GroupClient:
         if self._group_client is None:
-            AtlanClient.set_current_client(self)
             self._group_client = GroupClient(client=self)
         return self._group_client
 
     @property
     def role(self) -> RoleClient:
         if self._role_client is None:
-            AtlanClient.set_current_client(self)
             self._role_client = RoleClient(client=self)
         return self._role_client
 
     @property
     def asset(self) -> AssetClient:
         if self._asset_client is None:
-            AtlanClient.set_current_client(self)
             self._asset_client = AssetClient(client=self)
         return self._asset_client
 
     @property
     def impersonate(self) -> ImpersonationClient:
         if self._impersonate_client is None:
-            AtlanClient.set_current_client(self)
             self._impersonate_client = ImpersonationClient(client=self)
         return self._impersonate_client
 
     @property
     def queries(self) -> QueryClient:
         if self._query_client is None:
-            AtlanClient.set_current_client(self)
             self._query_client = QueryClient(client=self)
         return self._query_client
 
     @property
     def token(self) -> TokenClient:
         if self._token_client is None:
-            AtlanClient.set_current_client(self)
             self._token_client = TokenClient(client=self)
         return self._token_client
 
     @property
     def typedef(self) -> TypeDefClient:
         if self._typedef_client is None:
-            AtlanClient.set_current_client(self)
             self._typedef_client = TypeDefClient(client=self)
         return self._typedef_client
 
     @property
     def user(self) -> UserClient:
         if self._user_client is None:
-            AtlanClient.set_current_client(self)
             self._user_client = UserClient(client=self)
         return self._user_client
 
     @property
     def tasks(self) -> TaskClient:
         if self._task_client is None:
-            AtlanClient.set_current_client(self)
             self._task_client = TaskClient(client=self)
         return self._task_client
 
     @property
     def sso(self) -> SSOClient:
         if self._sso_client is None:
-            AtlanClient.set_current_client(self)
             self._sso_client = SSOClient(client=self)
         return self._sso_client
 
     @property
     def open_lineage(self) -> OpenLineageClient:
         if self._open_lineage_client is None:
-            AtlanClient.set_current_client(self)
             self._open_lineage_client = OpenLineageClient(client=self)
         return self._open_lineage_client
 
     @property
     def files(self) -> FileClient:
         if self._file_client is None:
-            AtlanClient.set_current_client(self)
             self._file_client = FileClient(client=self)
         return self._file_client
 
     @property
     def contracts(self) -> ContractClient:
         if self._contract_client is None:
-            AtlanClient.set_current_client(self)
             self._contract_client = ContractClient(client=self)
         return self._contract_client
 
     @property
     def atlan_tag_cache(self) -> AtlanTagCache:
         if self._atlan_tag_cache is None:
-            AtlanClient.set_current_client(self)
             self._atlan_tag_cache = AtlanTagCache(client=self)
         return self._atlan_tag_cache
 
     @property
     def enum_cache(self) -> EnumCache:
         if self._enum_cache is None:
-            AtlanClient.set_current_client(self)
             self._enum_cache = EnumCache(client=self)
         return self._enum_cache
 
     @property
     def group_cache(self) -> GroupCache:
         if self._group_cache is None:
-            AtlanClient.set_current_client(self)
             self._group_cache = GroupCache(client=self)
         return self._group_cache
 
     @property
     def role_cache(self) -> RoleCache:
         if self._role_cache is None:
-            AtlanClient.set_current_client(self)
             self._role_cache = RoleCache(client=self)
         return self._role_cache
 
     @property
     def user_cache(self) -> UserCache:
         if self._user_cache is None:
-            AtlanClient.set_current_client(self)
             self._user_cache = UserCache(client=self)
         return self._user_cache
 
     @property
     def custom_metadata_cache(self) -> CustomMetadataCache:
         if self._custom_metadata_cache is None:
-            AtlanClient.set_current_client(self)
             self._custom_metadata_cache = CustomMetadataCache(client=self)
         return self._custom_metadata_cache
 
     @property
     def connection_cache(self) -> ConnectionCache:
         if self._connection_cache is None:
-            AtlanClient.set_current_client(self)
             self._connection_cache = ConnectionCache(client=self)
         return self._connection_cache
 
     @property
     def source_tag_cache(self) -> SourceTagCache:
         if self._source_tag_cache is None:
-            AtlanClient.set_current_client(self)
             self._source_tag_cache = SourceTagCache(client=self)
         return self._source_tag_cache
 
@@ -1845,31 +1796,29 @@ def max_retries(
 
 @contextlib.contextmanager
 def client_connection(
+    client: AtlanClient,
     base_url: Optional[HttpUrl] = None,
     api_key: Optional[str] = None,
     connect_timeout: float = 30.0,
     read_timeout: float = 120.0,
     retry: Retry = DEFAULT_RETRY,
 ) -> Generator[AtlanClient, None, None]:
     """
-    Creates a new client created with the given base_url and/api_key. The AtlanClient.default_client will
-    be set to the new client. AtlanClient.default_client will be reset to the current default_client before
-    exiting the context.
-    :param base_url: the base_url to be used for the new connection. If not specified the current value will be used
-    :param api_key: the api_key to be used for the new connection. If not specified the current value will be used
+    Creates a new client created with the given base_url and/api_key.
+
+    :param base_url: the base_url to be used for the new connection.
+    If not specified the current value will be used
+    :param api_key: the api_key to be used for the new connection.
+    If not specified the current value will be used
     """
-    current_client = AtlanClient.get_current_client()
     tmp_client = AtlanClient(
-        base_url=base_url or current_client.base_url,
-        api_key=api_key or current_client.api_key,
+        base_url=base_url or client.base_url,
+        api_key=api_key or client.api_key,
         connect_timeout=connect_timeout,
         read_timeout=read_timeout,
         retry=retry,
     )
-    try:
-        yield tmp_client
-    finally:
-        AtlanClient.set_current_client(current_client)
+    yield tmp_client
 
 
 from pyatlan.model.keycloak_events import (  # noqa: E402

pyatlan/client/user.py

@@ -414,7 +414,7 @@ def _add_as(
             )
 
         token_user = self.get_current().username or ""
-        with client_connection(api_key=impersonation_token) as tmp:
+        with client_connection(client=self._client, api_key=impersonation_token) as tmp:  # type: ignore[arg-type]
             request = (
                 FluentSearch()
                 .where(Asset.GUID.eq(asset_guid))

pyatlan/model/assets/core/referenceable.py

@@ -103,18 +103,6 @@ def can_be_archived(self) -> bool:
         """
         return True
 
-    # @property
-    # def atlan_tag_names(self) -> List[str]:
-    #     from pyatlan.client.atlan import AtlanClient
-    #     from pyatlan.model.constants import DELETED_
-
-    #     if self.classification_names:
-    #         return [
-    #             AtlanClient.get_current_client().atlan_tag_cache.get_name_for_id(tag_id)
-    #             or DELETED_
-    #             for tag_id in self.classification_names
-    #         ]
-    #     return []
     @property
     def atlan_tag_names(self) -> List[str]:
         return self.classification_names or []

tests/integration/purpose_test.py

@@ -299,8 +299,10 @@ def test_token_permissions(client: AtlanClient, token):
 
 @pytest.mark.skip(reason="Test failing with HekaException")
 @pytest.mark.order(after="test_token_permissions")
-def test_run_query_with_policy(assign_tag_to_asset, token, query):
-    with client_connection(api_key=token.attributes.access_token) as redacted:
+def test_run_query_with_policy(assign_tag_to_asset, token, query, client: AtlanClient):
+    with client_connection(
+        client=client, api_key=token.attributes.access_token
+    ) as redacted:
         # The policy will take some time to go into effect
         # start by waiting a reasonable set amount of time
         # (limit the same query re-running multiple times on data store)

tests/integration/test_index_search.py

@@ -843,7 +843,9 @@ def test_purpose_search(client: AtlanClient, known_issues_purpose: Purpose):
 
 def test_read_timeout(client: AtlanClient):
     request = (FluentSearch().select()).to_request()
-    with client_connection(read_timeout=0.1, retry=Retry(total=0)) as timed_client:
+    with client_connection(
+        client=client, read_timeout=0.1, retry=Retry(total=0)
+    ) as timed_client:
         with pytest.raises(
             requests.exceptions.ReadTimeout,
             match=".Read timed out\. \(read timeout=0\.1\)",  # noqa W605
@@ -854,7 +856,7 @@ def test_read_timeout(client: AtlanClient):
 def test_connect_timeout(client: AtlanClient):
     request = (FluentSearch().select()).to_request()
     with client_connection(
-        connect_timeout=0.0001, retry=Retry(total=0)
+        client=client, connect_timeout=0.0001, retry=Retry(total=0)
     ) as timed_client:
         with pytest.raises(
             requests.exceptions.ConnectionError,