Merge pull request #50 from atlanhq/ACTIV-557

Adds shortcuts to create custom metadata properties

Author: ErnestoLoma

Diff for: .github/workflows/pytest.yaml

@@ -29,7 +29,9 @@ jobs:
           flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
       - name: Test with pytest
         env: # Or as an environment variable
-          ATLAN_API_KEY: ${{ secrets.MARK_ATLAN_API_KEY }}
-          ATLAN_BASE_URL: https://mark.atlan.com
+          ATLAN_API_KEY: ${{ secrets.ATLAN_API_KEY }}
+          ATLAN_BASE_URL: ${{ secrets.ATLAN_BASE_URL }}
+          MARK_API_KEY: ${{ secrets.MARK_ATLAN_API_KEY }}
+          MARK_BASE_URL: https://mark.atlan.com
         run: |
           pytest

Diff for: pyatlan/cache/classification_cache.py

@@ -7,7 +7,6 @@
 
 
 class ClassificationCache:
-
     cache_by_id: dict[str, ClassificationDef] = dict()
     map_id_to_name: dict[str, str] = dict()
     map_name_to_id: dict[str, str] = dict()

Diff for: pyatlan/cache/custom_metadata_cache.py

@@ -24,7 +24,6 @@ def __get__(self, instance, owner):
 
 
 class CustomMetadataCache:
-
     cache_by_id: dict[str, CustomMetadataDef] = dict()
     map_id_to_name: dict[str, str] = dict()
     map_id_to_type: dict[str, type] = dict()
@@ -66,25 +65,26 @@ def refresh_cache(cls) -> None:
                 applicable_types: set[str] = set()
                 if cm.attribute_defs:
                     for attr in cm.attribute_defs:
-                        if attr.options.custom_applicable_entity_types:
+                        if attr.options and attr.options.custom_applicable_entity_types:
                             applicable_types.update(
                                 json.loads(attr.options.custom_applicable_entity_types)
                             )
-                        attr_id = attr.name
-                        attr_name = attr.display_name
-                        cls.map_attr_id_to_name[type_id][attr_id] = attr_name
+                        attr_id = str(attr.name)
+                        attr_name = str(attr.display_name)
+                        # Use a renamed attribute everywhere
+                        attr_renamed = to_snake_case(attr_name.replace(" ", ""))
+                        cls.map_attr_id_to_name[type_id][attr_id] = attr_renamed
                         if attr.options and attr.options.is_archived:
-                            cls.archived_attr_ids[attr_id] = attr_name
-                        elif attr_name in cls.map_attr_name_to_id[type_id]:
+                            cls.archived_attr_ids[attr_id] = attr_renamed
+                        elif attr_renamed in cls.map_attr_name_to_id[type_id]:
                             raise LogicError(
-                                f"Multiple custom attributes with exactly the same name ({attr_name}) "
+                                f"Multiple custom attributes with exactly the same name ({attr_renamed}) "
                                 f"found for: {type_name}",
                                 code="ATLAN-PYTHON-500-100",
                             )
                         else:
-                            attr_name = to_snake_case(attr_name.replace(" ", ""))
-                            setattr(attrib_type, attr_name, Synonym(attr_id))
-                            cls.map_attr_name_to_id[type_id][attr_name] = attr_id
+                            setattr(attrib_type, attr_renamed, Synonym(attr_id))
+                            cls.map_attr_name_to_id[type_id][attr_renamed] = attr_id
                     for asset_type in applicable_types:
                         if asset_type not in cls.types_by_asset:
                             cls.types_by_asset[asset_type] = set()
@@ -230,3 +230,16 @@ def get_custom_metadata(
                 else ba_type()
             )
         raise ValueError(f"Custom metadata {name} is not applicable to {type_name}")
