Merge pull request #39 from atlanhq/change_custom_metadata

Made changes to updating custom meta data

Author: ErnestoLoma

.gitignore

@@ -50,3 +50,4 @@ wheels/
 .installed.cfg
 *.egg
 *.DS_Store
+examples/

pyatlan/cache/custom_metadata_cache.py

@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 Atlan Pte. Ltd.
 import json
-from typing import Optional
+from typing import Any, Optional
 
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.error import LogicError, NotFoundError
+from pyatlan.model.core import CustomMetadata
 from pyatlan.model.enums import AtlanTypeCategory
 from pyatlan.model.typedef import AttributeDef, CustomMetadataDef
 
@@ -35,7 +36,7 @@ class CustomMetadataCache:
 
     @classmethod
     def _refresh_cache(cls) -> None:
-        from pyatlan.model.core import BusinessAttributes, to_snake_case
+        from pyatlan.model.core import CustomMetadata, to_snake_case
 
         client = AtlanClient.get_default_client()
         if client is None:
@@ -58,7 +59,7 @@ def _refresh_cache(cls) -> None:
                 cls.map_attr_name_to_id[type_id] = {}
                 meta_name = cm.display_name.replace(" ", "")
                 attribute_class_name = f"Attributes_{meta_name}"
-                attrib_type = type(attribute_class_name, (BusinessAttributes,), {})
+                attrib_type = type(attribute_class_name, (CustomMetadata,), {})
                 attrib_type._meta_data_type_id = type_id  # type: ignore
                 attrib_type._meta_data_type_name = type_name  # type: ignore
                 cls.map_id_to_type[type_id] = attrib_type
@@ -202,3 +203,30 @@ def get_attributes_for_search_results(cls, set_name: str) -> Optional[list[str]]
             cls._refresh_cache()
             return cls._get_attributes_for_search_results(set_id)
         return None
+
+    @classmethod
+    def get_custom_metadata(
+        cls,
+        name: str,
+        asset_type: type,
+        business_attributes: Optional[dict[str, Any]] = None,
+    ) -> CustomMetadata:
+        type_name = asset_type.__name__
+        ba_id = cls.get_id_for_name(name)
+        if ba_id is None:
+            raise ValueError(f"No custom metadata with the name: {name} exist")
+        for a_type in CustomMetadataCache.types_by_asset[type_name]:
+            if (
+                hasattr(a_type, "_meta_data_type_name")
+                and a_type._meta_data_type_name == name
+            ):
+                break
+        else:
+            raise ValueError(f"Custom metadata {name} is not applicable to {type_name}")
+        if ba_type := CustomMetadataCache.get_type_for_id(ba_id):
+            return (
+                ba_type(business_attributes[ba_id])
+                if business_attributes and ba_id in business_attributes
+                else ba_type()
+            )
+        raise ValueError(f"Custom metadata {name} is not applicable to {type_name}")

pyatlan/client/atlan.py

@@ -20,6 +20,7 @@
 from urllib3.util.retry import Retry
 
 from pyatlan.client.constants import (
+    ADD_BUSINESS_ATTRIBUTE_BY_ID,
     BULK_UPDATE,
     CREATE_TYPE_DEFS,
     DELETE_ENTITY_BY_ATTRIBUTE,
@@ -57,6 +58,8 @@
     Classification,
     ClassificationName,
     Classifications,
+    CustomMetadata,
+    CustomMetadataReqest,
 )
 from pyatlan.model.enums import AtlanDeleteType, AtlanTypeCategory, CertificateStatus
 from pyatlan.model.response import AssetMutationResponse
@@ -377,11 +380,12 @@ def upsert(
         entity: Union[Asset, list[Asset]],
         replace_classifications: bool = False,
         replace_custom_metadata: bool = False,
+        overwrite_custom_metadata: bool = False,
     ) -> AssetMutationResponse:
         query_params = {
             "replaceClassifications": replace_classifications,
             "replaceBusinessAttributes": replace_custom_metadata,
-            "overwriteBusinessAttribute": replace_custom_metadata,
+            "overwriteBusinessAttributes": overwrite_custom_metadata,
         }
         entities: list[Asset] = []
         if isinstance(entity, list):
