Adopt installation package from databricks-labs-blueprint (#860)

1. The `pyproject.toml` file, which is used to manage the dependencies
and metadata of the project, has been updated to upgrade the version of
the `databricks-labs-blueprint` dependency from `0.2.2` to `0.2.4`.
2. The pull request also introduces changes to the `account.py` file in
the `databricks.labs.ucx` package. The changes include the addition of a
new `Installation` class from the
`databricks.labs.blueprint.installation` module and the removal of the
`InstallationManager` class. The `AccountWorkspaces` class has been
updated to accept an `Installation` object instead of an `AccountConfig`
object, and the `sync_workspace_info` method has been updated to use the
`Installation.save` method instead of writing the workspace info to a
JSON file.
3. The `aws.py` and `azure.py` files in the
`databricks.labs.ucx.assessment` package have also been updated to use
the `Installation` object instead of the `WorkspaceClient` object. The
`AzureResourcePermissions` and `AWSResourcePermissions` classes have
been updated to accept an `Installation` object and a `WorkspaceClient`
object as arguments, and the `save_spn_permissions` and
`save_instance_profile_permissions` methods have been updated to use the
`Installation.save` method instead of writing the permission info to a
CSV file.
4. The `cli.py` file in the `databricks.labs.ucx` package has been
updated to use the `Installation` object instead of the
`InstallationManager` object. The `installations` command has been
updated to use the `Installation.existing` method instead of the
`InstallationManager.user_installations` method. The `skip` command has
been updated to use the `TableMapping.current` method instead of the
`TableMapping.for_cli` method. The `sync_workspace_info` command has
been updated to use the `AccountWorkspaces.sync_workspace_info` method
instead of the `AccountWorkspaces.sync_workspace_info` method. The
`create_table_mapping` command has been updated to use the
`TableMapping.current` method instead of the `TableMapping.for_cli`
method. The `validate_external_locations` command has been updated to
use the `ExternalLocations.save_as_terraform_definitions_on_workspace`
method instead of the
`ExternalLocations.save_as_terraform_definitions_on_workspace` method.
The `ensure_assessment_run` command has been updated to use the
`WorkspaceInstallation.validate_and_run` method instead of the
`WorkspaceInstaller.validate_and_run` method. The `repair_run` command
has been updated to use the `WorkspaceInstallation.repair_run` method
instead of the `WorkspaceInstaller.repair_run` method. The
`validate_groups_membership` command has been updated to use the
`GroupManager.validate_group_membership` method instead of the
`GroupManager.validate_group_membership` method. The
`revert_migrated_tables` command has been updated to use the
`TablesMigrate.revert_migrated_tables` method instead of the
`TablesMigrate.revert_migrated_tables` method. The `move` and `alias`
commands have been updated to use the `TableMove.for_cli` method instead
of the `TableMove.for_cli` method. The `save_azure_storage_accounts`
command has been updated to use the `AzureResourcePermissions.for_cli`
method instead of the `AzureResourcePermissions.for_cli` method. The
`save_aws_iam_profiles` command has been updated to use the
`AWSResourcePermissions.for_cli` method instead of the
`AWSResourcePermissions.for_cli` method.
5. The `config.py` file in the `databricks.labs.ucx` package has been
updated to add a new `WorkspaceConfig` class and remove the
`AccountConfig` class. The `WorkspaceConfig` class has been updated to
include a new `connect` field, which is used to store the configuration
of the Databricks workspace. The `AccountConfig.from_dict` and
`WorkspaceConfig.from_dict` methods have been updated to migrate the
configuration from the old format to the new format. The
`AccountConfig.to_account_client` and
`WorkspaceConfig.to_workspace_client` methods have been removed, as they
are no longer needed.
6. The `crawlers.py` file in the `databricks.labs.ucx.framework` package
has been updated to add a new `DatabricksError` class, which is used to
raise exceptions when an error occurs during the execution of a SQL
statement. The `RuntimeBackend` and `SqlBackend` classes have been
updated to accept a new `debug_truncate_bytes` argument, which is used
to specify the maximum number of bytes to be printed when logging a SQL
statement. The `SqlBackend.execute` and `SqlBackend.fetch` methods have
been updated to raise a `DatabricksError` exception if an error occurs
during the execution of a SQL statement. The `CrawlerBase` class has
been updated to accept a new `SqlBackend` argument, which is used to
execute SQL statements.
7. The `dashboards.py` file in the `databricks.labs.ucx.framework`
package has been updated to add a new `Task` class, which is used to
define the tasks that can be executed by the `trigger` command. The
`Task` class has been updated to include a new `fn` field, which is used
to specify the function that should be executed when the task is run.
The `trigger` command has been updated to use the `Task.fn` field
instead of the `Task.fn` field.
8. The `locations.py` file in the `databricks.labs.ucx.hive_metastore`
package has been updated to add a new `Installation` class from the
`databricks.labs.blueprint.installation` module and remove the
`WorkspaceInfo` class. The `ExternalLocations` class has been updated to
accept an `Installation` object instead of a `WorkspaceClient` object.
The `save_as_terraform_definitions_on_workspace` method has been updated
to use the `Installation.upload` method instead of writing the external
location definitions to a file.
9. The `mapping.py` file in the `databricks.labs.ucx.hive_metastore`
package has been updated to add a new `Installation` class from the
`databricks.labs.blueprint.installation` module and remove the
`WorkspaceInfo` class. The `TableMapping` class has been updated to
accept an `Installation` object and a `WorkspaceClient` object as
arguments. The `save` method has been updated to use the
`Installation.save` method instead of writing the mapping to a CSV file.
The `load` method has been updated to use the `Installation.load` method
instead of reading the mapping from a CSV file. The `skip_table` and
`skip_schema` methods have been updated to use the
`Installation.execute` method instead of executing SQL statements
directly. The `get_tables_to_migrate` method has been updated to use the
`Installation.execute` method instead of executing SQL statements
directly.

In summary, this pull request introduces several changes to the
Databricks Labs project, including updates to the `labs.yml` and
`pyproject.toml` files, the addition of a new `Installation` class, the
removal of the `InstallationManager` class, the update of several
classes to use the `Installation` object instead of the
`WorkspaceClient` object, and the addition of a new `Task` class. These
changes aim to improve the usability and maintainability of the project.

Author: nfx

labs.yml

@@ -22,8 +22,8 @@ commands:
   - name: installations
     description: Show installations by different users on the same workspace
     table_template: |-
-      User\tDatabase\tWarehouse
-      {{range .}}{{.user_name}}\t{{.database}}\t{{.warehouse_id}}
+      Path\tDatabase\tWarehouse
+      {{range .}}{{.path}}\t{{.database}}\t{{.warehouse_id}}
       {{end}}
 
   - name: skip

pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
     "Programming Language :: Python :: Implementation :: CPython",
 ]
 dependencies = ["databricks-sdk~=0.18.0",
-                "databricks-labs-blueprint~=0.2.2",
+                "databricks-labs-blueprint~=0.2.4",
                 "PyYAML>=6.0.0,<7.0.0"]
 
 [project.entry-points.databricks]

