DBFS Root Support for HMS Federation (#3425)

Closes #3406 

Add external location for DBFS root location

Author: FastLee

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

@@ -208,7 +208,7 @@ def _get_role_access_task(self, arn: str, role_name: str):
         return policy_actions
 
     def _identify_missing_paths(self):
-        external_locations = self._locations.snapshot()
+        external_locations = self._locations.external_locations_with_root()
         compatible_roles = self.load_uc_compatible_roles()
         missing_paths = set()
         for external_location in external_locations:
@@ -226,13 +226,14 @@ def get_roles_to_migrate(self) -> list[AWSCredentialCandidate]:
         """
         Identify the roles that need to be migrated to UC from the UC compatible roles list.
         """
-        external_locations = self._locations.snapshot()
+        external_locations = self._locations.external_locations_with_root()
+        logger.info(f"Found {len(external_locations)} external locations")
         compatible_roles = self.load_uc_compatible_roles()
         roles: dict[str, AWSCredentialCandidate] = {}
         for external_location in external_locations:
             path = PurePath(external_location.location)
             for role in compatible_roles:
-                if not (path.match(role.resource_path) or path.match(role.resource_path + "/*")):
+                if not (PurePath(role.resource_path) in path.parents or path.match(role.resource_path)):
                     continue
                 if role.role_arn not in roles:
                     roles[role.role_arn] = AWSCredentialCandidate(
@@ -323,7 +324,7 @@ def _create_uber_instance_profile(self, iam_role_name: str, iam_policy_name: str
 
     def create_uber_principal(self, prompts: Prompts):
         config = self._installation.load(WorkspaceConfig)
-        s3_paths = {loc.location for loc in self._locations.snapshot()}
+        s3_paths = {loc.location for loc in self._locations.external_locations_with_root()}
         if len(s3_paths) == 0:
             logger.info("No S3 paths to migrate found")
             return
@@ -374,11 +375,16 @@ def create_uber_principal(self, prompts: Prompts):
             logger.error(f"Failed to assign instance profile to cluster policy {iam_role_name}")
             self._aws_resources.delete_instance_profile(iam_role_name, iam_role_name)
 
+    @classmethod
+    def _clean_external_location_name(cls, location: str) -> str:
+        # Remove leading s3:// s3a:// and trailing /
+        return location.replace("s3://", "").replace("s3a://", "").replace("/", "_").replace(":", "_").replace(".", "_")
+
     def _generate_role_name(self, single_role: bool, role_name: str, location: str) -> str:
         if single_role:
             metastore_id = self._ws.metastores.current().as_dict()["metastore_id"]
             return f"{role_name}_{metastore_id}"
-        return f"{role_name}_{location[5:]}"
+        return f"{role_name}_{self._clean_external_location_name(location)}"
 
     def delete_uc_role(self, role_name: str):
         self._aws_resources.delete_role(role_name)

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

@@ -29,7 +29,7 @@ def __init__(
         self._principal_acl = principal_acl
         # When HMS federation is enabled, the fallback bit is set for all the
         # locations which are created by UCX.
-        self._enable_fallback_mode = enable_hms_federation
+        self._enable_hms_federation = enable_hms_federation
 
     def run(self) -> None:
         """
@@ -38,7 +38,7 @@ def run(self) -> None:
         Create external location for the path using the credential identified
         """
         credential_dict = self._get_existing_credentials_dict()
-        external_locations = self._external_locations.snapshot()
+        external_locations = list(self._external_locations.external_locations_with_root())
         existing_external_locations = self._ws.external_locations.list()
         existing_paths = []
         for external_location in existing_external_locations:
@@ -56,7 +56,7 @@ def run(self) -> None:
                 path,
                 credential_dict[role_arn],
                 skip_validation=True,
-                fallback=self._enable_fallback_mode,
+                fallback=self._enable_hms_federation,
             )
         self._principal_acl.apply_location_acl()
 
@@ -91,7 +91,7 @@ def _identify_missing_external_locations(
                 path = role.resource_path
                 if path.endswith("/*"):
                     path = path[:-2]
-                if new_path.match(path + "/*") or new_path.match(path):
+                if PurePath(path) in new_path.parents or new_path.match(path):
                     matching_role = role.role_arn
                     continue
             if matching_role:

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

@@ -690,7 +690,7 @@ def load(self):
         return self._installation.load(list[StoragePermissionMapping], filename=self.FILENAME)
 
     def _get_storage_accounts(self) -> list[StorageAccount]:
-        external_locations = self._locations.snapshot()
+        external_locations = self._locations.external_locations_with_root()
         used_storage_accounts = []
         for location in external_locations:
             if location.location.startswith("abfss://"):

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

@@ -27,7 +27,7 @@ def __init__(
         self._resource_permissions = resource_permissions
         self._azurerm = azurerm
         self._principal_acl = principal_acl
-        self._enable_fallback_mode = enable_hms_federation
+        self._enable_hms_federation = enable_hms_federation
 
     def _app_id_credential_name_mapping(self) -> tuple[dict[str, str], dict[str, str]]:
         # list all storage credentials.
@@ -128,7 +128,7 @@ def _create_external_location_helper(
                 comment=comment,
                 read_only=read_only,
                 skip_validation=skip_validation,
-                fallback=self._enable_fallback_mode,
+                fallback=self._enable_hms_federation,
             )
             return url
         except InvalidParameterValue as invalid:

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

@@ -397,7 +397,6 @@ def mounts_crawler(self) -> MountsCrawler:
             self.sql_backend,
             self.workspace_client,
             self.inventory_database,
-            self.config.enable_hms_federation,
         )
 
     @cached_property
@@ -416,6 +415,7 @@ def external_locations(self) -> ExternalLocations:
             self.inventory_database,
             self.tables_crawler,
             self.mounts_crawler,
+            self.config.enable_hms_federation,
         )
 
     @cached_property

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

@@ -85,8 +85,8 @@ def _get_authorized_paths(self) -> str:
         current_user = self._workspace_client.current_user.me()
         if not current_user.user_name:
             raise NotFound('Current user not found')
-        for external_location_info in self._external_locations.snapshot():
-            location = external_location_info.location.rstrip('/').replace('s3a://', 's3://')
+        for external_location_info in self._external_locations.external_locations_with_root():
+            location = ExternalLocations.clean_location(external_location_info.location)
             existing_location = existing.get(location)
             if not existing_location:
                 logger.warning(f'External location {location} not found')

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

@@ -69,7 +69,8 @@ def location(self) -> str | None:
     def _parse_location(cls, location: str | None) -> list[str]:
         if not location:
             return []
-        parse_result = cls._parse_url(location.rstrip("/"))
+        location = ExternalLocations.clean_location(location)
+        parse_result = cls._parse_url(location)
         if not parse_result:
             return []
         parts = [parse_result.scheme, parse_result.netloc]
@@ -154,17 +155,66 @@ def __init__(
         schema: str,
         tables_crawler: TablesCrawler,
         mounts_crawler: 'MountsCrawler',
+        enable_hms_federation: bool = False,
     ):
         super().__init__(sql_backend, "hive_metastore", schema, "external_locations", ExternalLocation)
         self._ws = ws
         self._tables_crawler = tables_crawler
         self._mounts_crawler = mounts_crawler
+        self._enable_hms_federation = enable_hms_federation
 
     @cached_property
     def _mounts_snapshot(self) -> list['Mount']:
         """Returns all mounts, sorted by longest prefixes first."""
         return sorted(self._mounts_crawler.snapshot(), key=lambda _: (len(_.name), _.name), reverse=True)
 
+    @staticmethod
+    def clean_location(location: str) -> str:
+        # remove the s3a scheme and replace it with s3 as these can be considered the same and will be treated as such
+        # Having s3a and s3 as separate locations will cause issues when trying to find overlapping locations
+        return re.sub(r"^s3a:/", r"s3:/", location).rstrip("/")
+
+    def external_locations_with_root(self) -> list[ExternalLocation]:
+        """
+        Produces a list of external locations with the DBFS root location appended to the list.
+        Utilizes the snapshot method.
+        Used for HMS Federation.
+
+        Returns:
+                List of ExternalLocation objects
+        """
+
+        external_locations = list(self.snapshot())
+        dbfs_root = self._get_dbfs_root()
+        if dbfs_root:
+            external_locations.append(dbfs_root)
+        return external_locations
+
+    def _get_dbfs_root(self) -> ExternalLocation | None:
+        """
+        Get the root location of the DBFS only if HMS Fed is enabled.
+        Utilizes an undocumented Databricks API call
+
+        Returns:
+            Cloud storage root location for dbfs
+
+        """
+        if not self._enable_hms_federation:
+            return None
+        logger.debug("Retrieving DBFS root location")
+        try:
+            response = self._ws.api_client.do("GET", "/api/2.0/dbfs/resolve-path", query={"path": "dbfs:/"})
+            if isinstance(response, dict):
+                resolved_path = response.get("resolved_path")
+                if resolved_path:
+                    path = f"{self.clean_location(resolved_path)}/user/hive/warehouse"
+                    return ExternalLocation(path, 0)
+        except NotFound:
+            # Couldn't retrieve the DBFS root location
+            logger.warning("DBFS root location not found")
+            return None
+        return None
+
     def _external_locations(self) -> Iterable[ExternalLocation]:
         trie = LocationTrie()
         for table in self._tables_crawler.snapshot():
@@ -356,11 +406,9 @@ def __init__(
         sql_backend: SqlBackend,
         ws: WorkspaceClient,
         inventory_database: str,
-        enable_hms_federation: bool = False,
     ):
         super().__init__(sql_backend, "hive_metastore", inventory_database, "mounts", Mount)
         self._dbutils = ws.dbutils
-        self._enable_hms_federation = enable_hms_federation
 
     @staticmethod
     def _deduplicate_mounts(mounts: list) -> list:
@@ -389,6 +437,7 @@ def _jvm(self):
             return None
 
     def _resolve_dbfs_root(self) -> Mount | None:
+        # TODO: Consider deprecating this method and rely on the new API call
         # pylint: disable=broad-exception-caught,too-many-try-statements
         try:
             jvm = self._jvm
@@ -412,12 +461,6 @@ def _crawl(self) -> Iterable[Mount]:
         try:
             for mount_point, source, _ in self._dbutils.fs.mounts():
                 mounts.append(Mount(mount_point, source))
-            if self._enable_hms_federation:
-                root_mount = self._resolve_dbfs_root()
-                if root_mount:
-                    # filter out DatabricksRoot, otherwise ExternalLocations.resolve_mount() won't work
-                    mounts = list(filter(lambda _: _.source != 'DatabricksRoot', mounts))
-                    mounts.append(root_mount)
         except Exception as error:  # pylint: disable=broad-except
             if "com.databricks.backend.daemon.dbutils.DBUtilsCore.mounts() is not whitelisted" in str(error):
                 logger.warning(

tests/unit/aws/test_access.py

@@ -389,7 +389,7 @@ def test_create_uber_principal_no_storage(mock_ws, mock_installation, locations)
     )
     mock_ws.cluster_policies.get.return_value = cluster_policy
     external_locations = create_autospec(ExternalLocations)
-    external_locations.snapshot.return_value = []
+    external_locations.external_locations_with_root.return_value = []
     prompts = MockPrompts({})
     aws = create_autospec(AWSResources)
     aws_resource_permissions = AWSResourcePermissions(
@@ -459,7 +459,7 @@ def test_create_uc_role_multiple_raises_error(mock_ws, installation_single_role,
 
 def test_create_uc_no_roles(installation_no_roles, mock_ws, caplog):
     external_locations = create_autospec(ExternalLocations)
-    external_locations.snapshot.return_value = []
+    external_locations.external_locations_with_root.return_value = []
     aws = create_autospec(AWSResources)
     aws_resource_permissions = AWSResourcePermissions(
         installation_no_roles,
@@ -867,15 +867,15 @@ def command_call(_: str):
     aws = AWSResources("profile", command_call)
 
     external_locations = create_autospec(ExternalLocations)
-    external_locations.snapshot.return_value = [
+    external_locations.external_locations_with_root.return_value = [
         ExternalLocation("s3://BUCKET1", 1),
         ExternalLocation("s3://BUCKET2/Folder1", 1),
     ]
     resource_permissions = AWSResourcePermissions(installation_multiple_roles, mock_ws, aws, external_locations)
     roles = resource_permissions.get_roles_to_migrate()
     assert len(roles) == 1
     assert len(roles[0].paths) == 2
-    external_locations.snapshot.assert_called_once()
+    external_locations.external_locations_with_root.assert_called_once()
 
 
 def test_delete_uc_roles(mock_ws, installation_multiple_roles, backend, locations):
@@ -939,7 +939,7 @@ def command_call(cmd: str):
 
     aws = AWSResources("profile", command_call)
     external_locations = create_autospec(ExternalLocations)
-    external_locations.snapshot.return_value = []
+    external_locations.external_locations_with_root.return_value = []
     resource_permissions = AWSResourcePermissions(installation_no_roles, mock_ws, aws, external_locations)
     resource_permissions.delete_uc_role("uc_role_1")
     assert '/path/aws iam delete-role --role-name uc_role_1 --profile profile --output json' in command_calls