Merge branch 'prototype_sdk' of https://github.com/atlanhq/atlan-python into prototype_sdk

Merging Chris's changes

Author: ErnestoLoma

pyatlan/cache/__init__.py

@@ -0,0 +1,3 @@
+from pyatlan.cache.role_cache import RoleCache
+from pyatlan.cache.classification_cache import ClassificationCache
+from pyatlan.cache.custom_metadata_cache import CustomMetadataCache

pyatlan/cache/classification_cache.py

@@ -0,0 +1,63 @@
+
+from pyatlan.client.atlan import AtlanClient
+from pyatlan.client.typedef import TypeDefClient
+from pyatlan.model.typedef import ClassificationDef
+from pyatlan.model.enums import AtlanTypeCategory
+
+from typing import Optional
+
+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()
+    deleted_ids: set[str] = set()
+    deleted_names: set[str] = set()
+
+    @classmethod
+    def _refresh_cache(cls) -> None:
+        response = TypeDefClient(AtlanClient()).get_typedefs(type=AtlanTypeCategory.CLASSIFICATION)
+        if response is not None:
+            cls.cache_by_id = dict()
+            cls.map_id_to_name = dict()
+            cls.map_name_to_id = dict()
+            for classification in response.classification_defs:
+                classification_id = classification.name
+                classification_name = classification.display_name
+                cls.cache_by_id[classification_id] = classification
+                cls.map_id_to_name[classification_id] = classification_name
+                cls.map_name_to_id[classification_name] = classification_id
+
+    @classmethod
+    def get_id_for_name(cls, name: str) -> Optional[str]:
+        """
+        Translate the provided human-readable classification name to its Atlan-internal ID string.
+        """
+        cls_id = cls.map_name_to_id.get(name)
+        if not cls_id and not name in cls.deleted_names:
+            # If not found, refresh the cache and look again (could be stale)
+            cls._refresh_cache()
+            cls_id = cls.map_name_to_id.get(name)
+            if not cls_id:
+                # If still not found after refresh, mark it as deleted (could be
+                # an entry in an audit log that refers to a classification that
+                # no longer exists)
+                cls.deleted_names.add(name)
+        return cls_id
+
+    @classmethod
+    def get_name_for_id(cls, idstr: str) -> Optional[str]:
+        """
+        Translate the provided Atlan-internal classification ID string to the human-readable classification name.
+        """
+        cls_name = cls.map_id_to_name.get(idstr)
+        if not cls_name and not idstr in cls.deleted_ids:
+            # If not found, refresh the cache and look again (could be stale)
+            cls._refresh_cache()
+            cls_name = cls.map_id_to_name.get(idstr)
+            if not cls_name:
+                # If still not found after refresh, mark it as deleted (could be
+                # an entry in an audit log that refers to a classification that
+                # no longer exists)
+                cls.deleted_ids.add(idstr)
+        return cls_name

pyatlan/cache/custom_metadata_cache.py

