Add Support for Migrating Table ACL of Interactive clusters using SPN

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

@@ -6,7 +6,7 @@
 from databricks.labs.lsql.backends import SqlBackend
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import NotFound
-from databricks.sdk.service.compute import ClusterSource, Policy
+from databricks.sdk.service.compute import ClusterSource, DataSecurityMode, Policy
 
 from databricks.labs.ucx.assessment.crawlers import azure_sp_conf_present_check, logger
 from databricks.labs.ucx.assessment.jobs import JobsMixin
@@ -30,6 +30,15 @@ class AzureServicePrincipalInfo:
     storage_account: str | None = None
 
 
+@dataclass
+class ServicePrincipalClusterMapping:
+    # this class is created separately as we need cluster to spn mapping
+    # Cluster id where the spn is used
+    cluster_id: str
+    # spn info data class
+    spn_info: set[AzureServicePrincipalInfo]
+
+
 class AzureServicePrincipalCrawler(CrawlerBase[AzureServicePrincipalInfo], JobsMixin, SecretsMixin):
     def __init__(self, ws: WorkspaceClient, sbe: SqlBackend, schema):
         super().__init__(sbe, "hive_metastore", schema, "azure_service_principals", AzureServicePrincipalInfo)
@@ -171,3 +180,16 @@ def _get_azure_spn_from_config(self, config: dict) -> set[AzureServicePrincipalI
                 )
             )
         return set_service_principals
+
+    def get_cluster_to_storage_mapping(self):
+        # 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]
+            ):
+                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

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

@@ -4,12 +4,27 @@
 from dataclasses import dataclass
 from functools import partial
 
+from databricks.labs.blueprint.installation import Installation
 from databricks.labs.blueprint.parallel import ManyError, Threads
-from databricks.sdk.service.catalog import SchemaInfo, TableInfo
-
+from databricks.labs.lsql.backends import SqlBackend, StatementExecutionBackend
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import ResourceDoesNotExist
+from databricks.sdk.service.catalog import ExternalLocationInfo, SchemaInfo, TableInfo
+
+from databricks.labs.ucx.assessment.azure import (
+    AzureServicePrincipalCrawler,
+    AzureServicePrincipalInfo,
+)
+from databricks.labs.ucx.azure.access import (
+    AzureResourcePermissions,
+    StoragePermissionMapping,
+)
+from databricks.labs.ucx.azure.resources import AzureAPIClient, AzureResources
+from databricks.labs.ucx.config import WorkspaceConfig
 from databricks.labs.ucx.framework.crawlers import CrawlerBase
 from databricks.labs.ucx.framework.utils import escape_sql_identifier
-from databricks.labs.ucx.hive_metastore.tables import TablesCrawler
+from databricks.labs.ucx.hive_metastore.locations import ExternalLocations
+from databricks.labs.ucx.hive_metastore.tables import Table, TablesCrawler
 from databricks.labs.ucx.hive_metastore.udfs import UdfsCrawler
 
 logger = logging.getLogger(__name__)
@@ -127,6 +142,7 @@
             ("TABLE", "SELECT"): self._uc_action("SELECT"),
             ("TABLE", "MODIFY"): self._uc_action("MODIFY"),
             ("TABLE", "READ_METADATA"): self._uc_action("BROWSE"),
+            ("TABLE", "ALL PRIVILEGES"): self._uc_action("ALL PRIVILEGES"),
             ("TABLE", "OWN"): self._set_owner_sql,
             ("VIEW", "SELECT"): self._uc_action("SELECT"),
             ("VIEW", "READ_METADATA"): self._uc_action("BROWSE"),
@@ -307,3 +323,193 @@
             # TODO: https://github.com/databrickslabs/ucx/issues/406
             logger.error(f"Couldn't fetch grants for object {on_type} {key}: {e}")
             return []
