Merge pull request #635 from atlanhq/APP-6703

APP-6703: Move translation logic out of pydantic validators

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(