@@ -0,0 +1,150 @@
+
+from pyatlan.client.atlan import AtlanClient
+from pyatlan.client.typedef import TypeDefClient
+from pyatlan.model.typedef import CustomMetadataDef, AttributeDef
+from pyatlan.model.enums import AtlanTypeCategory
+from pyatlan.error import LogicError
+
+from typing import Optional
+
+class CustomMetadataCache:
+
+    cache_by_id: dict[str, CustomMetadataDef] = dict()
+    map_id_to_name: dict[str, str] = dict()
+    map_name_to_id: dict[str, str] = dict()
+    map_attr_id_to_name: dict[str, dict[str, str]] = dict()
+    map_attr_name_to_id: dict[str, dict[str, str]] = dict()
+    archived_attr_ids: dict[str, str] = dict()
+
+    @classmethod
+    def _refresh_cache(cls) -> None:
+        response = TypeDefClient(AtlanClient()).get_typedefs(type=AtlanTypeCategory.CUSTOM_METADATA)
+        if response is not None:
+            cls.cache_by_id = dict()
+            cls.map_id_to_name = dict()
+            cls.map_name_to_id = dict()
+            cls.map_attr_id_to_name = dict()
+            cls.map_attr_name_to_id = dict()
+            cls.archived_attr_ids = dict()
+            for cm in response.custom_metadata_defs:
+                type_id = cm.name
+                type_name = cm.display_name
+                cls.cache_by_id[type_id] = cm
+                cls.map_id_to_name[type_id] = type_name
+                cls.map_name_to_id[type_name] = type_id
+                cls.map_attr_id_to_name[type_id] = dict()
+                cls.map_attr_name_to_id[type_id] = dict()
+                if cm.attribute_defs:
+                    for attr in cm.attribute_defs:
+                        attr_id = attr.name
+                        attr_name = attr.display_name
+                        cls.map_attr_id_to_name[type_id][attr_id] = attr_name
+                        if attr.options and attr.options.is_archived:
+                            cls.archived_attr_ids[attr_id] = attr_name
+                        else:
+                            if attr_name in cls.map_attr_name_to_id[type_id]:
+                                raise LogicError(
+                                    "Multiple custom attributes with exactly the same name (" + attr_name + ") found for: " + type_name,
+                                    code="ATLAN-PYTHON-500-100"
+                                )
+                            cls.map_attr_name_to_id[type_id][attr_name] = attr_id
+
+
+    @classmethod
+    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.
+        """
+        cm_id = cls.map_name_to_id.get(name)
+        if cm_id:
+            return cm_id
+        else:
+            # If not found, refresh the cache and look again (could be stale)
+            cls._refresh_cache()
+            return cls.map_name_to_id.get(name)
+
+    @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.
+        """
+        cm_name = cls.map_id_to_name.get(idstr)
+        if cm_name:
+            return cm_name
+        else:
+            # If not found, refresh the cache and look again (could be stale)
+            cls._refresh_cache()
+            return cls.map_id_to_name.get(idstr)
+
+    @classmethod
+    def get_all_custom_attributes(cls, include_deleted: bool=False, force_refresh: bool=False) -> dict[str, list[AttributeDef]]:
+        """
+        Retrieve all the custom metadata attributes. The map will be keyed by custom metadata set
+        name, and the value will be a listing of all the attributes within that set (with all the details
+        of each of those attributes).
+        """
+        if len(cls.cache_by_id) == 0 or force_refresh:
+            cls._refresh_cache()
+        m = {}
+        for type_id, cm in cls.cache_by_id.items():
+            type_name = cls.get_name_for_id(type_id)
+            attribute_defs = cm.attribute_defs
+            if include_deleted:
+                to_include = attribute_defs
+            else:
+                to_include = []
+                if attribute_defs:
+                    for attr in attribute_defs:
+                        if not attr.options or not attr.options.is_archived:
+                            to_include.append(attr)
+            m[type_name] = to_include
+        return m
+
+    @classmethod
+    def get_attr_id_for_name(cls, set_name: str, attr_name: str) -> Optional[str]:
+        """
+        Translate the provided human-readable custom metadata set and attribute names to the Atlan-internal ID string
+        for the attribute.
+        """
+        attr_id = None
+        set_id = cls.get_id_for_name(set_name)
+        if set_id:
+            sub_map = cls.map_attr_name_to_id.get(set_id)
+            if sub_map:
+                attr_id = sub_map.get(attr_name)
+            if attr_id:
+                # If found, return straight away
+                return attr_id
+            else:
+                # Otherwise, refresh the cache and look again (could be stale)
+                cls._refresh_cache()
+                sub_map = cls.map_attr_name_to_id.get(set_id)
+                if sub_map:
+                    return sub_map.get(attr_name)
+        return None
+
+    @classmethod
+    def _get_attributes_for_search_results(cls, set_id: str) -> Optional[list[str]]:
+        sub_map = cls.map_attr_name_to_id.get(set_id)
+        if sub_map:
+            attr_ids = sub_map.values()
+            dot_names = []
+            for idstr in attr_ids:
+                dot_names.append(set_id + "." + idstr)
+            return dot_names
+        return None
+
+    @classmethod
+    def get_attributes_for_search_results(cls, set_name: str) -> Optional[list[str]]:
+        """
+        Retrieve the full set of custom attributes to include on search results.
+        """
+        set_id = cls.get_id_for_name(set_name)
+        if set_id:
+            dot_names = cls._get_attributes_for_search_results(set_id)
+            if dot_names:
+                return dot_names
+            else:
+                cls._refresh_cache()
+                return cls._get_attributes_for_search_results(set_id)
+        return None