src/databricks/labs/ucx/account.py

@@ -1,17 +1,14 @@
-import json
 import logging
 from typing import ClassVar
 
 import requests
+from databricks.labs.blueprint.installation import Installation
 from databricks.labs.blueprint.tui import Prompts
-from databricks.sdk import WorkspaceClient
+from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.errors import NotFound
 from databricks.sdk.service.provisioning import Workspace
-from databricks.sdk.service.workspace import ImportFormat
 
 from databricks.labs.ucx.__about__ import __version__
-from databricks.labs.ucx.config import AccountConfig
-from databricks.labs.ucx.installer import InstallationManager
 
 logger = logging.getLogger(__name__)
 
@@ -25,24 +22,12 @@ class AccountWorkspaces:
 
     SYNC_FILE_NAME: ClassVar[str] = "workspaces.json"
 
-    def __init__(
-        self, cfg: AccountConfig, new_workspace_client=WorkspaceClient, new_installation_manager=InstallationManager
-    ):
+    def __init__(self, ac: AccountClient, new_workspace_client=WorkspaceClient):
         self._new_workspace_client = new_workspace_client
-        self._new_installation_manager = new_installation_manager
-        self._ac = cfg.to_account_client()
-        self._cfg = cfg
-
-    def _configured_workspaces(self):
-        for workspace in self._ac.workspaces.list():
-            if self._cfg.include_workspace_names:
-                if workspace.workspace_name not in self._cfg.include_workspace_names:
-                    logger.debug(
-                        f"skipping {workspace.workspace_name} ({workspace.workspace_id} because "
-                        f"its not explicitly included"
-                    )
-                    continue
-            yield workspace
+        self._ac = ac
+
+    def _workspaces(self):
+        return self._ac.workspaces.list()
 
     def _get_cloud(self) -> str:
         if self._ac.config.is_azure:
