Add support for migrating Table ACL for SQL Warehouse cluster in AWS using Instance Profile and Azure using SPN (#2258)

This PR contains the following changes for migrating Table ACL for SQL
warehouse cluster used through Instance Profile in AWS scenario and spn
info for Azure

Identifies all the SQL Warehouse
identifies all the instance profiles used by the data access config for
aws and spn info for azure
identify the list of principals who have access to the warehouse (user,
group, spn)
get the prefix these instance role arn or spn have access and the type
of permission
get the list of external locations matching the list of prefix
get the list of tables having the matching external location prefix from
the table mapping CSV

Resolves #2238

Author: HariGS-DB

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

@@ -181,15 +181,27 @@ def _get_azure_spn_from_config(self, config: dict) -> set[AzureServicePrincipalI
             )
         return set_service_principals
 
-    def get_cluster_to_storage_mapping(self):
+    def get_cluster_to_storage_mapping(self) -> list[ServicePrincipalClusterMapping]:
         # this function gives a mapping between an interactive cluster and the spn used by it
         # either directly or through a cluster policy.
         set_service_principals = set[AzureServicePrincipalInfo]()
         spn_cluster_mapping = []
         for cluster in self._ws.clusters.list():
             if cluster.cluster_source != ClusterSource.JOB and (
                 cluster.data_security_mode in [DataSecurityMode.LEGACY_SINGLE_USER, DataSecurityMode.NONE]
+                and cluster.cluster_id is not None
             ):
                 set_service_principals = self._get_azure_spn_from_cluster_config(cluster)
                 spn_cluster_mapping.append(ServicePrincipalClusterMapping(cluster.cluster_id, set_service_principals))
         return spn_cluster_mapping
+
+    def get_warehouse_to_storage_mapping(self) -> list[ServicePrincipalClusterMapping]:
+        # this function gives a mapping between a sql warehouse and the spn used by it
+        spn_warehouse_mapping = []
+        set_service_principals = self._list_all_spn_in_sql_warehouses_spark_conf()
+        if len(set_service_principals) == 0:
+            return []
+        for warehouse in self._ws.warehouses.list():
+            if warehouse.id is not None:
+                spn_warehouse_mapping.append(ServicePrincipalClusterMapping(warehouse.id, set_service_principals))
+        return spn_warehouse_mapping

src/databricks/labs/ucx/hive_metastore/grants.py

@@ -43,9 +43,10 @@
 
 
 @dataclass
-class LocationACL:
-    location_name: str
-    principal: str
+class ComputeLocations:
+    compute_id: str
+    locations: dict
+    compute_type: str
 
 
 @dataclass(frozen=True)
@@ -386,6 +387,7 @@ def _get_cluster_to_instance_profile_mapping(self) -> dict[str, str]:
         # this function gives a mapping between an interactive cluster and the instance profile used by it
         # either directly or through a cluster policy.
         cluster_instance_profiles = {}
+
         for cluster in self._ws.clusters.list():
             if (
                 cluster.cluster_id is None
@@ -403,14 +405,26 @@ def _get_cluster_to_instance_profile_mapping(self) -> dict[str, str]:
 
         return cluster_instance_profiles
 
-    def get_eligible_locations_principals(self) -> dict[str, dict]:
-        cluster_locations = {}
-        eligible_locations = {}
+    def _update_warehouse_to_instance_profile_mapping(
+        self,
+    ) -> dict[str, str]:
+        warehouse_instance_profiles = {}
+        sql_config = self._ws.warehouses.get_workspace_warehouse_config()
+        if sql_config.instance_profile_arn is not None:
+            role_name = sql_config.instance_profile_arn
+            for warehouse in self._ws.warehouses.list():
+                if warehouse.id is not None:
+                    warehouse_instance_profiles[warehouse.id] = role_name
+        return warehouse_instance_profiles
+
+    def get_eligible_locations_principals(self) -> list[ComputeLocations]:
         cluster_instance_profiles = self._get_cluster_to_instance_profile_mapping()
-        if len(cluster_instance_profiles) == 0:
-            # if there are no interactive clusters , then return empty grants
-            logger.info("No interactive cluster found with instance profiles configured")
-            return {}
+        warehouse_instance_profiles = self._update_warehouse_to_instance_profile_mapping()
+        compute_locations = []
+        if len(cluster_instance_profiles) == 0 and len(warehouse_instance_profiles) == 0:
+            # if there are no interactive clusters or warehouse with instance profile , then return empty grants
+            logger.info("No interactive cluster or sql warehouse found with instance profiles configured")
+            return []
         external_locations = list(self._ws.external_locations.list())
         if len(external_locations) == 0:
             # if there are no external locations, then throw an error to run migrate_locations cli command
@@ -434,12 +448,17 @@ def get_eligible_locations_principals(self) -> dict[str, dict]:
             logger.error(msg)
             raise ResourceDoesNotExist(msg) from None
 
-        for cluster_id, role_name in cluster_instance_profiles.items():
-            eligible_locations.update(self._get_external_locations(role_name, external_locations, permission_mappings))
+        for cluster_id, role_compute in cluster_instance_profiles.items():
+            eligible_locations = self._get_external_locations(role_compute, external_locations, permission_mappings)
             if len(eligible_locations) == 0:
                 continue
-            cluster_locations[cluster_id] = eligible_locations
-        return cluster_locations
+            compute_locations.append(ComputeLocations(cluster_id, eligible_locations, "clusters"))
+        for warehouse_id, role_compute in warehouse_instance_profiles.items():
+            eligible_locations = self._get_external_locations(role_compute, external_locations, permission_mappings)
+            if len(eligible_locations) == 0:
+                continue
+            compute_locations.append(ComputeLocations(warehouse_id, eligible_locations, "warehouses"))
+        return compute_locations
 
     @staticmethod
     def _get_external_locations(
@@ -475,14 +494,14 @@ def __init__(
         self._spn_crawler = spn_crawler
         self._installation = installation
 
-    def get_eligible_locations_principals(self) -> dict[str, dict]:
-        cluster_locations = {}
-        eligible_locations = {}
+    def get_eligible_locations_principals(self) -> list[ComputeLocations]:
+        compute_locations = []
         spn_cluster_mapping = self._spn_crawler.get_cluster_to_storage_mapping()
-        if len(spn_cluster_mapping) == 0:
+        spn_warehouse_mapping = self._spn_crawler.get_warehouse_to_storage_mapping()
+        if len(spn_cluster_mapping) == 0 and len(spn_warehouse_mapping) == 0:
             # if there are no interactive clusters , then return empty grants
             logger.info("No interactive cluster found with spn configured")
-            return {}
+            return []
         external_locations = list(self._ws.external_locations.list())
         if len(external_locations) == 0:
             # if there are no external locations, then throw an error to run migrate_locations cli command
@@ -507,10 +526,18 @@ def get_eligible_locations_principals(self) -> dict[str, dict]:
             raise ResourceDoesNotExist(msg) from None
 
         for cluster_spn in spn_cluster_mapping:
+            eligible_locations = {}
             for spn in cluster_spn.spn_info:
                 eligible_locations.update(self._get_external_locations(spn, external_locations, permission_mappings))
-            cluster_locations[cluster_spn.cluster_id] = eligible_locations
-        return cluster_locations
+            compute_locations.append(ComputeLocations(cluster_spn.cluster_id, eligible_locations, "clusters"))
+
+        for warehouse_spn in spn_warehouse_mapping:
+            eligible_locations = {}
+            for spn in warehouse_spn.spn_info:
+                eligible_locations.update(self._get_external_locations(spn, external_locations, permission_mappings))
+            compute_locations.append(ComputeLocations(warehouse_spn.cluster_id, eligible_locations, "warehouses"))
+
+        return compute_locations
 
     def _get_external_locations(
         self,
@@ -543,25 +570,25 @@ def __init__(
         installation: Installation,
         tables_crawler: TablesCrawler,
         mounts_crawler: Mounts,
-        cluster_locations: dict[str, dict],
+        cluster_locations: list[ComputeLocations],
     ):
         self._backend = backend
         self._ws = ws
         self._installation = installation
         self._tables_crawler = tables_crawler
         self._mounts_crawler = mounts_crawler
-        self._cluster_locations = cluster_locations
+        self._compute_locations = cluster_locations
 
     def get_interactive_cluster_grants(self) -> list[Grant]:
         tables = self._tables_crawler.snapshot()
         mounts = list(self._mounts_crawler.snapshot())
         grants: set[Grant] = set()
 
-        for cluster_id, locations in self._cluster_locations.items():
-            principals = self._get_cluster_principal_mapping(cluster_id)
+        for compute_location in self._compute_locations:
+            principals = self._get_cluster_principal_mapping(compute_location.compute_id, compute_location.compute_type)
             if len(principals) == 0:
                 continue
-            cluster_usage = self._get_grants(locations, principals, tables, mounts)
+            cluster_usage = self._get_grants(compute_location.locations, principals, tables, mounts)
             grants.update(cluster_usage)
         return list(grants)
 
@@ -628,11 +655,11 @@ def _get_grants(
 
         return grants
 
-    def _get_cluster_principal_mapping(self, cluster_id: str) -> list[str]:
+    def _get_cluster_principal_mapping(self, cluster_id: str, object_type: str) -> list[str]:
         # gets all the users,groups,spn which have access to the clusters and returns a dataclass of that mapping
         principal_list = []
         try:
-            cluster_permission = self._ws.permissions.get("clusters", cluster_id)
+            cluster_permission = self._ws.permissions.get(object_type, cluster_id)
         except ResourceDoesNotExist:
             return []
         if cluster_permission.access_control_list is None:
@@ -661,12 +688,12 @@ def apply_location_acl(self):
             "CREATE EXTERNAL VOLUME and READ_FILES for existing eligible interactive cluster users"
         )
         # get the eligible location mapped for each interactive cluster
-        for cluster_id, locations in self._cluster_locations.items():
+        for compute_location in self._compute_locations:
             # get interactive cluster users
-            principals = self._get_cluster_principal_mapping(cluster_id)
+            principals = self._get_cluster_principal_mapping(compute_location.compute_id, compute_location.compute_type)
             if len(principals) == 0:
                 continue
-            for location_url in locations.keys():
+            for location_url in compute_location.locations.keys():
                 # get the location name for the given url
                 location_name = self._get_location_name(location_url)
                 if location_name is None:

tests/integration/hive_metastore/test_migrate.py

@@ -7,7 +7,6 @@
 from databricks.sdk.service.compute import DataSecurityMode, AwsAttributes
 from databricks.sdk.service.catalog import Privilege, SecurableType, TableInfo, TableType
 from databricks.sdk.service.iam import PermissionLevel
-
 from databricks.labs.ucx.config import WorkspaceConfig
 from databricks.labs.ucx.hive_metastore.mapping import Rule, TableMapping
 from databricks.labs.ucx.hive_metastore.tables import AclMigrationWhat, Table, What
@@ -564,6 +563,33 @@ def test_migrate_external_tables_with_principal_acl_aws(
     assert match
 
 
+@retried(on=[NotFound], timeout=timedelta(minutes=3))
+def test_migrate_external_tables_with_principal_acl_aws_warehouse(
+    ws, make_user, prepared_principal_acl, make_warehouse_permissions, make_warehouse, env_or_skip
+):
+    if not ws.config.is_aws:
+        pytest.skip("temporary: only works in aws test env")
+    ctx, table_full_name, _, _ = prepared_principal_acl
+    ctx.with_dummy_resource_permission()
+    warehouse = make_warehouse()
+    table_migrate = ctx.tables_migrator
+    user = make_user()
+    make_warehouse_permissions(
+        object_id=warehouse.id,
+        permission_level=PermissionLevel.CAN_USE,
+        user_name=user.user_name,
+    )
+    table_migrate.migrate_tables(what=What.EXTERNAL_SYNC, acl_strategy=[AclMigrationWhat.PRINCIPAL])
+
+    target_table_grants = ws.grants.get(SecurableType.TABLE, table_full_name)
+    match = False
+    for _ in target_table_grants.privilege_assignments:
+        if _.principal == user.user_name and _.privileges == [Privilege.ALL_PRIVILEGES]:
+            match = True
+            break
+    assert match
+
+
 def test_migrate_table_in_mount(
     ws,
     sql_backend,

tests/unit/assessment/test_azure.py

@@ -7,7 +7,7 @@
     ClusterSource,
     DataSecurityMode,
 )
-
+from databricks.sdk.service.sql import EndpointInfo
 from databricks.labs.ucx.assessment.azure import (
     AzureServicePrincipalCrawler,
     AzureServicePrincipalInfo,
@@ -219,6 +219,13 @@ def test_get_cluster_to_storage_mapping_no_cluster_return_empty():
     assert not crawler.get_cluster_to_storage_mapping()
 
 
+def test_get_warehouse_to_storage_mapping_no_warehouse_return_empty():
+    ws = create_autospec(WorkspaceClient)
+    ws.warehouses.list.return_value = []
+    crawler = AzureServicePrincipalCrawler(ws, MockBackend(), "ucx")
+    assert not crawler.get_warehouse_to_storage_mapping()
+
+
 def test_get_cluster_to_storage_mapping_no_interactive_cluster_return_empty():
     ws = mock_workspace_client(cluster_ids=['azure-spn-secret'])
     ws.clusters.list.return_value = [
@@ -229,6 +236,15 @@ def test_get_cluster_to_storage_mapping_no_interactive_cluster_return_empty():
     assert not crawler.get_cluster_to_storage_mapping()
 
 
+def test_get_warehouse_to_storage_mapping_no_spn_info_return_empty():
+    ws = mock_workspace_client()
+    ws.warehouses.list.return_value = [
+        EndpointInfo(id="123", name="warehouse1"),
+    ]
+    crawler = AzureServicePrincipalCrawler(ws, MockBackend(), "ucx")
+    assert not crawler.get_warehouse_to_storage_mapping()
+
+
 def test_get_cluster_to_storage_mapping_interactive_cluster_no_spn_return_empty():
     ws = mock_workspace_client(cluster_ids=['azure-spn-secret-interactive-multiple-spn'])
 
@@ -254,3 +270,29 @@ def test_get_cluster_to_storage_mapping_interactive_cluster_no_spn_return_empty(
     assert cluster_spn_info[0].cluster_id == "azure-spn-secret-interactive"
     assert len(cluster_spn_info[0].spn_info) == 2
     assert cluster_spn_info[0].spn_info == spn_info
+
+
+def test_get_warehouse_to_storage_mapping_spn():
+    ws = mock_workspace_client(warehouse_config="spn-secret-config")
+    ws.warehouses.list.return_value = [
+        EndpointInfo(id="azure-spn-secret-warehouse", name="warehouse1"),
+    ]
+
+    crawler = AzureServicePrincipalCrawler(ws, MockBackend(), "ucx")
+    cluster_spn_info = crawler.get_warehouse_to_storage_mapping()
+    spn_info = {
+        AzureServicePrincipalInfo(
+            application_id='Hello, World!',
+            tenant_id='dummy_tenant_id2',
+            storage_account='xyz',
+        ),
+        AzureServicePrincipalInfo(
+            application_id='dummy_application_id',
+            tenant_id='dummy_tenant_id',
+            storage_account='abcde',
+        ),
+    }
+
+    assert cluster_spn_info[0].cluster_id == "azure-spn-secret-warehouse"
+    assert len(cluster_spn_info[0].spn_info) == 2
+    assert cluster_spn_info[0].spn_info == spn_info