enhancement: Adding filter to quicksight delete_all methods (#1913)

* Adding filter to delete_all methods

* renaming arg

* adding delete all datasets tests

* adding test with unmatching and matching filter

* adding data to glue table

* adding randomly generated suffixes to avoid conflicts

* fixing syntax

* consolidating tests

* fixing test

Author: malachi-constant

awswrangler/quicksight/_delete.py

@@ -1,6 +1,7 @@
 """Amazon QuickSight Delete Module."""
 
 import logging
+import re
 from typing import Any, Callable, Dict, Optional
 
 import boto3
@@ -231,13 +232,17 @@ def delete_template(
     _delete(**args)
 
 
-def delete_all_dashboards(account_id: Optional[str] = None, boto3_session: Optional[boto3.Session] = None) -> None:
+def delete_all_dashboards(
+    account_id: Optional[str] = None, regex_filter: Optional[str] = None, boto3_session: Optional[boto3.Session] = None
+) -> None:
     """Delete all dashboards.
 
     Parameters
     ----------
     account_id : str, optional
         If None, the account ID will be inferred from your boto3 session.
+    regex_filter : str, optional
+        Regex regex_filter that will delete all dashboards with a match in their ``Name``
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
 
@@ -255,16 +260,23 @@ def delete_all_dashboards(account_id: Optional[str] = None, boto3_session: Optio
     if account_id is None:
         account_id = sts.get_account_id(boto3_session=session)
     for dashboard in list_dashboards(account_id=account_id, boto3_session=session):
+        if regex_filter:
+            if not re.search(f"^{regex_filter}$", dashboard["Name"]):
+                continue
         delete_dashboard(dashboard_id=dashboard["DashboardId"], account_id=account_id, boto3_session=session)
 
 
-def delete_all_datasets(account_id: Optional[str] = None, boto3_session: Optional[boto3.Session] = None) -> None:
+def delete_all_datasets(
+    account_id: Optional[str] = None, regex_filter: Optional[str] = None, boto3_session: Optional[boto3.Session] = None
+) -> None:
     """Delete all datasets.
 
     Parameters
     ----------
     account_id : str, optional
         If None, the account ID will be inferred from your boto3 session.
+    regex_filter : str, optional
+        Regex regex_filter that will delete all datasets with a match in their ``Name``
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
 
@@ -282,16 +294,23 @@ def delete_all_datasets(account_id: Optional[str] = None, boto3_session: Optiona
     if account_id is None:
         account_id = sts.get_account_id(boto3_session=session)
     for dataset in list_datasets(account_id=account_id, boto3_session=session):
+        if regex_filter:
+            if not re.search(f"^{regex_filter}$", dataset["Name"]):
+                continue
         delete_dataset(dataset_id=dataset["DataSetId"], account_id=account_id, boto3_session=session)
 
 
-def delete_all_data_sources(account_id: Optional[str] = None, boto3_session: Optional[boto3.Session] = None) -> None:
+def delete_all_data_sources(
+    account_id: Optional[str] = None, regex_filter: Optional[str] = None, boto3_session: Optional[boto3.Session] = None
+) -> None:
     """Delete all data sources.
 
     Parameters
     ----------
     account_id : str, optional
         If None, the account ID will be inferred from your boto3 session.
+    regex_filter : str, optional
+        Regex regex_filter that will delete all data sources with a match in their ``Name``
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
 
@@ -309,16 +328,23 @@ def delete_all_data_sources(account_id: Optional[str] = None, boto3_session: Opt
     if account_id is None:
         account_id = sts.get_account_id(boto3_session=session)
     for data_source in list_data_sources(account_id=account_id, boto3_session=session):
+        if regex_filter:
+            if not re.search(f"^{regex_filter}$", data_source["Name"]):
+                continue
         delete_data_source(data_source_id=data_source["DataSourceId"], account_id=account_id, boto3_session=session)
 
 
-def delete_all_templates(account_id: Optional[str] = None, boto3_session: Optional[boto3.Session] = None) -> None:
+def delete_all_templates(
+    account_id: Optional[str] = None, regex_filter: Optional[str] = None, boto3_session: Optional[boto3.Session] = None
+) -> None:
     """Delete all templates.
 
     Parameters
     ----------
     account_id : str, optional
         If None, the account ID will be inferred from your boto3 session.
+    regex_filter : str, optional
+        Regex regex_filter that will delete all templates with a match in their ``Name``
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
 
@@ -336,4 +362,7 @@ def delete_all_templates(account_id: Optional[str] = None, boto3_session: Option
     if account_id is None:
         account_id = sts.get_account_id(boto3_session=session)
     for template in list_templates(account_id=account_id, boto3_session=session):
+        if regex_filter:
+            if not re.search(f"^{regex_filter}$", template["Name"]):
+                continue
         delete_template(template_id=template["TemplateId"], account_id=account_id, boto3_session=session)

tests/test_quicksight.py

@@ -1,4 +1,5 @@
 import logging
+import uuid
 
 import boto3
 import pytest
@@ -18,39 +19,40 @@ def test_quicksight(path, glue_database, glue_table):
     wr.s3.to_parquet(
         df=df, path=path, dataset=True, database=glue_database, table=glue_table, partition_cols=["par0", "par1"]
     )
+    resource_name = f"test{str(uuid.uuid4())[:8]}"
 
-    wr.quicksight.delete_all_dashboards()
-    wr.quicksight.delete_all_datasets()
-    wr.quicksight.delete_all_data_sources()
-    wr.quicksight.delete_all_templates()
+    wr.quicksight.delete_all_dashboards(regex_filter="test.*")
+    wr.quicksight.delete_all_datasets(regex_filter="test.*")
+    wr.quicksight.delete_all_data_sources(regex_filter="test.*")
+    wr.quicksight.delete_all_templates(regex_filter="test.*")
 
     wr.quicksight.create_athena_data_source(
-        name="test", allowed_to_manage=[wr.sts.get_current_identity_name()], tags={"Env": "aws-sdk-pandas"}
+        name=resource_name, allowed_to_manage=[wr.sts.get_current_identity_name()], tags={"Env": "aws-sdk-pandas"}
     )
-    assert wr.quicksight.describe_data_source("test")["Name"] == "test"
+    assert wr.quicksight.describe_data_source(resource_name)["Name"] == resource_name
     assert (
-        wr.quicksight.describe_data_source_permissions("test")[0]["Principal"].endswith(
+        wr.quicksight.describe_data_source_permissions(resource_name)[0]["Principal"].endswith(
             wr.sts.get_current_identity_name()
         )
         is True
     )
-
+    dataset_name = f"{resource_name}-table"
     wr.quicksight.create_athena_dataset(
-        name="test-table",
+        name=dataset_name,
         database=glue_database,
         table=glue_table,
-        data_source_name="test",
+        data_source_name=resource_name,
         allowed_to_manage=[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"}}]},
     )
-    assert wr.quicksight.describe_dataset("test-table")["Name"] == "test-table"
+    assert wr.quicksight.describe_dataset(dataset_name)["Name"] == dataset_name
 
     wr.quicksight.create_athena_dataset(
-        name="test-sql",
+        name=f"{resource_name}-sql",
         sql=f"SELECT * FROM {glue_database}.{glue_table}",
-        data_source_name="test",
+        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()],
@@ -60,23 +62,67 @@ def test_quicksight(path, glue_database, glue_table):
         tags={"foo": "boo"},
     )
 
-    ingestion_id = wr.quicksight.create_ingestion("test-sql")
+    ingestion_id = wr.quicksight.create_ingestion(f"{resource_name}-sql")
     status = None
     while status not in ["FAILED", "COMPLETED", "CANCELLED"]:
-        status = wr.quicksight.describe_ingestion(ingestion_id, "test-sql")["IngestionStatus"]
+        status = wr.quicksight.describe_ingestion(ingestion_id, f"{resource_name}-sql")["IngestionStatus"]
     assert status == "COMPLETED"
 
-    ingestion_id = wr.quicksight.create_ingestion("test-sql")
-    wr.quicksight.cancel_ingestion(ingestion_id, "test-sql")
-    assert len(wr.quicksight.list_ingestions("test-sql")) == 3
+    ingestion_id = wr.quicksight.create_ingestion(f"{resource_name}-sql")
+    wr.quicksight.cancel_ingestion(ingestion_id, f"{resource_name}-sql")
+    assert len(wr.quicksight.list_ingestions(f"{resource_name}-sql")) == 3
 
     wr.quicksight.list_groups()
     wr.quicksight.list_iam_policy_assignments()
     wr.quicksight.list_iam_policy_assignments_for_user(wr.sts.get_current_identity_name())
     wr.quicksight.list_user_groups(wr.sts.get_current_identity_name())
     wr.quicksight.list_users()
-    wr.quicksight.get_dataset_ids("test-sql")
+    wr.quicksight.get_dataset_ids(f"{resource_name}-sql")
     wr.quicksight.get_data_source_ids("test")
 
-    wr.quicksight.delete_all_datasets()
-    wr.quicksight.delete_all_data_sources()
+    wr.quicksight.delete_all_datasets(regex_filter="test.*")
+    wr.quicksight.delete_all_data_sources(regex_filter="test.*")
+
+
+def test_quicksight_delete_all_datasources_filter():
+    wr.quicksight.delete_all_data_sources(regex_filter="test.*")
+    resource_name = "test-delete"
+    wr.quicksight.create_athena_data_source(
+        name=resource_name, allowed_to_manage=[wr.sts.get_current_identity_name()], tags={"Env": "aws-sdk-pandas"}
+    )
+    wr.quicksight.delete_all_data_sources(regex_filter="test-no-delete")
+
+    assert len(wr.quicksight.get_data_source_ids(resource_name)) == 1
+
+    wr.quicksight.delete_all_data_sources(regex_filter="test-delete.*")
+    assert len(wr.quicksight.get_data_source_ids(resource_name)) == 0
+
+
+def test_quicksight_delete_all_datasets(path, glue_database, glue_table):
+    df = get_df_quicksight()
+    wr.s3.to_parquet(
+        df=df, path=path, dataset=True, database=glue_database, table=glue_table, partition_cols=["par0", "par1"]
+    )
+    wr.quicksight.delete_all_datasets(regex_filter="test.*")
+    wr.quicksight.delete_all_data_sources(regex_filter="test.*")
+
+    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"}
+    )
+    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()],
+        rename_columns={"iint16": "new_col"},
+        cast_columns_types={"new_col": "STRING"},
+        tag_columns={"string": [{"ColumnGeographicRole": "CITY"}, {"ColumnDescription": {"Text": "some description"}}]},
+        tags={"foo": "boo"},
+    )
+    wr.quicksight.delete_all_datasets(regex_filter="test.*")
+    wr.quicksight.delete_all_data_sources(regex_filter="test.*")
+
+    assert len(wr.quicksight.get_dataset_ids(resource_name)) == 0