databricks-labs-ucx 0.17.0__py3-none-any.whl

databricks/__init__.py

@@ -0,0 +1,3 @@
+# DO NOT ADD ANYTHING ELSE TO THIS FILE FOR COMPATIBILITY WITH OTHER databricks.* PACKAGES
+# SEE https://packaging.python.org/guides/packaging-namespace-packages/#pkgutil-style-namespace-packages
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

databricks/labs/__init__.py

@@ -0,0 +1,3 @@
+# DO NOT ADD ANYTHING ELSE TO THIS FILE FOR COMPATIBILITY WITH OTHER databricks.* PACKAGES
+# SEE https://packaging.python.org/guides/packaging-namespace-packages/#pkgutil-style-namespace-packages
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

databricks/labs/ucx/__about__.py

@@ -0,0 +1,2 @@
+# DO NOT MODIFY THIS FILE
+__version__ = "0.17.0"

databricks/labs/ucx/__init__.py

@@ -0,0 +1,3 @@
+from databricks.labs.blueprint.logger import install_logger
+
+install_logger()

databricks/labs/ucx/account.py

@@ -0,0 +1,253 @@
+import logging
+from typing import ClassVar
+
+from databricks.labs.blueprint.installation import Installation
+from databricks.labs.blueprint.tui import Prompts
+from databricks.sdk import AccountClient, WorkspaceClient
+from databricks.sdk.errors import NotFound, ResourceConflict
+from databricks.sdk.service.iam import ComplexValue, Group, Patch, PatchOp, PatchSchema
+from databricks.sdk.service.provisioning import Workspace
+
+logger = logging.getLogger(__name__)
+
+
+class AccountWorkspaces:
+    _tlds: ClassVar[dict[str, str]] = {
+        "aws": "cloud.databricks.com",
+        "azure": "azuredatabricks.net",
+        "gcp": "gcp.databricks.com",
+    }
+
+    SYNC_FILE_NAME: ClassVar[str] = "workspaces.json"
+
+    def __init__(self, account_client: AccountClient, new_workspace_client=WorkspaceClient):
+        self._new_workspace_client = new_workspace_client
+        self._ac = account_client
+
+    def _workspaces(self):
+        return self._ac.workspaces.list()
+
+    def _get_cloud(self) -> str:
+        if self._ac.config.is_azure:
+            return "azure"
+        if self._ac.config.is_gcp:
+            return "gcp"
+        return "aws"
+
+    def client_for(self, workspace: Workspace) -> WorkspaceClient:
+        return self._ac.get_workspace_client(workspace)
+
+    def workspace_clients(self) -> list[WorkspaceClient]:
+        """
+        Return a list of WorkspaceClient for each configured workspace in the account
+        :return: list[WorkspaceClient]
+        """
+        clients = []
+        for workspace in self._workspaces():
+            ws = self.client_for(workspace)
+            clients.append(ws)
+        return clients
+
+    def sync_workspace_info(self):
+        """
+        Create a json dump for each Workspace in account
+        For each user that has ucx installed in their workspace,
+        upload the json dump of workspace info in the .ucx folder
+        """
+        workspaces = []
+        for workspace in self._workspaces():
+            workspaces.append(workspace)
+        for ws in self.workspace_clients():
+            for installation in Installation.existing(ws, "ucx"):
+                installation.save(workspaces, filename=self.SYNC_FILE_NAME)
+
+    def create_account_level_groups(self, prompts: Prompts, workspace_ids: list[int] | None = None):
+        acc_groups = self._get_account_groups()
+        workspace_ids = self._get_valid_workspaces_ids(workspace_ids)
+        all_valid_workspace_groups = self._get_valid_workspaces_groups(prompts, workspace_ids)
+
+        for group_name, valid_group in all_valid_workspace_groups.items():
+            acc_group = self._try_create_account_groups(group_name, acc_groups)
+
+            if not acc_group or not valid_group.members or not acc_group.id:
+                continue
+            if len(valid_group.members) > 0:
+                self._add_members_to_acc_group(self._ac, acc_group.id, group_name, valid_group)
+            logger.info(f"Group {group_name} created in the account")
+
+    def _try_create_account_groups(
+        self, group_name: str, acc_groups: dict[str | None, list[ComplexValue] | None]
+    ) -> Group | None:
+        try:
+            if group_name in acc_groups:
+                logger.info(f"Group {group_name} already exist in the account, ignoring")
+                return None
+            return self._ac.groups.create(display_name=group_name)
+        except ResourceConflict:
+            logger.info(f"Group {group_name} already exist in the account, ignoring")
+            return None
+
+    def _get_valid_workspaces_ids(self, workspace_ids: list[int] | None = None) -> list[int]:
+        if not workspace_ids:
+            logger.info("No workspace ids provided, using current workspace instead")
+            return [self._new_workspace_client().get_workspace_id()]
+
+        all_workspace_ids = [workspace.workspace_id for workspace in self._workspaces()]
+
+        valid_workspace_ids = []
+        for workspace_id in workspace_ids:
+            if workspace_id in all_workspace_ids:
+                valid_workspace_ids.append(workspace_id)
+            else:
+                logger.info(f"Workspace id {workspace_id} not found on the account")
+
+        if not valid_workspace_ids:
+            raise ValueError("No workspace ids provided in the configuration found in the account")
+
+        workspace_ids_str = ','.join(str(x) for x in valid_workspace_ids)
+        logger.info(f"Creating account groups for workspaces IDs : {workspace_ids_str}")
+        return valid_workspace_ids
+
+    def _add_members_to_acc_group(
+        self, acc_client: AccountClient, acc_group_id: str, group_name: str, valid_group: Group
+    ):
+        for chunk in self._chunks(valid_group.members, 20):
+            logger.debug(f"Adding {len(chunk)} members to acc group {group_name}")
+            acc_client.groups.patch(
+                acc_group_id,
+                operations=[Patch(op=PatchOp.ADD, path="members", value=[x.as_dict() for x in chunk])],
+                schemas=[PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP],
+            )
+
+    def _chunks(self, lst, chunk_size):
+        """Yield successive n-sized chunks from lst."""
+        for i in range(0, len(lst), chunk_size):
+            yield lst[i : i + chunk_size]
+
+    def _get_valid_workspaces_groups(self, prompts: Prompts, workspace_ids: list[int]) -> dict[str, Group]:
+        all_workspaces_groups: dict[str, Group] = {}
+
+        for workspace in self._workspaces():
+            if workspace.workspace_id not in workspace_ids:
+                continue
+            self._load_workspace_groups(prompts, workspace, all_workspaces_groups)
+
+        return all_workspaces_groups
+
+    def _load_workspace_groups(self, prompts, workspace, all_workspaces_groups):
+        client = self.client_for(workspace)
+        logger.info(f"Crawling groups in workspace {client.config.host}")
+        ws_group_ids = client.groups.list(attributes="id")
+        for group_id in ws_group_ids:
+            full_workspace_group = self._safe_groups_get(client, group_id.id)
+            if not full_workspace_group:
+                continue
+            group_name = full_workspace_group.display_name
+            if self._is_group_out_of_scope(full_workspace_group):
+                continue
+            if not group_name:
+                continue
+            if group_name in all_workspaces_groups:
+                if self._has_same_members(all_workspaces_groups[group_name], full_workspace_group):
+                    logger.info(f"Workspace group {group_name} already found, ignoring")
+                    continue
+                if prompts.confirm(
+                    f"Group {group_name} does not have the same amount of members "
+                    f"in workspace {client.config.host} than previous workspaces which contains the same group name,"
+                    f"it will be created at the account with name : {workspace.workspace_name}_{group_name}"
+                ):
+                    all_workspaces_groups[f"{workspace.workspace_name}_{group_name}"] = full_workspace_group
+                    continue
+            logger.info(f"Found new group {group_name}")
+            all_workspaces_groups[group_name] = full_workspace_group
+        logger.info(f"Found a total of {len(all_workspaces_groups)} groups to migrate to the account")
+
+    def _is_group_out_of_scope(self, group: Group) -> bool:
+        if group.display_name in {"users", "admins", "account users"}:
+            logger.debug(f"Group {group.display_name} is a system group, ignoring")
+            return True
+        meta = group.meta
+        if not meta:
+            return False
+        if meta.resource_type != "WorkspaceGroup":
+            logger.debug(f"Group {group.display_name} is an account group, ignoring")
+            return True
+        return False
+
+    def _has_same_members(self, group_1: Group, group_2: Group) -> bool:
+        ws_members_set_1 = set([m.display for m in group_1.members] if group_1.members else [])
+        ws_members_set_2 = set([m.display for m in group_2.members] if group_2.members else [])
+        return not bool((ws_members_set_1 - ws_members_set_2).union(ws_members_set_2 - ws_members_set_1))
+
+    def _get_account_groups(self) -> dict[str | None, list[ComplexValue] | None]:
+        logger.debug("Listing groups in account")
+        acc_groups = {}
+        for acc_grp_id in self._ac.groups.list(attributes="id"):
+            if not acc_grp_id.id:
+                continue
+            full_account_group = self._safe_groups_get(self._ac, acc_grp_id.id)
+            if not full_account_group:
+                continue
+            logger.debug(f"Found account group {full_account_group.display_name}")
+            acc_groups[full_account_group.display_name] = full_account_group.members
+
+        logger.info(f"{len(acc_groups)} account groups found")
+        return acc_groups
+
+    def _safe_groups_get(self, interface, group_id) -> Group | None:
+        try:
+            if not group_id:
+                return None
+            return interface.groups.get(group_id)
+        except NotFound:
+            logger.info(f"Group {group_id} has been deleted")
+            return None
+
+
+class WorkspaceInfo:
+    def __init__(self, installation: Installation, ws: WorkspaceClient):
+        self._installation = installation
+        self._ws = ws
+
+    def _load_workspace_info(self) -> dict[int, Workspace]:
+        try:
+            id_to_workspace = {}
+            for workspace in self._installation.load(list[Workspace], filename=AccountWorkspaces.SYNC_FILE_NAME):
+                assert workspace.workspace_id is not None
+                id_to_workspace[workspace.workspace_id] = workspace
+            return id_to_workspace
+        except NotFound:
+            msg = "Please run as account-admin: databricks labs ucx sync-workspace-info"
+            raise ValueError(msg) from None
+
+    def current(self) -> str:
+        workspace_id = self._ws.get_workspace_id()
+        workspaces = self._load_workspace_info()
+        if workspace_id not in workspaces:
+            msg = f"Current workspace is not known: {workspace_id}"
+            raise KeyError(msg) from None
+        workspace = workspaces[workspace_id]
+        assert workspace.workspace_name, "workspace name undefined"
+        return workspace.workspace_name
+
+    def manual_workspace_info(self, prompts: Prompts):
+        logger.warning(
+            'You are strongly recommended to run "databricks labs ucx sync-workspace-info" by account admin,'
+            ' otherwise there is a significant risk of inconsistencies between different workspaces. This '
+            'command will overwrite all UCX installations on this given workspace. Result may be consistent '
+            f'only within {self._ws.config.host}'
+        )
+        workspaces = []
+        workspace_id = self._ws.get_workspace_id()
+        while workspace_id:
+            workspace_name = prompts.question(
+                f"Workspace name for {workspace_id}", default=f"workspace-{workspace_id}", valid_regex=r"^[\w-]+$"
+            )
+            workspaces.append(Workspace(workspace_id=int(workspace_id), workspace_name=workspace_name))
+            answer = prompts.question("Next workspace id", valid_number=True, default="stop")
+            if answer == "stop":
+                break
+            workspace_id = int(answer)
+        for installation in Installation.existing(self._ws, 'ucx'):
+            installation.save(workspaces, filename=AccountWorkspaces.SYNC_FILE_NAME)
+        logger.info("Synchronised workspace id mapping for installations on current workspace")