@@ -66,7 +51,7 @@ def workspace_clients(self) -> list[WorkspaceClient]:
         :return: list[WorkspaceClient]
         """
         clients = []
-        for workspace in self._configured_workspaces():
+        for workspace in self._workspaces():
             ws = self.client_for(workspace)
             clients.append(ws)
         return clients
@@ -78,23 +63,17 @@ def sync_workspace_info(self):
         upload the json dump of workspace info in the .ucx folder
         """
         workspaces = []
-        for workspace in self._configured_workspaces():
-            workspaces.append(workspace.as_dict())
-        info = json.dumps(workspaces, indent=2).encode("utf8")
+        for workspace in self._workspaces():
+            workspaces.append(workspace)
         for ws in self.workspace_clients():
-            installation_manager = self._new_installation_manager(ws)
-            for installation in installation_manager.user_installations():
-                path = f"{installation.path}/{self.SYNC_FILE_NAME}"
-                ws.workspace.upload(path, info, overwrite=True, format=ImportFormat.AUTO)
+            for installation in Installation.existing(ws, "ucx"):
+                installation.save(workspaces, filename=self.SYNC_FILE_NAME)
 
 
 class WorkspaceInfo:
-    def __init__(self, ws: WorkspaceClient, folder: str | None = None, new_installation_manager=InstallationManager):
-        if not folder:
-            folder = f"/Users/{ws.current_user.me().user_name}/.ucx"
+    def __init__(self, installation: Installation, ws: WorkspaceClient):
+        self._installation = installation
         self._ws = ws
-        self._folder = folder
-        self._new_installation_manager = new_installation_manager
 
     def _current_workspace_id(self) -> int:
         headers = self._ws.config.authenticate()
@@ -109,12 +88,10 @@ def _current_workspace_id(self) -> int:
     def _load_workspace_info(self) -> dict[int, Workspace]:
         try:
             id_to_workspace = {}
-            with self._ws.workspace.download(f"{self._folder}/{AccountWorkspaces.SYNC_FILE_NAME}") as f:
-                for workspace_metadata in json.loads(f.read()):
-                    workspace = Workspace.from_dict(workspace_metadata)
-                    assert workspace.workspace_id is not None
-                    id_to_workspace[workspace.workspace_id] = workspace
-                return id_to_workspace
+            for workspace in self._installation.load(list[Workspace], filename=AccountWorkspaces.SYNC_FILE_NAME):
+                assert workspace.workspace_id is not None
+                id_to_workspace[workspace.workspace_id] = workspace
+            return id_to_workspace
         except NotFound:
             msg = "Please run as account-admin: databricks labs ucx sync-workspace-info"
             raise ValueError(msg) from None
