Added ability to run create-missing-principals command as collection (#2675)

This PR enables create-missing-principals command to run as a
collection.

Linked issues
#1782

Resolves #2606 

### Functionality

- [ ] modified existing command: `databricks labs ucx ...`

### Tests
<!-- How is this tested? Please see the checklist below and also
describe any other relevant tests -->

- [ ] added unit tests
- [ ] verified on staging environment (screenshot attached)

Author: HariGS-DB

labs.yml

@@ -157,6 +157,8 @@ commands:
         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 S3 locations. (default:False)
+      - name: run-as-collection
+        description: (Optional) boolean flag to indicate to run the cmd as a collection. Default is False.
 
   - name: delete-missing-principals
     description: For AWS, this command identifies all the UC roles that are created through the create-missing-principals

src/databricks/labs/ucx/cli.py

@@ -371,6 +371,8 @@ def create_missing_principals(
     w: WorkspaceClient,
     prompts: Prompts,
     ctx: WorkspaceContext | None = None,
+    run_as_collection: bool = False,
+    a: AccountClient | None = None,
     single_role: bool = False,
     role_name="UC_ROLE",
     policy_name="UC_POLICY",
@@ -380,11 +382,17 @@ def create_missing_principals(
     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  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, role_name=role_name, policy_name=policy_name)
-    raise ValueError("Unsupported cloud provider")
+    workspace_contexts = _get_workspace_contexts(w, a, run_as_collection, **named_parameters)
+    if ctx:
+        workspace_contexts = [ctx]
+    if w.config.is_aws:
+        for workspace_ctx in workspace_contexts:
+            logger.info(f"Running cmd for workspace {workspace_ctx.workspace_client.get_workspace_id()}")
+            workspace_ctx.iam_role_creation.run(
+                prompts, single_role=single_role, role_name=role_name, policy_name=policy_name
+            )
+    else:
+        raise ValueError("Unsupported cloud provider")
 
 
 @ucx.command

tests/unit/test_cli.py

@@ -966,27 +966,30 @@ def test_migrate_tables_calls_external_tables_ctas_job_run_now(ws) -> None:
     ws.jobs.wait_get_run_job_terminated_or_skipped.call_count = 2
 
 
-def test_create_missing_principal_aws(ws):
+def test_create_missing_principal_aws(ws, acc_client):
+    ws.config.is_aws = True
     aws_resource_permissions = create_autospec(AWSResourcePermissions)
     ctx = WorkspaceContext(ws).replace(is_aws=True, is_azure=False, aws_resource_permissions=aws_resource_permissions)
     prompts = MockPrompts({'.*': 'yes'})
-    create_missing_principals(ws, prompts=prompts, ctx=ctx)
+    create_missing_principals(ws, prompts=prompts, ctx=ctx, a=acc_client)
     aws_resource_permissions.create_uc_roles.assert_called_once()
 
 
-def test_create_missing_principal_aws_not_approved(ws):
+def test_create_missing_principal_aws_not_approved(ws, acc_client):
+    ws.config.is_aws = True
     aws_resource_permissions = create_autospec(AWSResourcePermissions)
     ctx = WorkspaceContext(ws).replace(is_aws=True, is_azure=False, aws_resource_permissions=aws_resource_permissions)
     prompts = MockPrompts({'.*': 'No'})
-    create_missing_principals(ws, prompts=prompts, ctx=ctx)
+    create_missing_principals(ws, prompts=prompts, ctx=ctx, a=acc_client)
     aws_resource_permissions.create_uc_roles.assert_not_called()
 
 
-def test_create_missing_principal_azure(ws, caplog):
+def test_create_missing_principal_azure(ws, caplog, acc_client):
+    ws.config.is_aws = False
     ctx = WorkspaceContext(ws).replace(is_aws=False, is_azure=True)
     prompts = MockPrompts({'.*': 'yes'})
     with pytest.raises(ValueError) as failure:
-        create_missing_principals(ws, prompts=prompts, ctx=ctx)
+        create_missing_principals(ws, prompts=prompts, ctx=ctx, a=acc_client)
     assert str(failure.value) == "Unsupported cloud provider"