Fixed External Location widget to return hostname for jdbc: locations (#621)

pritishpai · web-flow · commit 9981fac91a14 · 2023-11-28T12:30:13.000+01:00
The UCX tool fetches the 'location' from the table properties while for an external location with jdbc, the host information is stored in 'storage properties'. Add a field for Table class and fetch Storage Properties from the table. Using the host field from Storage properties, create the full jdbc url Older PR closed due to Git issues: #582
diff --git a/src/databricks/labs/ucx/hive_metastore/data_objects.py b/src/databricks/labs/ucx/hive_metastore/data_objects.py
@@ -1,4 +1,6 @@
+import logging
 import os
+import re
 import typing
 from dataclasses import dataclass
 
@@ -8,6 +10,8 @@
 from databricks.labs.ucx.hive_metastore.mounts import Mounts
 from databricks.labs.ucx.mixins.sql import Row
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class ExternalLocation:
@@ -32,8 +36,10 @@ def _external_locations(self, tables: list[Row], mounts) -> list[ExternalLocatio
                         if location[5:].startswith(mount.name):
                             location = location[5:].replace(mount.name, mount.source)
                             break
-                if not location.startswith("dbfs") and (
-                    self._prefix_size[0] < location.find(":/") < self._prefix_size[1]
+                if (
+                    not location.startswith("dbfs")
+                    and (self._prefix_size[0] < location.find(":/") < self._prefix_size[1])
+                    and not location.startswith("jdbc")
                 ):
                     dupe = False
                     loc = 0
@@ -50,10 +56,47 @@ def _external_locations(self, tables: list[Row], mounts) -> list[ExternalLocatio
                         loc += 1
                     if not dupe:
                         external_locations.append(ExternalLocation(os.path.dirname(location) + "/"))
+                if location.startswith("jdbc"):
+                    pattern = r"(\w+)=(.*?)(?=\s*,|\s*\])"
+
+                    # Find all matches in the input string
+                    # Storage properties is of the format
+                    # "[personalAccessToken=*********(redacted), \
+                    #  httpPath=/sql/1.0/warehouses/65b52fb5bd86a7be, host=dbc-test1-aa11.cloud.databricks.com, \
+                    #  dbtable=samples.nyctaxi.trips]"
+                    matches = re.findall(pattern, table.storage_properties)
+
+                    # Create a dictionary from the matches
+                    result_dict = dict(matches)
+
+                    # Fetch the value of host from the newly created dict
+                    host = result_dict.get("host", "")
+                    port = result_dict.get("port", "")
+                    database = result_dict.get("database", "")
+                    httppath = result_dict.get("httpPath", "")
+                    provider = result_dict.get("provider", "")
+                    # dbtable = result_dict.get("dbtable", "")
+
+                    # 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}"
+                    external_locations.append(ExternalLocation(jdbc_location))
+
         return external_locations
 
     def _external_location_list(self):
-        tables = list(self._backend.fetch(f"SELECT location FROM {self._schema}.tables WHERE location IS NOT NULL"))
+        tables = list(
+            self._backend.fetch(
+                f"SELECT location, storage_properties FROM {self._schema}.tables WHERE location IS NOT NULL"
+            )
+        )
         mounts = Mounts(self._backend, self._ws, self._schema).snapshot()
         return self._external_locations(list(tables), list(mounts))
 
diff --git a/src/databricks/labs/ucx/hive_metastore/tables.py b/src/databricks/labs/ucx/hive_metastore/tables.py
@@ -26,6 +26,8 @@ class Table:
     view_text: str = None
     upgraded_to: str = None
 
+    storage_properties: str = None
+
     @property
     def is_delta(self) -> bool:
         if self.table_format is None:
@@ -162,6 +164,7 @@ def _describe(self, catalog: str, database: str, table: str) -> Table | None:
                 location=describe.get("Location", None),
                 view_text=describe.get("View Text", None),
                 upgraded_to=self._parse_table_props(describe.get("Table Properties", "")).get("upgraded_to", None),
+                storage_properties=self._parse_table_props(describe.get("Storage Properties", "")),
             )
         except Exception as e:
             # TODO: https://github.com/databrickslabs/ucx/issues/406
diff --git a/src/databricks/labs/ucx/hive_metastore/tables.scala b/src/databricks/labs/ucx/hive_metastore/tables.scala
@@ -8,7 +8,7 @@ import org.apache.spark.sql.DataFrame
 
 // must follow the same structure as databricks.labs.ucx.hive_metastore.tables.Table
 case class TableDetails(catalog: String, database: String, name: String, object_type: String,
-                        table_format: String, location: String, view_text: String, upgraded_to: String)
+                        table_format: String, location: String, view_text: String, upgraded_to: String, storage_properties: String)
 
 // recording error log in the database
 case class TableError(catalog: String, database: String, name: String, error: String)
@@ -37,10 +37,22 @@ def metadataForAllTables(databases: Seq[String], queue: ConcurrentLinkedQueue[Ta
           failures.add(TableError("hive_metastore", databaseName, tableName, s"result is null"))
           None
         } else {
-          val upgraded_to=table.properties.get("upgraded_to")
+          val upgraded_to = table.properties.get("upgraded_to")
+          val redactedKey = "*********"
+
+          val formattedString = table.storage.properties.map {
+            case (key, value) =>
+              if (key == "personalAccessToken")
+                s"$key=$redactedKey(redacted)"
+              else if (key.equalsIgnoreCase("password"))
+                s"$key=$redactedKey(redacted)"
+              else
+                s"$key=$value"
+          }.mkString("[", ", ", "]")
+
           Some(TableDetails("hive_metastore", databaseName, tableName, table.tableType.name, table.provider.orNull,
             table.storage.locationUri.map(_.toString).orNull, table.viewText.orNull,
-              upgraded_to match {case Some(target) => target case None => null}))
+            upgraded_to match { case Some(target) => target case None => null }, formattedString))
         }
       } catch {
         case err: Throwable =>
diff --git a/tests/integration/hive_metastore/test_external_locations.py b/tests/integration/hive_metastore/test_external_locations.py
@@ -13,11 +13,60 @@ def test_external_locations(ws, sql_backend, inventory_schema, env_or_skip):
         Table("hive_metastore", "foo", "bar", "MANAGED", "delta", location="s3://test_location/test1/table1"),
         Table("hive_metastore", "foo", "bar", "EXTERNAL", "delta", location="s3://test_location/test2/table2"),
         Table("hive_metastore", "foo", "bar", "EXTERNAL", "delta", location="dbfs:/mnt/foo/test3/table3"),
+        Table(
+            "hive_metastore",
+            "foo",
+            "bar",
+            "EXTERNAL",
+            "delta",
+            location="jdbc://databricks/",
+            storage_properties="[personalAccessToken=*********(redacted), \
+            httpPath=/sql/1.0/warehouses/65b52fb5bd86a7be, host=dbc-test1-aa11.cloud.databricks.com, \
+            dbtable=samples.nyctaxi.trips]",
+        ),
+        Table(
+            "hive_metastore",
+            "foo",
+            "bar",
+            "EXTERNAL",
+            "delta",
+            location="jdbc://MYSQL/",
+            storage_properties="[database=test_db, host=somemysql.us-east-1.rds.amazonaws.com, \
+            port=3306, dbtable=movies, user=*********(redacted), password=*********(redacted)]",
+        ),
+        Table(
+            "hive_metastore",
+            "foo",
+            "bar",
+            "EXTERNAL",
+            "delta",
+            location="jdbc://providerknown/",
+            storage_properties="[database=test_db, host=somedb.us-east-1.rds.amazonaws.com, \
+            port=1234, dbtable=sometable, user=*********(redacted), password=*********(redacted), \
+            provider=providerknown]",
+        ),
+        Table(
+            "hive_metastore",
+            "foo",
+            "bar",
+            "EXTERNAL",
+            "delta",
+            location="jdbc://providerunknown/",
+            storage_properties="[database=test_db, host=somedb.us-east-1.rds.amazonaws.com, \
+            port=1234, dbtable=sometable, user=*********(redacted), password=*********(redacted)]",
+        ),
     ]
     sql_backend.save_table(f"{inventory_schema}.tables", tables, Table)
     sql_backend.save_table(f"{inventory_schema}.mounts", [Mount("/mnt/foo", "s3://bar")], Mount)
 
     crawler = ExternalLocationCrawler(ws, sql_backend, inventory_schema)
     results = crawler.snapshot()
