Add databricks labs ucx delete-credential cmd to delete the UC roles created by UCX (#2504)

## Changes
<!-- Summary of your changes that are easy to understand. Add
screenshots when necessary -->

### Linked issues
<!-- DOC: Link issue with a keyword: close, closes, closed, fix, fixes,
fixed, resolve, resolves, resolved. See
https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword
-->

Resolves #2359 


https://github.com/user-attachments/assets/b0e036b5-aca9-4805-ae22-922fe6c97c1c

Author: HariGS-DB

labs.yml

@@ -151,6 +151,14 @@ commands:
       - name: single-role
         description: (Optional) Create a single role for all S3 locations. (default:False)
 
+  - name: delete-missing-principals
+    description: For AWS, this command identifies all the UC roles that are created through the create-missing-principals
+      cmd. It lists all the UC roles in aws and lets users select the roles to delete. It also validates if the selected
+      roles are used by any storage credential and prompts to confirm if roles should still be deleted.
+    flags:
+      - name: aws-profile
+        description: AWS Profile to use for authentication
+
   - name: create-uber-principal
     description: For azure cloud, creates a service principal and gives STORAGE BLOB READER access on all the storage account
       used by tables in the workspace and stores the spn info in the UCX cluster policy. For aws,

src/databricks/labs/ucx/assessment/aws.py

@@ -409,6 +409,12 @@ def delete_instance_profile(self, instance_profile_name: str, role_name: str):
             f" --role-name {role_name}"
         )
         self._run_json_command(f"iam delete-instance-profile --instance-profile-name {instance_profile_name}")
+        self.delete_role(role_name)
+
+    def delete_role(self, role_name: str):
+        role_policies = self.list_role_policies(role_name)
+        for policy in role_policies:
+            self._run_json_command(f"iam delete-role-policy --role-name {role_name} --policy-name {policy}")
         self._run_json_command(f"iam delete-role --role-name {role_name}")
 
     def add_role_to_instance_profile(self, instance_profile_name: str, role_name: str) -> bool:

src/databricks/labs/ucx/aws/access.py

@@ -367,3 +367,6 @@ def _generate_role_name(self, single_role: bool, role_name: str, location: str)
             metastore_id = self._ws.metastores.current().as_dict()["metastore_id"]
             return f"{role_name}_{metastore_id}"
         return f"{role_name}_{location[5:]}"
+
+    def delete_uc_role(self, role_name: str):
+        self._aws_resources.delete_role(role_name)

src/databricks/labs/ucx/aws/credentials.py

