Added DBFS Root resolution when HMS Federation is enabled (#2947)

This PR adds a DBFS Resolver which is used by HMS federation to resolve
the DBFS root locations.

TODO:

- [x] #2954
- [x] resolve special `/user/hive/metastore` as special case of DBFS
mount
- [x] use `LocationTrie` to come guess locations -
#2965

co-authored by @vihangk-db

Author: nfx

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

@@ -294,7 +294,12 @@ def table_move(self) -> TableMove:
 
     @cached_property
     def mounts_crawler(self) -> MountsCrawler:
-        return MountsCrawler(self.sql_backend, self.workspace_client, self.inventory_database)
+        return MountsCrawler(
+            self.sql_backend,
+            self.workspace_client,
+            self.inventory_database,
+            self.config.enable_hms_federation,
+        )
 
     @cached_property
     def azure_service_principal_crawler(self) -> AzureServicePrincipalCrawler:

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

@@ -341,14 +341,18 @@ class Mount:
     def as_scheme_prefix(self) -> str:
         return f'dbfs:{self.name}'  # dbfs:/mnt/mount-name
 
-    def as_fuse_prefix(self) -> str:
-        return f'/dbfs{self.name}'  # /dbfs/mnt/mount-name
-
 
 class MountsCrawler(CrawlerBase[Mount]):
-    def __init__(self, backend: SqlBackend, ws: WorkspaceClient, inventory_database: str):
+    def __init__(
+        self,
+        backend: SqlBackend,
+        ws: WorkspaceClient,
+        inventory_database: str,
+        enable_hms_federation: bool = False,
+    ):
         super().__init__(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:
@@ -364,11 +368,48 @@ def _deduplicate_mounts(mounts: list) -> list:
                 deduplicated_mounts.append(obj)
         return deduplicated_mounts
 
+    @cached_property
+    def _jvm(self):
+        # pylint: disable=import-error,import-outside-toplevel,broad-exception-caught
+        try:
+            from pyspark.sql.session import SparkSession  # type: ignore[import-not-found]
+
+            spark = SparkSession.builder.getOrCreate()
+            return spark._jvm  # pylint: disable=protected-access
+        except Exception as err:
+            logger.warning(f"Cannot create Py4j proxy: {err}")
+            return None
+
+    def _resolve_dbfs_root(self) -> Mount | None:
+        # pylint: disable=broad-exception-caught,too-many-try-statements
+        try:
+            jvm = self._jvm
+            if not jvm:
+                return None
+            uri = jvm.java.net.URI
+            some = jvm.scala.Some
+            hms_fed_dbfs_utils = jvm.com.databricks.sql.managedcatalog.connections.HmsFedDbfsUtils
+            root_location_opt = hms_fed_dbfs_utils.resolveDbfsPath(some(uri("dbfs:/user/hive/warehouse")))
+            if root_location_opt.isDefined():
+                source: str = root_location_opt.get().toString()
+                source = source.removesuffix('user/hive/warehouse')
+                return Mount("/", source)
+            return None
+        except Exception as err:
+            logger.warning(f"Failed to resolve DBFS root location: {err}")
+            return None
+
     def _crawl(self) -> Iterable[Mount]:
         mounts = []
         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/hive_metastore/test_locations.py

@@ -717,3 +717,20 @@ def my_side_effect(path, **_):
     assert results[0].location == "adls://bucket/table1"
     assert results[1].location == "dbfs:/mnt/test_mount/table2"
     assert results[2].location is None
+
+
+def test_resolve_dbfs_root_in_hms_federation():
+    jvm = Mock()
+    sql_backend = MockBackend()
+    client = create_autospec(WorkspaceClient)
+    client.dbutils.fs.mounts.return_value = [MountInfo('/', 'DatabricksRoot', '')]
+
+    mounts_crawler = MountsCrawler(sql_backend, client, "test", enable_hms_federation=True)
+    mounts_crawler.__dict__['_jvm'] = jvm
+
+    hms_fed_dbfs_utils = jvm.com.databricks.sql.managedcatalog.connections.HmsFedDbfsUtils
+    hms_fed_dbfs_utils.resolveDbfsPath().get().toString.return_value = 's3://original/bucket/user/hive/warehouse'
+
+    mounts = mounts_crawler.snapshot()
+
+    assert [Mount("/", 's3://original/bucket/')] == mounts