Added databricks labs ucx validate-external-locations command for cli (#715)

HariGS-DB · web-flow · commit 46ea0f16f164 · 2023-12-21T16:47:50.000+01:00
Description:
This change contains steps to validate external location defined in the
external tables with the ones already available in the unity external
location. It prints the count of tables that can be migrated for
external location already present. Also shares the list of external
location that needs to be created for the missing ones. It generates a
tf file

Changes:
labs.yml - to add the new command
cli.py - to add code for invoking the new command
locations.py - add a field which gives count of tables using the
external location. also added logic to identify duplicate for jdbc
connection objects, functionality to compare locations and print details
and generate tf file
updated test cases to locations
added new test cases for mapping
diff --git a/labs.yml b/labs.yml
@@ -43,3 +43,6 @@ commands:
 
   - name: create-table-mapping
     description: create initial table mapping for review
+
+  - name: validate-external-locations
+    description: validates and provides mapping to external table to external location and shared generation tf scripts
diff --git a/src/databricks/labs/ucx/cli.py b/src/databricks/labs/ucx/cli.py
@@ -10,7 +10,7 @@
 from databricks.labs.ucx.config import AccountConfig, ConnectConfig
 from databricks.labs.ucx.framework.crawlers import StatementExecutionBackend
 from databricks.labs.ucx.framework.tui import Prompts
-from databricks.labs.ucx.hive_metastore import TablesCrawler
+from databricks.labs.ucx.hive_metastore import ExternalLocations, TablesCrawler
 from databricks.labs.ucx.hive_metastore.mapping import TableMapping
 from databricks.labs.ucx.install import WorkspaceInstaller
 from databricks.labs.ucx.installer import InstallationManager
@@ -90,13 +90,26 @@ def create_table_mapping():
     webbrowser.open(f"{ws.config.host}/#workspace{path}")
 
 
+def validate_external_locations():
+    ws = WorkspaceClient()
+    prompts = Prompts()
+    installation_manager = InstallationManager(ws)
+    installation = installation_manager.for_user(ws.current_user.me())
+    sql_backend = StatementExecutionBackend(ws, installation.config.warehouse_id)
+    location_crawler = ExternalLocations(ws, sql_backend, installation.config.inventory_database)
+    path = location_crawler.save_as_terraform_definitions_on_workspace()
+    if len(path) > 0 and prompts.confirm(f"external_locations.tf file written to {path}. Do you want to open it?"):
+        webbrowser.open(f"{ws.config.host}/#workspace{path}")
+
+
 MAPPING = {
     "open-remote-config": open_remote_config,
     "installations": list_installations,
     "workflows": workflows,
     "sync-workspace-info": sync_workspace_info,
     "manual-workspace-info": manual_workspace_info,
     "create-table-mapping": create_table_mapping,
+    "validate-external-locations": validate_external_locations,
     "skip": skip,
 }
 
diff --git a/src/databricks/labs/ucx/hive_metastore/locations.py b/src/databricks/labs/ucx/hive_metastore/locations.py
@@ -1,3 +1,4 @@
+import io
 import logging
 import os
 import re
@@ -6,6 +7,7 @@
 from typing import ClassVar
 
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.workspace import ImportFormat
 
 from databricks.labs.ucx.framework.crawlers import CrawlerBase, SqlBackend
 from databricks.labs.ucx.mixins.sql import Row
@@ -16,6 +18,7 @@
 @dataclass
 class ExternalLocation:
     location: str
+    table_count: int
 
 
 @dataclass
@@ -30,6 +33,7 @@ class ExternalLocations(CrawlerBase[ExternalLocation]):
     def __init__(self, ws: WorkspaceClient, sbe: SqlBackend, schema):
         super().__init__(sbe, "hive_metastore", schema, "external_locations", ExternalLocation)
         self._ws = ws
+        self._folder = f"/Users/{ws.current_user.me().user_name}/.ucx"
 
     def _external_locations(self, tables: list[Row], mounts) -> Iterable[ExternalLocation]:
         min_slash = 2
@@ -57,12 +61,14 @@ def _external_locations(self, tables: list[Row], mounts) -> Iterable[ExternalLoc
                             + "/"
                         )
                         if common.count("/") > min_slash:
-                            external_locations[loc] = ExternalLocation(common)
+                            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) + "/"))
+                        external_locations.append(ExternalLocation(os.path.dirname(location) + "/", 1))
                 if location.startswith("jdbc"):
+                    dupe = False
                     pattern = r"(\w+)=(.*?)(?=\s*,|\s*\])"
 
                     # Find all matches in the input string