databricks/labs/ucx/assessment/aws.py

@@ -0,0 +1,391 @@
+import json
+import logging
+import re
+import shutil
+import subprocess
+import typing
+from collections.abc import Callable
+from dataclasses import dataclass
+from functools import lru_cache
+
+from databricks.sdk.service.catalog import Privilege
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AWSRole:
+    path: str
+    role_name: str
+    role_id: str
+    arn: str
+
+
+@dataclass
+class AWSPolicyAction:
+    resource_type: str
+    privilege: str
+    resource_path: str
+
+
+@dataclass
+class AWSRoleAction:
+    role_arn: str
+    resource_type: str
+    privilege: str
+    resource_path: str
+
+    @property
+    def role_name(self):
+        role_match = re.match(AWSInstanceProfile.ROLE_NAME_REGEX, self.role_arn)
+        return role_match.group(1)
+
+
+@dataclass
+class AWSInstanceProfile:
+    instance_profile_arn: str
+    iam_role_arn: str | None = None
+
+    ROLE_NAME_REGEX = r"arn:aws:iam::[0-9]+:(?:instance-profile|role)\/([a-zA-Z0-9+=,.@_-]*)$"
+
+    @property
+    def role_name(self) -> str | None:
+        if self.iam_role_arn:
+            arn = self.iam_role_arn
+        else:
+            arn = self.instance_profile_arn
+        role_match = re.match(self.ROLE_NAME_REGEX, arn)
+        if not role_match:
+            logger.error(f"Role ARN is mismatched {self.iam_role_arn}")
+            return None
+        return role_match.group(1)
+
+
+@lru_cache(maxsize=1024)
+def run_command(command):
+    logger.info(f"Invoking Command {command}")
+    with subprocess.Popen(command.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE) as process:
+        output, error = process.communicate()
+        return process.returncode, output.decode("utf-8"), error.decode("utf-8")
+
+
+class AWSResources:
+    S3_ACTIONS: typing.ClassVar[set[str]] = {"s3:PutObject", "s3:GetObject", "s3:DeleteObject", "s3:PutObjectAcl"}
+    S3_READONLY: typing.ClassVar[str] = "s3:GetObject"
+    S3_REGEX: typing.ClassVar[str] = r"arn:aws:s3:::([a-zA-Z0-9\/+=,.@_-]*)\/\*$"
+    S3_PREFIX: typing.ClassVar[str] = "arn:aws:s3:::"
+    S3_PATH_REGEX: typing.ClassVar[str] = r"((s3:\/\/)|(s3a:\/\/))(.*)"
+    UC_MASTER_ROLES_ARN: typing.ClassVar[list[str]] = [
+        "arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL",
+        "arn:aws:iam::707343435239:role/unity-catalog-dev-UCMasterRole-G3MMN8SP21FO",
+    ]
+
+    def __init__(self, profile: str, command_runner: Callable[[str], tuple[int, str, str]] = run_command):
+        self._profile = profile
+        self._command_runner = command_runner
+
+    def validate_connection(self):
+        validate_command = f"sts get-caller-identity --profile {self._profile}"
+        result = self._run_json_command(validate_command)
+        if result:
+            logger.info(result)
+            return result
+        return None
+
+    def list_role_policies(self, role_name: str):
+        list_policies_cmd = f"iam list-role-policies --profile {self._profile} --role-name {role_name}"
+        policies = self._run_json_command(list_policies_cmd)
+        if not policies:
+            return []
+        return policies.get("PolicyNames", [])
+
+    def list_attached_policies_in_role(self, role_name: str):
+        list_attached_policies_cmd = (
+            f"iam list-attached-role-policies --profile {self._profile} --role-name {role_name}"
+        )
+        policies = self._run_json_command(list_attached_policies_cmd)
+        if not policies:
+            return []
+        attached_policies = []
+        for policy in policies.get("AttachedPolicies", []):
+            attached_policies.append(policy.get("PolicyArn"))
+        return attached_policies
+
+    def list_all_uc_roles(self) -> list[AWSRole]:
+        roles = self._run_json_command(f"iam list-roles --profile {self._profile}")
+        uc_roles: list[AWSRole] = []
+        roles = roles.get("Roles")
+        if not roles:
+            logger.warning("list-roles couldn't find any roles")
+            return uc_roles
+        for role in roles:
+            policy_document = role.get("AssumeRolePolicyDocument")
+            if not policy_document:
+                continue
+            for statement in policy_document["Statement"]:
+                effect = statement.get("Effect")
+                action = statement.get("Action")
+                principal = statement.get("Principal")
+                if not (effect and action and principal):
+                    continue
+                if effect != "Allow":
+                    continue
+                if action != "sts:AssumeRole":
+                    continue
+                principal = principal.get("AWS")
+                if not principal:
+                    continue
+                if not self._is_uc_principal(principal):
+                    continue
+                uc_roles.append(
+                    AWSRole(
+                        role_id=role.get("RoleId"),
+                        role_name=role.get("RoleName"),
+                        arn=role.get("Arn"),
+                        path=role.get("Path"),
+                    )
+                )
+
+        return uc_roles
+
+    def _is_uc_principal(self, principal):
+        if isinstance(principal, list):
+            for single_principal in principal:
+                if single_principal in self.UC_MASTER_ROLES_ARN:
+                    return True
+            return False
+        return principal in self.UC_MASTER_ROLES_ARN
+
+    def get_role_policy(self, role_name, policy_name: str | None = None, attached_policy_arn: str | None = None):
+        if policy_name:
+            get_policy = (
+                f"iam get-role-policy --profile {self._profile} --role-name {role_name} " f"--policy-name {policy_name}"
+            )
+        elif attached_policy_arn:
+            get_attached_policy = f"iam get-policy --profile {self._profile} --policy-arn {attached_policy_arn}"
+            attached_policy = self._run_json_command(get_attached_policy)
+            if not attached_policy:
+                return []
+            policy_version = attached_policy["Policy"]["DefaultVersionId"]
+            get_policy = (
+                f"iam get-policy-version --profile {self._profile} --policy-arn {attached_policy_arn} "
+                f"--version-id {policy_version}"
+            )
+        else:
+            logger.error("Failed to retrieve role. No role name or attached role ARN specified.")
+            return []
+        policy = self._run_json_command(get_policy)
+        if not policy:
+            return []
+        if policy_name:
+            actions = policy["PolicyDocument"].get("Statement", [])
+        else:
+            actions = policy["PolicyVersion"]["Document"].get("Statement", [])
+        return self._policy_actions(actions)
+
+    def _policy_actions(self, actions):
+        policy_actions = []
+        for action in actions:
+            if action.get("Effect", "Deny") != "Allow":
+                continue
+            actions = action["Action"]
+            s3_actions = self._s3_actions(actions)
+            if not s3_actions or self.S3_READONLY not in s3_actions:
+                continue
+            privilege = Privilege.WRITE_FILES.value
+            for s3_action_type in self.S3_ACTIONS:
+                if s3_action_type not in s3_actions:
+                    privilege = Privilege.READ_FILES.value
+                    continue
+            for resource in action.get("Resource", []):
+                match = re.match(self.S3_REGEX, resource)
+                if match:
+                    policy_actions.append(AWSPolicyAction("s3", privilege, f"s3://{match.group(1)}"))
+                    policy_actions.append(AWSPolicyAction("s3", privilege, f"s3a://{match.group(1)}"))
+        return policy_actions
+
+    def _s3_actions(self, actions):
+        s3_actions = []
+        if isinstance(actions, list):
+            for single_action in actions:
+                if single_action in self.S3_ACTIONS:
+                    s3_actions.append(single_action)
+                    continue
+        elif actions in self.S3_ACTIONS:
+            s3_actions = [actions]
+        return s3_actions
+
+    def _aws_role_trust_doc(self, external_id="0000"):
+        return self._get_json_for_cli(
+            {
+                "Version": "2012-10-17",
+                "Statement": [
+                    {
+                        "Effect": "Allow",
+                        "Principal": {
+                            "AWS": "arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL"
+                        },
+                        "Action": "sts:AssumeRole",
+                        "Condition": {"StringEquals": {"sts:ExternalId": external_id}},
+                    }
+                ],
+            }
+        )
+
+    def _aws_s3_policy(self, s3_prefixes, account_id, role_name, kms_key=None):
+        """
+        Create the UC IAM policy for the given S3 prefixes, account ID, role name, and KMS key.
+        """
+        s3_prefixes_strip = set()
+        for path in s3_prefixes:
+            match = re.match(AWSResources.S3_PATH_REGEX, path)
+            if match:
+                s3_prefixes_strip.add(match.group(4))
+
+        s3_prefixes_enriched = sorted([self.S3_PREFIX + s3_prefix for s3_prefix in s3_prefixes_strip])
+        statement = [
+            {
+                "Action": ["s3:GetObject", "s3:PutObject", "s3:DeleteObject", "s3:ListBucket", "s3:GetBucketLocation"],
+                "Resource": s3_prefixes_enriched,
+                "Effect": "Allow",
+            },
+            {
+                "Action": ["sts:AssumeRole"],
+                "Resource": [f"arn:aws:iam::{account_id}:role/{role_name}"],
+                "Effect": "Allow",
+            },
+        ]
+        if kms_key:
+            statement.append(
+                {
+                    "Action": ["kms:Decrypt", "kms:Encrypt", "kms:GenerateDataKey*"],
+                    "Resource": [f"arn:aws:kms:{kms_key}"],
+                    "Effect": "Allow",
+                }
+            )
+        return self._get_json_for_cli(
+            {
+                "Version": "2012-10-17",
+                "Statement": statement,
+            }
+        )
+
+    def _create_role(self, role_name: str, assume_role_json: str) -> str | None:
+        """
+        Create an AWS role with the given name and assume role policy document.
+        """
+        add_role = self._run_json_command(
+            f"iam create-role --role-name {role_name} --assume-role-policy-document {assume_role_json}"
+        )
+        if not add_role:
+            return None
+        return add_role["Role"]["Arn"]
+
+    def create_uc_role(self, role_name: str) -> str | None:
+        """
+        Create an IAM role for Unity Catalog to access the S3 buckets.
+        the AssumeRole condition will be modified later with the external ID captured from the UC credential.
+        https://docs.databricks.com/en/connect/unity-catalog/storage-credentials.html
+        """
+        return self._create_role(role_name, self._aws_role_trust_doc())
+
+    def update_uc_trust_role(self, role_name: str, external_id: str = "0000") -> str | None:
+        """
+        Modify an existing IAM role for Unity Catalog to access the S3 buckets with the external ID
+        captured from the UC credential.
+        https://docs.databricks.com/en/connect/unity-catalog/storage-credentials.html
+        """
+        update_role = self._run_json_command(
+            f"iam update-assume-role-policy --role-name {role_name} --policy-document {self._aws_role_trust_doc(external_id)}"
+        )
+        if not update_role:
+            return None
+        return update_role["Role"]["Arn"]
+
+    def put_role_policy(
+        self, role_name: str, policy_name: str, s3_prefixes: set[str], account_id: str, kms_key=None
+    ) -> bool:
+        if not self._run_command(
+            f"iam put-role-policy --role-name {role_name} "
+            f"--policy-name {policy_name} "
+            f"--policy-document {self._aws_s3_policy(s3_prefixes, account_id, role_name, kms_key)}"
+        ):
+            return False
+        return True
+
+    def create_migration_role(self, role_name: str) -> str | None:
+        aws_role_trust_doc = {
+            "Version": "2012-10-17",
+            "Statement": [
+                {"Effect": "Allow", "Principal": {"Service": "ec2.amazonaws.com"}, "Action": "sts:AssumeRole"}
+            ],
+        }
+        assume_role_json = self._get_json_for_cli(aws_role_trust_doc)
+        return self._create_role(role_name, assume_role_json)
+
+    def get_instance_profile(self, instance_profile_name: str) -> str | None:
+        instance_profile = self._run_json_command(
+            f"iam get-instance-profile --instance-profile-name {instance_profile_name}"
+        )
+
+        if not instance_profile:
+            return None
+
+        return instance_profile["InstanceProfile"]["Arn"]
+
+    def create_instance_profile(self, instance_profile_name: str) -> str | None:
+        instance_profile = self._run_json_command(
+            f"iam create-instance-profile --instance-profile-name {instance_profile_name}"
+        )
+
+        if not instance_profile:
+            return None
+
+        return instance_profile["InstanceProfile"]["Arn"]
+
+    def delete_instance_profile(self, instance_profile_name: str, role_name: str):
+        self._run_json_command(
+            f"iam remove-role-from-instance-profile --instance-profile-name {instance_profile_name}"
+            f" --role-name {role_name}"
+        )
+        self._run_json_command(f"iam delete-instance-profile --instance-profile-name {instance_profile_name}")
+        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):
+        # there can only be one role associated with iam instance profile
+        self._run_command(
+            f"iam add-role-to-instance-profile --instance-profile-name {instance_profile_name} --role-name {role_name}"
+        )
+
+    def role_exists(self, role_name: str) -> bool:
+        """
+        Check if the given role exists in the AWS account.
+        """
+        result = self._run_json_command("iam list-roles")
+        roles = result.get("Roles", [])
+        for role in roles:
+            if role["RoleName"] == role_name:
+                return True
+        return False
+
+    def _run_json_command(self, command: str):
+        aws_cmd = shutil.which("aws")
+        code, output, error = self._command_runner(f"{aws_cmd} {command} --output json")
+        if code != 0:
+            logger.error(error)
+            return None
+        if output == "":
+            return {}
+        return json.loads(output)
+
+    def _run_command(self, command: str):
+        aws_cmd = shutil.which("aws")
+        code, _, error = self._command_runner(f"{aws_cmd} {command} --output json")
+        if code != 0:
+            logger.error(error)
+            return False
+        return True
+
+    @staticmethod
+    def _get_json_for_cli(input_json: dict) -> str:
+        return json.dumps(input_json).replace('\n', '').replace(" ", "")