@@ -142,17 +119,11 @@ def manual_workspace_info(self, prompts: Prompts):
             workspace_name = prompts.question(
                 f"Workspace name for {workspace_id}", default=f"workspace-{workspace_id}", valid_regex=r"^[\w-]+$"
             )
-            workspace = Workspace(workspace_id=int(workspace_id), workspace_name=workspace_name)
-            workspaces.append(workspace.as_dict())
+            workspaces.append(Workspace(workspace_id=int(workspace_id), workspace_name=workspace_name))
             answer = prompts.question("Next workspace id", valid_number=True, default="stop")
             if answer == "stop":
                 break
             workspace_id = int(answer)
-        info = json.dumps(workspaces, indent=2).encode("utf8")
-        installation_manager = self._new_installation_manager(self._ws)
-        logger.info("Detecting UCX installations on current workspace...")
-        for installation in installation_manager.user_installations():
-            path = f"{installation.path}/{AccountWorkspaces.SYNC_FILE_NAME}"
-            logger.info(f"Overwriting {path}")
-            self._ws.workspace.upload(path, info, overwrite=True, format=ImportFormat.AUTO)  # type: ignore[arg-type]
+        for installation in Installation.existing(self._ws, 'ucx'):
+            installation.save(workspaces, filename=AccountWorkspaces.SYNC_FILE_NAME)
         logger.info("Synchronised workspace id mapping for installations on current workspace")

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

@@ -1,6 +1,3 @@
-import csv
-import dataclasses
-import io
 import json
 import logging
 import re
@@ -11,10 +8,10 @@
 from dataclasses import dataclass
 from functools import lru_cache, partial
 
+from databricks.labs.blueprint.installation import Installation
 from databricks.labs.blueprint.parallel import Threads
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.catalog import Privilege
-from databricks.sdk.service.workspace import ImportFormat
 
 logger = logging.getLogger(__name__)
 
@@ -172,18 +169,18 @@ def _run_json_command(self, command: str):
 
 
 class AWSResourcePermissions:
-    def __init__(
-        self,
-        ws: WorkspaceClient,
-        aws_resources: AWSResources,
-        folder: str | None = None,
-    ):
-        if not folder:
-            folder = f"/Users/{ws.current_user.me().user_name}/.ucx"
-        self._folder = folder
+    def __init__(self, installation: Installation, ws: WorkspaceClient, aws_resources: AWSResources):
+        self._installation = installation
         self._aws_resources = aws_resources
         self._ws = ws
-        self._field_names = [_.name for _ in dataclasses.fields(AWSInstanceProfileAction)]
+
+    @classmethod
+    def for_cli(cls, ws: WorkspaceClient, aws_profile, product='ucx'):
+        installation = Installation.current(ws, product)
+        aws = AWSResources(aws_profile)
+        if not aws.validate_connection():
+            raise ResourceWarning("AWS CLI is not configured properly.")
+        return cls(installation, ws, aws)
 
     def _get_instance_profiles(self) -> Iterable[AWSInstanceProfile]:
         instance_profiles = self._ws.instance_profiles.list()
@@ -240,18 +237,4 @@ def save_instance_profile_permissions(self) -> str | None:
         if len(instance_profile_access) == 0:
             logger.warning("No Mapping Was Generated.")
             return None
-        return self._save(instance_profile_access)
-
-    def _overwrite_mapping(self, buffer) -> str:
-        path = f"{self._folder}/aws_instance_profile_info.csv"
-        self._ws.workspace.upload(path, buffer, overwrite=True, format=ImportFormat.AUTO)
-        return path
-
-    def _save(self, instance_profile_actions: list[AWSInstanceProfileAction]) -> str:
-        buffer = io.StringIO()
-        writer = csv.DictWriter(buffer, self._field_names)
-        writer.writeheader()
-        for instance_profile_action in instance_profile_actions:
-            writer.writerow(dataclasses.asdict(instance_profile_action))
-        buffer.seek(0)
-        return self._overwrite_mapping(buffer)
+        return self._installation.save(instance_profile_access, filename='aws_instance_profile_info.csv')

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

