Add download cli command to download file from a collection of workspaces (#2523)

## Changes
Add a `download` cli command to download file from a collection of
workspaces.

### Linked issues

Progresses #1783

### Functionality

- [x] added relevant user documentation
- [x] added new CLI command: `download`

### Tests

- [x] manually tested
- [x] added unit tests
- [ ] ~added integration tests~ : Covering after #2507

Author: JCZuurmond

README.md

@@ -101,6 +101,7 @@ See [contributing instructions](CONTRIBUTING.md) to help improve this project.
   * [`cluster-remap` command](#cluster-remap-command)
   * [`revert-cluster-remap` command](#revert-cluster-remap-command)
   * [`upload` command](#upload-command)
+  * [`download` command](#download-command)
 * [Common Challenges and the Solutions](#common-challenges-and-the-solutions)
     * [Network Connectivity Issues](#network-connectivity-issues)
     * [Insufficient Privileges](#insufficient-privileges)
@@ -1421,6 +1422,16 @@ $ databricks labs ucx upload --file <file_path> --run-as-collection True
 Upload a file to a single workspace (`--run-as-collection False`) or a collection of workspaces
 (`--run-as-collection True`). This command is especially useful when uploading the same file to multiple workspace.
 
+## `download` command
+
+```text
+$ databricks labs ucx download --file <file_path> --run-as-collection True
+21:31:29 INFO [d.labs.ucx] Finished downloading: <file_path>
+```
+
+Download a csv file from a single workspace (`--run-as-collection False`) or a collection of workspaces
+(`--run-as-collection True`). This command is especially useful when downloading the same file from multiple workspace.
+
 # Common Challenges and the Solutions
 Users might encounter some challenges while installing and executing UCX. Please find the listing of some common challenges and the solutions below.
 

labs.yml

@@ -272,3 +272,11 @@ commands:
         description: The file to upload
       - name: run-as-collection
         description: Run the command for the collection of workspace with ucx installed. Default is False.
+
+  - name: download
+    description: download file from all workspaces in the account where ucx is installed
+    flags:
+      - name: file
+        description: The file to download
+      - name: run-as-collection
+        description: Run the command for the collection of workspace with ucx installed. Default is False.

src/databricks/labs/ucx/cli.py

@@ -1,3 +1,4 @@
+from io import BytesIO
 import json
 import webbrowser
 from pathlib import Path
@@ -8,6 +9,7 @@
 from databricks.labs.blueprint.tui import Prompts
 from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.errors import NotFound
+from databricks.sdk.service.workspace import ExportFormat
 from databricks.labs.ucx.__about__ import __version__
 
 from databricks.labs.ucx.config import WorkspaceConfig
@@ -550,6 +552,48 @@ def upload(
         logger.info(f"Finished uploading {file}")
 
 
+@ucx.command
+def download(
+    file: Path | str,
+    w: WorkspaceClient,
+    run_as_collection: bool = False,
+    a: AccountClient | None = None,  # Only used while testing
+):
+    """Download and merge a CSV file from the ucx installation in a (collection of) workspace(s)"""
+    file = Path(file)
+    if file.suffix != ".csv":
+        raise ValueError("Command only supported for CSV files")
+    contexts = _get_workspace_contexts(w, run_as_collection=run_as_collection, a=a)
+    csv_header = None
+    with file.open("wb") as output:
+        for ctx in contexts:
+            remote_file_name = f"{ctx.installation.install_folder()}/{file.name}"
+            try:
+                # Installation does not have a download method
+                data = ctx.workspace_client.workspace.download(remote_file_name, format=ExportFormat.AUTO).read()
+            except NotFound:
+                logger.warning(f"File not found for {ctx.workspace_client.config.host}: {remote_file_name}")
+                continue
+            input_ = BytesIO()  # BytesIO supports .readline() to read the header, where StreamingResponse does not
+            input_.write(data.rstrip(b"\n"))
+            input_.seek(0)  # Go back to the beginning of the file
+            csv_header_next = input_.readline()
+            if csv_header is None:
+                csv_header = csv_header_next
+                output.write(csv_header)
+            elif csv_header == csv_header_next:
+                output.write(b"\n")
+            else:
+                raise ValueError("CSV files have different headers")
+            output.write(input_.read())
+    if csv_header is None:
+        logger.warning("No file(s) to download found")
+    if file.is_file() and file.stat().st_size == 0:
+        file.unlink()
+    else:
+        logger.info(f"Finished downloading {file}")
+
+
 @ucx.command
 def lint_local_code(
     w: WorkspaceClient, prompts: Prompts, path: str | None = None, ctx: LocalCheckoutContext | None = None

tests/unit/test_cli.py

@@ -1,5 +1,6 @@
 import io
 import json
+import logging
 import time
 from pathlib import Path
 from unittest.mock import create_autospec, patch, Mock
@@ -16,7 +17,7 @@
 from databricks.sdk.service.iam import ComplexValue, User
 from databricks.sdk.service.jobs import Run, RunResultState, RunState
 from databricks.sdk.service.provisioning import Workspace
-from databricks.sdk.service.workspace import ImportFormat, ObjectInfo, ObjectType
+from databricks.sdk.service.workspace import ExportFormat, ImportFormat, ObjectInfo, ObjectType
 
 from databricks.labs.ucx.assessment.aws import AWSResources, AWSRoleAction
 from databricks.labs.ucx.aws.access import AWSResourcePermissions
@@ -31,6 +32,7 @@
     create_missing_principals,
     create_table_mapping,
     create_uber_principal,
+    download,
     ensure_assessment_run,
     installations,
     join_collection,
@@ -106,7 +108,7 @@ def create_workspace_client_mock(workspace_id: int) -> WorkspaceClient:
 """,
     }
 
-    def download(path: str) -> io.StringIO | io.BytesIO:
+    def mock_download(path: str, **_) -> io.StringIO | io.BytesIO:
         if path not in state:
             raise NotFound(path)
         if ".csv" in path or ".log" in path:
@@ -117,7 +119,7 @@ def download(path: str) -> io.StringIO | io.BytesIO:
     workspace_client.get_workspace_id.return_value = workspace_id
     workspace_client.config.host = 'https://localhost'
     workspace_client.current_user.me.return_value = User(user_name="foo", groups=[ComplexValue(display="admins")])
-    workspace_client.workspace.download = download
+    workspace_client.workspace.download.side_effect = mock_download
     workspace_client.statement_execution.execute_statement.return_value = sql.StatementResponse(
         status=sql.StatementStatus(state=sql.StatementState.SUCCEEDED),
         manifest=sql.ResultManifest(schema=sql.ResultSchema()),
@@ -798,3 +800,72 @@ def test_join_collection():
     w.workspace.download.return_value = io.StringIO(json.dumps([{"workspace_id": 123, "workspace_name": "some"}]))
     join_collection(a, "123")
     w.workspace.download.assert_not_called()
+
+
+def test_download_raises_value_error_if_not_downloading_a_csv(ws1):
+    with pytest.raises(ValueError) as e:
+        download(Path("test.txt"), ws1)
+    assert "Command only supported for CSV files" in str(e)
+
+
+@pytest.mark.parametrize("run_as_collection", [False, True])
+def test_download_calls_workspace_download(tmp_path, workspace_clients, acc_client, run_as_collection):
+    if not run_as_collection:
+        workspace_clients = [workspace_clients[0]]
+
+    download(
+        tmp_path / "test.csv",
+        workspace_clients[0],
+        run_as_collection=run_as_collection,
+        a=acc_client,
+    )
+
+    for ws in workspace_clients:
+        ws.workspace.download.assert_called_with(
+            "/Users/foo/.ucx/test.csv",
+            format=ExportFormat.AUTO,
+        )
+
+
+def test_download_warns_if_file_not_found(caplog, ws1, acc_client):
+    ws1.workspace.download.side_effect = NotFound("test.csv")
+    with caplog.at_level(logging.WARNING, logger="databricks.labs.ucx.cli"):
+        download(
+            Path("test.csv"),
+            ws1,
+            run_as_collection=False,
+            a=acc_client,
+        )
+    assert "File not found for https://localhost: /Users/foo/.ucx/test.csv" in caplog.messages
+    assert "No file(s) to download found" in caplog.messages
+
+
+def test_download_deletes_empty_file(tmp_path, ws1, acc_client):
+    ws1.workspace.download.side_effect = NotFound("test.csv")
+    mapping_path = tmp_path / "mapping.csv"
+    download(
+        mapping_path,
+        ws1,
+        run_as_collection=False,
+        a=acc_client,
+    )
+    assert not mapping_path.is_file()
+
+
+def test_download_has_expected_content(tmp_path, workspace_clients, acc_client):
+    expected = (
+        "workspace_name,catalog_name,src_schema,dst_schema,src_table,dst_table"
+        "\ntest,test,test,test,test,test"
+        "\ntest,test,test,test,test,test"
+    )
+    mapping_path = tmp_path / "mapping.csv"
+
+    download(
+        mapping_path,
+        workspace_clients[0],
+        run_as_collection=True,
+        a=acc_client,
+    )
+
+    content = mapping_path.read_text()
+    assert content == expected