Added databricks labs ucx cluster-remap command to remap legacy cluster configurations to UC-compatible

labs.yml

@@ -143,3 +143,9 @@ commands:
   - name: create-catalogs-schemas
     description: Create UC external catalogs and schemas based on the destinations created from create_table_mapping command.
       This command is supposed to be run before migrating tables to UC.
+
+  - name: cluster-remap
+    description: Re-mapping the cluster to UC
+
+  - name: revert-cluster-remap
+    description: Reverting the Re-mapping of the  cluster from  UC

src/databricks/labs/ucx/cli.py

@@ -26,6 +26,7 @@
 from databricks.labs.ucx.hive_metastore.table_migrate import TablesMigrate
 from databricks.labs.ucx.hive_metastore.table_move import TableMove
 from databricks.labs.ucx.install import WorkspaceInstallation
+from databricks.labs.ucx.workspace_access.clusters import ClusterAccess
 from databricks.labs.ucx.workspace_access.groups import GroupManager
 
 ucx = App(__file__)
@@ -504,5 +505,43 @@
     catalog_schema.create_catalog_schema()
 
 
+@ucx.command
+def cluster_remap(w: WorkspaceClient, prompts: Prompts):
+    """Re-mapping the cluster to UC"""
+    logger.info("Remapping the Clusters to UC")
+    installation = Installation.current(w, 'ucx')
+    cluster = ClusterAccess(installation, w, prompts)
+    cluster_list = cluster.list_cluster()
+    print("Cluster Name\tCluster Id")
+    for name, cluster_id in cluster_list.items():
+        print(f"{name}\t{cluster_id}")
+    cluster_ids = prompts.question(
+        "Please provide the cluster id's as comma separated value from the above list", default="<ALL>"
+    )
+    cluster.map_cluster_to_uc(cluster_ids)
+
+
+@ucx.command
+def revert_cluster_remap(w: WorkspaceClient, prompts: Prompts):
+    """Reverting Re-mapping of  clusters from UC"""
+    logger.info("Reverting the Remapping of the Clusters from UC")
+    installation = Installation.current(w, 'ucx')
+    cluster_ids = [
+        cluster_files.path.split("/")[-1].split(".")[0]
+        for cluster_files in installation.files()
+        if cluster_files.path is not None and cluster_files.path.find("backup/clusters") > 0
+    ]
+    if not cluster_ids:
+        logger.info("There is no cluster files in the backup folder.Skipping the reverting process")
+        return
+    for cluster in cluster_ids:
+        print(cluster)
+    cluster_list = prompts.question(
+        "Please provide the cluster id's as comma separated value from the above list", default="<ALL>"
+    )
+    cluster_class = ClusterAccess(installation, w, prompts)
+    cluster_class.revert_cluster_remap(cluster_list, cluster_ids)
+
+
 if __name__ == "__main__":
     ucx()

src/databricks/labs/ucx/workspace_access/clusters.py

