Added databricks labs ucx create-uber-principal command to create Azure Service Principal for migration

labs.yml

@@ -108,6 +108,13 @@ commands:
       - name: aws-profile
         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.
+    flags:
+      - name: subscription-id
+        description: Subscription to scan storage account in
+
   - name: validate-groups-membership
     description: Validate groups to check if the groups at account level and workspace level have different memberships
     table_template: |-

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

@@ -1,11 +1,20 @@
+import json
+import uuid
 from dataclasses import dataclass
 
 from databricks.labs.blueprint.installation import Installation
+from databricks.labs.blueprint.tui import Prompts
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import NotFound
 from databricks.sdk.service.catalog import Privilege
 
 from databricks.labs.ucx.assessment.crawlers import logger
-from databricks.labs.ucx.azure.resources import AzureResource, AzureResources
+from databricks.labs.ucx.azure.resources import (
+    AzureAPIClient,
+    AzureResource,
+    AzureResources,
+    PrincipalSecret,
+)
 from databricks.labs.ucx.config import WorkspaceConfig
 from databricks.labs.ucx.framework.crawlers import StatementExecutionBackend
 from databricks.labs.ucx.hive_metastore.locations import ExternalLocations
@@ -46,7 +55,12 @@ def for_cli(cls, ws: WorkspaceClient, product='ucx', include_subscriptions=None)
         installation = Installation.current(ws, product)
         config = installation.load(WorkspaceConfig)
         sql_backend = StatementExecutionBackend(ws, config.warehouse_id)
-        azurerm = AzureResources(ws, include_subscriptions=include_subscriptions)
+        azure_mgmt_client = AzureAPIClient(
+            ws.config.arm_environment.resource_manager_endpoint,
+            ws.config.arm_environment.service_management_endpoint,
+        )
+        graph_client = AzureAPIClient("https://graph.microsoft.com", "https://graph.microsoft.com")
+        azurerm = AzureResources(azure_mgmt_client, graph_client, include_subscriptions)
         locations = ExternalLocations(ws, sql_backend, config.inventory_database)
         return cls(installation, ws, azurerm, locations)
 
@@ -91,6 +105,98 @@ def save_spn_permissions(self) -> str | None:
             return None
         return self._installation.save(storage_account_infos, filename=self._filename)
 
