Add create method to Badge

Author: ErnestoLoma

pyatlan/cache/custom_metadata_cache.py

@@ -4,7 +4,7 @@
 from typing import Any, Optional
 
 from pyatlan.client.atlan import AtlanClient
-from pyatlan.error import LogicError, NotFoundError
+from pyatlan.error import InvalidRequestError, LogicError, NotFoundError
 from pyatlan.model.core import CustomMetadata
 from pyatlan.model.enums import AtlanTypeCategory
 from pyatlan.model.typedef import AttributeDef, CustomMetadataDef
@@ -95,22 +95,44 @@ def get_id_for_name(cls, name: str) -> Optional[str]:
         """
         Translate the provided human-readable custom metadata set name to its Atlan-internal ID string.
         """
+        if name is None or not name.strip():
+            raise InvalidRequestError(
+                message="No name was provided when attempting to retrieve custom metadata.",
+                code="ATLAN-PYTHON-404-008",
+                param="",
+            )
         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()
-        return cls.map_name_to_id.get(name)
+        if cm_id := cls.map_name_to_id.get(name):
+            return cm_id
+        raise NotFoundError(
+            message=f"Custom metadata with name {name} does not exist.",
+            code="ATLAN-PYTHON-404-009",
+        )
 
     @classmethod
     def get_name_for_id(cls, idstr: str) -> Optional[str]:
         """
         Translate the provided Atlan-internal custom metadata ID string to the human-readable custom metadata set name.
         """
+        if idstr is None or not idstr.strip():
+            raise InvalidRequestError(
+                message="No ID was provided when attempting to retrieve custom metadata.",
+                code="ATLAN-PYTHON-404-008",
+                param="",
+            )
         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()
-        return cls.map_id_to_name.get(idstr)
+        if cm_name := cls.map_id_to_name.get(idstr):
+            return cm_name
+        raise NotFoundError(
+            message=f"Custom metadata with ID {idstr} does not exist.",
+            code="ATLAN-PYTHON-404-009",
+        )
 
     @classmethod
     def get_type_for_id(cls, idstr: str) -> Optional[type]:

pyatlan/exceptions.py

@@ -48,6 +48,7 @@ def __init__(
         code: Optional[str],
         cause: Optional[Exception] = None,
     ):
+        super().__init__(message)
         self.message = message
         self.code = code
         self.status_code = status_code

pyatlan/generator/templates/entity.jinja2

@@ -50,7 +50,10 @@ from pyatlan.model.structs import (
     GoogleTag,
     Histogram,
     KafkaTopicConsumption,
+    MCRuleComparison,
+    MCRuleSchedule,
     PopularityInsights,
+    SourceTagAttribute,
 )
 from pyatlan.utils import next_id
 from urllib.parse import quote, unquote
@@ -361,7 +364,17 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
             raise ValueError('must be {{ entity_def.name }}')
         return v
 
-    {%- if entity_def.name == "Readme" %}
+    {%- if entity_def.name == "Badge" %}
+    @classmethod
+    # @validate_arguments()
+    def create(cls, *, name: StrictStr, cm_name: str, cm_attribute: str) -> Badge:
+        validate_required_fields(["name", "cm_name", "cm_attribute"], [name])
+        return cls(
+            attributes=Badge.Attributes.create(
+                name=name, cm_name=cm_name, cm_attribute=cm_attribute
+            )
+        )
+    {%- elif entity_def.name == "Readme" %}
     @classmethod
     # @validate_arguments()
     def create(
@@ -523,6 +536,26 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
                 inputs=inputs,
                 outputs=outputs,
             )
+        {%- elif entity_def.name == "Badge"  %}
+        @classmethod
+        # @validate_arguments()
+        def create(
+            cls, *, name: StrictStr, cm_name: str, cm_attribute: str
+        ) -> Badge.Attributes:
+            validate_required_fields(
+                ["name", "cm_name", "cm_attribute"], [name, cm_name, cm_attribute]
+            )
+            from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
+
+            cm_id = CustomMetadataCache.get_id_for_name(cm_name)
+            cm_attr_id = CustomMetadataCache.get_attr_id_for_name(
+                set_name=cm_name, attr_name=cm_attribute
+            )
+            return Badge.Attributes(
+                name=name,
+                qualified_name=f"badges/global/{cm_id}.{cm_attr_id}",
+                badge_metadata_attribute=f"{cm_id}.{cm_attr_id}",
+            )
         {%- elif entity_def.name == "Readme"  %}
         @classmethod
         # @validate_arguments()