feat: Allow user groups to be passed in allowed_to_use and allowed_to_manage when creating QuickSight resources (#2278)

Author: LeonLuttenberger

awswrangler/quicksight/_create.py

@@ -2,16 +2,25 @@
 
 import logging
 import uuid
-from typing import Any, Dict, List, Literal, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Set, TypeVar, Union, cast
 
 import boto3
 
 from awswrangler import _utils, exceptions, sts
-from awswrangler.quicksight._get_list import get_data_source_arn, get_dataset_id, list_users
-from awswrangler.quicksight._utils import extract_athena_query_columns, extract_athena_table_columns
+from awswrangler.quicksight._get_list import get_data_source_arn, get_dataset_id, list_groups, list_users
+from awswrangler.quicksight._utils import (
+    _QuicksightPrincipalList,
+    extract_athena_query_columns,
+    extract_athena_table_columns,
+)
+
+if TYPE_CHECKING:
+    from mypy_boto3_quicksight.type_defs import GroupTypeDef, UserTypeDef
+
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
+
 _ALLOWED_ACTIONS: Dict[str, Dict[str, List[str]]] = {
     "data_source": {
         "allowed_to_use": [
@@ -52,41 +61,51 @@
 }
 
 
-def _usernames_to_arns(user_names: List[str], all_users: List[Dict[str, Any]]) -> List[str]:
-    return [cast(str, u["Arn"]) for u in all_users if u.get("UserName") in user_names]
+def _groupnames_to_arns(group_names: Set[str], all_groups: List["GroupTypeDef"]) -> List[str]:
+    return [u["Arn"] for u in all_groups if u.get("GroupName") in group_names]
 
 
-def _generate_permissions(
+def _usernames_to_arns(user_names: Set[str], all_users: List["UserTypeDef"]) -> List[str]:
+    return [u["Arn"] for u in all_users if u.get("UserName") in user_names]
+
+
+_PrincipalTypeDef = TypeVar("_PrincipalTypeDef", "UserTypeDef", "GroupTypeDef")
+
+
+def _generate_permissions_base(
     resource: str,
     namespace: str,
     account_id: str,
     boto3_session: Optional[boto3.Session],
-    allowed_to_use: Optional[List[str]] = None,
-    allowed_to_manage: Optional[List[str]] = None,
+    allowed_to_use: Optional[List[str]],
+    allowed_to_manage: Optional[List[str]],
+    principal_names_to_arns_func: Callable[[Set[str], List[_PrincipalTypeDef]], List[str]],
+    list_principals: Callable[[str, str, Optional[boto3.Session]], List[Dict[str, Any]]],
 ) -> List[Dict[str, Union[str, List[str]]]]:
     permissions: List[Dict[str, Union[str, List[str]]]] = []
     if (allowed_to_use is None) and (allowed_to_manage is None):
         return permissions
 
-    # Forcing same user not be in both lists at the same time.
-    if (allowed_to_use is not None) and (allowed_to_manage is not None):
-        allowed_to_use = list(set(allowed_to_use) - set(allowed_to_manage))
+    allowed_to_use_set = set(allowed_to_use) if allowed_to_use else None
+    allowed_to_manage_set = set(allowed_to_manage) if allowed_to_manage else None
 
-    all_users: List[Dict[str, Any]] = list_users(
-        namespace=namespace, account_id=account_id, boto3_session=boto3_session
-    )
+    # Forcing same principal not be in both lists at the same time.
+    if (allowed_to_use_set is not None) and (allowed_to_manage_set is not None):
+        allowed_to_use_set = allowed_to_use_set - allowed_to_manage_set
+
+    all_principals = cast(List[_PrincipalTypeDef], list_principals(namespace, account_id, boto3_session))
 
-    if allowed_to_use is not None:
-        allowed_arns: List[str] = _usernames_to_arns(user_names=allowed_to_use, all_users=all_users)
+    if allowed_to_use_set is not None:
+        allowed_arns: List[str] = principal_names_to_arns_func(allowed_to_use_set, all_principals)
         permissions += [
             {
                 "Principal": arn,
                 "Actions": _ALLOWED_ACTIONS[resource]["allowed_to_use"],
             }
             for arn in allowed_arns
         ]
-    if allowed_to_manage is not None:
-        allowed_arns = _usernames_to_arns(user_names=allowed_to_manage, all_users=all_users)
+    if allowed_to_manage_set is not None:
+        allowed_arns = principal_names_to_arns_func(allowed_to_manage_set, all_principals)
         permissions += [
             {
                 "Principal": arn,
@@ -97,6 +116,41 @@ def _generate_permissions(
     return permissions
 
 
+def _generate_permissions(
+    resource: str,
+    namespace: str,
+    account_id: str,
+    boto3_session: Optional[boto3.Session],
+    allowed_users_to_use: Optional[List[str]] = None,
+    allowed_groups_to_use: Optional[List[str]] = None,
+    allowed_users_to_manage: Optional[List[str]] = None,
+    allowed_groups_to_manage: Optional[List[str]] = None,
+) -> List[Dict[str, Union[str, List[str]]]]:
+    permissions_users = _generate_permissions_base(
+        resource=resource,
+        namespace=namespace,
+        account_id=account_id,
+        boto3_session=boto3_session,
+        allowed_to_use=allowed_users_to_use,
+        allowed_to_manage=allowed_users_to_manage,
+        principal_names_to_arns_func=_usernames_to_arns,
+        list_principals=list_users,
+    )
+
+    permissions_groups = _generate_permissions_base(
+        resource=resource,
+        namespace=namespace,
+        account_id=account_id,
+        boto3_session=boto3_session,
+        allowed_to_use=allowed_groups_to_use,
+        allowed_to_manage=allowed_groups_to_manage,
+        principal_names_to_arns_func=_groupnames_to_arns,
+        list_principals=list_groups,
+    )
+
+    return permissions_users + permissions_groups
+
+
 def _generate_transformations(
     rename_columns: Optional[Dict[str, str]],
     cast_columns_types: Optional[Dict[str, str]],
@@ -115,11 +169,27 @@ def _generate_transformations(
     return trans
 
 
+_AllowedType = Optional[Union[List[str], _QuicksightPrincipalList]]
+
+
+def _get_principal_names(principals: _AllowedType, type: Literal["users", "groups"]) -> Optional[List[str]]:
+    if principals is None:
+        return None
+
+    if isinstance(principals, list):
+        if type == "users":
+            return principals
+        else:
+            return None
+
+    return principals.get(type)
+
+
 def create_athena_data_source(
     name: str,
     workgroup: str = "primary",
-    allowed_to_use: Optional[List[str]] = None,
-    allowed_to_manage: Optional[List[str]] = None,
+    allowed_to_use: _AllowedType = None,
+    allowed_to_manage: _AllowedType = None,
     tags: Optional[Dict[str, str]] = None,
     account_id: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
@@ -140,13 +210,19 @@ def create_athena_data_source(
         Athena workgroup.
     tags : Dict[str, str], optional
         Key/Value collection to put on the Cluster.
-        e.g. {"foo": "boo", "bar": "xoo"})
-    allowed_to_use : optional
-        List of principals that will be allowed to see and use the data source.
-        e.g. ["John"]
-    allowed_to_manage : optional
-        List of principals that will be allowed to see, use, update and delete the data source.
-        e.g. ["Mary"]
+        e.g. ```{"foo": "boo", "bar": "xoo"})```
+    allowed_to_use: dict["users" | "groups", list[str]], optional
+        Dictionary containing usernames and groups that will be allowed to see and
+        use the data.
+        e.g. ```{"users": ["john", "Mary"], "groups": ["engineering", "customers"]}```
+        Alternatively, if a list of string is passed,
+        it will be interpreted as a list of usernames only.
+    allowed_to_manage: dict["users" | "groups", list[str]], optional
+        Dictionary containing usernames and groups that will be allowed to see, use,
+        update and delete the data source.
+        e.g. ```{"users": ["Mary"], "groups": ["engineering"]}```
+        Alternatively, if a list of string is passed,
+        it will be interpreted as a list of usernames only.
     account_id : str, optional
         If None, the account ID will be inferred from your boto3 session.
     boto3_session : boto3.Session(), optional
@@ -180,11 +256,13 @@ def create_athena_data_source(
     }
     permissions: List[Dict[str, Union[str, List[str]]]] = _generate_permissions(
         resource="data_source",
+        namespace=namespace,
         account_id=account_id,
         boto3_session=boto3_session,
-        allowed_to_use=allowed_to_use,
-        allowed_to_manage=allowed_to_manage,
-        namespace=namespace,
+        allowed_users_to_use=_get_principal_names(allowed_to_use, "users"),
+        allowed_users_to_manage=_get_principal_names(allowed_to_manage, "users"),
+        allowed_groups_to_use=_get_principal_names(allowed_to_use, "groups"),
+        allowed_groups_to_manage=_get_principal_names(allowed_to_manage, "groups"),
     )
     if permissions:
         args["Permissions"] = permissions
@@ -203,8 +281,8 @@ def create_athena_dataset(
     data_source_name: Optional[str] = None,
     data_source_arn: Optional[str] = None,
     import_mode: Literal["SPICE", "DIRECT_QUERY"] = "DIRECT_QUERY",
-    allowed_to_use: Optional[List[str]] = None,
-    allowed_to_manage: Optional[List[str]] = None,
+    allowed_to_use: _AllowedType = None,
+    allowed_to_manage: _AllowedType = None,
     logical_table_alias: str = "LogicalTable",
     rename_columns: Optional[Dict[str, str]] = None,
     cast_columns_types: Optional[Dict[str, str]] = None,
@@ -251,12 +329,18 @@ def create_athena_dataset(
     tags : Dict[str, str], optional
         Key/Value collection to put on the Cluster.
         e.g. {"foo": "boo", "bar": "xoo"}
-    allowed_to_use : optional
-        List of usernames that will be allowed to see and use the data source.
-        e.g. ["john", "Mary"]
-    allowed_to_manage : optional
-        List of usernames that will be allowed to see, use, update and delete the data source.
-        e.g. ["Mary"]
+    allowed_to_use: dict["users" | "groups", list[str]], optional
+        Dictionary containing usernames and groups that will be allowed to see and
+        use the data.
+        e.g. ```{"users": ["john", "Mary"], "groups": ["engineering", "customers"]}```
+        Alternatively, if a list of string is passed,
+        it will be interpreted as a list of usernames only.
+    allowed_to_manage: dict["users" | "groups", list[str]], optional
+        Dictionary containing usernames and groups that will be allowed to see, use,
+        update and delete the data source.
+        e.g. ```{"users": ["Mary"], "groups": ["engineering"]}```
+        Alternatively, if a list of string is passed,
+        it will be interpreted as a list of usernames only.
     logical_table_alias : str
         A display name for the logical table.
     rename_columns : Dict[str, str], optional
@@ -347,13 +431,16 @@ def create_athena_dataset(
     )
     if trans:
         args["LogicalTableMap"][table_uuid]["DataTransforms"] = trans
+
     permissions: List[Dict[str, Union[str, List[str]]]] = _generate_permissions(
         resource="dataset",
+        namespace=namespace,
         account_id=account_id,
         boto3_session=boto3_session,
-        allowed_to_use=allowed_to_use,
-        allowed_to_manage=allowed_to_manage,
-        namespace=namespace,
+        allowed_users_to_use=_get_principal_names(allowed_to_use, "users"),
+        allowed_users_to_manage=_get_principal_names(allowed_to_manage, "users"),
+        allowed_groups_to_use=_get_principal_names(allowed_to_use, "groups"),
+        allowed_groups_to_manage=_get_principal_names(allowed_to_manage, "groups"),
     )
     if permissions:
         args["Permissions"] = permissions

awswrangler/quicksight/_utils.py

@@ -1,16 +1,22 @@
 """Internal (private) Amazon QuickSight Utilities Module."""
 
 import logging
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, TypedDict
 
 import boto3
+from typing_extensions import NotRequired
 
 from awswrangler import _data_types, athena, catalog, exceptions
 from awswrangler.quicksight._get_list import list_data_sources
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
+class _QuicksightPrincipalList(TypedDict):
+    users: NotRequired[List[str]]
+    groups: NotRequired[List[str]]
+
+
 def extract_athena_table_columns(
     database: str, table: str, boto3_session: Optional[boto3.Session]
 ) -> List[Dict[str, str]]:

tests/unit/test_quicksight.py

@@ -50,8 +50,8 @@ def test_quicksight(path, quicksight_datasource, quicksight_dataset, glue_databa
         sql=f"SELECT * FROM {glue_database}.{glue_table}",
         data_source_name=quicksight_datasource,
         import_mode="SPICE",
-        allowed_to_use=[wr.sts.get_current_identity_name()],
-        allowed_to_manage=[wr.sts.get_current_identity_name()],
+        allowed_to_use={"users": [wr.sts.get_current_identity_name()]},
+        allowed_to_manage={"users": [wr.sts.get_current_identity_name()]},
         rename_columns={"iint16": "new_col"},
         cast_columns_types={"new_col": "STRING"},
         tag_columns={"string": [{"ColumnGeographicRole": "CITY"}, {"ColumnDescription": {"Text": "some description"}}]},
@@ -84,7 +84,9 @@ def test_quicksight_delete_all_datasources_filter():
     wr.quicksight.delete_all_data_sources(regex_filter="test.*")
     resource_name = f"test-delete-{uuid.uuid4()}"
     wr.quicksight.create_athena_data_source(
-        name=resource_name, allowed_to_manage=[wr.sts.get_current_identity_name()], tags={"Env": "aws-sdk-pandas"}
+        name=resource_name,
+        allowed_to_manage={"users": [wr.sts.get_current_identity_name()]},
+        tags={"Env": "aws-sdk-pandas"},
     )
     wr.quicksight.delete_all_data_sources(regex_filter="test-no-delete")
 
@@ -104,15 +106,17 @@ def test_quicksight_delete_all_datasets(path, glue_database, glue_table):
 
     resource_name = f"test{str(uuid.uuid4())[:8]}"
     wr.quicksight.create_athena_data_source(
-        name=resource_name, allowed_to_manage=[wr.sts.get_current_identity_name()], tags={"Env": "aws-sdk-pandas"}
+        name=resource_name,
+        allowed_to_manage={"users": [wr.sts.get_current_identity_name()]},
+        tags={"Env": "aws-sdk-pandas"},
     )
     wr.quicksight.create_athena_dataset(
         name=f"{resource_name}-sql",
         sql=f"SELECT * FROM {glue_database}.{glue_table}",
         data_source_name=resource_name,
         import_mode="SPICE",
-        allowed_to_use=[wr.sts.get_current_identity_name()],
-        allowed_to_manage=[wr.sts.get_current_identity_name()],
+        allowed_to_use={"users": [wr.sts.get_current_identity_name()]},
+        allowed_to_manage={"users": [wr.sts.get_current_identity_name()]},
         rename_columns={"iint16": "new_col"},
         cast_columns_types={"new_col": "STRING"},
         tag_columns={"string": [{"ColumnGeographicRole": "CITY"}, {"ColumnDescription": {"Text": "some description"}}]},