+    def _update_cluster_policy_definition(
+        self, policy_definition: str, storage_accounts: list[AzureResource], uber_principal: PrincipalSecret
+    ) -> str:
+        policy_dict = json.loads(policy_definition)
+        tenant_id = self._azurerm.tenant_id()
+        endpoint = f"https://login.microsoftonline.com/{tenant_id}/oauth2/token"
+        for storage in storage_accounts:
+            policy_dict[
+                f"spark_conf.fs.azure.account.oauth2.client.id.{storage.storage_account}.dfs.core.windows.net"
+            ] = self._policy_config(uber_principal.client.client_id)
+            policy_dict[
+                f"spark_conf.fs.azure.account.oauth.provider.type.{storage.storage_account}.dfs.core.windows.net"
+            ] = self._policy_config("org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider")
+            policy_dict[
+                f"spark_conf.fs.azure.account.oauth2.client.endpoint.{storage.storage_account}.dfs.core.windows.net"
+            ] = self._policy_config(endpoint)
+            policy_dict[f"spark_conf.fs.azure.account.auth.type.{storage.storage_account}.dfs.core.windows.net"] = (
+                self._policy_config("OAuth")
+            )
+            if uber_principal.secret is not None:
+                policy_dict[
+                    f"spark_conf.fs.azure.account.oauth2.client.secret.{storage.storage_account}.dfs.core.windows.net"
+                ] = self._policy_config(uber_principal.secret)
+        return json.dumps(policy_dict)
+
+    @staticmethod
+    def _policy_config(value: str):
+        return {"type": "fixed", "value": value}
+
+    def _update_cluster_policy_with_spn(
+        self, policy_id: str, storage_accounts: list[AzureResource], uber_principal: PrincipalSecret
+    ):
+        try:
+            policy_definition = ""
+            cluster_policy = self._ws.cluster_policies.get(policy_id)
+            self._installation.save(cluster_policy, filename="policy-backup.json")
+            if cluster_policy.definition is not None:
+                policy_definition = self._update_cluster_policy_definition(
+                    cluster_policy.definition, storage_accounts, uber_principal
+                )
+            if cluster_policy.name is not None:
+                self._ws.cluster_policies.edit(policy_id, cluster_policy.name, definition=policy_definition)
+        except NotFound:
+            msg = f"cluster policy {policy_id} not found, please run UCX installation to create UCX cluster policy"
+            raise NotFound(msg) from None
+
+    def create_uber_principal(self, prompts: Prompts):
+        config = self._installation.load(WorkspaceConfig)
+        display_name = f"unity-catalog-migration-{config.inventory_database}-{self._ws.get_workspace_id()}"
+        uber_principal_name = prompts.question(
+            "Enter a name for the uber service principal to be created", default=display_name
+        )
+        policy_id = config.policy_id
+        if policy_id is None:
+            msg = "UCX cluster policy not found in config. Please run latest UCX installation to set cluster policy"
+            logger.error(msg)
+            raise ValueError(msg) from None
+        if config.uber_spn_id is not None:
+            logger.error("Uber service principal already created for this workspace.")
+            return
+        used_storage_accounts = self._get_storage_accounts()
+        if len(used_storage_accounts) == 0:
+            logger.warning(
+                "There are no external table present with azure storage account. "
+                "Please check if assessment job is run"
+            )
+            return
+        storage_account_info = []
+        for storage in self._azurerm.storage_accounts():
+            if storage.storage_account in used_storage_accounts:
+                storage_account_info.append(storage)
+        logger.info("Creating service principal")
+        uber_principal = self._azurerm.create_service_principal(uber_principal_name)
+        config = self._installation.load(WorkspaceConfig)
+        config.uber_spn_id = uber_principal.client.client_id
+        self._installation.save(config)
+        logger.info(
+            f"Created service principal of client_id {config.uber_spn_id}. " f"Applying permission on storage accounts"
+        )
+        for storage in storage_account_info:
+            role_name = str(uuid.uuid4())
+            self._azurerm.apply_storage_permission(
+                uber_principal.client.object_id, storage, "STORAGE_BLOB_DATA_READER", role_name
+            )
+            logger.debug(
+                f"Storage Data Blob Reader permission applied for spn {config.uber_spn_id} "
+                f"to storage account {storage.storage_account}"
+            )
+        self._update_cluster_policy_with_spn(policy_id, storage_account_info, uber_principal)
+
+        logger.info(f"Update UCX cluster policy {policy_id} with spn connection details for storage accounts")
+
     def load(self):
         return self._installation.load(list[StoragePermissionMapping], filename=self._filename)
 

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

@@ -18,7 +18,7 @@
     AzureResourcePermissions,
     StoragePermissionMapping,
 )
-from databricks.labs.ucx.azure.resources import AzureResources
+from databricks.labs.ucx.azure.resources import AzureAPIClient, AzureResources
 from databricks.labs.ucx.config import WorkspaceConfig
 from databricks.labs.ucx.framework.crawlers import StatementExecutionBackend
 from databricks.labs.ucx.hive_metastore.locations import ExternalLocations
@@ -171,7 +171,12 @@ def for_cli(cls, ws: WorkspaceClient, installation: Installation, prompts: Promp
 
         config = installation.load(WorkspaceConfig)
         sql_backend = StatementExecutionBackend(ws, config.warehouse_id)
-        azurerm = AzureResources(ws)
+        azure_mgmt_client = AzureAPIClient(
+            ws.config.arm_environment.resource_manager_endpoint,
+            ws.config.arm_environment.service_management_endpoint,
+        )
+        graph_client = AzureAPIClient("https://graph.microsoft.com", "https://graph.microsoft.com")
+        azurerm = AzureResources(azure_mgmt_client, graph_client)
         locations = ExternalLocations(ws, sql_backend, config.inventory_database)
 
         resource_permissions = AzureResourcePermissions(installation, ws, azurerm, locations)