@@ -0,0 +1,117 @@
+import logging
+
+from databricks.labs.blueprint.installation import Installation
+from databricks.labs.blueprint.tui import Prompts
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import InvalidParameterValue
+from databricks.sdk.service.compute import ClusterDetails, DataSecurityMode
+
+logger = logging.getLogger(__name__)
+
+
+class ClusterAccess:
+    def __init__(self, installation: Installation, ws: WorkspaceClient, prompts: Prompts):
+        self._ws = ws
+        self._prompts = prompts
+        self._installation = installation
+
+    def list_cluster(self):
+        cluster_list = {}
+        clusters = self._ws.clusters.list()
+        for cluster in clusters:
+            if cluster.cluster_source is not None and cluster.cluster_source.name != "JOB":
+                cluster_list[cluster.cluster_name] = cluster.cluster_id
+        return cluster_list
+
+    def _get_cluster_id(self, cluster_id: str):
+        if cluster_id != "<ALL>":
+            return [x.strip() for x in cluster_id.split(",")]
+        cluster_list = []
+        clusters = self._ws.clusters.list()
+        for cluster in clusters:
+            if cluster.cluster_source is not None and cluster.cluster_source.name != "JOB":
+                cluster_list.append(cluster.cluster_id)
+        return cluster_list
+
+    def _get_access_mode(self, access_mode: str):
+        if access_mode in {"LEGACY_SINGLE_USER", "SINGLE_USER"}:
+            return DataSecurityMode.SINGLE_USER
+        return DataSecurityMode.USER_ISOLATION
+
+    def map_cluster_to_uc(self, cluster_id: str):
+
+        cluster_id_list = self._get_cluster_id(cluster_id)
+        spark_version = self._ws.clusters.select_spark_version(latest=True)
+        for cluster in cluster_id_list:
+            try:
+                cluster_details = self._ws.clusters.get(cluster)
+                if cluster_details.data_security_mode is None:
+                    logger.info(f"Data security Mode is None. Skipping the remapping for the cluster: {cluster}")
+                    continue
+                access_mode = self._get_access_mode(cluster_details.data_security_mode.name)
+                self._installation.save(cluster_details, filename=f'backup/clusters/{cluster_details.cluster_id}.json')
+                self._ws.clusters.edit(
+                    cluster_id=cluster,
+                    cluster_name=cluster_details.cluster_name,
+                    spark_version=spark_version,
+                    num_workers=cluster_details.num_workers,
+                    spark_conf=cluster_details.spark_conf,
+                    spark_env_vars=cluster_details.spark_env_vars,
+                    data_security_mode=access_mode,
+                    node_type_id=cluster_details.node_type_id,
+                    autoscale=cluster_details.autoscale,
+                    policy_id=cluster_details.policy_id,
+                    autotermination_minutes=cluster_details.autotermination_minutes,
+                    custom_tags=cluster_details.custom_tags,
+                    init_scripts=cluster_details.init_scripts,
+                    cluster_log_conf=cluster_details.cluster_log_conf,
+                    aws_attributes=cluster_details.aws_attributes,
+                    ssh_public_keys=cluster_details.ssh_public_keys,
+                    enable_elastic_disk=cluster_details.enable_elastic_disk,
+                    cluster_source=cluster_details.cluster_source,
+                    instance_pool_id=cluster_details.instance_pool_id,
+                    enable_local_disk_encryption=cluster_details.enable_local_disk_encryption,
+                    driver_instance_pool_id=cluster_details.driver_instance_pool_id,
+                )
+            except InvalidParameterValue as e:
+                logger.warning(f"skipping cluster remapping: {e}")
+
+    def revert_cluster_remap(self, cluster_ids: str, total_cluster_ids: list):
+        if cluster_ids != "<ALL>":
+            cluster_list = [x.strip() for x in cluster_ids.split(",")]
+        else:
+            cluster_list = total_cluster_ids
+        logger.info(f"Reverting the configurations for the cluster {cluster_list}")
+        for cluster in cluster_list:
+            try:
+                cluster_details = self._installation.load(ClusterDetails, filename=f"/backup/clusters/{cluster}.json")
+                if cluster_details.spark_version is None:
+                    raise InvalidParameterValue("cluster does not have spark version")
+                if cluster_details.cluster_id is None:
+                    raise InvalidParameterValue("cluster Id is not present in the config file")
+                num_workers = cluster_details.num_workers if cluster_details.num_workers else 0
+                self._ws.clusters.edit(
+                    cluster_id=cluster_details.cluster_id,
+                    cluster_name=cluster_details.cluster_name,
+                    spark_version=cluster_details.spark_version,
+                    num_workers=num_workers,
+                    spark_conf=cluster_details.spark_conf,
+                    spark_env_vars=cluster_details.spark_env_vars,
+                    data_security_mode=cluster_details.data_security_mode,
+                    node_type_id=cluster_details.node_type_id,
+                    autoscale=cluster_details.autoscale,
+                    policy_id=cluster_details.policy_id,
+                    autotermination_minutes=cluster_details.autotermination_minutes,
+                    custom_tags=cluster_details.custom_tags,
+                    init_scripts=cluster_details.init_scripts,
+                    cluster_log_conf=cluster_details.cluster_log_conf,
+                    aws_attributes=cluster_details.aws_attributes,
+                    ssh_public_keys=cluster_details.ssh_public_keys,
+                    enable_elastic_disk=cluster_details.enable_elastic_disk,
+                    cluster_source=cluster_details.cluster_source,
+                    instance_pool_id=cluster_details.instance_pool_id,
+                    enable_local_disk_encryption=cluster_details.enable_local_disk_encryption,
+                    driver_instance_pool_id=cluster_details.driver_instance_pool_id,
+                )
+            except InvalidParameterValue as e:
+                logger.warning(f"skipping cluster remapping: {e}")

tests/unit/test_cli.py

@@ -5,16 +5,20 @@
 
 import pytest
 import yaml
