Let create-uber-principal command run on collection of workspaces (#2640)

## Changes
Let `create-uber-principal` command run on a collection of workspaces.

### Linked issues

Resolves #2605

### Functionality

- [x] modified existing command: `databricks labs ucx
create-uber-principal`

### Tests

- [x] manually tested
- [x] added unit tests
- [ ] ~added integration tests~ : Covering after #2507

Author: JCZuurmond

labs.yml

@@ -167,14 +167,17 @@ commands:
         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,
-      it identifies all s3 buckets used by the Instance Profiles configured in the workspace.
+    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 service principal information in the UCX cluster policy.
+      For aws, indentify all s3 buckets used by the Instance Profiles configured in the workspace.
     flags:
       - name: subscription-id
         description: Subscription to scan storage account in
       - name: aws-profile
         description: AWS Profile to use for authentication
+      - name: run-as-collection
+        description: Run the command for the collection of workspaces with ucx installed. Default is False.
 
   - name: validate-groups-membership
     description: Validate groups to check if the groups at account level and workspace level have different memberships

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

@@ -373,7 +373,7 @@ def _remove_service_principal_configuration_from_workspace_warehouse_config(
             )
             raise error
 
-    def create_uber_principal(self, prompts: Prompts):
+    def create_uber_principal(self, prompts: Prompts) -> None:
         config = self._installation.load(WorkspaceConfig)
         inventory_database = config.inventory_database
         display_name = f"unity-catalog-migration-{inventory_database}-{self._ws.get_workspace_id()}"

src/databricks/labs/ucx/cli.py

@@ -313,19 +313,25 @@ def create_uber_principal(
     w: WorkspaceClient,
     prompts: Prompts,
     ctx: WorkspaceContext | None = None,
+    run_as_collection: bool = False,
+    a: AccountClient | None = None,
     **named_parameters,
 ):
     """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,
     it identifies all s3 buckets used by the Instance Profiles configured in the workspace.
     Pass subscription_id for azure and aws_profile for aws."""
-    if not ctx:
-        ctx = WorkspaceContext(w, named_parameters)
-    if ctx.is_azure:
-        return ctx.azure_resource_permissions.create_uber_principal(prompts)
-    if ctx.is_aws:
-        return ctx.aws_resource_permissions.create_uber_principal(prompts)
-    raise ValueError("Unsupported cloud provider")
+    if ctx:
+        workspace_contexts = [ctx]
+    else:
+        workspace_contexts = _get_workspace_contexts(w, a, run_as_collection, **named_parameters)
+    for workspace_context in workspace_contexts:
+        if workspace_context.is_azure:
+            workspace_context.azure_resource_permissions.create_uber_principal(prompts)
+        elif workspace_context.is_aws:
+            workspace_context.aws_resource_permissions.create_uber_principal(prompts)
+        else:
+            raise ValueError("Unsupported cloud provider")
 
 
 @ucx.command

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

@@ -119,8 +119,6 @@ def connect_config(self) -> core.Config:
 
     @cached_property
     def is_azure(self) -> bool:
-        if self.is_aws:
-            return False
         return self.connect_config.is_azure
 
     @cached_property

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

@@ -132,7 +132,7 @@ def aws_cli_run_command(self):
         return run_command
 
     @cached_property
-    def aws_profile(self):
+    def aws_profile(self) -> str:
         aws_profile = self.named_parameters.get("aws_profile")
         if not aws_profile:
             aws_profile = os.getenv("AWS_DEFAULT_PROFILE")

tests/unit/test_cli.py

@@ -87,6 +87,10 @@ def create_workspace_client_mock(workspace_id: int) -> WorkspaceClient:
                     'token': 'bar',
                 },
                 'installed_workspace_ids': installed_workspace_ids,
+                'policy_id': '01234567A8BCDEF9',
+                # Exit Azure's `create_uber_principal` early by setting the uber service principal id
+                # to isolate cli testing as much as possible to the cli commands and not the invoked ucx functionality.
+                'uber_spn_id': '0123456789',
             }
         ),
         '/Users/foo/.ucx/state.json': json.dumps(
@@ -593,30 +597,90 @@ def test_migrate_credentials_limit_aws(ws, acc_client):
         migrate_credentials(ws, prompts, ctx=ctx, a=acc_client)
 
 
-def test_create_master_principal_not_azure(ws):
-    ws.config.is_azure = False
-    ws.config.is_aws = False
+def test_create_uber_principal_raises_value_error_for_unsupported_cloud(ws) -> None:
+    ctx = WorkspaceContext(ws).replace(
+        is_azure=False,
+        is_aws=False,
+    )
     prompts = MockPrompts({})
-    ctx = WorkspaceContext(ws)
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, match="Unsupported cloud provider"):
         create_uber_principal(ws, prompts, ctx=ctx)
 
 