@@ -93,7 +99,13 @@ def _external_locations(self, tables: list[Row], mounts) -> Iterable[ExternalLoc
                         jdbc_location = f"jdbc:{provider.lower()}://{host}:{port}/{database}"
                     else:
                         jdbc_location = f"{location.lower()}/{host}:{port}/{database}"
-                    external_locations.append(ExternalLocation(jdbc_location))
+                    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 external_locations
 
@@ -113,6 +125,81 @@ def _try_fetch(self) -> Iterable[ExternalLocation]:
         for row in self._fetch(f"SELECT * FROM {self._schema}.{self._table}"):
             yield ExternalLocation(*row)
 
+    def _get_ext_location_definitions(self, missing_locations: list[ExternalLocation]) -> list:
+        tf_script = []
+        cnt = 1
+        for loc in missing_locations:
+            if loc.location.startswith("s3://"):
+                res_name = loc.location[5:].rstrip("/").replace("/", "_")
+            elif loc.location.startswith("gcs://"):
+                res_name = loc.location[6:].rstrip("/").replace("/", "_")
+            elif loc.location.startswith("abfss://"):
+                container_name = loc.location[8 : loc.location.index("@")]
+                res_name = (
+                    loc.location[loc.location.index("@") + 1 :]
+                    .replace(".dfs.core.windows.net", "")
+                    .rstrip("/")
+                    .replace("/", "_")
+                )
+                res_name = f"{container_name}_{res_name}"
+            else:
+                # if the cloud storage url doesn't match the above condition or incorrect (example wasb://)
+                # dont generate tf script and ignore
+                logger.warning(f"unsupported storage format {loc.location}")
+                continue
+            script = f'resource "databricks_external_location" "{res_name}" {{ \n'
+            script += f'    name = "{res_name}"\n'
+            script += f'    url  = "{loc.location.rstrip("/")}"\n'
+            script += "    credential_name = databricks_storage_credential.<storage_credential_reference>.id\n"
+            script += "}\n"
+            tf_script.append(script)
+            cnt += 1
+        return tf_script
+
+    def _match_table_external_locations(self) -> tuple[list[list], list[ExternalLocation]]:
+        external_locations = list(self._ws.external_locations.list())
+        location_path = [_.url.lower() for _ in external_locations]
+        table_locations = self.snapshot()
+        matching_locations = []
+        missing_locations = []
+        for loc in table_locations:
+            # external_location.list returns url without trailing "/" but ExternalLocation.snapshot
+            # does so removing the trailing slash before comparing
+            if loc.location.rstrip("/").lower() in location_path:
+                # identify the index of the matching external_locations
+                iloc = location_path.index(loc.location.rstrip("/"))
+                matching_locations.append([external_locations[iloc].name, loc.table_count])
+                continue
+            missing_locations.append(loc)
+        return matching_locations, missing_locations
+
+    def save_as_terraform_definitions_on_workspace(self, folder: str | None = None) -> str:
+        if folder:
+            self._folder = folder
+        matching_locations, missing_locations = self._match_table_external_locations()
+        if len(matching_locations) > 0:
+            logger.info("following external locations are already configured.")
+            logger.info("sharing details of # tables that can be migrated for each location")
+            for _ in matching_locations:
+                logger.info(f"{_[1]} tables can be migrated using external location {_[0]}.")
+        if len(missing_locations) > 0:
+            logger.info("following external location need to be created.")
+            for _ in missing_locations:
+                logger.info(f"{_.table_count} tables can be migrated using external location {_.location}.")
+            buffer = io.StringIO()
+            for script in self._get_ext_location_definitions(missing_locations):
+                buffer.write(script)
+            buffer.seek(0)
+            return self._overwrite_mapping(buffer)
+        else:
+            logger.info("no additional external location to be created.")
+            return ""
+
+    def _overwrite_mapping(self, buffer) -> str:
+        path = f"{self._folder}/external_locations.tf"
+        self._ws.workspace.upload(path, buffer, overwrite=True, format=ImportFormat.AUTO)
+        return path
+
 
 class Mounts(CrawlerBase[Mount]):
     def __init__(self, backend: SqlBackend, ws: WorkspaceClient, inventory_database: str):
diff --git a/tests/integration/hive_metastore/test_external_locations.py b/tests/integration/hive_metastore/test_external_locations.py
@@ -1,6 +1,10 @@
 import logging
 
-from databricks.labs.ucx.hive_metastore.locations import ExternalLocations, Mount
+from databricks.labs.ucx.hive_metastore.locations import (
+    ExternalLocation,
+    ExternalLocations,
+    Mount,
+)
 from databricks.labs.ucx.hive_metastore.tables import Table
 
 logger = logging.getLogger(__name__)
@@ -44,6 +48,17 @@ def test_external_locations(ws, sql_backend, inventory_schema, env_or_skip):
             port=1234, dbtable=sometable, user=*********(redacted), password=*********(redacted), \
             provider=providerknown]",
         ),
