Merge pull request #53 from atlanhq/ACTIV-559

Add support for creating Badges

Author: ErnestoLoma

.github/workflows/pytest.yaml

@@ -4,7 +4,7 @@ on: [push, workflow_dispatch]
 
 jobs:
   build:
-
+    concurrency: integration_tests
     runs-on: ubuntu-latest
     strategy:
       matrix:

HISTORY.md

@@ -1,3 +1,8 @@
+## 0.0.32 (May 18, 2023)
+
+* Add create method to Badge
+* Add integrations tests for CustomMetaData
+
 ## 0.0.31 (May 15, 2023)
 
 * Added the following classes to support lineage retrieval

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
@@ -72,7 +72,7 @@ def refresh_cache(cls) -> None:
                         attr_id = str(attr.name)
                         attr_name = str(attr.display_name)
                         # Use a renamed attribute everywhere
-                        attr_renamed = to_snake_case(attr_name.replace(" ", ""))
+                        attr_renamed = to_snake_case(attr_name)
                         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_renamed
@@ -91,26 +91,48 @@ def refresh_cache(cls) -> None:
                         cls.types_by_asset[asset_type].add(attrib_type)
 
     @classmethod
-    def get_id_for_name(cls, name: str) -> Optional[str]:
+    def get_id_for_name(cls, name: str) -> 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]:
+    def get_name_for_id(cls, idstr: str) -> 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]:
@@ -152,23 +174,32 @@ def get_all_custom_attributes(
         return m
 
     @classmethod
-    def get_attr_id_for_name(cls, set_name: str, attr_name: str) -> Optional[str]:
+    def get_attr_id_for_name(cls, set_name: str, attr_name: str) -> str:
         """
         Translate the provided human-readable custom metadata set and attribute names to the Atlan-internal ID string
         for the attribute.
         """
-        attr_id = None
-        if set_id := cls.get_id_for_name(set_name):
-            if sub_map := cls.map_attr_name_to_id.get(set_id):
-                attr_id = sub_map.get(attr_name)
+        set_id = cls.get_id_for_name(set_name)
+        if sub_map := cls.map_attr_name_to_id.get(set_id):
+            attr_id = sub_map.get(attr_name)
             if attr_id:
                 # If found, return straight away
                 return attr_id
-            # Otherwise, refresh the cache and look again (could be stale)
-            cls.refresh_cache()
-            if sub_map := cls.map_attr_name_to_id.get(set_id):
-                return sub_map.get(attr_name)
-        return None
+        # Otherwise, refresh the cache and look again (could be stale)
+        cls.refresh_cache()
+        if sub_map := cls.map_attr_name_to_id.get(set_id):
+            attr_id = sub_map.get(attr_name)
+            if attr_id:
+                # If found, return straight away
+                return attr_id
+            raise NotFoundError(
+                message=f"Custom metadata property with name {attr_name} does not exist in custom metadata {set_name}.",
+                code="ATLAN-PYTHON-404-009",
+            )
+        raise NotFoundError(
+            message=f"Custom metadata with ID {set_id} does not exist.",
+            code="ATLAN-PYTHON-404-009",
+        )
 
     @classmethod
     def get_attr_name_for_id(cls, set_id: str, attr_id: str) -> Optional[str]:
@@ -177,8 +208,7 @@ def get_attr_name_for_id(cls, set_id: str, attr_id: str) -> Optional[str]:
         for the attribute.
         """
         if sub_map := cls.map_attr_id_to_name.get(set_id):
-            attr_name = sub_map.get(attr_id)
-            if attr_name:
+            if attr_name := sub_map.get(attr_id):
                 return attr_name
             cls.refresh_cache()
             if sub_map := cls.map_attr_id_to_name.get(set_id):

pyatlan/cache/group_cache.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2022 Atlan Pte. Ltd.
+from typing import Optional
+
+from pyatlan.client.atlan import AtlanClient
+
+
+class GroupCache:
+    map_id_to_name: dict[str, str] = dict()
+    map_name_to_id: dict[str, str] = dict()
+    map_alias_to_id: dict[str, str] = dict()
+
+    @classmethod
+    def _refresh_cache(cls) -> None:
+        client = AtlanClient.get_default_client()
+        if client is None:
+            client = AtlanClient()
+        groups = client.get_all_groups()
+        if groups is not None:
+            cls.map_id_to_name = {}
+            cls.map_name_to_id = {}
+            cls.map_alias_to_id = {}
+            for group in groups:
+                group_id = str(group.id)
+                group_name = str(group.name)
+                group_alias = str(group.alias)
+                cls.map_id_to_name[group_id] = group_name
+                cls.map_name_to_id[group_name] = group_id
+                cls.map_alias_to_id[group_alias] = group_id
+
+    @classmethod
+    def get_id_for_name(cls, name: str) -> Optional[str]:
+        """
+        Translate the provided human-readable group name to its GUID.
+        """
+        if group_id := cls.map_name_to_id.get(name):
+            return group_id
+        cls._refresh_cache()
+        return cls.map_name_to_id.get(name)
+
+    @classmethod
+    def get_id_for_alias(cls, alias: str) -> Optional[str]:
+        """
+        Translate the provided alias to its GUID.
+        """
+        if group_id := cls.map_alias_to_id.get(alias):
+            return group_id
+        cls._refresh_cache()
+        return cls.map_alias_to_id.get(alias)
+
+    @classmethod
+    def get_name_for_id(cls, idstr: str) -> Optional[str]:
+        """
+        Translate the provided group GUID to the human-readable group name.
+        """
+        if group_name := cls.map_id_to_name.get(idstr):
+            return group_name
+        cls._refresh_cache()
+        return cls.map_id_to_name.get(idstr)