+
+
+class PrincipalACL:
+    def __init__(
+        self,
+        ws: WorkspaceClient,
+        backend: SqlBackend,
+        installation: Installation,
+        table_crawler: TablesCrawler,
+        spn_crawler: AzureServicePrincipalCrawler | None = None,
+        resource_permission: AzureResourcePermissions | None = None,
+    ):
+        self._backend = backend
+        self._ws = ws
+        self._spn_crawler = spn_crawler
+        self._installation = installation
+        self._resource_permission = resource_permission
+        self._table_crawler = table_crawler
+
+    @classmethod
+    def for_cli(cls, ws: WorkspaceClient, installation: Installation):
+        config = installation.load(WorkspaceConfig)
+        sql_backend = StatementExecutionBackend(ws, config.warehouse_id)
+        locations = ExternalLocations(ws, sql_backend, config.inventory_database)
+        table_crawler = TablesCrawler(sql_backend, config.inventory_database)
+        if ws.config.is_azure:
+            azure_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_client, graph_client)
+            resource_permissions = AzureResourcePermissions(installation, ws, azurerm, locations)
+            spn_crawler = AzureServicePrincipalCrawler(ws, sql_backend, config.inventory_database)
+            return cls(ws, sql_backend, installation, table_crawler, spn_crawler, resource_permissions)
+        if ws.config.is_aws:
+            return None
+        if ws.config.is_gcp:
+            logger.error("UCX is not supported for GCP yet. Please run it on azure or aws")
+            return None
+        return None
+
+    def get_interactive_cluster_grants(self) -> list[Grant]:
+        if self._ws.config.is_azure:
+            return self._get_azure_grants()
+        return []
+
+    def _get_azure_grants(self) -> list[Grant]:
+        assert self._spn_crawler is not None
+        assert self._resource_permission is not None
+        spn_cluster_mapping = self._spn_crawler.get_cluster_to_storage_mapping()
+        if len(spn_cluster_mapping) == 0:
+            # if there are no interactive clusters , then return empty grants
+            logger.info("No interactive cluster found with spn 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
+            msg = (
+                "No external location found, If hive metastore tables are created in external storage, "
+                "ensure migrate_locations cli cmd is run to create the required locations."
+            )
+            logger.error(msg)
+            raise ResourceDoesNotExist(msg) from None
+
+        permission_mappings = self._resource_permission.load()
+        if len(permission_mappings) == 0:
+            # if permission mapping is empty, raise an error to run principal_prefix cmd
+            msg = "No storage permission file found. Please ensure principal_prefix_access cli cmd is run to create the access permission file."
+            logger.error(msg)
+            raise ResourceDoesNotExist(msg) from None
+        tables = self._table_crawler.snapshot()
+        grants: list[Grant] = []
+
+        for cluster_spn in spn_cluster_mapping:
+            principals = self._get_cluster_principal_mapping(cluster_spn.cluster_id)
+            if len(principals) == 0:
+                continue
+            for spn in cluster_spn.spn_info:
+                eligible_locations = self._get_external_location(spn, external_locations, permission_mappings)
+                if len(eligible_locations) == 0:
+                    continue
+                grant = self._get_grants(eligible_locations, principals, tables)
+                grants.extend(grant)
+        catalog_grants = [Grant(principal, "USE", "hive_metastore") for principal in principals]
+        grants.extend(catalog_grants)
+
+        return list(set(grants))
+
+    def _get_aws_grants(self) -> list[Grant]:
+        # TODO
+        return []
+
+    def _get_privilege(self, table: Table, locations: dict[str, str]):
+        if table.view_text is not None:
+            # return nothing for view so that it goes to the separate view logic
+            return None
+        if table.location is None:
+            return "WRITE_FILES"
+        if table.location.startswith('dbfs:/') or table.location.startswith('/dbfs/'):
+            return "WRITE_FILES"
+
+        for loc, privilege in locations.items():
+            if loc is not None and table.location.startswith(loc):
+                return privilege
+        return None
+
+    def _get_database_grants(self, tables: list[Table], principals: list[str]) -> list[Grant]:
+        databases = []
+        for table in tables:
+            if table.database not in databases:
+                databases.append(table.database)
+        return [
+            Grant(principal, "USE", "hive_metastore", database) for database in databases for principal in principals
+        ]
+
+    def _get_grants(
+        self,
+        locations: dict[str, str],
+        principals: list[str],
+        tables: list[Table],
+    ) -> list[Grant]:
+        grants = []
+        filtered_tables = []
+        for table in tables:
+            privilege = self._get_privilege(table, locations)
+            if privilege == "READ_FILES":
+                grants.extend(
+                    [Grant(principal, "SELECT", table.catalog, table.database, table.name) for principal in principals]
+                )
+                filtered_tables.append(table)
+            if privilege == "WRITE_FILES":
+                grants.extend(
+                    [
+                        Grant(principal, "ALL PRIVILEGES", table.catalog, table.database, table.name)
+                        for principal in principals
+                    ]
+                )
+                filtered_tables.append(table)
+            if table.view_text is not None:
+                grants.extend(
+                    [
+                        Grant(principal, "ALL PRIVILEGES", table.catalog, table.database, view=table.name)
+                        for principal in principals
+                    ]
+                )
+                filtered_tables.append(table)
+
+        database_grants = self._get_database_grants(filtered_tables, principals)
+
+        grants.extend(database_grants)
+
+        return grants
+
+    def _get_external_location(
+        self,
+        spn: AzureServicePrincipalInfo,
+        external_locations: list[ExternalLocationInfo],
+        permission_mappings: list[StoragePermissionMapping],
+    ) -> dict[str, str]:
+        matching_location = {}
+        for location in external_locations:
+            if location.url is None:
+                continue
+            for permission_mapping in permission_mappings:
+                if (
+                    location.url.startswith(permission_mapping.prefix)
+                    and permission_mapping.client_id == spn.application_id
+                    and spn.storage_account is not None
+                    and spn.storage_account in permission_mapping.prefix
+                ):
+                    matching_location[location.url] = permission_mapping.privilege
+        return matching_location
+
+    def _get_cluster_principal_mapping(self, cluster_id: str) -> list[str]:
+        # gets all the users,groups,spn which have access to the clusters and returns a dataclass of that mapping
+        principal_list = []
+        cluster_permission = self._ws.permissions.get("clusters", cluster_id)
+        if cluster_permission.access_control_list is None:
+            return []
+        for acl in cluster_permission.access_control_list:
+            if acl.user_name is not None:
+                principal_list.append(acl.user_name)
+            if acl.group_name is not None:
+                if acl.group_name == "admins":
+                    continue
+                principal_list.append(acl.group_name)
+            if acl.service_principal_name is not None:
+                principal_list.append(acl.service_principal_name)
+        return principal_list

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

@@ -15,7 +15,7 @@
 from databricks.labs.ucx.framework.crawlers import CrawlerBase
 from databricks.labs.ucx.framework.utils import escape_sql_identifier
 from databricks.labs.ucx.hive_metastore import GrantsCrawler, TablesCrawler
-from databricks.labs.ucx.hive_metastore.grants import Grant
+from databricks.labs.ucx.hive_metastore.grants import Grant, PrincipalACL
 from databricks.labs.ucx.hive_metastore.mapping import Rule, TableMapping
 from databricks.labs.ucx.hive_metastore.tables import (
     AclMigrationWhat,
@@ -52,6 +52,7 @@ def __init__(
         table_mapping: TableMapping,
         group_manager: GroupManager,
         migration_status_refresher,
+        principal_grants: PrincipalACL,
     ):
         self._tc = table_crawler
         self._gc = grant_crawler
@@ -61,6 +62,7 @@ def __init__(
         self._group = group_manager
         self._migration_status_refresher = migration_status_refresher
         self._seen_tables: dict[str, str] = {}
+        self._principal_grants = principal_grants
 
     @classmethod
     def for_cli(cls, ws: WorkspaceClient, product='ucx'):
@@ -72,9 +74,17 @@ def for_cli(cls, ws: WorkspaceClient, product='ucx'):
         grants_crawler = GrantsCrawler(table_crawler, udfs_crawler)
         table_mapping = TableMapping(installation, ws, sql_backend)
         group_manager = GroupManager(sql_backend, ws, config.inventory_database)
+        principal_grants = PrincipalACL.for_cli(ws, installation)
         migration_status_refresher = MigrationStatusRefresher(ws, sql_backend, config.inventory_database, table_crawler)
         return cls(
-            table_crawler, grants_crawler, ws, sql_backend, table_mapping, group_manager, migration_status_refresher
+            table_crawler,
+            grants_crawler,
+            ws,
+            sql_backend,
+            table_mapping,
+            group_manager,
+            migration_status_refresher,
+            principal_grants,
         )
 
     def index(self):
@@ -87,6 +97,7 @@ def migrate_tables(self, *, what: What | None = None, acl_strategy: list[AclMigr
         if acl_strategy is not None:
             grants_to_migrate = self._gc.snapshot()
             migrated_groups = self._group.snapshot()
+            principal_grants = self._principal_grants.get_interactive_cluster_grants()
         else:
             acl_strategy = []
         for table in tables_to_migrate:
@@ -95,6 +106,8 @@ def migrate_tables(self, *, what: What | None = None, acl_strategy: list[AclMigr
                 continue
             if AclMigrationWhat.LEGACY_TACL in acl_strategy:
                 grants.extend(self._match_grants(table.src, grants_to_migrate, migrated_groups))
+            if AclMigrationWhat.PRINCIPAL in acl_strategy:
+                grants.extend(self._match_grants(table.src, principal_grants, migrated_groups))
             tasks.append(partial(self._migrate_table, table.src, table.rule, grants))
         Threads.strict("migrate tables", tasks)
 

src/databricks/labs/ucx/runtime.py

@@ -19,13 +19,14 @@
     Mounts,
     TablesCrawler,
 )
+from databricks.labs.ucx.hive_metastore.grants import PrincipalACL
 from databricks.labs.ucx.hive_metastore.mapping import TableMapping
 from databricks.labs.ucx.hive_metastore.table_migrate import (
     MigrationStatusRefresher,
     TablesMigrate,
 )
 from databricks.labs.ucx.hive_metastore.table_size import TableSizeCrawler
-from databricks.labs.ucx.hive_metastore.tables import AclMigrationWhat, What
+from databricks.labs.ucx.hive_metastore.tables import What
 from databricks.labs.ucx.hive_metastore.udfs import UdfsCrawler
 from databricks.labs.ucx.hive_metastore.verification import VerifyHasMetastore
 from databricks.labs.ucx.workspace_access.generic import WorkspaceListing
@@ -434,11 +435,19 @@ def migrate_external_tables_sync(
     table_crawler = TablesCrawler(sql_backend, cfg.inventory_database)
     udf_crawler = UdfsCrawler(sql_backend, cfg.inventory_database)
     grant_crawler = GrantsCrawler(table_crawler, udf_crawler)
-    table_mapping = TableMapping(install, ws, sql_backend)
+    table_mappings = TableMapping(install, ws, sql_backend)
     migration_status_refresher = MigrationStatusRefresher(ws, sql_backend, cfg.inventory_database, table_crawler)
     group_manager = GroupManager(sql_backend, ws, cfg.inventory_database)
+    interactive_grants = PrincipalACL.for_cli(ws, install)
     TablesMigrate(
-        table_crawler, grant_crawler, ws, sql_backend, table_mapping, group_manager, migration_status_refresher
+        table_crawler,
+        grant_crawler,
+        ws,
+        sql_backend,
+        table_mappings,
+        group_manager,
+        migration_status_refresher,
+        interactive_grants,
     ).migrate_tables(what=What.EXTERNAL_SYNC)
 
 
@@ -454,11 +463,19 @@ def migrate_dbfs_root_delta_tables(
     table_crawler = TablesCrawler(sql_backend, cfg.inventory_database)
     udf_crawler = UdfsCrawler(sql_backend, cfg.inventory_database)
     grant_crawler = GrantsCrawler(table_crawler, udf_crawler)
-    table_mapping = TableMapping(install, ws, sql_backend)
+    table_mappings = TableMapping(install, ws, sql_backend)
     migration_status_refresher = MigrationStatusRefresher(ws, sql_backend, cfg.inventory_database, table_crawler)
     group_manager = GroupManager(sql_backend, ws, cfg.inventory_database)
+    interactive_grants = PrincipalACL.for_cli(ws, install)
     TablesMigrate(
-        table_crawler, grant_crawler, ws, sql_backend, table_mapping, group_manager, migration_status_refresher
+        table_crawler,
+        grant_crawler,
+        ws,
+        sql_backend,
+        table_mappings,
+        group_manager,
+        migration_status_refresher,
+        interactive_grants,
     ).migrate_tables(what=What.DBFS_ROOT_DELTA)