Added export CLI functionality for assessment results

README.md

@@ -10,8 +10,10 @@ so that you'll be able to [scope the migration](docs/assessment.md) and execute
 The [README notebook](#readme-notebook), which can be found in the installation folder contains further instructions and explanations of the different ucx workflows & dashboards.
 Once the migration is scoped, you can start with the [table migration process](#Table-Migration).
 
+
 More workflows, like notebook code migration are coming in future releases.
 
+
 UCX also provides a number of command line utilities accessible via `databricks labs ucx`.
 
 For questions, troubleshooting or bug fixes, please see our [troubleshooting guide](docs/troubleshooting.md) or submit [an issue](https://github.com/databrickslabs/ucx/issues).
@@ -88,6 +90,7 @@ See [contributing instructions](CONTRIBUTING.md) to help improve this project.
   * [`open-remote-config` command](#open-remote-config-command)
   * [`installations` command](#installations-command)
   * [`report-account-compatibility` command](#report-account-compatibility-command)
+  * [`export-assessment` command](#export-assessment-command)
 * [Metastore related commands](#metastore-related-commands)
   * [`show-all-metastores` command](#show-all-metastores-command)
   * [`assign-metastore` command](#assign-metastore-command)
@@ -1109,6 +1112,29 @@ databricks labs ucx report-account-compatibility --profile labs-azure-account
 
 [[back to top](#databricks-labs-ucx)]
 
+## `export-assessment` command
+
+```commandline
+databricks labs ucx export-assessment
+```
+The `export-assessment` command exports UCX results to a specified location and prompts for details on the path and report type. If not specified, it defaults to exporting the main results to the current directory as `export_to_zipped_csv.zip`.
+
+### Prompts
+
+- **Choose a path to save the UCX Assessment results:**
+    - **Description:** Specify the path where the results should be saved. If not provided, results will be saved in the current directory.
+
+- **Choose which assessment results to export:**
+    - **Description:** Select the type of results to export. Options include:
+        - `azure`
+        - `estimates`
+        - `interactive`
+        - `main`
+    - **Default:** `main`
+
+[[back to top](#databricks-labs-ucx)]
+
+
 # Metastore related commands
 
 These commands are used to assign a Unity Catalog metastore to a workspace. The metastore assignment is a pre-requisite

labs.yml

@@ -290,3 +290,6 @@ commands:
         description: The file to download
       - name: run-as-collection
         description: Run the command for the collection of workspaces with ucx installed. Default is False.
+
+  - name: export-assessment
+    description: Export UCX results to a specified location

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

@@ -0,0 +1,51 @@
+import logging
+from pathlib import Path
+
+from databricks.labs.blueprint.tui import Prompts
+
+from databricks.labs.ucx.config import WorkspaceConfig
+from databricks.labs.lsql.backends import SqlBackend
+from databricks.labs.lsql.dashboards import DashboardMetadata
+
+logger = logging.getLogger(__name__)
+
+
+class AssessmentExporter:
+    # File and Path Constants
+    _EXPORT_FILE_NAME = "ucx_assessment_results.zip"
+
+    def __init__(self, sql_backend: SqlBackend, config: WorkspaceConfig):
+        self._sql_backend = sql_backend
+        self._config = config
+
+    def export_results(self, prompts: Prompts):
+        """Main method to export results to CSV files inside a ZIP archive."""
+        project_root = Path(__file__).resolve().parents[3]
+        queries_path_root = project_root / f"labs/ucx/queries/assessment"
+        valid_queries = {subdir.name for subdir in queries_path_root.iterdir() if subdir.is_dir()}
+
+        export_path = Path(
+            prompts.question(
+                "Choose a path to save the UCX Assessment results",
+                default=Path.cwd().as_posix(),
+                validate=lambda p_: Path(p_).exists(),
+            )
+        )
+
+        query_choice = prompts.question(
+            "Choose which assessment results to export",
+            default="main",
+            validate=lambda q: q in valid_queries,
+        )
+
+        queries_path = queries_path_root / query_choice
+        display_name = "UCX Assessment Results"
+        dashboard_metadata = DashboardMetadata(display_name)
+        assessment_results = (
+            DashboardMetadata.from_path(queries_path)
+                .replace_database(catalog="hive_metastore", database=self._config.inventory_database)
+        )
+
+        logger.info("Exporting assessment results....")
+        results_path = assessment_results.export_to_zipped_csv(self._sql_backend, export_path)
+        logger.info(f"Results exported to {results_path}")

src/databricks/labs/ucx/cli.py

@@ -19,6 +19,8 @@
 from databricks.labs.ucx.install import AccountInstaller
 from databricks.labs.ucx.source_code.linters.files import LocalCodeLinter
 
+from databricks.labs.ucx.assessment.export import AssessmentExporter
+
 ucx = App(__file__)
 logger = get_logger(__file__)
 
@@ -566,7 +568,6 @@ def join_collection(a: AccountClient, workspace_ids: str):
     w_ids = [int(_.strip()) for _ in workspace_ids.split(",") if _]
     account_installer.join_collection(w_ids)
 
-
 @ucx.command
 def upload(
     file: Path | str,
@@ -637,5 +638,13 @@ def lint_local_code(
     linter.lint(prompts, None if path is None else Path(path))
 
 
+@ucx.command
+def export_assessment(w: WorkspaceClient, prompts: Prompts):
+    """Export the UCX assessment queries to a zip file."""
+    ctx: WorkspaceContext = WorkspaceContext(w)
+    exporter : AssessmentExporter = ctx.assessment_exporter
+    exporter.export_results(prompts)
+
+
 if __name__ == "__main__":
     ucx()

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

@@ -22,6 +22,7 @@
 
 from databricks.labs.ucx.account.workspaces import WorkspaceInfo
 from databricks.labs.ucx.assessment.azure import AzureServicePrincipalCrawler
+from databricks.labs.ucx.assessment.export import AssessmentExporter
 from databricks.labs.ucx.aws.credentials import CredentialManager
 from databricks.labs.ucx.config import WorkspaceConfig
 from databricks.labs.ucx.hive_metastore import ExternalLocations, Mounts, TablesCrawler
@@ -250,6 +251,10 @@ def tables_migrator(self):
             self.migrate_grants,
         )
 
+    @cached_property
+    def assessment_exporter(self):
+        return AssessmentExporter(self.sql_backend, self.config)
+
     @cached_property
     def acl_migrator(self):
         return ACLMigrator(

src/databricks/labs/ucx/installer/workflows.py

@@ -112,6 +112,111 @@
      f'--parent_run_id=' + dbutils.widgets.get('parent_run_id'))
 """
 
+EXPORT_TO_EXCEL_NOTEBOOK = """
+# Databricks notebook source
+# MAGIC %md
+# MAGIC ##### Exporter of UCX assessment results
+# MAGIC ##### Instructions:
+# MAGIC 1. Execute using an all-purpose cluster with Databricks Runtime 14 or higher.
+# MAGIC 1. Hit **Run all** button and wait for completion.
+# MAGIC 1. Go to the bottom of the notebook and click the Download UCX Results button.
+# MAGIC
+# MAGIC ##### Important:
+# MAGIC Please note that this is only meant to serve as example code.
+# MAGIC
+# MAGIC Example code developed by **Databricks Shared Technical Services team**.
+
+# COMMAND ----------
+
+# DBTITLE 1,Installing Packages
+# MAGIC %pip install {remote_wheel} -qqq
+# MAGIC %pip install xlsxwriter -qqq
+# MAGIC dbutils.library.restartPython()
+
+# COMMAND ----------
+
+# DBTITLE 1,Libraries Import and Setting UCX
+import os
+import shutil
+import logging, threading
+from functools import partial
+from threading import Lock
+
+# third party Libraries imports
+import pandas as pd
+import xlsxwriter
+
+# Databricks imports
+from databricks.sdk.config import with_user_agent_extra
+from databricks.labs.blueprint.logger import install_logger
+from databricks.labs.blueprint.parallel import Threads
+from databricks.labs.lsql.dashboards import QueryTile
+from databricks.labs.ucx.contexts.workflow_task import RuntimeContext
+
+# ctx
+install_logger()
+with_user_agent_extra("cmd", "export-assessment")
+named_parameters = dict(config="/Workspace{config_file}")
+ctx = RuntimeContext(named_parameters)
+lock = Lock()
+
+# COMMAND ----------
+
+# DBTITLE 1,Assessment Export
+
+# File and Path Constants
+FILE_NAME = "assessment_results.xlsx"
+TMP_PATH = f'/tmp' + ctx.installation.install_folder() + '/excel-export'
+DOWNLOAD_PATH = "/dbfs/FileStore/excel-export"
+
+def _cleanup() -> None:
+    \"\"\"Move the temporary results file to the download path and clean up the temp directory.\"\"\"
+    shutil.move(
+        os.path.join(TMP_PATH, FILE_NAME),
+        os.path.join(DOWNLOAD_PATH, FILE_NAME),
+    )
+    shutil.rmtree(TMP_PATH)
+
+def _prepare_directories() -> None:
+    '''Ensure that the necessary directories exist.'''
+    os.makedirs(TMP_PATH, exist_ok=True)
+    os.makedirs(DOWNLOAD_PATH, exist_ok=True)
+
+def _to_excel(tile: QueryTile, writer: ...) -> None:
+    '''Execute a SQL query and write the result to an Excel sheet.'''
+    sdf = spark.sql(tile.content)
+    df = sdf.toPandas()
+    with lock:
+        df.to_excel(writer, sheet_name=tile.metadata.title, index=False)
+
+def _render_export() -> None:
+    '''Render an HTML link for downloading the results.'''
+    html_content = f'''
+            <style>@font-face{{font-family:'DM Sans';src:url(https://cdn.bfldr.com/9AYANS2F/at/p9qfs3vgsvnp5c7txz583vgs/dm-sans-regular.ttf?auto=webp&format=ttf) format('truetype');font-weight:400;font-style:normal}}body{{font-family:'DM Sans',Arial,sans-serif}}.export-container{{text-align:center;margin-top:20px}}.export-container h2{{color:#1B3139;font-size:24px;margin-bottom:20px}}.export-container a{{display:inline-block;padding:12px 25px;background-color:#1B3139;color:#fff;text-decoration:none;border-radius:4px;font-size:18px;font-weight:500;transition:background-color 0.3s ease,transform 0.3s ease}}.export-container a:hover{{background-color:#FF3621;transform:translateY(-2px)}}</style><div class="export-container"><h2>Export Results</h2><a href='{workspace_host}files/excel-export/assessment_results.xlsx?o={workspace_id}' target='_blank' download>Download UCX Results </a></div>
+    '''
+    displayHTML(html_content)
+
+def export_results() -> None:
+    '''Main method to export results to an Excel file.'''
+    _prepare_directories()
+    try:
+        target = TMP_PATH + '/assessment_results.xlsx'
+        with pd.ExcelWriter(target, engine="xlsxwriter") as writer:
+            tasks = []
+            for query_tile in ctx.assessment_exporter.queries():
+                tasks.append(partial(_to_excel, query_tile, writer))
+                Threads.strict("exporting", tasks)
+        _cleanup()
+        _render_export()
+    except Exception as e:
+        print(f"Error exporting results ", e)
+
+# COMMAND ----------
+
+# DBTITLE 1,Data Export
+export_results()
+"""
+
 
 class DeployedWorkflows:
     def __init__(self, ws: WorkspaceClient, install_state: InstallState):
@@ -792,6 +897,15 @@ def _create_debug(self, remote_wheels: list[str]):
         ).encode("utf8")
         self._installation.upload('DEBUG.py', content)
 
+    def _create_export(self, remote_wheels: list[str]):
+        content = EXPORT_TO_EXCEL_NOTEBOOK.format(
+            remote_wheel=remote_wheels,
+            config_file=self._config_file,
+            workspace_host=self._ws.config.host,
+            workspace_id=self._ws.get_workspace_id(),
+        ).encode("utf8")
+        self._installation.upload('EXPORT_TO_EXCEL_NOTEBOOK.py', content)
+
 
 class MaxedStreamHandler(logging.StreamHandler):
 

tests/unit/assessment/test_export.py

@@ -0,0 +1,39 @@
+from databricks.labs.ucx.config import WorkspaceConfig
+from databricks.labs.ucx.assessment.export import AssessmentExport
+from databricks.labs.lsql.backends import MockBackend
+from databricks.labs.blueprint.tui import MockPrompts
+from databricks.labs.lsql.core import Row
+
+
+def test_export(tmp_path):
+    """Test the export_results method of the AssessmentExport class."""
+    query = {
+        "SELECT\n  one\nFROM ucx.external_locations": [
+            Row(location="s3://bucket1/folder1", table_count=1),
+            Row(location="abfss://container1@storage1.dfs.core.windows.net/folder1", table_count=1),
+            Row(location="gcp://folder1", table_count=2),
+        ]
+    }
+
+    # Setup workspace configuration
+    config = WorkspaceConfig(inventory_database="ucx")
+
+    # Prepare temporary paths and files
+    export_path = tmp_path / "export"
+    export_path.mkdir(parents=True, exist_ok=True)
+
+    # Mock backend and prompts
+    mock_backend = MockBackend(rows=query)
+    mock_prompts = MockPrompts(
+        {
+            "Choose a path to save the UCX Assessment results": export_path.as_posix(),
+            "Choose which assessment results to export": "main",
+        }
+    )
+
+    # Execute export process
+    export = AssessmentExporter(mock_backend, config)
+    exported = export.export_results(mock_prompts)
+
+    # Assertions
+    assert exported == export_path / "export_to_zipped_csv.zip"

tests/unit/test_cli.py

@@ -59,6 +59,7 @@
     validate_groups_membership,
     workflows,
     delete_missing_principals,
+    export_assessment,
 )
 from databricks.labs.ucx.contexts.account_cli import AccountContext
 from databricks.labs.ucx.contexts.workspace_cli import WorkspaceContext
@@ -887,3 +888,12 @@ def test_delete_principals(ws):
     prompts = MockPrompts({"Select the list of roles *": "0"})
     delete_missing_principals(ws, prompts, ctx)
     role_creation.delete_uc_roles.assert_called_once()
+
+def test_export_assessment(ws, tmp_path):
+    mock_prompts = MockPrompts({
+            "Choose a path to save the UCX Assessment results": tmp_path.as_posix(),
+            "Choose which assessment results to export": "main",
+        })
+
+    export_assessment(ws, mock_prompts)
+    assert len(list(tmp_path.glob("export_to_zipped_csv.zip"))) == 1