Merge pull request #54 from atlanhq/ACTIV-623

Adds user management, caching and related tests

Author: ErnestoLoma

pyatlan/cache/user_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 UserCache:
+    map_id_to_name: dict[str, str] = dict()
+    map_name_to_id: dict[str, str] = dict()
+    map_email_to_id: dict[str, str] = dict()
+
+    @classmethod
+    def _refresh_cache(cls) -> None:
+        client = AtlanClient.get_default_client()
+        if client is None:
+            client = AtlanClient()
+        users = client.get_all_users()
+        if users is not None:
+            cls.map_id_to_name = {}
+            cls.map_name_to_id = {}
+            cls.map_email_to_id = {}
+            for user in users:
+                user_id = str(user.id)
+                username = str(user.username)
+                user_email = str(user.email)
+                cls.map_id_to_name[user_id] = username
+                cls.map_name_to_id[username] = user_id
+                cls.map_email_to_id[user_email] = user_id
+
+    @classmethod
+    def get_id_for_name(cls, name: str) -> Optional[str]:
+        """
+        Translate the provided human-readable username to its GUID.
+        """
+        if user_id := cls.map_name_to_id.get(name):
+            return user_id
+        cls._refresh_cache()
+        return cls.map_name_to_id.get(name)
+
+    @classmethod
+    def get_id_for_email(cls, email: str) -> Optional[str]:
+        """
+        Translate the provided email to its GUID.
+        """
+        if user_id := cls.map_email_to_id.get(email):
+            return user_id
+        cls._refresh_cache()
+        return cls.map_email_to_id.get(email)
+
+    @classmethod
+    def get_name_for_id(cls, idstr: str) -> Optional[str]:
+        """
+        Translate the provided user GUID to the human-readable username.
+        """
+        if username := cls.map_id_to_name.get(idstr):
+            return username
+        cls._refresh_cache()
+        return cls.map_id_to_name.get(idstr)

pyatlan/client/atlan.py

@@ -39,9 +39,20 @@
     CREATE_GROUP,
     DELETE_GROUP,
     UPDATE_GROUP,
+    GET_GROUP_MEMBERS,
+    REMOVE_USERS_FROM_GROUP,
+    GET_USERS,
+    CREATE_USERS,
+    UPDATE_USER,
+    DELETE_USER,
+    GET_USER_GROUPS,
+    ADD_USER_TO_GROUPS,
+    CHANGE_USER_ROLE,
+    GET_CURRENT_USER,
 )
 from pyatlan.error import AtlanError, NotFoundError
 from pyatlan.exceptions import AtlanServiceException, InvalidRequestException
+from pyatlan.model import group
 from pyatlan.model.assets import (
     Asset,
     AtlasGlossary,
@@ -78,6 +89,7 @@
     AtlanGroup,
     CreateGroupResponse,
     CreateGroupRequest,
+    RemoveFromGroupRequest,
 )
 from pyatlan.model.lineage import LineageRequest, LineageResponse
 from pyatlan.model.response import AssetMutationResponse
@@ -90,6 +102,14 @@
     TypeDef,
     TypeDefResponse,
 )
+from pyatlan.model.user import (
+    AtlanUser,
+    CreateUserRequest,
+    UserMinimalResponse,
+    AddToGroupsRequest,
+    ChangeRoleRequest,
+    UserResponse,
+)
 from pyatlan.utils import HTTPStatus, get_logger
 
 LOGGER = get_logger()
@@ -371,7 +391,9 @@ def create_group(
         group: AtlanGroup,
         user_ids: Optional[list[str]] = None,
     ) -> CreateGroupResponse:
-        payload = CreateGroupRequest(group=group, user_ids=user_ids)
+        payload = CreateGroupRequest(group=group)
+        if user_ids:
+            payload.users = user_ids
         raw_json = self._call_api(CREATE_GROUP, request_obj=payload, exclude_unset=True)
         return CreateGroupResponse(**raw_json)
 