+
+    @classmethod
+    def get_custom_metadata_def(cls, name: str) -> CustomMetadataDef:
+        """
+        Retrieve the full custom metadata structure definition.
+        """
+        ba_id = cls.get_id_for_name(name)
+        if ba_id is None:
+            raise ValueError(f"No custom metadata with the name: {name} exist")
+        if typedef := cls.cache_by_id.get(ba_id):
+            return typedef
+        else:
+            raise ValueError(f"No custom metadata with the name: {name} found")

Diff for: pyatlan/cache/enum_cache.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2023 Atlan Pte. Ltd.
+from typing import Optional
+
+from pyatlan.model.enums import AtlanTypeCategory
+from pyatlan.model.typedef import EnumDef
+
+
+class EnumCache:
+    cache_by_name: dict[str, EnumDef] = dict()
+
+    @classmethod
+    def refresh_cache(cls) -> None:
+        from pyatlan.client.atlan import AtlanClient
+
+        client = AtlanClient.get_default_client()
+        if client is None:
+            client = AtlanClient()
+        response = client.get_typedefs(type_category=AtlanTypeCategory.ENUM)
+        cls.cache_by_name = {}
+        if response is not None:
+            for enum in response.enum_defs:
+                type_name = enum.name
+                cls.cache_by_name[type_name] = enum
+
+    @classmethod
+    def get_by_name(cls, name: str) -> Optional[EnumDef]:
+        """
+        Retrieve the enumeration definition by its name.
+        """
+        if name:
+            if enum_def := cls.cache_by_name.get(name):
+                return enum_def
+            cls.refresh_cache()
+            return cls.cache_by_name.get(name)
+        return None

Diff for: pyatlan/cache/role_cache.py

@@ -7,7 +7,6 @@
 
 
 class RoleCache:
-
     cache_by_id: dict[str, AtlanRole] = dict()
     map_id_to_name: dict[str, str] = dict()
     map_name_to_id: dict[str, str] = dict()

Diff for: pyatlan/client/atlan.py

@@ -23,6 +23,7 @@
     ADD_BUSINESS_ATTRIBUTE_BY_ID,
     BULK_UPDATE,
     CREATE_TYPE_DEFS,
+    UPDATE_TYPE_DEFS,
     DELETE_ENTITY_BY_ATTRIBUTE,
     DELETE_ENTITY_BY_GUID,
     DELETE_TYPE_DEF_BY_NAME,
@@ -122,6 +123,61 @@ def get_session():
     return session
 
 
+def _build_typdef_request(typedef: TypeDef) -> TypeDefResponse:
+    if isinstance(typedef, ClassificationDef):
+        # Set up the request payload...
+        payload = TypeDefResponse(
+            classification_defs=[typedef],
+            enum_defs=[],
+            struct_defs=[],
+            entity_defs=[],
+            relationship_defs=[],
+            custom_metadata_defs=[],
+        )
+    elif isinstance(typedef, CustomMetadataDef):
+        # Set up the request payload...
+        payload = TypeDefResponse(
+            classification_defs=[],
+            enum_defs=[],
+            struct_defs=[],
+            entity_defs=[],
+            relationship_defs=[],
+            custom_metadata_defs=[typedef],
+        )
+    elif isinstance(typedef, EnumDef):
+        # Set up the request payload...
+        payload = TypeDefResponse(
+            classification_defs=[],
+            enum_defs=[typedef],
+            struct_defs=[],
+            entity_defs=[],
+            relationship_defs=[],
+            custom_metadata_defs=[],
+        )
+    else:
+        raise InvalidRequestException(
+            "Unable to update type definitions of category: " + typedef.category.value,
+            param="category",
+        )
+        # Throw an invalid request exception
+    return payload
+
+
+def _refresh_caches(typedef: TypeDef) -> None:
+    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()
+    if isinstance(typedef, EnumDef):
+        from pyatlan.cache.enum_cache import EnumCache
+
+        EnumCache.refresh_cache()
+
+
 class AtlanClient(BaseSettings):
     _default_client: "ClassVar[Optional[AtlanClient]]" = None
     base_url: HttpUrl