@@ -587,3 +591,14 @@ def remove_announcement(
         asset.name = name
         asset.remove_announcement()
         return self._update_asset_by_attribute(asset, asset_type, qualified_name)
+
+    def replace_custom_metadata(self, guid: str, custom_metadata: CustomMetadata):
+        # TODO: This endpoint is not currently functioning correctly on the server
+        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,
+        )

pyatlan/client/constants.py

@@ -122,7 +122,12 @@
 
 # Business Attributes APIs
 ADD_BUSINESS_ATTRIBUTE = API(
-    ENTITY_API + "guid/{entity_guid}/businessmetadata",
+    ENTITY_API + "guid/{entity_guid}/businessmetadata/",
+    HTTPMethod.POST,
+    HTTPStatus.NO_CONTENT,
+)
+ADD_BUSINESS_ATTRIBUTE_BY_ID = API(
+    ENTITY_API + "guid/{entity_guid}/businessmetadata/{bm_id}",
     HTTPMethod.POST,
     HTTPStatus.NO_CONTENT,
 )

pyatlan/generator/templates/entity.jinja2

@@ -9,7 +9,7 @@ from typing import Any, ClassVar, Dict, List, Optional, TypeVar
 
 from pydantic import Field, StrictStr, root_validator, validator
 
-from pyatlan.model.core import Announcement, AtlanObject, BusinessAttributes, Classification, Meaning
+from pyatlan.model.core import Announcement, AtlanObject, CustomMetadata, Classification, Meaning
 from pyatlan.model.enums import (
     ADLSAccessTier,
     ADLSAccountStatus,
@@ -226,7 +226,7 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
         if not self.create_time or self.created_by:
             self.attributes.validate_required()
 
-    def get_business_attributes(self, name: str) -> BusinessAttributes:
+    def get_custom_metadata(self, name: str) -> CustomMetadata:
         from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
 
         ba_id = CustomMetadataCache.get_id_for_name(name)
@@ -240,7 +240,7 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
                 break
         else:
             raise ValueError(
-                f"Business attributes {name} are not applicable to {self.type_name}"
+                f"Custom metadata attributes {name} are not applicable to {self.type_name}"
             )
         if ba_type := CustomMetadataCache.get_type_for_id(ba_id):
             return (
@@ -250,27 +250,27 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
             )
         else:
             raise ValueError(
-                f"Business attributes {name} are not applicable to {self.type_name}"
+                f"Custom metadata attributes {name} are not applicable to {self.type_name}"
             )
 
-    def set_business_attribute(self, business_attributes: BusinessAttributes) -> None:
+    def set_custom_metadata(self, custom_metadata: CustomMetadata) -> None:
         from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
 
-        if not isinstance(business_attributes, BusinessAttributes):
+        if not isinstance(custom_metadata, CustomMetadata):
             raise ValueError(
-                "business_attributes must be an instance of BusinessAttributes"
+                "business_attributes must be an instance of CustomMetadata"
             )
         if (
-            type(business_attributes)
+            type(custom_metadata)
             not in CustomMetadataCache.types_by_asset[self.type_name]
         ):
             raise ValueError(
-                f"Business attributes {business_attributes._meta_data_type_name} are not applicable to {self.type_name}"
+                f"Business attributes {custom_metadata._meta_data_type_name} are not applicable to {self.type_name}"
             )
-        ba_dict = dict(business_attributes)
+        ba_dict = dict(custom_metadata)
         if not self.business_attributes:
             self.business_attributes = {}
-        self.business_attributes[business_attributes._meta_data_type_id] = ba_dict
+        self.business_attributes[custom_metadata._meta_data_type_id] = ba_dict
 
 {%- else %}
     {%- if entity_def.name == "Asset"  %}
@@ -536,7 +536,7 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
                 )
             else:
                 raise ValueError(
-                    f"parent_type must be either Table, View or MaterializeView"
+                    "parent_type must be either Table, View or MaterializeView"
                 )
             return ret_value