Assessment report (#157)

First steps, since I want to merge Liran's work and continue there.

---------

Co-authored-by: FastLee <[email protected]>
Co-authored-by: Serge Smertin <[email protected]>

Author: 3 people

src/databricks/labs/ucx/assessment/__init__.py

@@ -0,0 +1,32 @@
+import org.apache.spark.sql.catalyst.catalog.{CatalogTable, CatalogTableType};
+
+import org.apache.spark.sql.catalyst.TableIdentifier;
+import java.io.{FileWriter, BufferedWriter, File};
+
+val schema = "SCHEMA"
+val bw = new BufferedWriter(new FileWriter(new File("/tmp/metastore_schema.csv"), true));
+
+bw.write("db,table,format,type,table_location,created_version,created_time,last_access,lib,inputformat,outputformat\n")
+val dbs = spark.sharedState.externalCatalog.listDatabases()
+for( db <- dbs) {
+  //println(s"database: ${db}")
+  val tables = spark.sharedState.externalCatalog.listTables(db)
+  for (t <- tables) {
+    try {
+      //println(s"table: ${t}")
+      val table: CatalogTable = spark.sharedState.externalCatalog.getTable(db = db, table = t)
+      val row = s"${db},${t},${table.provider.getOrElse("Unknown")},${table.tableType.name},${table.storage.locationUri.getOrElse("None")},${table.createVersion},${table.createTime},${table.lastAccessTime},${table.storage.serde.getOrElse("Unknown")},${table.storage.inputFormat.getOrElse("Unknown")},${table.storage.outputFormat.getOrElse("Unknown")}\n"
+      bw.write(row)
+    } catch {
+      case e: Exception => bw.write(s"${db},${t},Unknown,Unknown,NONE,,,,,,,\n")
+    }
+  }
+
+}
+
+bw.close;
+
+dbutils.fs.cp("file:///tmp/metastore_schema.csv","dbfs:/tmp/uc_assessment/hms/")
+spark.sql(s"create database if not exists ${schema}")
+val tables_df=spark.read.option("header","true").csv("/tmp/uc_assessment/hms/metastore_schema.csv")
+tables_df.write.mode("overwrite").saveAsTable(s"${schema}.hms_tables");

src/databricks/labs/ucx/assessment/commands/create_table_inventory.scala

@@ -1,3 +1,4 @@
+import os
 from pathlib import Path
 from typing import Annotated
 
@@ -23,3 +24,21 @@ def migrate_groups(config_file: Annotated[Path, typer.Argument(help="Path to con
     toolkit.apply_permissions_to_account_groups()
     toolkit.delete_backup_groups()
     toolkit.cleanup_inventory_table()
+
+
+@app.command()
+def generate_assessment_report():
+    from databricks.sdk import WorkspaceClient
+
+    from databricks.labs.ucx.toolkits.assessment import AssessmentToolkit
+
+    cluster_id = os.getenv("DATABRICKS_CLUSTER_ID")
+    catalog = os.getenv("INVENTORY_CATALOG", "ucx")
+    schema = os.getenv("INVENTORY_SCHEMA", "ucx")
+    ws = WorkspaceClient()
+    toolkit = AssessmentToolkit(ws=ws, cluster_id=cluster_id, inventory_catalog=catalog, inventory_schema=schema)
+    toolkit.compile_report()
+
+
+if __name__ == "__main__":
+    app()

src/databricks/labs/ucx/cli/app.py

@@ -79,7 +79,7 @@ def run(self, code):
 
         ctx = self._running_command_context()
         result = self._commands.execute(
-            cluster_id=self._cluster_id, language=compute.Language.PYTHON, context_id=ctx.id, command=code
+            cluster_id=self._cluster_id, language=self._language, context_id=ctx.id, command=code
         ).result()
 
         results = result.results

src/databricks/labs/ucx/providers/mixins/compute.py

@@ -0,0 +1,79 @@
+import logging
+import re
+from importlib import resources
+
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.compute import Language
+
+from databricks.labs.ucx.assessment import commands
+from databricks.labs.ucx.providers.mixins.compute import CommandExecutor
+
+logger = logging.getLogger(__name__)
+
+
+class AssessmentToolkit:
+    def __init__(self, ws: WorkspaceClient, cluster_id, inventory_catalog, inventory_schema, warehouse_id=None):
+        self._ws = ws
+        self._inventory_catalog = inventory_catalog
+        self._inventory_schema = inventory_schema
+        self._warehouse_id = warehouse_id
+        self._cluster_id = cluster_id
+        self._command_executor = None
+        self._managed_executor = False
+
+    def _get_command_executor(self, executor: CommandExecutor | None = None, language=None):
+        ce = executor
+        if ce is None:
+            if language:
+                ce = CommandExecutor(self._ws, language=language, cluster_id=self._cluster_id)
+            else:
+                ce = CommandExecutor(self._ws, cluster_id=self._cluster_id)
+            self._managed_executor = True
+        self._command_executor = ce
+        return ce
+
+    def _remove_command_executor(self):
+        if self._managed_executor:
+            self._command_executor = None
+            self._managed_executor = False
+
+    @staticmethod
+    def _load_command_code(name):
+        cmd_file = resources.files(commands) / name
+        with cmd_file.open("rt") as f:
+            cmd_code = f.read()
+        return cmd_code
+
+    def _get_command(self, name, params: dict | None = None):
+        cmd_code = self._load_command_code(name)
+        if params:
+            for pattern, replace in params.items():
+                p = re.compile(pattern)
+                cmd_code = p.sub(replace, cmd_code)
+        return cmd_code
+
+    @staticmethod
+    def _verify_ws_client(w: WorkspaceClient):
+        _me = w.current_user.me()
+        is_workspace_admin = any(g.display == "admins" for g in _me.groups)
+        if not is_workspace_admin:
+            msg = "Current user is not a workspace admin"
+            raise RuntimeError(msg)
+
+    def table_inventory(self, executor: CommandExecutor | None = None):
+        logger.info("Started dataset inventorization...")
+        ce = self._get_command_executor(executor, language=Language.SCALA)
+        params = {"SCHEMA": self._inventory_schema}
+        cmd_code = self._get_command("create_table_inventory.scala", params=params)
+        command_output = ce.run(cmd_code)
+        logger.debug(command_output)
+        if executor is None:
+            self._remove_command_executor()
+        logger.info("Completed dataset inventorization...")
+
+    def compile_report(self):
+        logger.info("Started report compilation...")
+        ce = self._get_command_executor(None, language=Language.SCALA)
+        self.table_inventory(ce)
+        self._remove_command_executor()
+        logger.info("Completed report compilation...")

src/databricks/labs/ucx/toolkits/assessment.py

@@ -0,0 +1,6 @@
+from databricks.labs.ucx.toolkits.assessment import AssessmentToolkit
+
+
+def test_table_inventory(ws, make_catalog, make_schema):
+    assess = AssessmentToolkit(ws, make_catalog(), make_schema())
+    assess.table_inventory()

tests/integration/test_assessment.py

@@ -0,0 +1,14 @@
+from unittest.mock import Mock
+
+from databricks.labs.ucx.toolkits.assessment import AssessmentToolkit
+
+
+def test_get_command():
+    def _mock_load(name):
+        return f"This is a test! Given name: {name}"
+
+    ws = Mock()
+    at = AssessmentToolkit(ws, cluster_id=1, inventory_catalog="foo", inventory_schema="bar")
+    at._load_command_code = _mock_load
+    c = at._get_command("ignored", {"test": "WIN"})
+    assert c.startswith("This is a WIN!")