+        Table(
+            "hive_metastore",
+            "foo",
+            "bar2",
+            "EXTERNAL",
+            "delta",
+            location="jdbc://providerknown/",
+            storage_properties="[database=test_db, host=somedb.us-east-1.rds.amazonaws.com, \
+            port=1234, dbtable=sometable2, user=*********(redacted), password=*********(redacted), \
+            provider=providerknown]",
+        ),
         Table(
             "hive_metastore",
             "foo",
@@ -68,4 +83,17 @@ def test_external_locations(ws, sql_backend, inventory_schema, env_or_skip):
     )
     assert results[3].location == "jdbc:mysql://somemysql.us-east-1.rds.amazonaws.com:3306/test_db"
     assert results[4].location == "jdbc:providerknown://somedb.us-east-1.rds.amazonaws.com:1234/test_db"
+    assert results[4].table_count == 2
     assert results[5].location == "jdbc://providerunknown//somedb.us-east-1.rds.amazonaws.com:1234/test_db"
+
+
+def test_save_external_location_mapping_missing_location(ws, sql_backend, inventory_schema):
+    logger.info("setting up fixtures")
+    locations = [
+        ExternalLocation("abfss://cont1@storage123/test_location", 2),
+        ExternalLocation("abfss://cont1@storage456/test_location2", 1),
+    ]
+    sql_backend.save_table(f"{inventory_schema}.external_locations", locations, ExternalLocation)
+    location_crawler = ExternalLocations(ws, sql_backend, inventory_schema)
+    path = location_crawler.save_as_terraform_definitions_on_workspace()
+    assert ws.workspace.get_status(path)
diff --git a/tests/unit/hive_metastore/test_locations.py b/tests/unit/hive_metastore/test_locations.py
@@ -1,6 +1,7 @@
 from unittest.mock import MagicMock, Mock
 
 from databricks.sdk.dbutils import MountInfo
+from databricks.sdk.service.catalog import ExternalLocationInfo
 
 from databricks.labs.ucx.hive_metastore.locations import (
     ExternalLocations,
@@ -79,6 +80,14 @@ def test_external_locations():
             provider=providerknown]",
             ]
         ),
+        row_factory(
+            [
+                "jdbc:providerknown:/",
+                "[database=test_db, host=somedb.us-east-1.rds.amazonaws.com, \
+            port=1234, dbtable=sometable2, user=*********(redacted), password=*********(redacted), \
+            provider=providerknown]",
+            ]
+        ),
         row_factory(
             [
                 "jdbc:providerunknown:/",
@@ -91,6 +100,7 @@ def test_external_locations():
     result_set = crawler._external_locations(sample_locations, sample_mounts)
     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
@@ -99,3 +109,61 @@ def test_external_locations():
     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"
+
+
+def make_row(data, columns):
+    row = Row(data)
+    row.__columns__ = columns
+    return row
+
+
+def test_save_external_location_mapping_missing_location():
+    ws = MagicMock()
+    select_cols = ["location", "storage_properties"]
+    sbe = MockBackend(
+        rows={
+            "SELECT location, storage_properties FROM test.tables WHERE location IS NOT NULL": [
+                make_row(("s3://test_location/test1/table1", ""), select_cols),
+                make_row(("gcs://test_location2/test2/table2", ""), select_cols),
+                make_row(("abfss://cont1@storagetest1.dfs.core.windows.net/test2/table3", ""), select_cols),
+            ],
+        }
+    )
+    location_crawler = ExternalLocations(ws, sbe, "test")
+    ws.external_locations.list.return_value = [ExternalLocationInfo(name="loc1", url="s3://test_location/test11")]
+    location_crawler.save_as_terraform_definitions_on_workspace("~/.ucx")
+    (path, content), _ = ws.workspace.upload.call_args
+    assert "~/.ucx/external_locations.tf" == path
+    assert (
+        'resource "databricks_external_location" "test_location_test1" { \n'
+        '    name = "test_location_test1"\n'
+        '    url  = "s3://test_location/test1"\n'
+        "    credential_name = databricks_storage_credential.<storage_credential_reference>.id\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"
+        '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"
+    ) == content.read()
+
+
+def test_save_external_location_mapping_no_missing_location():
+    ws = MagicMock()
+    select_cols = ["location", "storage_properties"]
+    sbe = MockBackend(
+        rows={
+            "SELECT location, storage_properties FROM test.tables WHERE location IS NOT NULL": [
+                make_row(("s3://test_location/test1/table1", ""), select_cols),
+            ],
+        }
+    )
+    location_crawler = ExternalLocations(ws, sbe, "test")
+    ws.external_locations.list.return_value = [ExternalLocationInfo(name="loc1", url="s3://test_location/test1")]
+    location_crawler.save_as_terraform_definitions_on_workspace("~/.ucx")
+    ws.workspace.upload.assert_not_called()