pyatlan/cache/role_cache.py

@@ -0,0 +1,50 @@
+
+from pyatlan.client.atlan import AtlanClient
+from pyatlan.client.role import RoleClient
+from pyatlan.model.role import AtlanRole
+
+from typing import Optional
+
+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()
+
+    @classmethod
+    def _refresh_cache(cls) -> None:
+        response = RoleClient(AtlanClient()).get_all_roles()
+        if response is not None:
+            cls.cache_by_id = dict()
+            cls.map_id_to_name = dict()
+            cls.map_name_to_id = dict()
+            for role in response.records:
+                role_id = role.id
+                role_name = role.name
+                cls.cache_by_id[role_id] = role
+                cls.map_id_to_name[role_id] = role_name
+                cls.map_name_to_id[role_name] = role_id
+
+    @classmethod
+    def get_id_for_name(cls, name: str) -> Optional[str]:
+        """
+        Translate the provided human-readable role name to its GUID.
+        """
+        role_id = cls.map_name_to_id.get(name)
+        if role_id:
+            return role_id
+        else:
+            cls._refresh_cache()
+            return cls.map_name_to_id.get(name)
+
+    @classmethod
+    def get_name_for_id(cls, idstr: str) -> Optional[str]:
+        """
+        Translate the provided role GUID to the human-readable role name.
+        """
+        role_name = cls.map_id_to_name.get(idstr)
+        if role_name:
+            return role_name
+        else:
+            cls._refresh_cache()
+            return cls.map_id_to_name.get(idstr)

pyatlan/client/atlan.py

@@ -85,7 +85,7 @@ def call_api(self, api, query_params=None, request_obj=None):
             return None
         elif response.status_code == api.expected_status:
             try:
-                if response.content is None:
+                if response.content is None or response.status_code == HTTPStatus.NO_CONTENT:
                     return None
                 if LOGGER.isEnabledFor(logging.DEBUG):
                     LOGGER.debug(

pyatlan/client/role.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env/python
+# Copyright 2022 Atlan Pte, Ltd
+# Copyright [2015-2021] The Apache Software Foundation
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from pyatlan.client.atlan import AtlanClient
+from typing import Optional
+from pyatlan.model.role import (
+    RoleResponse
+)
+from pyatlan.utils import (
+    API,
+    ADMIN_URI,
+    HTTPMethod,
+    HTTPStatus,
+)
+
+
+class RoleClient:
+    ROLE_API = ADMIN_URI + "roles"
+    
+    # Role APIs
+    GET_ROLES = API(ROLE_API, HTTPMethod.GET, HTTPStatus.OK)
+
+    def __init__(self, client: AtlanClient):
+        self.client = client
+
+
+    def get_roles(self, limit: int, filter: Optional[str]=None, sort: Optional[str]=None, count: bool=True, offset: int=0) -> RoleResponse:
+        if filter is None:
+            filter = ""
+        if sort is None:
+            sort = ""
+        query_params = {
+            "filter": filter,
+            "sort": sort,
+            "count": count,
+            "offset": offset,
+            "limit": limit
+        }
+        raw_json = self.client.call_api(RoleClient.GET_ROLES.format_path_with_params(), query_params)
+        response = RoleResponse(**raw_json)
+        return response
+
+
+    def get_all_roles(self) -> RoleResponse:
+        """
+        Retrieve all roles defined in Atlan.
+        """
+        raw_json = self.client.call_api(RoleClient.GET_ROLES.format_path_with_params())
+        response = RoleResponse(**raw_json)
+        return response