-    assert len(results) == 2
+    assert len(results) == 6
     assert results[1].location == "s3://bar/test3/"
+    assert (
+        results[2].location
+        == "jdbc:databricks://dbc-test1-aa11.cloud.databricks.com;httpPath=/sql/1.0/warehouses/65b52fb5bd86a7be"
+    )
+    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[5].location == "jdbc://providerunknown//somedb.us-east-1.rds.amazonaws.com:1234/test_db"
diff --git a/tests/unit/assessment/test_assessment.py b/tests/unit/assessment/test_assessment.py
@@ -65,22 +65,58 @@ def test_spark_version_compatibility():
 
 def test_external_locations():
     crawler = ExternalLocationCrawler(Mock(), MockBackend(), "test")
-    row_factory = type("Row", (Row,), {"__columns__": ["location"]})
+    row_factory = type("Row", (Row,), {"__columns__": ["location", "storage_properties"]})
     sample_locations = [
-        row_factory(["s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-1/Location/Table"]),
-        row_factory(["s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-1/Location/Table2"]),
-        row_factory(["s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-23/testloc/Table3"]),
-        row_factory(["s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-23/anotherloc/Table4"]),
-        row_factory(["dbfs:/mnt/ucx/database1/table1"]),
-        row_factory(["dbfs:/mnt/ucx/database2/table2"]),
-        row_factory(["DatabricksRootmntDatabricksRoot"]),
+        row_factory(["s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-1/Location/Table", ""]),
+        row_factory(["s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-1/Location/Table2", ""]),
+        row_factory(["s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-23/testloc/Table3", ""]),
+        row_factory(["s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-23/anotherloc/Table4", ""]),
+        row_factory(["dbfs:/mnt/ucx/database1/table1", ""]),
+        row_factory(["dbfs:/mnt/ucx/database2/table2", ""]),
+        row_factory(["DatabricksRootmntDatabricksRoot", ""]),
+        row_factory(
+            [
+                "jdbc:databricks://",
+                "[personalAccessToken=*********(redacted), \
+        httpPath=/sql/1.0/warehouses/65b52fb5bd86a7be, host=dbc-test1-aa11.cloud.databricks.com, \
+        dbtable=samples.nyctaxi.trips]",
+            ]
+        ),
+        row_factory(
+            [
+                "jdbc:/MYSQL",
+                "[database=test_db, host=somemysql.us-east-1.rds.amazonaws.com, \
+            port=3306, dbtable=movies, user=*********(redacted), password=*********(redacted)]",
+            ]
+        ),
+        row_factory(
+            [
+                "jdbc:providerknown:/",
+                "[database=test_db, host=somedb.us-east-1.rds.amazonaws.com, \
+            port=1234, dbtable=sometable, user=*********(redacted), password=*********(redacted), \
+            provider=providerknown]",
+            ]
+        ),
+        row_factory(
+            [
+                "jdbc:providerunknown:/",
+                "[database=test_db, host=somedb.us-east-1.rds.amazonaws.com, \
+            port=1234, dbtable=sometable, user=*********(redacted), password=*********(redacted)]",
+            ]
+        ),
     ]
     sample_mounts = [Mount("/mnt/ucx", "s3://us-east-1-ucx-container")]
     result_set = crawler._external_locations(sample_locations, sample_mounts)
-    assert len(result_set) == 3
+    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[1].location == "s3://us-east-1-dev-account-staging-uc-ext-loc-bucket-23/"
-    assert result_set[2].location == "s3://us-east-1-ucx-container/"
+    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"
 
 
 def test_job_assessment():
