Merge branch 'main' into get_lineage

Author: ErnestoLoma

pyatlan/cache/classification_cache.py

@@ -15,7 +15,7 @@ class ClassificationCache:
     deleted_names: set[str] = set()
 
     @classmethod
-    def _refresh_cache(cls) -> None:
+    def refresh_cache(cls) -> None:
         from pyatlan.client.atlan import AtlanClient
 
         client = AtlanClient.get_default_client()
@@ -41,7 +41,7 @@ def get_id_for_name(cls, name: str) -> Optional[str]:
         cls_id = cls.map_name_to_id.get(name)
         if not cls_id and name not in cls.deleted_names:
             # If not found, refresh the cache and look again (could be stale)
-            cls._refresh_cache()
+            cls.refresh_cache()
             cls_id = cls.map_name_to_id.get(name)
             if not cls_id:
                 # If still not found after refresh, mark it as deleted (could be
@@ -58,7 +58,7 @@ def get_name_for_id(cls, idstr: str) -> Optional[str]:
         cls_name = cls.map_id_to_name.get(idstr)
         if not cls_name and idstr not in cls.deleted_ids:
             # If not found, refresh the cache and look again (could be stale)
-            cls._refresh_cache()
+            cls.refresh_cache()
             cls_name = cls.map_id_to_name.get(idstr)
             if not cls_name:
                 # If still not found after refresh, mark it as deleted (could be

pyatlan/cache/custom_metadata_cache.py

@@ -35,7 +35,7 @@ class CustomMetadataCache:
     types_by_asset: dict[str, set[type]] = dict()
 
     @classmethod
-    def _refresh_cache(cls) -> None:
+    def refresh_cache(cls) -> None:
         from pyatlan.model.core import CustomMetadata, to_snake_case
 
         client = AtlanClient.get_default_client()
@@ -98,7 +98,7 @@ def get_id_for_name(cls, name: str) -> Optional[str]:
         if cm_id := cls.map_name_to_id.get(name):
             return cm_id
         # If not found, refresh the cache and look again (could be stale)
-        cls._refresh_cache()
+        cls.refresh_cache()
         return cls.map_name_to_id.get(name)
 
     @classmethod
@@ -109,14 +109,14 @@ def get_name_for_id(cls, idstr: str) -> Optional[str]:
         if cm_name := cls.map_id_to_name.get(idstr):
             return cm_name
         # If not found, refresh the cache and look again (could be stale)
-        cls._refresh_cache()
+        cls.refresh_cache()
         return cls.map_id_to_name.get(idstr)
 
     @classmethod
     def get_type_for_id(cls, idstr: str) -> Optional[type]:
         if cm_type := cls.map_id_to_type.get(idstr):
             return cm_type
-        cls._refresh_cache()
+        cls.refresh_cache()
         return cls.map_id_to_type.get(idstr)
 
     @classmethod
@@ -129,7 +129,7 @@ def get_all_custom_attributes(
         of each of those attributes).
         """
         if len(cls.cache_by_id) == 0 or force_refresh:
-            cls._refresh_cache()
+            cls.refresh_cache()
         m = {}
         for type_id, cm in cls.cache_by_id.items():
             type_name = cls.get_name_for_id(type_id)
@@ -165,7 +165,7 @@ def get_attr_id_for_name(cls, set_name: str, attr_name: str) -> Optional[str]:
                 # If found, return straight away
                 return attr_id
             # Otherwise, refresh the cache and look again (could be stale)
-            cls._refresh_cache()
+            cls.refresh_cache()
             if sub_map := cls.map_attr_name_to_id.get(set_id):
                 return sub_map.get(attr_name)
         return None
@@ -180,7 +180,7 @@ def get_attr_name_for_id(cls, set_id: str, attr_id: str) -> Optional[str]:
             attr_name = sub_map.get(attr_id)
             if attr_name:
                 return attr_name
-            cls._refresh_cache()
+            cls.refresh_cache()
             if sub_map := cls.map_attr_id_to_name.get(set_id):
                 return sub_map.get(attr_id)
         return None
@@ -200,7 +200,7 @@ def get_attributes_for_search_results(cls, set_name: str) -> Optional[list[str]]
         if set_id := cls.get_id_for_name(set_name):
             if dot_names := cls._get_attributes_for_search_results(set_id):
                 return dot_names
-            cls._refresh_cache()
+            cls.refresh_cache()
             return cls._get_attributes_for_search_results(set_id)
         return None
 

pyatlan/client/atlan.py

@@ -495,10 +495,22 @@ def create_typedef(self, typedef: TypeDef) -> TypeDefResponse:
         raw_json = self._call_api(
             CREATE_TYPE_DEFS, request_obj=payload, exclude_unset=False
         )
+        if isinstance(typedef, ClassificationDef):
+            from pyatlan.cache.classification_cache import ClassificationCache
+            ClassificationCache.refresh_cache()
+        if isinstance(typedef, CustomMetadataDef):
+            from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
+            CustomMetadataCache.refresh_cache()
         return TypeDefResponse(**raw_json)
 
     def purge_typedef(self, internal_name: str) -> None:
         self._call_api(DELETE_TYPE_DEF_BY_NAME.format_path_with_params(internal_name))
+        # TODO: if we know which kind of typedef is being purged, we only need
+        #  to refresh that particular cache
+        from pyatlan.cache.classification_cache import ClassificationCache
+        from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
+        ClassificationCache.refresh_cache()
+        CustomMetadataCache.refresh_cache()
 
     @validate_arguments()
     def add_classifications(