@@ -230,3 +230,38 @@ def run(self, prompts: Prompts, *, single_role=False, role_name="UC_ROLE", polic
 
         self._resource_permissions.create_uc_roles(iam_list)
         self._resource_permissions.save_uc_compatible_roles()
+
+    def delete_uc_roles(self, prompts: Prompts) -> None:
+        uc_roles = self._resource_permissions.load_uc_compatible_roles()
+        if len(uc_roles) == 0:
+            self._resource_permissions.save_uc_compatible_roles()
+            uc_roles = self._resource_permissions.load_uc_compatible_roles()
+        storage_credentials = self._ws.storage_credentials.list()
+
+        uc_role_mapping = {role.role_arn: role.role_name for role in uc_roles}
+        selected_roles = prompts.multiple_choice_from_dict("Select the list of roles created by UCX", uc_role_mapping)
+        if len(selected_roles) == 0:
+            logger.info("No roles selected...")
+            return
+        matching_credentials = []
+        for storage_credential in storage_credentials:
+            if (
+                storage_credential.aws_iam_role is not None
+                and uc_role_mapping.get(storage_credential.aws_iam_role.role_arn) in selected_roles
+            ):
+                matching_credentials.append(storage_credential)
+
+        for credential in matching_credentials:
+            if credential.aws_iam_role is not None:
+                logger.info(f"Storage credential: {credential.name} IAM Role: {credential.aws_iam_role.role_arn}")
+        if len(matching_credentials) == 0:
+            logger.info("No storage credential using the selected UC roles, proceeding to delete.")
+        if len(matching_credentials) == 0 or prompts.confirm(
+            "The above storage credential will be impacted on deleting the selected IAM roles,"
+            " Are you sure you want to confirm"
+        ):
+
+            logger.info("Deleting UCX created roles...")
+            for role_name in selected_roles:
+                logger.info(f"Deleting role {role_name}.")
+                self._resource_permissions.delete_uc_role(role_name)

src/databricks/labs/ucx/cli.py

@@ -347,6 +347,25 @@ def create_missing_principals(
     raise ValueError("Unsupported cloud provider")
 
 
+@ucx.command
+def delete_missing_principals(
+    w: WorkspaceClient,
+    prompts: Prompts,
+    ctx: WorkspaceContext | None = None,
+    **named_parameters,
+):
+    """Not supported for Azure.
+    For AWS, this command identifies all the UC roles that are created through the create-missing-principals cmd.
+    It lists all the UC roles in aws and lets users select the roles to delete. It also validates if the selected roles
+    are used by any storage credential and prompts to confirm if roles should still be deleted.
+    """
+    if not ctx:
+        ctx = WorkspaceContext(w, named_parameters)
+    if ctx.is_aws:
+        return ctx.iam_role_creation.delete_uc_roles(prompts)
+    raise ValueError("Unsupported cloud provider")
+
+
 @ucx.command
 def migrate_credentials(w: WorkspaceClient, prompts: Prompts, ctx: WorkspaceContext | None = None, **named_parameters):
     """For Azure, this command prompts to i) create UC storage credentials for the access connectors with a

tests/unit/aws/test_access.py

@@ -864,3 +864,69 @@ def command_call(_: str):
     assert len(roles) == 1
     assert len(roles[0].paths) == 2
     external_locations.snapshot.assert_called_once()
+
+
+def test_delete_uc_roles(mock_ws, installation_multiple_roles, backend, locations):
+    aws = create_autospec(AWSResources)
+    aws.validate_connection.return_value = {}
+    aws_resource_permissions = AWSResourcePermissions(installation_multiple_roles, mock_ws, aws, locations)
+    mock_ws.storage_credentials.list.return_value = [
+        StorageCredentialInfo(
+            id="1",
+            name="cred1",
+            aws_iam_role=AwsIamRoleResponse("arn:aws:iam::12345:role/uc-role1"),
+        )
+    ]
+    role_creation = IamRoleCreation(installation_multiple_roles, mock_ws, aws_resource_permissions)
+    prompts = MockPrompts({"Select the list of roles *": "1", "The above storage credential will be impacted *": "Yes"})
+    role_creation.delete_uc_roles(prompts)
+    calls = [call("uc-role1"), call("uc-rolex")]
+    assert aws.delete_role.mock_calls == calls
+
+
+def test_delete_uc_roles_not_present(mock_ws, installation_no_roles, backend, locations):
+    aws = create_autospec(AWSResources)
+    aws.validate_connection.return_value = {}
+    aws.delete_role.return_value = []
+    aws_resource_permissions = AWSResourcePermissions(installation_no_roles, mock_ws, aws, locations)
+    mock_ws.storage_credentials.list.return_value = [
+        StorageCredentialInfo(
+            id="1",
+            name="cred1",
+            aws_iam_role=AwsIamRoleResponse("arn:aws:iam::12345:role/uc-role1"),
+        )
+    ]
+    role_creation = IamRoleCreation(installation_no_roles, mock_ws, aws_resource_permissions)
+    aws.list_all_uc_roles.return_value = [AWSRole("", "uc-role1", "123", "arn:aws:iam::12345:role/uc-role1")]
+    aws.get_role_policy.side_effect = [
+        [
+            AWSPolicyAction(
+                resource_type="s3",
+                privilege="READ_FILES",
+                resource_path="s3://bucket1",
+            )
+        ]
+    ]
+    aws.list_role_policies.return_value = ["Policy1"]
+    aws.list_all_uc_roles.return_value = [
+        AWSRole(path='/', role_name='uc-role1', role_id='12345', arn='arn:aws:iam::12345:role/uc-role1')
+    ]
+    prompts = MockPrompts({"Select the list of roles *": "1", "The above storage credential will be impacted *": "Yes"})
+    role_creation.delete_uc_roles(prompts)
+    calls = [call("uc-role1")]
+    assert aws.delete_role.mock_calls == calls
+
+
+def test_delete_role(mock_ws, installation_no_roles, backend, mocker):
+    command_calls = []
+    mocker.patch("shutil.which", return_value="/path/aws")
+
+    def command_call(cmd: str):
+        command_calls.append(cmd)
+        return 0, '{"account":"1234"}', ""
+
+    aws = AWSResources("profile", command_call)
+    external_locations = ExternalLocations(mock_ws, backend, 'ucx')
+    resource_permissions = AWSResourcePermissions(installation_no_roles, mock_ws, aws, external_locations)
+    resource_permissions.delete_uc_role("uc_role_1")
+    assert '/path/aws iam delete-role --role-name uc_role_1 --profile profile --output json' in command_calls

tests/unit/test_cli.py

@@ -7,6 +7,7 @@
 import pytest
 import yaml
 from databricks.labs.blueprint.tui import MockPrompts
+from databricks.labs.ucx.aws.credentials import IamRoleCreation
 from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.errors import NotFound
 from databricks.sdk.errors.platform import BadRequest
@@ -53,6 +54,7 @@
     validate_external_locations,
     validate_groups_membership,
     workflows,
+    delete_missing_principals,
 )
 from databricks.labs.ucx.contexts.account_cli import AccountContext
 from databricks.labs.ucx.contexts.workspace_cli import WorkspaceContext
@@ -742,3 +744,13 @@ def test_join_collection():
     w.workspace.download.return_value = io.StringIO(json.dumps([{"workspace_id": 123, "workspace_name": "some"}]))
     join_collection(a, "123")
     w.workspace.download.assert_not_called()
+
+
+def test_delete_principals(ws):
+    ws.config.is_azure = False
+    ws.config.is_aws = True
+    role_creation = create_autospec(IamRoleCreation)
+    ctx = WorkspaceContext(ws).replace(iam_role_creation=role_creation, workspace_client=ws)
+    prompts = MockPrompts({"Select the list of roles *": "0"})
+    delete_missing_principals(ws, prompts, ctx)
+    role_creation.delete_uc_roles.assert_called_once()