Automatically retrieved aws_account_id from aws profile instead of prompting

labs.yml

@@ -123,13 +123,19 @@ commands:
 
   - name: create-missing-principals
     description: For AWS, this command identifies all the S3 locations that are missing a UC compatible role and 
-      creates them. It takes single-role optional parameter. 
-        If set to True, it will create a single role for all the S3 locations.
+      creates them. It accepts a number of optional parameters, i.e. KMS Key, Role Name, Policy Name, and whether to 
+      create a single role for all the S3 locations.
     flags:
       - name: aws-profile
         description: AWS Profile to use for authentication
+      - name: kms-key
+        description: (Optional) KMS Key to be specified for the UC roles.
+      - name: role-name
+        description: (Optional) IAM Role name to be specified for the UC roles. (default:UC_ROLE)
+      - name: policy-name
+        description: (Optional) IAM policy Name to be specified for the UC roles. (default:UC_POLICY)
       - name: single-role
-        description: (Optional) Create a single role for all the S3 locations (default:True)
+        description: (Optional) Create a single role for all S3 locations. (default:False)
 
   - name: create-uber-principal
     description: For azure cloud, creates a service principal and gives STORAGE BLOB READER access on all the storage account 

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

@@ -330,7 +330,12 @@ def update_uc_trust_role(self, role_name: str, external_id: str = "0000") -> str
         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
+        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} --policy-name {policy_name} "

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

@@ -33,14 +33,13 @@ def __init__(
         ws: WorkspaceClient,
         aws_resources: AWSResources,
         external_locations: ExternalLocations,
-        aws_account_id=None,
         kms_key=None,
     ):
         self._installation = installation
         self._aws_resources = aws_resources
         self._ws = ws
         self._locations = external_locations
-        self._aws_account_id = aws_account_id
+        self._aws_account_id = aws_resources.validate_connection().get("Account")
         self._kms_key = kms_key
 
     def list_uc_roles(self, *, single_role=True, role_name="UC_ROLE", policy_name="UC_POLICY"):

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

@@ -205,10 +205,12 @@ def _print_action_plan(iam_list: list[AWSUCRoleCandidate]):
     def save(self, migration_results: list[CredentialValidationResult]) -> str:
         return self._installation.save(migration_results, filename=self._output_file)
 
-    def run(self, prompts: Prompts, *, single_role=True, role_name="UC_ROLE", policy_name="UC_POLICY"):
+    def run(self, prompts: Prompts, *, single_role=False, role_name="UC_ROLE", policy_name="UC_POLICY"):
 
         iam_list = self._resource_permissions.list_uc_roles(
-            single_role=single_role, role_name=role_name, policy_name=policy_name
+            single_role=single_role,
+            role_name=role_name,
+            policy_name=policy_name,
         )
         if not iam_list:
             logger.info("No IAM Role created")

src/databricks/labs/ucx/cli.py

@@ -293,17 +293,19 @@ def create_missing_principals(
     w: WorkspaceClient,
     prompts: Prompts,
     ctx: WorkspaceContext | None = None,
-    single_role: bool = True,
+    single_role: bool = False,
+    role_name="UC_ROLE",
+    policy_name="UC_POLICY",
     **named_parameters,
 ):
     """Not supported for Azure.
     For AWS, this command identifies all the S3 locations that are missing a UC compatible role and creates them.
-    By default, it will create a single role for all S3. Set the optional single_role parameter to False, to create one role per S3 location.
+    By default, it will create a  role per S3 location. Set the optional single_role parameter to True to create a single role for all S3 locations.
     """
     if not ctx:
         ctx = WorkspaceContext(w, named_parameters)
     if ctx.is_aws:
-        return ctx.iam_role_creation.run(prompts, single_role=single_role)
+        return ctx.iam_role_creation.run(prompts, single_role=single_role, role_name=role_name, policy_name=policy_name)
     raise ValueError("Unsupported cloud provider")
 
 

src/databricks/labs/ucx/contexts/workspace_cli.py

@@ -151,7 +151,6 @@ def aws_resource_permissions(self):
             self.workspace_client,
             self.aws_resources,
             self.external_locations,
-            self.named_parameters.get("aws_account_id"),
             self.named_parameters.get("kms_key"),
         )
 

tests/integration/aws/test_access.py