@@ -409,6 +465,44 @@ def upsert(
         raw_json = self._call_api(BULK_UPDATE, query_params, request)
         return AssetMutationResponse(**raw_json)
 
+    def upsert_merging_cm(
+        self, entity: Union[Asset, list[Asset]], replace_classifications: bool = False
+    ) -> AssetMutationResponse:
+        query_params = {
+            "replaceClassifications": replace_classifications,
+            "replaceBusinessAttributes": True,
+            "overwriteBusinessAttributes": False,
+        }
+        entities: list[Asset] = []
+        if isinstance(entity, list):
+            entities.extend(entity)
+        else:
+            entities.append(entity)
+        for asset in entities:
+            asset.validate_required()
+        request = BulkRequest[Asset](entities=entities)
+        raw_json = self._call_api(BULK_UPDATE, query_params, request)
+        return AssetMutationResponse(**raw_json)
+
+    def upsert_replacing_cm(
+        self, entity: Union[Asset, list[Asset]], replace_classifications: bool = False
+    ) -> AssetMutationResponse:
+        query_params = {
+            "replaceClassifications": replace_classifications,
+            "replaceBusinessAttributes": True,
+            "overwriteBusinessAttributes": True,
+        }
+        entities: list[Asset] = []
+        if isinstance(entity, list):
+            entities.extend(entity)
+        else:
+            entities.append(entity)
+        for asset in entities:
+            asset.validate_required()
+        request = BulkRequest[Asset](entities=entities)
+        raw_json = self._call_api(BULK_UPDATE, query_params, request)
+        return AssetMutationResponse(**raw_json)
+
     def purge_entity_by_guid(self, guid) -> AssetMutationResponse:
         raw_json = self._call_api(
             DELETE_ENTITY_BY_GUID.format_path_with_params(guid),
@@ -455,54 +549,19 @@ def get_typedefs(self, type_category: AtlanTypeCategory) -> TypeDefResponse:
         return TypeDefResponse(**raw_json)
 
     def create_typedef(self, typedef: TypeDef) -> TypeDefResponse:
-        if isinstance(typedef, ClassificationDef):
-            # Set up the request payload...
-            payload = TypeDefResponse(
-                classification_defs=[typedef],
-                enum_defs=[],
-                struct_defs=[],
-                entity_defs=[],
-                relationship_defs=[],
-                custom_metadata_defs=[],
-            )
-        elif isinstance(typedef, CustomMetadataDef):
-            # Set up the request payload...
-            payload = TypeDefResponse(
-                classification_defs=[],
-                enum_defs=[],
-                struct_defs=[],
-                entity_defs=[],
-                relationship_defs=[],
-                custom_metadata_defs=[typedef],
-            )
-        elif isinstance(typedef, EnumDef):
-            # Set up the request payload...
-            payload = TypeDefResponse(
-                classification_defs=[],
-                enum_defs=[typedef],
-                struct_defs=[],
-                entity_defs=[],
-                relationship_defs=[],
-                custom_metadata_defs=[],
-            )
-        else:
-            raise InvalidRequestException(
-                "Unable to create new type definitions of category: "
-                + typedef.category.value,
-                param="category",
-            )
-            # Throw an invalid request exception
+        payload = _build_typdef_request(typedef)
         raw_json = self._call_api(
-            CREATE_TYPE_DEFS, request_obj=payload, exclude_unset=False
+            CREATE_TYPE_DEFS, request_obj=payload, exclude_unset=True
         )
-        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
+        _refresh_caches(typedef)
+        return TypeDefResponse(**raw_json)
 
-            CustomMetadataCache.refresh_cache()
+    def update_typedef(self, typedef: TypeDef) -> TypeDefResponse:
+        payload = _build_typdef_request(typedef)
+        raw_json = self._call_api(
+            UPDATE_TYPE_DEFS, request_obj=payload, exclude_unset=True
+        )
+        _refresh_caches(typedef)
         return TypeDefResponse(**raw_json)
 
     def purge_typedef(self, internal_name: str) -> None:
@@ -511,9 +570,11 @@ def purge_typedef(self, internal_name: str) -> None:
         #  to refresh that particular cache
         from pyatlan.cache.classification_cache import ClassificationCache
         from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
+        from pyatlan.cache.enum_cache import EnumCache
 
         ClassificationCache.refresh_cache()
         CustomMetadataCache.refresh_cache()
+        EnumCache.refresh_cache()
 
     @validate_arguments()
     def add_classifications(
@@ -628,8 +689,35 @@ def remove_announcement(
         asset.remove_announcement()
         return self._update_asset_by_attribute(asset, asset_type, qualified_name)
 
+    def update_custom_metadata_attributes(
+        self, guid: str, custom_metadata: CustomMetadata
+    ):
+        custom_metadata_request = CustomMetadataReqest(__root__=custom_metadata)
+        self._call_api(
+            ADD_BUSINESS_ATTRIBUTE_BY_ID.format_path(
+                {"entity_guid": guid, "bm_id": custom_metadata._meta_data_type_id}
+            ),
+            None,
+            custom_metadata_request,
+        )
+
     def replace_custom_metadata(self, guid: str, custom_metadata: CustomMetadata):
-        # TODO: This endpoint is not currently functioning correctly on the server
+        from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
+
+        # Iterate through the custom metadata provided and explicitly set every
+        # single attribute, so that they are all serialized out (forcing removal
+        # of any empty ones)
+        for k, v in custom_metadata.items():
+            # Need to translate the hashed-string key here back to an attribute name
+            attr_name = str(
+                CustomMetadataCache.get_attr_name_for_id(
+                    set_id=custom_metadata._meta_data_type_id, attr_id=k
+                )
+            )
+            if not v:
+                setattr(custom_metadata, attr_name, None)
+            else:
+                setattr(custom_metadata, attr_name, v)
         custom_metadata_request = CustomMetadataReqest(__root__=custom_metadata)
         self._call_api(
             ADD_BUSINESS_ATTRIBUTE_BY_ID.format_path(
@@ -639,6 +727,24 @@ def replace_custom_metadata(self, guid: str, custom_metadata: CustomMetadata):
             custom_metadata_request,
         )
 
+    def remove_custom_metadata(self, guid: str, cm_name: str):
+        from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
+
+        # Ensure the custom metadata exists first - let this throw an error if not
+        if cm_id := CustomMetadataCache.get_id_for_name(cm_name):
+            # Initialize a dict of empty attributes for the custom metadata, and then
+            # send that so that they are removed accordingly
+            if cm_type := CustomMetadataCache.get_type_for_id(cm_id):
+                custom_metadata = cm_type()
+                custom_metadata_request = CustomMetadataReqest(__root__=custom_metadata)
+                self._call_api(
+                    ADD_BUSINESS_ATTRIBUTE_BY_ID.format_path(
+                        {"entity_guid": guid, "bm_id": cm_id}
+                    ),
+                    None,
+                    custom_metadata_request,
+                )
+
     @validate_arguments()
     def append_terms(
         self,

Diff for: pyatlan/error.py

@@ -66,7 +66,6 @@ def __init__(
         status_code: Optional[int] = None,
         should_retry: bool = False,
     ):
-
         super(APIConnectionError, self).__init__(message, code, status_code)
         self.should_retry = should_retry
 

Diff for: pyatlan/generator/generate_from_typdefs.py

@@ -53,7 +53,7 @@
 def get_type(type_: str):
     ret_value = type_
 
-    for (field, replacement) in TYPE_REPLACEMENTS:
+    for field, replacement in TYPE_REPLACEMENTS:
         ret_value = ret_value.replace(field, replacement)
     return ret_value