Adds more efficient cache management for typedefs, and group CRUD

Signed-off-by: Christopher Grote <[email protected]>

Author: cmgrote

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)

pyatlan/client/atlan.py

@@ -32,9 +32,13 @@
     GET_ENTITY_BY_UNIQUE_ATTRIBUTE,
     GET_LINEAGE,
     GET_ROLES,
+    GET_GROUPS,
     INDEX_SEARCH,
     PARTIAL_UPDATE_ENTITY_BY_ATTRIBUTE,
     UPDATE_ENTITY_BY_ATTRIBUTE,
+    CREATE_GROUP,
+    DELETE_GROUP,
+    UPDATE_GROUP,
 )
 from pyatlan.error import AtlanError, NotFoundError
 from pyatlan.exceptions import AtlanServiceException, InvalidRequestException
@@ -69,6 +73,12 @@
     AtlanTypeCategory,
     CertificateStatus,
 )
+from pyatlan.model.group import (
+    GroupResponse,
+    AtlanGroup,
+    CreateGroupResponse,
+    CreateGroupRequest,
+)
 from pyatlan.model.lineage import LineageRequest, LineageResponse
 from pyatlan.model.response import AssetMutationResponse
 from pyatlan.model.role import RoleResponse
@@ -123,7 +133,7 @@ def get_session():
     return session
 
 