@@ -76,9 +76,10 @@ def test_create_uber_instance_profile(
 ):
     profile = env_or_skip("AWS_DEFAULT_PROFILE")
     aws = AWSResources(profile)
-    account_id = aws.validate_connection().get("Account")
     sql_backend.save_table(
-        f"{inventory_schema}.external_locations", [ExternalLocation("s3://bucket1/FOLDER1", 1)], ExternalLocation
+        f"{inventory_schema}.external_locations",
+        [ExternalLocation("s3://bucket1/FOLDER1", 1)],
+        ExternalLocation,
     )
     installation = Installation(ws, make_random(4))
     # create a new policy
@@ -91,7 +92,6 @@ def test_create_uber_instance_profile(
         aws,
         ExternalLocations(ws, sql_backend, inventory_schema),
         aws_cli_ctx.principal_acl,
-        account_id,
     )
     aws_permissions.create_uber_principal(
         MockPrompts(

tests/unit/aws/test_access.py

@@ -199,6 +199,7 @@ def test_create_uber_principal_existing_role_in_policy(mock_ws, mock_installatio
     )
     mock_ws.cluster_policies.get.return_value = cluster_policy
     aws = create_autospec(AWSResources)
+    aws.validate_connection.return_value = {}
     aws.get_instance_profile_arn.return_value = instance_profile_arn
     locations = ExternalLocations(mock_ws, backend, "ucx")
     prompts = MockPrompts({"We have identified existing UCX migration role *": "yes"})
@@ -310,10 +311,11 @@ def test_create_uber_principal_no_storage(mock_ws, mock_installation, locations)
 
 def test_create_uc_role_single(mock_ws, installation_single_role, backend, locations):
     aws = create_autospec(AWSResources)
+    aws.validate_connection.return_value = {}
     aws_resource_permissions = AWSResourcePermissions(installation_single_role, mock_ws, aws, locations)
     role_creation = IamRoleCreation(installation_single_role, mock_ws, aws_resource_permissions)
     aws.list_all_uc_roles.return_value = []
-    role_creation.run(MockPrompts({"Above *": "yes"}))
+    role_creation.run(MockPrompts({"Above *": "yes"}), single_role=True)
     assert aws.create_uc_role.assert_called
     assert (
         call('UC_ROLE', 'UC_POLICY', {'s3://BUCKET1/FOLDER1', 's3://BUCKET2/FOLDER2'}, None, None)
@@ -323,6 +325,7 @@ def test_create_uc_role_single(mock_ws, installation_single_role, backend, locat
 
 def test_create_uc_role_multiple(mock_ws, installation_single_role, backend, locations):
     aws = create_autospec(AWSResources)
+    aws.validate_connection.return_value = {}
     aws_resource_permissions = AWSResourcePermissions(installation_single_role, mock_ws, aws, locations)
     role_creation = IamRoleCreation(installation_single_role, mock_ws, aws_resource_permissions)
     aws.list_all_uc_roles.return_value = []

tests/unit/test_cli.py

@@ -388,7 +388,23 @@ def test_migrate_locations_azure(ws):
 
 
 def test_migrate_locations_aws(ws, caplog):
-    ctx = WorkspaceContext(ws).replace(is_aws=True, is_azure=False, aws_profile="profile")
+    successful_return = """
+    {
+        "UserId": "[email protected]",
+        "Account": "1234",
+        "Arn": "arn:aws:sts::1234:assumed-role/AWSVIEW/[email protected]"
+    }
+    """
+
+    def successful_call(_):
+        return 0, successful_return, ""
+
+    ctx = WorkspaceContext(ws).replace(
+        is_aws=True,
+        is_azure=False,
+        aws_profile="profile",
+        aws_cli_run_command=successful_call,
+    )
     migrate_locations(ws, ctx=ctx)
     ws.external_locations.list.assert_called()
 
@@ -542,3 +558,10 @@ def test_migrate_dbsql_dashboards(ws, caplog):
 def test_revert_dbsql_dashboards(ws, caplog):
     revert_dbsql_dashboards(ws)
     ws.dashboards.list.assert_called_once()
+
+
+def test_cli_missing_awscli(ws, mocker, caplog):
+    mocker.patch("shutil.which", side_effect=ValueError("Couldn't find AWS CLI in path"))
+    with pytest.raises(ValueError):
+        ctx = WorkspaceContext(ws).replace(is_aws=True, is_azure=False, aws_profile="profile")
+        migrate_locations(ws, ctx)