Use LocationTrie to infer a list of UC external locations (#2965)

nfx · web-flow · commit c385d773a46a · 2024-10-16T15:35:47.000+02:00
This PR removes duplicate code for determining overlapping storage
prefixes. this also lays down some groundwork for federated SQL
connections
diff --git a/src/databricks/labs/ucx/hive_metastore/locations.py b/src/databricks/labs/ucx/hive_metastore/locations.py
@@ -6,7 +6,7 @@
 from dataclasses import dataclass
 from functools import cached_property
 from typing import ClassVar, Optional
-from urllib.parse import urlparse
+from urllib.parse import urlparse, ParseResult
 
 from databricks.labs.blueprint.installation import Installation
 from databricks.labs.lsql.backends import SqlBackend
@@ -42,9 +42,9 @@ class LocationTrie:
     """
 
     key: str = ""
-    parent: OptionalLocationTrie = None
+    parent: OptionalLocationTrie = dataclasses.field(repr=False, default=None)
     children: dict[str, "LocationTrie"] = dataclasses.field(default_factory=dict)
-    tables: list[Table] = dataclasses.field(default_factory=list)
+    tables: list[Table] = dataclasses.field(repr=False, default_factory=list)
 
     @cached_property
     def _path(self) -> list[str]:
@@ -57,19 +57,47 @@ def _path(self) -> list[str]:
         return list(reversed(parts))[1:]
 
     @property
-    def location(self):
-        scheme, netloc, *path = self._path
-        return f"{scheme}://{netloc}/{'/'.join(path)}"
+    def location(self) -> str | None:
+        if not self.is_valid():
+            return None
+        try:
+            scheme, netloc, *path = self._path
+            return f"{scheme}://{netloc}/{'/'.join(path)}".rstrip("/")
+        except ValueError:
+            return None
 
-    @staticmethod
-    def _parse_location(location: str | None) -> list[str]:
+    @classmethod
+    def _parse_location(cls, location: str | None) -> list[str]:
         if not location:
             return []
-        parse_result = urlparse(location)
+        parse_result = cls._parse_url(location.rstrip("/"))
+        if not parse_result:
+            return []
         parts = [parse_result.scheme, parse_result.netloc]
-        parts.extend(parse_result.path.strip("/").split("/"))
+        for part in parse_result.path.split("/"):
+            if not part:
+                continue  # remove empty strings
+            parts.append(part)
         return parts
 
+    @staticmethod
+    def _parse_url(location: str) -> ParseResult | None:
+        parse_result = urlparse(location)
+        if parse_result.scheme == 'jdbc':
+            jdbc_path = parse_result.path.split('://')
+            if len(jdbc_path) != 2:
+                return None
+            netloc, path = jdbc_path[1].split('/', 1)
+            parse_result = ParseResult(
+                scheme=f'{parse_result.scheme}:{jdbc_path[0]}',
+                netloc=netloc,
+                path=path,
+                params='',
+                query='',
+                fragment='',
+            )
+        return parse_result
+
     def insert(self, table: Table) -> None:
         current = self
         for part in self._parse_location(table.location):
@@ -91,11 +119,22 @@ def find(self, table: Table) -> OptionalLocationTrie:
 
     def is_valid(self) -> bool:
         """A valid location has a scheme and netloc; the path is optional."""
-        if len(self._path) < 3:
+        if len(self._path) < 2:
             return False
         scheme, netloc, *_ = self._path
+        if scheme.startswith('jdbc:') and len(netloc) > 0:
+            return True
         return scheme in _EXTERNAL_FILE_LOCATION_SCHEMES and len(netloc) > 0
 
+    def is_jdbc(self) -> bool:
+        if not self.is_valid():
+            return False
+        return self._path[0].startswith('jdbc:')
+
+    def all_tables(self) -> Iterable[Table]:
+        for node in self:
+            yield from node.tables
+
     def has_children(self):
         return len(self.children) > 0
 
@@ -125,64 +164,59 @@ def __init__(
     @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), reverse=True)
+        return sorted(self._mounts_crawler.snapshot(), key=lambda _: (len(_.name), _.name), reverse=True)
 
     def _external_locations(self) -> Iterable[ExternalLocation]:
-        min_slash = 2
-        external_locations: list[ExternalLocation] = []
+        trie = LocationTrie()
         for table in self._tables_crawler.snapshot():
-            location = table.location
-            if not location:
+            table = self._resolve_location(table)
+            if not table.location:
                 continue
-            # TODO: refactor this with LocationTrie
-            if location.startswith("dbfs:/mnt"):
-                location = self.resolve_mount(location)
-            if not location:
+            trie.insert(table)
+        queue = list(trie.children.values())
+        external_locations = []
+        while queue:
+            curr = queue.pop()
+            num_children = len(curr.children)  # 0 - take parent
+            if curr.location and (num_children > 1 or num_children == 0):
+                if curr.parent and num_children == 0 and not curr.is_jdbc():  # one table having the prefix
+                    curr = curr.parent
+                assert curr.location is not None
+                external_location = ExternalLocation(curr.location, len(list(curr.all_tables())))
+                external_locations.append(external_location)
                 continue
-            if (
-                not location.startswith("dbfs")
-                and (self._prefix_size[0] < location.find(":/") < self._prefix_size[1])
-                and not location.startswith("jdbc")
-            ):
-                self._dbfs_locations(external_locations, location, min_slash)
-            if location.startswith("jdbc"):
-                self._add_jdbc_location(external_locations, location, table)
-        return external_locations
+            queue.extend(curr.children.values())
+        return sorted(external_locations, key=lambda _: _.location)
+
+    def _resolve_location(self, table: Table) -> Table:
+        location = table.location
+        if not location:
+            return table
+        location = self._resolve_jdbc(table)
+        location = self.resolve_mount(location)
+        return dataclasses.replace(table, location=location)
 
     def resolve_mount(self, location: str | None) -> str | None:
         if not location:
             return None
+        if location.startswith('/dbfs'):
+            location = 'dbfs:' + location[5:]  # convert FUSE path to DBFS path
+        if not location.startswith('dbfs:'):
+            return location  # not a mount, save some cycles
         for mount in self._mounts_snapshot:
-            for prefix in (mount.as_scheme_prefix(), mount.as_fuse_prefix()):
-                if not location.startswith(prefix):
-                    continue
-                logger.debug(f"Replacing location {prefix} with {mount.source} in {location}")
-                location = location.replace(prefix, mount.source)
-                return location
+            prefix = mount.as_scheme_prefix()
+            if not location.startswith(prefix):
+                continue
+            logger.debug(f"Replacing location {prefix} with {mount.source} in {location}")
+            location = location.replace(prefix, mount.source)
+            return location
         logger.debug(f"Mount not found for location {location}. Skipping replacement.")
         return location
 
-    @staticmethod
-    def _dbfs_locations(external_locations, location, min_slash):
-        dupe = False
-        loc = 0
-        while loc < len(external_locations) and not dupe:
-            common = (
-                os.path.commonpath([external_locations[loc].location, os.path.dirname(location) + "/"]).replace(
-                    ":/", "://"
-                )
-                + "/"
-            )
-            if common.count("/") > min_slash:
-                table_count = external_locations[loc].table_count
-                external_locations[loc] = ExternalLocation(common, table_count + 1)
-                dupe = True
-            loc += 1
-        if not dupe:
-            external_locations.append(ExternalLocation(os.path.dirname(location) + "/", 1))
-
-    def _add_jdbc_location(self, external_locations, location, table):
-        dupe = False
+    def _resolve_jdbc(self, table: Table) -> str | None:
+        location = table.location
+        if not location or not table.storage_properties or not location.startswith('jdbc:'):
+            return location
         pattern = r"(\w+)=(.*?)(?=\s*,|\s*\])"
         # Find all matches in the input string
         # Storage properties is of the format
@@ -201,20 +235,12 @@ def _add_jdbc_location(self, external_locations, location, table):
         # currently supporting databricks and mysql external tables
         # add other jdbc types
         if "databricks" in location.lower():
-            jdbc_location = f"jdbc:databricks://{host};httpPath={httppath}"
-        elif "mysql" in location.lower():
-            jdbc_location = f"jdbc:mysql://{host}:{port}/{database}"
-        elif not provider == "":
-            jdbc_location = f"jdbc:{provider.lower()}://{host}:{port}/{database}"
-        else:
-            jdbc_location = f"{location.lower()}/{host}:{port}/{database}"
-        for ext_loc in external_locations:
-            if ext_loc.location == jdbc_location:
-                ext_loc.table_count += 1
-                dupe = True
-                break
-        if not dupe:
-            external_locations.append(ExternalLocation(jdbc_location, 1))
+            return f"jdbc:databricks://{host};httpPath={httppath}"
+        if "mysql" in location.lower():
+            return f"jdbc:mysql://{host}:{port}/{database}"
+        if not provider == "":
+            return f"jdbc:{provider.lower()}://{host}:{port}/{database}"
+        return f"{location.lower()}/{host}:{port}/{database}"
 
     def _crawl(self) -> Iterable[ExternalLocation]:
         return self._external_locations()
diff --git a/tests/integration/hive_metastore/test_external_locations.py b/tests/integration/hive_metastore/test_external_locations.py
@@ -68,7 +68,7 @@ def test_external_locations(ws, sql_backend, inventory_schema, env_or_skip):
             "bar",
             "EXTERNAL",
             "delta",
-            location="jdbc://providerunknown/",
+            location="jdbc:providerunknown:/",
             storage_properties="[database=test_db, host=somedb.us-east-1.rds.amazonaws.com, \
             port=1234, dbtable=sometable, user=*********(redacted), password=*********(redacted)]",
         ),
@@ -81,14 +81,14 @@ def test_external_locations(ws, sql_backend, inventory_schema, env_or_skip):
     crawler = ExternalLocations(ws, sql_backend, inventory_schema, tables_crawler, mounts_crawler)
     results = crawler.snapshot()
     assert results == [
-        ExternalLocation('s3://test_location/', 2),
-        ExternalLocation('s3://bar/test3/', 1),
         ExternalLocation(
             'jdbc:databricks://dbc-test1-aa11.cloud.databricks.com;httpPath=/sql/1.0/warehouses/65b52fb5bd86a7be', 1
         ),
         ExternalLocation('jdbc:mysql://somemysql.us-east-1.rds.amazonaws.com:3306/test_db', 1),
         ExternalLocation('jdbc:providerknown://somedb.us-east-1.rds.amazonaws.com:1234/test_db', table_count=2),
-        ExternalLocation('jdbc://providerunknown//somedb.us-east-1.rds.amazonaws.com:1234/test_db', 1),
+        ExternalLocation('jdbc:providerunknown://somedb.us-east-1.rds.amazonaws.com:1234/test_db', 1),
+        ExternalLocation('s3://bar/test3', 1),
+        ExternalLocation('s3://test_location', 2),
     ]
 
 
diff --git a/tests/unit/hive_metastore/test_locations.py b/tests/unit/hive_metastore/test_locations.py
@@ -22,7 +22,7 @@
 @pytest.mark.parametrize(
     "location",
     [
-        "s3://databricks-e2demofieldengwest/b169/b50"
+        "s3://databricks-e2demofieldengwest/b169/b50",
         "s3a://databricks-datasets-oregon/delta-sharing/share/open-datasets.share",
         "s3n://bucket-name/path-to-file-in-bucket",
         "gcs://test_location2/test2/table2",
@@ -166,8 +166,9 @@ def test_external_locations():
         table_factory(["s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-1/Location/Table2", ""]),
         table_factory(["s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-23/testloc/Table3", ""]),
         table_factory(["s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-23/anotherloc/Table4", ""]),
+        table_factory(["gcs://test_location2/a/b/table2", ""]),
         table_factory(["dbfs:/mnt/ucx/database1/table1", ""]),
-        table_factory(["dbfs:/mnt/ucx/database2/table2", ""]),
+        table_factory(["/dbfs/mnt/ucx/database2/table2", ""]),
         table_factory(["DatabricksRootmntDatabricksRoot", ""]),
         table_factory(
             [
@@ -212,18 +213,18 @@ def test_external_locations():
     mounts_crawler.snapshot.return_value = [Mount("/mnt/ucx", "s3://us-east-1-ucx-container")]
     sql_backend = MockBackend()
     crawler = ExternalLocations(Mock(), sql_backend, "test", tables_crawler, mounts_crawler)
-    result_set = crawler.snapshot()
-    assert len(result_set) == 7
-    assert result_set[0].location == "s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-1/Location/"
-    assert result_set[0].table_count == 2
-    assert result_set[1].location == "s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-23/"
-    assert (
-        result_set[3].location
-        == "jdbc:databricks://dbc-test1-aa11.cloud.databricks.com;httpPath=/sql/1.0/warehouses/65b52fb5bd86a7be"
-    )
-    assert result_set[4].location == "jdbc:mysql://somemysql.us-east-1.rds.amazonaws.com:3306/test_db"
-    assert result_set[5].location == "jdbc:providerknown://somedb.us-east-1.rds.amazonaws.com:1234/test_db"
-    assert result_set[6].location == "jdbc:providerunknown://somedb.us-east-1.rds.amazonaws.com:1234/test_db"
+    assert crawler.snapshot() == [
+        ExternalLocation('gcs://test_location2/a/b', 1),
+        ExternalLocation(
+            'jdbc:databricks://dbc-test1-aa11.cloud.databricks.com;httpPath=/sql/1.0/warehouses/65b52fb5bd86a7be', 1
+        ),
+        ExternalLocation('jdbc:mysql://somemysql.us-east-1.rds.amazonaws.com:3306/test_db', 1),
+        ExternalLocation('jdbc:providerknown://somedb.us-east-1.rds.amazonaws.com:1234/test_db', 2),
+        ExternalLocation('jdbc:providerunknown://somedb.us-east-1.rds.amazonaws.com:1234/test_db', 1),
+        ExternalLocation('s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-1/Location', 2),
+        ExternalLocation('s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-23', 2),
+        ExternalLocation('s3://us-east-1-ucx-container', 2),
+    ]
 
 
 LOCATION_STORAGE = MockBackend.rows("location", "storage_properties")
@@ -237,10 +238,7 @@ def test_save_external_location_mapping_missing_location():
     tables_crawler = create_autospec(TablesCrawler)
     tables_crawler.snapshot.return_value = [
         table_factory(["s3://test_location/test1/table1", ""]),
-        table_factory(["gcs://test_location2/test2/table2", ""]),
-        table_factory(["abfss://cont1@storagetest1.dfs.core.windows.net/test2/table3", ""]),
-        table_factory(["s3a://test_location_2/test1/table1", ""]),
-        table_factory(["s3n://test_location_3/test1/table1", ""]),
+        table_factory(["s3://test_location/test1/table2", ""]),
     ]
     mounts_crawler = create_autospec(MountsCrawler)
     mounts_crawler.snapshot.return_value = []
@@ -257,26 +255,6 @@ def test_save_external_location_mapping_missing_location():
             '    name = "test_location_test1"\n'
             '    url  = "s3://test_location/test1"\n'
             "    credential_name = databricks_storage_credential.<storage_credential_reference>.id\n"
-            "}\n\n"
-            'resource "databricks_external_location" "test_location2_test2" { \n'
-            '    name = "test_location2_test2"\n'
-            '    url  = "gcs://test_location2/test2"\n'
-            "    credential_name = databricks_storage_credential.<storage_credential_reference>.id\n"
-            "}\n\n"
-            'resource "databricks_external_location" "cont1_storagetest1_test2" { \n'
-            '    name = "cont1_storagetest1_test2"\n'
-            '    url  = "abfss://cont1@storagetest1.dfs.core.windows.net/test2"\n'
-            "    credential_name = databricks_storage_credential.<storage_credential_reference>.id\n"
-            "}\n\n"
-            'resource "databricks_external_location" "test_location_2_test1" { \n'
-            '    name = "test_location_2_test1"\n'
-            '    url  = "s3a://test_location_2/test1"\n'
-            "    credential_name = databricks_storage_credential.<storage_credential_reference>.id\n"
-            "}\n\n"
-            'resource "databricks_external_location" "test_location_3_test1" { \n'
-            '    name = "test_location_3_test1"\n'
-            '    url  = "s3n://test_location_3/test1"\n'
-            "    credential_name = databricks_storage_credential.<storage_credential_reference>.id\n"
             "}\n"
         ).encode("utf8"),
     )
@@ -317,8 +295,10 @@ def test_match_table_external_locations():
     matching_locations, missing_locations = location_crawler.match_table_external_locations()
 
     assert len(matching_locations) == 1
-    assert ExternalLocation("gcs://test_location2/a/b/", 1) in missing_locations
-    assert ExternalLocation("abfss://cont1@storagetest1/a/", 2) in missing_locations
+    assert [
+        ExternalLocation("abfss://cont1@storagetest1/a", 2),
+        ExternalLocation("gcs://test_location2/a/b", 1),
+    ] == missing_locations
 
 
 def test_mount_listing_multiple_folders():