-def _build_typdef_request(typedef: TypeDef) -> TypeDefResponse:
+def _build_typedef_request(typedef: TypeDef) -> TypeDefResponse:
     if isinstance(typedef, ClassificationDef):
         # Set up the request payload...
         payload = TypeDefResponse(
@@ -263,6 +273,8 @@ def _call_api(
             try:
                 if (
                     response.content is None
+                    or response.content == "null"
+                    or len(response.content) == 0
                     or response.status_code == HTTPStatus.NO_CONTENT
                 ):
                     return None
@@ -354,6 +366,88 @@ def get_all_roles(self) -> RoleResponse:
         raw_json = self._call_api(GET_ROLES.format_path_with_params())
         return RoleResponse(**raw_json)
 
+    def create_group(
+        self,
+        group: AtlanGroup,
+        user_ids: Optional[list[str]] = None,
+    ) -> CreateGroupResponse:
+        payload = CreateGroupRequest(group=group, user_ids=user_ids)
+        raw_json = self._call_api(CREATE_GROUP, request_obj=payload, exclude_unset=True)
+        return CreateGroupResponse(**raw_json)
+
+    def update_group(
+        self,
+        group: AtlanGroup,
+    ) -> None:
+        self._call_api(
+            UPDATE_GROUP.format_path_with_params(group.id),
+            request_obj=group,
+            exclude_unset=True,
+        )
+
+    def purge_group(
+        self,
+        guid: str,
+    ) -> None:
+        self._call_api(DELETE_GROUP.format_path({"group_guid": guid}))
+
+    def get_groups(
+        self,
+        limit: Optional[int] = None,
+        post_filter: Optional[str] = None,
+        sort: Optional[str] = None,
+        count: bool = True,
+        offset: int = 0,
+    ) -> GroupResponse:
+        query_params: dict[str, str] = {
+            "count": str(count),
+            "offset": str(offset),
+        }
+        if limit is not None:
+            query_params["limit"] = str(limit)
+        if post_filter is not None:
+            query_params["filter"] = post_filter
+        if sort is not None:
+            query_params["sort"] = sort
+        raw_json = self._call_api(GET_GROUPS.format_path_with_params(), query_params)
+        return GroupResponse(**raw_json)
+
+    def get_all_groups(self) -> list[AtlanGroup]:
+        """
+        Retrieve all groups defined in Atlan.
+        """
+        groups: list[AtlanGroup] = []
+        offset = 0
+        limit = 100
+        response: Optional[GroupResponse] = self.get_groups(
+            offset=offset, limit=limit, sort="createdAt"
+        )
+        while response:
+            if page := response.records:
+                groups.extend(page)
+                offset += limit
+                response = self.get_groups(offset=offset, limit=limit, sort="createdAt")
+            else:
+                response = None
+        return groups
+
+    def get_group_by_name(
+        self, alias: str, limit: int = 100
+    ) -> Optional[list[AtlanGroup]]:
+        """
+        Retrieve all groups with a name that contains the provided string.
+        (This could include a complete group name, in which case there should be at most
+        a single item in the returned list, or could be a partial group name to retrieve
+        all groups with that naming convention.)
+        """
+        if response := self.get_groups(
+            offset=0,
+            limit=limit,
+            post_filter='{"$and":[{"alias":{"$ilike":"%' + alias + '%"}}]}',
+        ):
+            return response.records
+        return None
+
     @validate_arguments()
     def get_asset_by_qualified_name(
         self,
@@ -549,32 +643,59 @@ def get_typedefs(self, type_category: AtlanTypeCategory) -> TypeDefResponse:
         return TypeDefResponse(**raw_json)
 
     def create_typedef(self, typedef: TypeDef) -> TypeDefResponse:
-        payload = _build_typdef_request(typedef)
+        payload = _build_typedef_request(typedef)
         raw_json = self._call_api(
             CREATE_TYPE_DEFS, request_obj=payload, exclude_unset=True
         )
         _refresh_caches(typedef)
         return TypeDefResponse(**raw_json)
 
     def update_typedef(self, typedef: TypeDef) -> TypeDefResponse:
-        payload = _build_typdef_request(typedef)
+        payload = _build_typedef_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:
-        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
-        from pyatlan.cache.enum_cache import EnumCache
+    def purge_typedef(self, name: str, typedef_type: type) -> None:
+        if typedef_type == CustomMetadataDef:
+            from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
 
-        ClassificationCache.refresh_cache()
-        CustomMetadataCache.refresh_cache()
-        EnumCache.refresh_cache()
+            internal_name = CustomMetadataCache.get_id_for_name(name)
+        elif typedef_type == EnumDef:
+            internal_name = name
+        elif typedef_type == ClassificationDef:
+            from pyatlan.cache.classification_cache import ClassificationCache
+
+            internal_name = str(ClassificationCache.get_id_for_name(name))
+        else:
+            raise InvalidRequestException(
+                message=f"Unable to purge type definitions of type: {typedef_type}",
+            )
+            # Throw an invalid request exception
+        if internal_name:
+            self._call_api(
+                DELETE_TYPE_DEF_BY_NAME.format_path_with_params(internal_name)
+            )
+        else:
+            raise NotFoundError(
+                message=f"Unable to find {typedef_type} with name: {name}",
+                code="ATLAN-PYTHON-404-000",
+            )
+
+        if typedef_type == CustomMetadataDef:
+            from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
+
+            CustomMetadataCache.refresh_cache()
+        elif typedef_type == EnumDef:
+            from pyatlan.cache.enum_cache import EnumCache
+
+            EnumCache.refresh_cache()
+        elif typedef_type == ClassificationDef:
+            from pyatlan.cache.classification_cache import ClassificationCache
+
+            ClassificationCache.refresh_cache()
 
     @validate_arguments()
     def add_classifications(

pyatlan/client/constants.py

@@ -13,10 +13,17 @@
 )
 
 ROLE_API = f"{ADMIN_URI}roles"
+GROUP_API = f"{ADMIN_URI}groups"
 
 # Role APIs
 GET_ROLES = API(ROLE_API, HTTPMethod.GET, HTTPStatus.OK)
 
+# Group APIs
+GET_GROUPS = API(GROUP_API, HTTPMethod.GET, HTTPStatus.OK)
+CREATE_GROUP = API(GROUP_API, HTTPMethod.POST, HTTPStatus.OK)
+UPDATE_GROUP = API(GROUP_API, HTTPMethod.POST, HTTPStatus.OK)
+DELETE_GROUP = API(GROUP_API + "/{group_guid}/delete", HTTPMethod.POST, HTTPStatus.OK)
+
 ENTITY_API = f"{BASE_URI}entity/"
 PREFIX_ATTR = "attr:"
 PREFIX_ATTR_ = "attr_"