-def test_create_master_principal_no_subscription(ws):
-    ws.config.auth_type = "azure-cli"
-    ws.config.is_azure = True
+def test_create_azure_uber_principal_raises_value_error_if_subscription_id_is_missing(ws) -> None:
+    ctx = WorkspaceContext(ws).replace(
+        is_azure=True,
+        is_aws=False,
+        azure_cli_authenticated=True,
+    )
+    prompts = MockPrompts({"Enter a name for the uber service principal to be created": "test"})
+    with pytest.raises(ValueError, match="Please enter subscription id to scan storage accounts in."):
+        create_uber_principal(ws, prompts, ctx=ctx)
+
+
+def test_create_azure_uber_principal_calls_workspace_id(ws) -> None:
+    ctx = WorkspaceContext(ws).replace(
+        is_azure=True,
+        is_aws=False,
+        azure_cli_authenticated=True,
+        azure_subscription_id="id",
+    )
+    prompts = MockPrompts({"Enter a name for the uber service principal to be created": "test"})
+
+    create_uber_principal(ws, prompts, ctx=ctx)
+
+    ws.get_workspace_id.assert_called_once()
+
+
+def test_create_azure_uber_principal_runs_as_collection_requests_workspace_ids(workspace_clients, acc_client) -> None:
+    for workspace_client in workspace_clients:
+        # Setting the auth as follows as we (currently) do not support injecting multiple workspace contexts
+        workspace_client.config.auth_type = "azure-cli"
+    prompts = MockPrompts({"Enter a name for the uber service principal to be created": "test"})
+
+    create_uber_principal(
+        workspace_clients[0],
+        prompts,
+        run_as_collection=True,
+        a=acc_client,
+        subscription_id="test",
+    )
+
+    for workspace_client in workspace_clients:
+        workspace_client.get_workspace_id.assert_called()
+
+
+def test_create_aws_uber_principal_raises_value_error_if_aws_profile_is_missing(ws) -> None:
+    ctx = WorkspaceContext(ws).replace(
+        is_azure=False,
+        is_aws=True,
+    )
     prompts = MockPrompts({})
-    ctx = WorkspaceContext(ws)
-    with pytest.raises(ValueError):
-        create_uber_principal(ws, prompts, ctx=ctx, subscription_id="")
+    with pytest.raises(ValueError, match="AWS Profile is not specified. .*"):
+        create_uber_principal(ws, prompts, ctx=ctx)
 
 
-def test_create_uber_principal(ws):
-    ws.config.auth_type = "azure-cli"
-    ws.config.is_azure = True
+def successful_aws_cli_call(_):
+    successful_return = """
+    {
+        "UserId": "[email protected]",
+        "Account": "1234",
+        "Arn": "arn:aws:sts::1234:assumed-role/AWSVIEW/[email protected]"
+    }
+    """
+    return 0, successful_return, ""
+
+
+def test_create_aws_uber_principal_calls_dbutils_fs_mounts(ws) -> None:
+    ctx = WorkspaceContext(ws).replace(
+        is_azure=False,
+        is_aws=True,
+        aws_profile="test",
+        aws_cli_run_command=successful_aws_cli_call,
+    )
     prompts = MockPrompts({})
-    with pytest.raises(ValueError):
-        create_uber_principal(ws, prompts, subscription_id="12")
+    create_uber_principal(ws, prompts, ctx=ctx)
+    ws.dbutils.fs.mounts.assert_called_once()
 
 
 def test_migrate_locations_raises_value_error_for_unsupported_cloud_provider(ws) -> None:
@@ -669,22 +733,11 @@ def test_migrate_locations_azure_run_as_collection(workspace_clients, acc_client
 
 
 def test_migrate_locations_aws(ws, caplog) -> None:
-    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,
+        aws_cli_run_command=successful_aws_cli_call,
     )
 
     migrate_locations(ws, ctx=ctx)