Support multiple subscription ids for command line commands (#2647)

## Changes
Support multiple subscription ids for `create-uber-principal` command.
See linked issue for motivation

### Linked issues
Resolves #2645

### Functionality

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

### Tests

- [x] manually tested
- [x] modified unit tests

README.md

@@ -1174,7 +1174,7 @@ Once you're done with the table migration, proceed to the [code migration](#code
 ## `principal-prefix-access` command
 
 ```text
-databricks labs ucx principal-prefix-access [--subscription-id <Azure Subscription ID>] [--aws-profile <AWS CLI profile>]
+databricks labs ucx principal-prefix-access [--subscription-ids <Azure Subscription ID>] [--aws-profile <AWS CLI profile>]
 ```
 
 This command depends on results from the [assessment workflow](#assessment-workflow) and requires [AWS CLI](#access-for-aws-s3-buckets)
@@ -1210,7 +1210,7 @@ Once done, proceed to the [`migrate-credentials` command](#migrate-credentials-c
 ### Access for Azure Storage Accounts
 
 ```commandline
-databricks labs ucx principal-prefix-access --subscription-id test-subscription-id
+databricks labs ucx principal-prefix-access --subscription-ids test-subscription-id
 ```
 
 Use to identify all storage account used by tables, identify the relevant Azure service principals and their permissions
@@ -1253,7 +1253,7 @@ and asks for confirmation from user. Once confirmed, it deletes the role and its
 ## `create-uber-principal` command
 
 ```text
-databricks labs ucx create-uber-principal [--subscription-id X]
+databricks labs ucx create-uber-principal [--subscription-ids X]
 ```
 
 **Requires Cloud IAM admin privileges.**

labs.yml

@@ -135,8 +135,8 @@ commands:
       its access to all the S3 buckets, along with AWS roles that are set with UC access and its access to S3 buckets.
       The output is stored in the workspace install folder.
     flags:
-      - name: subscription-id
-        description: Subscription to scan storage account in
+      - name: subscription-ids
+        description: Comma separated list of subscriptions to scan storage account in.
       - name: aws-profile
         description: AWS Profile to use for authentication
       - name: run-as-collection
@@ -172,8 +172,8 @@ commands:
       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: subscription-ids
+        description: Comma separated list of subscriptions to scan storage account in.
       - name: aws-profile
         description: AWS Profile to use for authentication
       - name: run-as-collection
@@ -192,8 +192,8 @@ commands:
   - name: migrate-credentials
     description: Migrate credentials for storage access to UC storage credential
     flags:
-      - name: subscription-id
-        description: Subscription to scan storage account in
+      - name: subscription-ids
+        description: Comma separated list of subscriptions to scan storage account in.
       - name: aws-profile
         description: AWS Profile to use for authentication
       - name: run-as-collection
@@ -211,8 +211,8 @@ commands:
   - name: migrate-locations
     description: Create UC external locations based on the output of guess_external_locations assessment task.
     flags:
-      - name: subscription-id
-        description: Subscription to scan storage account in
+      - name: subscription-ids
+        description: Comma separated list of subscriptions to scan storage account in.
       - name: aws-profile
         description: AWS Profile to use for authentication
       - name: run-as-collection

src/databricks/labs/ucx/cli.py

@@ -320,7 +320,7 @@ def create_uber_principal(
     """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."""
+    Pass subscription ids for Azure and aws_profile for AWS."""
     if ctx:
         workspace_contexts = [ctx]
     else:
@@ -346,7 +346,7 @@ def principal_prefix_access(
     permission on each storage accounts. For aws, identifies all the Instance Profiles configured in the workspace and
     its access to all the S3 buckets, along with AWS roles that are set with UC access and its access to S3 buckets.
     The output is stored in the workspace install folder.
-    Pass subscription_id for azure and aws_profile for aws."""
+    Pass subscription ids for Azure and aws_profile for AWS."""
     workspace_contexts = _get_workspace_contexts(w, a, run_as_collection, **named_parameters)
     if ctx:
         workspace_contexts = [ctx]

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

@@ -68,18 +68,18 @@ def microsoft_graph_client(self):
         return AzureAPIClient("https://graph.microsoft.com", "https://graph.microsoft.com")
 
     @cached_property
-    def azure_subscription_id(self):
-        subscription_id = self.named_parameters.get("subscription_id")
-        if not subscription_id:
-            raise ValueError("Please enter subscription id to scan storage accounts in.")
-        return subscription_id
+    def azure_subscription_ids(self) -> list[str]:
+        subscription_ids = self.named_parameters.get("subscription_ids", "")
+        if not subscription_ids:
+            raise ValueError("Please enter subscription ids to scan storage accounts in.")
+        return subscription_ids.split(",")
 
     @cached_property
     def azure_resources(self):
         return AzureResources(
             self.azure_management_client,
             self.microsoft_graph_client,
-            [self.azure_subscription_id],
+            self.azure_subscription_ids,
         )
 
     @cached_property

tests/integration/conftest.py

@@ -707,7 +707,7 @@ def azure_cli_authenticated(self):
         return True
 
     @cached_property
-    def azure_subscription_id(self):
+    def azure_subscription_id(self) -> str:
         return self._env_or_skip("TEST_SUBSCRIPTION_ID")
 
 

tests/unit/test_cli.py

@@ -330,7 +330,10 @@ def test_revert_migrated_tables(ws, caplog):
     # test with no schema and no table, user confirm to not retry
     prompts = MockPrompts({'.*': 'no'})
     ctx = WorkspaceContext(ws).replace(
-        is_azure=True, azure_cli_authenticated=True, azure_subscription_id='test', is_gcp=False
+        is_azure=True,
+        azure_cli_authenticated=True,
+        azure_subscription_ids=["test"],
+        is_gcp=False,
     )
     assert revert_migrated_tables(ws, prompts, schema=None, table=None, ctx=ctx) is None
 
@@ -481,7 +484,7 @@ def test_migrate_credentials_azure(ws, acc_client):
     ctx = WorkspaceContext(ws).replace(
         is_azure=True,
         azure_cli_authenticated=True,
-        azure_subscription_id='test',
+        azure_subscription_ids=["test"],
         azure_resources=azure_resources,
     )
     migrate_credentials(ws, prompts, ctx=ctx, a=acc_client)
@@ -528,7 +531,7 @@ def test_migrate_credentials_raises_runtime_warning_when_hitting_storage_credent
     ctx = WorkspaceContext(ws).replace(
         is_azure=True,
         azure_cli_authenticated=True,
-        azure_subscription_id='test',
+        azure_subscription_ids=["test"],
         azure_resources=azure_resources,
         external_locations=external_locations,
     )
@@ -614,7 +617,7 @@ def test_create_azure_uber_principal_raises_value_error_if_subscription_id_is_mi
         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."):
+    with pytest.raises(ValueError, match="Please enter subscription ids to scan storage accounts in."):
         create_uber_principal(ws, prompts, ctx=ctx)
 
 
@@ -623,7 +626,7 @@ def test_create_azure_uber_principal_calls_workspace_id(ws) -> None:
         is_azure=True,
         is_aws=False,
         azure_cli_authenticated=True,
-        azure_subscription_id="id",
+        azure_subscription_ids=["id"],
     )
     prompts = MockPrompts({"Enter a name for the uber service principal to be created": "test"})
 
@@ -643,7 +646,7 @@ def test_create_azure_uber_principal_runs_as_collection_requests_workspace_ids(w
         prompts,
         run_as_collection=True,
         a=acc_client,
-        subscription_id="test",
+        subscription_ids="test",
     )
 
     for workspace_client in workspace_clients:
@@ -695,7 +698,7 @@ def test_migrate_locations_azure(ws) -> None:
         is_azure=True,
         is_aws=False,
         azure_cli_authenticated=True,
-        azure_subscription_id='test',
+        azure_subscription_ids=["test"],
         azure_resources=azurerm,
     )
 
@@ -725,7 +728,7 @@ def test_migrate_locations_azure_run_as_collection(workspace_clients, acc_client
             workspace_clients[0],
             run_as_collection=True,
             a=acc_client,
-            subscription_id="test",
+            subscription_ids=["test"],
         )
 
     for workspace_client in workspace_clients:

tests/unit/test_factories.py

@@ -58,7 +58,7 @@ def test_replace_installation():
     )
     ctx = WorkspaceContext(ws).replace(
         is_azure=True,
-        azure_subscription_id='foo',
+        azure_subscription_ids=["test"],
         installation=mock_installation,
         sql_backend=MockBackend(
             rows={