[refactor/change] Explicitly provide client to public methods and CM methods

Author: Aryamanz29

pyatlan/cache/source_tag_cache.py

@@ -154,7 +154,7 @@ def get_name(self, asset: Asset):
         if not isinstance(asset, Asset):
             return
         try:
-            source_tag_name = str(SourceTagName(asset))
+            source_tag_name = str(SourceTagName(client=self.client, tag=asset))
         except AtlanError as e:
             LOGGER.error(
                 "Unable to construct a source tag name for: %s", asset.qualified_name
@@ -174,24 +174,18 @@ class SourceTagName(AbstractAssetName):
     _TYPE_NAME = "SourceTagAttachment"
     _CONNECTION_DELIMITER = "@@"
 
-    def __init__(self, tag: Union[str, Asset]):
+    def __init__(self, client: AtlanClient, tag: Union[str, Asset]):
         self.connection = None
         self.partial_tag_name = None
 
         # NOTE: Checking if the first result is an "Asset" since in pyatlan,
         # "DbtTag" extends "Dbt" (unlike other tags like "SnowflakeTag" that extend the "Tag" model),
         # preventing Dbt tags from being excluded from caching:
         if isinstance(tag, Asset):
-            from pyatlan.client.atlan import AtlanClient
-
             source_tag_qn = tag.qualified_name or ""
             tokens = source_tag_qn.split("/")
             connection_qn = "/".join(tokens[:3]) if len(tokens) >= 3 else ""
-            conn = (
-                AtlanClient.get_current_client().connection_cache.get_by_qualified_name(
-                    connection_qn
-                )
-            )
+            conn = client.connection_cache.get_by_qualified_name(connection_qn)
             self.connection = ConnectionName(conn)
             self.partial_tag_name = source_tag_qn[len(connection_qn) + 1 :]  # noqa
 

pyatlan/client/asset.py

@@ -628,6 +628,7 @@ def save(
             entities.append(entity)
         for asset in entities:
             asset.validate_required()
+            asset.flush_custom_metadata(client=self._client)  # type: ignore[arg-type]
         request = BulkRequest[Asset](entities=entities)
         raw_json = self._client._call_api(BULK_UPDATE, query_params, request)
         response = AssetMutationResponse(**raw_json)
@@ -745,6 +746,7 @@ def save_replacing_cm(
             entities.append(entity)
         for asset in entities:
             asset.validate_required()
+            asset.flush_custom_metadata(client=self._client)  # type: ignore[arg-type]
         request = BulkRequest[Asset](entities=entities)
         raw_json = self._client._call_api(BULK_UPDATE, query_params, request)
         return AssetMutationResponse(**raw_json)
@@ -885,6 +887,7 @@ def _restore_asset(self, asset: Asset) -> AssetMutationResponse:
             "replaceBusinessAttributes": False,
             "overwriteBusinessAttributes": False,
         }
+        to_restore.flush_custom_metadata(self._client)  # type: ignore[arg-type]
         request = BulkRequest[Asset](entities=[to_restore])
         raw_json = self._client._call_api(BULK_UPDATE, query_params, request)
         return AssetMutationResponse(**raw_json)
@@ -919,7 +922,7 @@ def _modify_tags(
             )
 
         atlan_tag = [
-            AtlanTag(
+            AtlanTag(  # type: ignore[call-arg]
                 type_name=AtlanTagName(display_text=name),
                 propagate=propagate,
                 remove_propagations_on_entity_delete=remove_propagation_on_delete,
@@ -1091,6 +1094,7 @@ def _update_asset_by_attribute(
         self, asset: A, asset_type: Type[A], qualified_name: str
     ):
         query_params = {"attr:qualifiedName": qualified_name}
+        asset.flush_custom_metadata(client=self._client)  # type: ignore[arg-type]
         raw_json = self._client._call_api(
             PARTIAL_UPDATE_ENTITY_BY_ATTRIBUTE.format_path_with_params(
                 asset_type.__name__
@@ -1404,7 +1408,7 @@ def remove_custom_metadata(self, guid: str, cm_name: str):
         :param cm_name: human-readable name of the custom metadata to remove
         :raises AtlanError: on any API communication issue
         """
-        custom_metadata = CustomMetadataDict(name=cm_name)
+        custom_metadata = CustomMetadataDict(client=self._client, name=cm_name)  # type: ignore[arg-type]
         # invoke clear_all so all attributes are set to None and consequently removed
         custom_metadata.clear_all()
         custom_metadata_request = CustomMetadataRequest.create(

pyatlan/client/open_lineage.py

@@ -45,10 +45,6 @@ def create_connection(
         :param admin_groups:list of admin groups to associate with this connection
         :return: details of the connection created
         """
-        from pyatlan.client.atlan import AtlanClient
-
-        client = AtlanClient.get_current_client()
-
         create_credential = Credential()
         create_credential.auth_type = "atlan_api_key"
         create_credential.name = (
@@ -65,7 +61,7 @@ def create_connection(
             "events.topic": f"openlineage_{connector_type.value}",
             "events.urlPath": f"/events/openlineage/{connector_type.value}/api/v1/lineage",
         }
-        response = client.credentials.creator(credential=create_credential)
+        response = self._client.credentials.creator(credential=create_credential)  # type: ignore[attr-defined]
         connection = Connection.creator(
             name=name,
             connector_type=connector_type,
@@ -75,7 +71,7 @@ def create_connection(
         )
 
         connection.default_credential_guid = response.id
-        return client.asset.save(connection)
+        return self._client.asset.save(connection)  # type: ignore[attr-defined]
 
     @validate_arguments
     def send(

pyatlan/client/typedef.py

@@ -64,18 +64,6 @@ def _build_typedef_request(typedef: TypeDef) -> TypeDefResponse:
     return payload
 
 
-def _refresh_caches(typedef: TypeDef) -> None:
-    from pyatlan.client.atlan import AtlanClient
-
-    client = AtlanClient.get_current_client()
-    if isinstance(typedef, AtlanTagDef):
-        client.atlan_tag_cache.refresh_cache()
-    if isinstance(typedef, CustomMetadataDef):
-        client.custom_metadata_cache.refresh_cache()
-    if isinstance(typedef, EnumDef):
-        client.enum_cache.refresh_cache()
-
-
 class TypeDefFactory:
     @staticmethod
     def create(raw_json: dict) -> TypeDef:
@@ -117,6 +105,14 @@ def __init__(self, client: ApiCaller):
             )
         self._client = client
 
+    def _refresh_caches(self, typedef: TypeDef) -> None:
+        if isinstance(typedef, AtlanTagDef):
+            self._client.atlan_tag_cache.refresh_cache()  # type: ignore[attr-defined]
+        if isinstance(typedef, CustomMetadataDef):
+            self._client.custom_metadata_cache.refresh_cache()  # type: ignore[attr-defined]
+        if isinstance(typedef, EnumDef):
+            self._client.enum_cache.refresh_cache()  # type: ignore[attr-defined]
+
     def get_all(self) -> TypeDefResponse:
         """
         Retrieves a TypeDefResponse object that contains a list of all the type definitions in Atlan.
@@ -189,7 +185,7 @@ def create(self, typedef: TypeDef) -> TypeDefResponse:
         raw_json = self._client._call_api(
             CREATE_TYPE_DEFS, request_obj=payload, exclude_unset=True
         )
-        _refresh_caches(typedef)
+        self._refresh_caches(typedef)
         return TypeDefResponse(**raw_json)
 
     @validate_arguments
@@ -210,7 +206,7 @@ def update(self, typedef: TypeDef) -> TypeDefResponse:
         raw_json = self._client._call_api(
             UPDATE_TYPE_DEFS, request_obj=payload, exclude_unset=True
         )
-        _refresh_caches(typedef)
+        self._refresh_caches(typedef)
         return TypeDefResponse(**raw_json)
 
     @validate_arguments
@@ -226,29 +222,26 @@ def purge(self, name: str, typedef_type: type) -> None:
         :raises NotFoundError: if the typedef you are trying to delete cannot be found
         :raises AtlanError: on any API communication issue
         """
-        from pyatlan.client.atlan import AtlanClient
-
-        client = AtlanClient.get_current_client()
         if typedef_type == CustomMetadataDef:
-            internal_name = client.custom_metadata_cache.get_id_for_name(name)
+            internal_name = self._client.custom_metadata_cache.get_id_for_name(name)  # type: ignore[attr-defined]
         elif typedef_type == EnumDef:
             internal_name = name
         elif typedef_type == AtlanTagDef:
-            internal_name = str(client.atlan_tag_cache.get_id_for_name(name))
+            internal_name = str(self._client.atlan_tag_cache.get_id_for_name(name))  # type: ignore[attr-defined]
         else:
             raise ErrorCode.UNABLE_TO_PURGE_TYPEDEF_OF_TYPE.exception_with_parameters(
                 typedef_type
             )
         if internal_name:
-            self._client._call_api(
+            self._client._client._call_api(  # type: ignore[attr-defined]
                 DELETE_TYPE_DEF_BY_NAME.format_path_with_params(internal_name)
             )
         else:
             raise ErrorCode.TYPEDEF_NOT_FOUND_BY_NAME.exception_with_parameters(name)
 
         if typedef_type == CustomMetadataDef:
-            client.custom_metadata_cache.refresh_cache()
+            self._client.custom_metadata_cache.refresh_cache()  # type: ignore[attr-defined]
         elif typedef_type == EnumDef:
-            client.enum_cache.refresh_cache()
+            self._client.enum_cache.refresh_cache()  # type: ignore[attr-defined]
         elif typedef_type == AtlanTagDef:
-            client.atlan_tag_cache.refresh_cache()
+            self._client.atlan_tag_cache.refresh_cache()  # type: ignore[attr-defined]

pyatlan/client/user.py

@@ -57,15 +57,12 @@ def create(
         :raises AtlanError: on any API communication issue
         :returns: a UserResponse object which contains the list of details of created users if `return_info` is `True`, otherwise `None`
         """
-        from pyatlan.client.atlan import AtlanClient
 
         cur = CreateUserRequest(users=[])
         for user in users:
             role_name = str(user.workspace_role)
             if (
-                role_id := AtlanClient.get_current_client().role_cache.get_id_for_name(
-                    role_name
-                )
+                role_id := self._client.role_cache.get_id_for_name(role_name)  # type: ignore[attr-defined]
             ) and user.email:
                 to_create = CreateUserRequest.CreateUser(
                     email=user.email,

pyatlan/model/assets/badge.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from typing import ClassVar, List, Optional
+from typing import TYPE_CHECKING, ClassVar, List, Optional
 from warnings import warn
 
 from pydantic.v1 import Field, StrictStr, validator
@@ -16,6 +16,9 @@
 
 from .core.asset import Asset
 
+if TYPE_CHECKING:
+    from pyatlan.client.atlan import AtlanClient
+
 
 class Badge(Asset, type_name="Badge"):
     """Description"""
@@ -131,6 +134,7 @@ class Attributes(Asset.Attributes):
         def create(
             cls,
             *,
+            client: AtlanClient,
             name: StrictStr,
             cm_name: str,
             cm_attribute: str,
@@ -140,9 +144,6 @@ def create(
                 ["name", "cm_name", "cm_attribute", "badge_conditions"],
                 [name, cm_name, cm_attribute, badge_conditions],
             )
-            from pyatlan.client.atlan import AtlanClient
-
-            client = AtlanClient.get_current_client()
             cm_id = client.custom_metadata_cache.get_id_for_name(cm_name)
             cm_attr_id = client.custom_metadata_cache.get_attr_id_for_name(
                 set_name=cm_name, attr_name=cm_attribute

pyatlan/model/assets/connection.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import ClassVar, Dict, List, Optional, Set
+from typing import TYPE_CHECKING, ClassVar, Dict, List, Optional, Set
 from warnings import warn
 
 from pydantic.v1 import Field, validator
@@ -22,6 +22,9 @@
 
 from .core.asset import Asset
 
+if TYPE_CHECKING:
+    from pyatlan.client.atlan import AtlanClient
+
 
 class Connection(Asset, type_name="Connection"):
     """Description"""
@@ -31,6 +34,7 @@ class Connection(Asset, type_name="Connection"):
     def creator(
         cls,
         *,
+        client: AtlanClient,
         name: str,
         connector_type: AtlanConnectorType,
         admin_users: Optional[List[str]] = None,
@@ -50,6 +54,9 @@ def creator(
             connector_name=connector_type.value,
             category=connector_type.category.value,
         )
+        client.user_cache.validate_names(names=admin_users or [])
+        client.role_cache.validate_idstrs(idstrs=admin_roles or [])
+        client.group_cache.validate_aliases(aliases=admin_groups or [])
         attr.admin_users = set() if admin_users is None else set(admin_users)
         attr.admin_groups = set() if admin_groups is None else set(admin_groups)
         attr.admin_roles = set() if admin_roles is None else set(admin_roles)
@@ -62,6 +69,7 @@ def creator(
     def create(
         cls,
         *,
+        client: AtlanClient,
         name: str,
         connector_type: AtlanConnectorType,
         admin_users: Optional[List[str]] = None,
@@ -77,6 +85,7 @@ def create(
             stacklevel=2,
         )
         return cls.creator(
+            client=client,
             name=name,
             connector_type=connector_type,
             admin_users=admin_users,
@@ -724,36 +733,6 @@ class Attributes(Asset.Attributes):
 
         is_loaded: bool = Field(default=True)
 
-        @validator("admin_users")
-        def admin_users_valid(cls, admin_users, values):
-            from pyatlan.client.atlan import AtlanClient
-
-            if values.get("is_loaded", False):
-                AtlanClient.get_current_client().user_cache.validate_names(
-                    names=admin_users
-                )
-            return admin_users
-
-        @validator("admin_roles")
-        def admin_roles_valid(cls, admin_roles, values):
-            from pyatlan.client.atlan import AtlanClient
-
-            if values.get("is_loaded", False):
-                AtlanClient.get_current_client().role_cache.validate_idstrs(
-                    idstrs=admin_roles
-                )
-            return admin_roles
-
-        @validator("admin_groups")
-        def admin_groups_valid(cls, admin_groups, values):
-            from pyatlan.client.atlan import AtlanClient
-
-            if values.get("is_loaded", False):
-                AtlanClient.get_current_client().group_cache.validate_aliases(
-                    aliases=admin_groups
-                )
-                return admin_groups
-
     attributes: Connection.Attributes = Field(
         default_factory=lambda: Connection.Attributes(),
         description=(

pyatlan/model/assets/core/data_product.py

@@ -125,18 +125,15 @@ def create_for_modification(
             name=name,
         )
 
-    def get_assets(self, client: Optional[AtlanClient] = None):
+    def get_assets(self, client: AtlanClient):
         """
         Reterieves list of all assets linked to the provided data product.
 
-        :param client: connectivity to an Atlan tenant (optional). If not provided, the default client will be used.
+        :param client: connectivity to an Atlan tenant. If not provided, the default client will be used.
         :raises InvalidRequestError: if DataProduct asset DSL cannot be found (does not exist) in Atlan
         :raises AtlanError: if there is an issue interacting with the API
         :returns: instance of `IndexSearchResults` with list of all assets linked to the provided data product
         """
-        from pyatlan.client.atlan import AtlanClient
-
-        client = AtlanClient.get_current_client() if not client else client
         dp_dsl = self.data_product_assets_d_s_l
         if not dp_dsl:
             raise ErrorCode.MISSING_DATA_PRODUCT_ASSET_DSL.exception_with_parameters()