Let migrate-locations command run as collection (#2652)

## Changes
Let `validate-external-locations` command to run as collection

### Linked issues

Resolves #2608

### Functionality

- [x] modified existing command: `databricks labs ucx migrate-locations`

### Tests

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

Author: JCZuurmond

labs.yml

@@ -208,6 +208,8 @@ commands:
         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: create-catalogs-schemas
     description: Create UC external catalogs and schemas based on the destinations created from create_table_mapping command.

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

@@ -27,7 +27,7 @@ def __init__(
         self._aws_resource_permissions = aws_resource_permissions
         self._principal_acl = principal_acl
 
-    def run(self):
+    def run(self) -> None:
         """
         For each path find out the role that has access to it
         Find out the credential that is pointing to this path
@@ -36,7 +36,10 @@ def run(self):
         credential_dict = self._get_existing_credentials_dict()
         external_locations = self._external_locations.snapshot()
         existing_external_locations = self._ws.external_locations.list()
-        existing_paths = [external_location.url for external_location in existing_external_locations]
+        existing_paths = []
+        for external_location in existing_external_locations:
+            if external_location.url is not None:
+                existing_paths.append(external_location.url)
         compatible_roles = self._aws_resource_permissions.load_uc_compatible_roles()
         missing_paths = self._identify_missing_external_locations(external_locations, existing_paths, compatible_roles)
         for path, role_arn in missing_paths:

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

@@ -182,7 +182,7 @@ def _filter_unsupported_location(self, location_urls: list[str]) -> list[str]:
             logger.warning(f"Skip unsupported location: {url}")
         return supported_urls
 
-    def run(self):
+    def run(self) -> list[str]:
         # list missing external locations in UC
         _, missing_locations = self._hms_locations.match_table_external_locations()
         # Extract the location URLs from the missing locations

src/databricks/labs/ucx/cli.py

@@ -428,17 +428,27 @@ def migrate_credentials(
 
 
 @ucx.command
-def migrate_locations(w: WorkspaceClient, ctx: WorkspaceContext | None = None, **named_parameters):
+def migrate_locations(
+    w: WorkspaceClient,
+    ctx: WorkspaceContext | None = None,
+    run_as_collection: bool = False,
+    a: AccountClient | None = None,
+    **named_parameters,
+):
     """This command creates UC external locations. The candidate locations to be created are extracted from
     guess_external_locations task in the assessment job. You can run validate_external_locations command to check
     the candidate locations. Please make sure the credentials haven migrated before running this command. The command
     will only create the locations that have corresponded UC Storage Credentials.
     """
-    if not ctx:
-        ctx = WorkspaceContext(w, named_parameters)
-    if ctx.is_azure or ctx.is_aws:
-        return ctx.external_locations_migration.run()
-    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 or workspace_context.is_aws:
+            workspace_context.external_locations_migration.run()
+        else:
+            raise ValueError("Unsupported cloud provider")
 
 
 @ucx.command

tests/unit/test_cli.py

@@ -584,7 +584,13 @@ def test_create_uber_principal(ws):
         create_uber_principal(ws, prompts, subscription_id="12")
 
 
-def test_migrate_locations_azure(ws):
+def test_migrate_locations_raises_value_error_for_unsupported_cloud_provider(ws) -> None:
+    ctx = WorkspaceContext(ws).replace(is_azure=False, is_aws=False)
+    with pytest.raises(ValueError, match="Unsupported cloud provider"):
+        migrate_locations(ws, ctx=ctx)
+
+
+def test_migrate_locations_azure(ws) -> None:
     azurerm = create_autospec(AzureResources)
     ctx = WorkspaceContext(ws).replace(
         is_azure=True,
@@ -593,12 +599,41 @@ def test_migrate_locations_azure(ws):
         azure_subscription_id='test',
         azure_resources=azurerm,
     )
+
     migrate_locations(ws, ctx=ctx)
+
     ws.external_locations.list.assert_called()
     azurerm.storage_accounts.assert_called()
 
 
-def test_migrate_locations_aws(ws, caplog):
+@pytest.mark.xfail(reason="Currently not supported in unit tests see TODO")
+def test_migrate_locations_azure_run_as_collection(workspace_clients, acc_client) -> None:
+    """Test migrate locations for a collection of workspaces
+
+    The "run as collection" test should run the same as the test `test_migrate_locations_azure`, but the assert should
+    be called on **all** workspace clients.
+    """
+    for workspace_client in workspace_clients:
+        # Setting the auth as follows as we (currently) do not support injecting multiple workspace contexts
+        workspace_client.config.is_aws = False
+        workspace_client.config.auth_type = "azure-cli"
+
+    # TODO: Migrate locations fails in unit testing as we currently not support injecting multiple workspace contexts
+    # thus we cannot mock azure-cli login for multiple workspaces.
+    # If we support this in the future, the `pytest.raises` can be removed and the test should pass
+    with pytest.raises(OSError, match=".*The provided request must include a 'scope' input parameter.*"):
+        migrate_locations(
+            workspace_clients[0],
+            run_as_collection=True,
+            a=acc_client,
+            subscription_id="test",
+        )
+
+    for workspace_client in workspace_clients:
+        workspace_client.external_locations.list.assert_called()
+
+
+def test_migrate_locations_aws(ws, caplog) -> None:
     successful_return = """
     {
         "UserId": "[email protected]",
@@ -616,10 +651,38 @@ def successful_call(_):
         aws_profile="profile",
         aws_cli_run_command=successful_call,
     )
+
     migrate_locations(ws, ctx=ctx)
+
     ws.external_locations.list.assert_called()
 
 
+@pytest.mark.xfail(reason="Currently not supported in unit tests see TODO")
+def test_migrate_locations_run_as_collection(workspace_clients, acc_client) -> None:
+    """Test migrate locations for a collection of workspaces
+
+    The "run as collection" test should run the same as the test `test_migrate_locations_aws`, but the assert should
+    be called on **all** workspace clients.
+    """
+    for workspace_client in workspace_clients:
+        # Setting the auth as follows as we (currently) do not support injecting multiple workspace contexts
+        workspace_client.config.is_azure = False
+
+    # TODO: Migrate locations fails in unit testing as we currently not support injecting multiple workspace contexts
+    # thus we cannot mock AWS login for multiple workspaces.
+    # If we support this in the future, the `pytest.raises` can be removed and the test should pass
+    with pytest.raises(AttributeError, match="'NoneType' object has no attribute 'get'"):
+        migrate_locations(
+            workspace_clients[0],
+            run_as_collection=True,
+            a=acc_client,
+            aws_profile="profile",
+        )
+
+    for workspace_client in workspace_clients:
+        workspace_client.external_locations.list.assert_called()
+
+
 def test_migrate_locations_gcp(ws):
     ctx = WorkspaceContext(ws).replace(is_aws=False, is_azure=False)
     with pytest.raises(ValueError):