@@ -1,12 +1,10 @@
 import base64
-import csv
-import dataclasses
-import io
 import json
 import re
 from collections.abc import Iterable
 from dataclasses import dataclass
 
+from databricks.labs.blueprint.installation import Installation
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.core import (
     ApiClient,
@@ -17,7 +15,6 @@
 from databricks.sdk.errors import NotFound
 from databricks.sdk.service.catalog import Privilege
 from databricks.sdk.service.compute import ClusterSource, Policy
-from databricks.sdk.service.workspace import ImportFormat
 
 from databricks.labs.ucx.assessment.crawlers import (
     _CLIENT_ENDPOINT_LENGTH,
@@ -28,7 +25,12 @@
     logger,
 )
 from databricks.labs.ucx.assessment.jobs import JobsMixin
-from databricks.labs.ucx.framework.crawlers import CrawlerBase, SqlBackend
+from databricks.labs.ucx.config import WorkspaceConfig
+from databricks.labs.ucx.framework.crawlers import (
+    CrawlerBase,
+    SqlBackend,
+    StatementExecutionBackend,
+)
 from databricks.labs.ucx.hive_metastore.locations import ExternalLocations
 
 
@@ -461,20 +463,27 @@ class StoragePermissionMapping:
 
 
 class AzureResourcePermissions:
-    def __init__(self, ws: WorkspaceClient, azurerm: AzureResources, lc: ExternalLocations, folder: str | None = None):
+    def __init__(self, installation: Installation, ws: WorkspaceClient, azurerm: AzureResources, lc: ExternalLocations):
+        self._filename = 'azure_storage_account_info.csv'
+        self._installation = installation
         self._locations = lc
         self._azurerm = azurerm
         self._ws = ws
-        self._field_names = [_.name for _ in dataclasses.fields(StoragePermissionMapping)]
-        if not folder:
-            folder = f"/Users/{ws.current_user.me().user_name}/.ucx"
-        self._folder = folder
         self._levels = {
             "Storage Blob Data Contributor": Privilege.WRITE_FILES,
             "Storage Blob Data Owner": Privilege.WRITE_FILES,
             "Storage Blob Data Reader": Privilege.READ_FILES,
         }
 
+    @classmethod
+    def for_cli(cls, ws: WorkspaceClient, product='ucx'):
+        installation = Installation.current(ws, product)
+        config = installation.load(WorkspaceConfig)
+        sql_backend = StatementExecutionBackend(ws, config.warehouse_id)
+        azurerm = AzureResources(ws)
+        locations = ExternalLocations(ws, sql_backend, config.inventory_database)
+        return cls(installation, ws, azurerm, locations)
+
     def _map_storage(self, storage: AzureResource) -> list[StoragePermissionMapping]:
         logger.info(f"Fetching role assignment for {storage.storage_account}")
         out = []
@@ -512,21 +521,7 @@ def save_spn_permissions(self) -> str | None:
         if len(storage_account_infos) == 0:
             logger.error("No storage account found in current tenant with spn permission")
             return None
-        return self._save(storage_account_infos)
-
-    def _save(self, storage_infos: list[StoragePermissionMapping]) -> str:
-        buffer = io.StringIO()
-        writer = csv.DictWriter(buffer, self._field_names)
-        writer.writeheader()
-        for storage_info in storage_infos:
-            writer.writerow(dataclasses.asdict(storage_info))
-        buffer.seek(0)
-        return self._overwrite_mapping(buffer)
-
-    def _overwrite_mapping(self, buffer) -> str:
-        path = f"{self._folder}/azure_storage_account_info.csv"
-        self._ws.workspace.upload(path, buffer, overwrite=True, format=ImportFormat.AUTO)
-        return path
+        return self._installation.save(storage_account_infos, filename=self._filename)
 
     def _get_storage_accounts(self) -> list[str]:
         external_locations = self._locations.snapshot()