+from databricks.labs.blueprint.installation import Installation
 from databricks.labs.blueprint.tui import MockPrompts
 from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.errors import NotFound
 from databricks.sdk.service import iam, sql
+from databricks.sdk.service.compute import ClusterDetails
+from databricks.sdk.service.workspace import ObjectInfo
 
 from databricks.labs.ucx.assessment.aws import AWSResources
 from databricks.labs.ucx.aws.access import AWSResourcePermissions
 from databricks.labs.ucx.azure.access import AzureResourcePermissions
 from databricks.labs.ucx.cli import (
     alias,
+    cluster_remap,
     create_account_groups,
     create_catalogs_schemas,
     create_table_mapping,
@@ -28,6 +32,7 @@
     open_remote_config,
     principal_prefix_access,
     repair_run,
+    revert_cluster_remap,
     revert_migrated_tables,
     skip,
     sync_workspace_info,
@@ -439,3 +444,28 @@ def test_create_catalogs_schemas(ws):
     prompts = MockPrompts({'.*': 's3://test'})
     create_catalogs_schemas(ws, prompts)
     ws.catalogs.list.assert_called_once()
+
+
+def test_cluster_remap(ws, caplog):
+    prompts = MockPrompts({"Please provide the cluster id's as comma separated value from the above list.*": "1"})
+    ws = create_autospec(WorkspaceClient)
+    ws.clusters.get.return_value = ClusterDetails(cluster_id="123", cluster_name="test_cluster")
+    installation = create_autospec(Installation)
+    installation.save.return_value = "a/b/c"
+    cluster_remap(ws, prompts)
+    assert "Remapping the Clusters to UC" in caplog.messages
+
+
+def test_revert_cluster_remap(ws, caplog, mocker):
+    prompts = MockPrompts({"Please provide the cluster id's as comma separated value from the above list.*": "1"})
+    ws = create_autospec(WorkspaceClient)
+    ws.workspace.list.return_value = [ObjectInfo(path='/ucx/backup/clusters/123.json')]
+    with pytest.raises(TypeError):
+        revert_cluster_remap(ws, prompts)
+
+
+def test_revert_cluster_remap_empty(ws, caplog):
+    prompts = MockPrompts({"Please provide the cluster id's as comma separated value from the above list.*": "1"})
+    ws = create_autospec(WorkspaceClient)
+    revert_cluster_remap(ws, prompts)
+    assert "There is no cluster files in the backup folder.Skipping the reverting process" in caplog.messages

tests/unit/workspace_access/test_clusters.py

@@ -0,0 +1,41 @@
+from unittest.mock import create_autospec
+
+from databricks.labs.blueprint.installation import Installation
+from databricks.labs.blueprint.tui import MockPrompts
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.compute import ClusterDetails
+
+from databricks.labs.ucx.workspace_access.clusters import ClusterAccess
+
+
+def test_map_cluster_to_uc():
+    ws = create_autospec(WorkspaceClient)
+    ws.clusters.get.return_value = ClusterDetails(cluster_id="123", cluster_name="test_cluster")
+    prompts = MockPrompts({})
+    installation = create_autospec(Installation)
+    installation.save.return_value = "a/b/c"
+    cluster = ClusterAccess(installation, ws, prompts)
+    cluster.map_cluster_to_uc(cluster_id="123")
+
+
+def test_map_cluster_to_uc_error(caplog):
+    ws = create_autospec(WorkspaceClient)
+    ws.clusters.get.return_value = ClusterDetails(cluster_id="123", cluster_name="test_cluster")
+    prompts = MockPrompts({})
+    installation = create_autospec(Installation)
+    installation.save.return_value = "a/b/c"
+    cluster = ClusterAccess(installation, ws, prompts)
+    with caplog.at_level('INFO'):
+        cluster.map_cluster_to_uc("123")
+        assert 'Data security Mode is None. Skipping the remapping for the cluster: 123' in caplog.messages
+
+
+def test_revert_map_cluster_to_uc(caplog):
+    ws = create_autospec(WorkspaceClient)
+    installation = create_autospec(Installation)
+    prompts = MockPrompts({})
+    installation.load.return_value = ClusterDetails(
+        cluster_id="123", cluster_name="test_cluster", spark_version="13.3.x-cpu-ml-scala2.12"
+    )
+    cluster = ClusterAccess(installation, ws, prompts)
+    cluster.revert_cluster_remap(cluster_ids="123", total_cluster_ids=["123"])