@@ -448,6 +470,169 @@ def get_group_by_name(
             return response.records
         return None
 
+    def get_group_members(self, guid: str) -> UserResponse:
+        """
+        Retrieves the members (users) of a group.
+        """
+        raw_json = self._call_api(GET_GROUP_MEMBERS.format_path({"group_guid": guid}))
+        return UserResponse(**raw_json)
+
+    def remove_users_from_group(self, guid: str, user_ids=list[str]) -> None:
+        """
+        Remove one or more users from a group.
+        """
+        rfgr = RemoveFromGroupRequest(users=user_ids)
+        self._call_api(
+            REMOVE_USERS_FROM_GROUP.format_path({"group_guid": guid}),
+            request_obj=rfgr,
+            exclude_unset=True,
+        )
+
+    def create_users(
+        self,
+        users: list[AtlanUser],
+    ) -> None:
+        from pyatlan.cache.role_cache import RoleCache
+
+        cur = CreateUserRequest(users=[])
+        for user in users:
+            role_name = str(user.workspace_role)
+            if role_id := RoleCache.get_id_for_name(role_name):
+                to_create = CreateUserRequest.CreateUser(
+                    email=user.email,
+                    role_name=role_name,
+                    role_id=role_id,
+                )
+                cur.users.append(to_create)
+        self._call_api(CREATE_USERS, request_obj=cur, exclude_unset=True)
+
+    def update_user(
+        self,
+        guid: str,
+        user: AtlanUser,
+    ) -> UserMinimalResponse:
+        raw_json = self._call_api(
+            UPDATE_USER.format_path_with_params(guid),
+            request_obj=user,
+            exclude_unset=True,
+        )
+        return UserMinimalResponse(**raw_json)
+
+    def purge_user(
+        self,
+        guid: str,
+    ) -> None:
+        self._call_api(DELETE_USER.format_path({"user_guid": guid}))
+
+    def get_groups_for_user(
+        self,
+        guid: str,
+    ) -> GroupResponse:
+        raw_json = self._call_api(GET_USER_GROUPS.format_path({"user_guid": guid}))
+        return GroupResponse(**raw_json)
+
+    def add_user_to_groups(
+        self,
+        guid: str,
+        group_ids: list[str],
+    ) -> None:
+        atgr = AddToGroupsRequest(groups=group_ids)
+        self._call_api(
+            ADD_USER_TO_GROUPS.format_path({"user_guid": guid}),
+            request_obj=atgr,
+            exclude_unset=True,
+        )
+
+    def change_user_role(
+        self,
+        guid: str,
+        role_id: str,
+    ) -> None:
+        crr = ChangeRoleRequest(role_id=role_id)
+        self._call_api(
+            CHANGE_USER_ROLE.format_path({"user_guid": guid}),
+            request_obj=crr,
+            exclude_unset=True,
+        )
+
+    def get_current_user(
+        self,
+    ) -> UserMinimalResponse:
+        raw_json = self._call_api(GET_CURRENT_USER)
+        return UserMinimalResponse(**raw_json)
+
+    def get_users(
+        self,
+        limit: Optional[int] = None,
+        post_filter: Optional[str] = None,
+        sort: Optional[str] = None,
+        count: bool = True,
+        offset: int = 0,
+    ) -> UserResponse:
+        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_USERS.format_path_with_params(), query_params)
+        return UserResponse(**raw_json)
+
+    def get_all_users(self) -> list[AtlanUser]:
+        """
+        Retrieve all users defined in Atlan.
+        """
+        users: list[AtlanUser] = []
+        offset = 0
+        limit = 100
+        response: Optional[UserResponse] = self.get_users(
+            offset=offset, limit=limit, sort="username"
+        )
+        while response:
+            if page := response.records:
+                users.extend(page)
+                offset += limit
+                response = self.get_users(offset=offset, limit=limit, sort="username")
+            else:
+                response = None
+        return users
+
+    def get_users_by_email(
+        self, email: str, limit: int = 100
+    ) -> Optional[list[AtlanUser]]:
+        """
+        Retrieves all users with email addresses that contain the provided email.
+        (This could include a complete email address, in which case there should be at
+        most a single item in the returned list, or could be a partial email address
+        such as "@example.com" to retrieve all users with that domain in their email
+        address.)
+        """
+        if response := self.get_users(
+            offset=0,
+            limit=limit,
+            post_filter='{"email":{"$ilike":"%' + email + '%"}}',
+        ):
+            return response.records
+        return None
+
+    def get_user_by_username(self, username: str) -> Optional[AtlanUser]:
+        """
+        Retrieves a user based on the username. (This attempts an exact match on username
+        rather than a contains search.)
+        """
+        if response := self.get_users(
+            offset=0,
+            limit=5,
+            post_filter='{"username":"' + username + '"}',
+        ):
+            if response.records and len(response.records) >= 1:
+                return response.records[0]
+        return None
+
     @validate_arguments()
     def get_asset_by_qualified_name(
         self,

pyatlan/client/constants.py

@@ -14,6 +14,7 @@
 
 ROLE_API = f"{ADMIN_URI}roles"
 GROUP_API = f"{ADMIN_URI}groups"
+USER_API = f"{ADMIN_URI}users"
 
 # Role APIs
 GET_ROLES = API(ROLE_API, HTTPMethod.GET, HTTPStatus.OK)
@@ -23,6 +24,26 @@
 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)
+GET_GROUP_MEMBERS = API(
+    GROUP_API + "/{group_guid}/members", HTTPMethod.GET, HTTPStatus.OK
+)
+REMOVE_USERS_FROM_GROUP = API(
+    GROUP_API + "/{group_guid}/members/remove", HTTPMethod.POST, HTTPStatus.OK
+)
+
+# User APIs
+GET_USERS = API(USER_API, HTTPMethod.GET, HTTPStatus.OK)
+CREATE_USERS = API(USER_API, HTTPMethod.POST, HTTPStatus.OK)
+UPDATE_USER = API(USER_API, HTTPMethod.POST, HTTPStatus.OK)
+DELETE_USER = API(USER_API + "/{user_guid}/delete", HTTPMethod.POST, HTTPStatus.OK)
+GET_USER_GROUPS = API(USER_API + "/{user_guid}/groups", HTTPMethod.GET, HTTPStatus.OK)
+ADD_USER_TO_GROUPS = API(
+    USER_API + "/{user_guid}/groups", HTTPMethod.POST, HTTPStatus.OK
+)
+CHANGE_USER_ROLE = API(
+    USER_API + "/{user_guid}/roles/update", HTTPMethod.POST, HTTPStatus.OK
+)
+GET_CURRENT_USER = API(f"{USER_API}/current", HTTPMethod.GET, HTTPStatus.OK)
 
 ENTITY_API = f"{BASE_URI}entity/"
 PREFIX_ATTR = "attr:"

pyatlan/model/group.py

@@ -58,7 +58,7 @@ def is_default(self) -> bool:
         return (
             self.attributes is not None
             and self.attributes.is_default is not None
-            and self.attributes.is_default == "true"
+            and self.attributes.is_default == ["true"]
         )
 
     @staticmethod