Use ws.groups.patch to rename groups instead of creating backup groups (#450)

we can use patch to rename groups, so that we can simplify the local
group migration step.
<img width="646" alt="image"
src="https://github.com/databrickslabs/ucx/assets/259697/b0f54013-3366-4fd2-a1a4-236c7f033f6b">
<img width="325" alt="image"
src="https://github.com/databrickslabs/ucx/assets/259697/3b2ea9e1-5b4c-4b34-94b4-8e25267fd5ef">

Closes #218

---------

Co-authored-by: William Conti <[email protected]>
Co-authored-by: william-conti <[email protected]>

Author: nfx

src/databricks/labs/ucx/config.py

@@ -9,27 +9,7 @@
 
 from databricks.labs.ucx.__about__ import __version__
 
-__all__ = ["GroupsConfig", "AccountConfig", "WorkspaceConfig"]
-
-
-# TODO: flatten down
-@dataclass
-class GroupsConfig:
-    selected: list[str] | None = None
-    auto: bool | None = None
-    backup_group_prefix: str | None = "db-temp-"
-
-    def __post_init__(self):
-        if not self.selected and self.auto is None:
-            msg = "Either selected or auto must be set"
-            raise ValueError(msg)
-        if self.selected and self.auto is False:
-            msg = "No selected groups provided, but auto-collection is disabled"
-            raise ValueError(msg)
-
-    @classmethod
-    def from_dict(cls, raw: dict):
-        return cls(**raw)
+__all__ = ["AccountConfig", "WorkspaceConfig"]
 
 
 @dataclass
@@ -97,7 +77,8 @@ def from_dict(cls, raw: dict):
 
 
 # Used to set the right expectation about configuration file schema
-_CONFIG_VERSION = 1
+# TODO: config versions migrate
+_CONFIG_VERSION = 2
 
 T = TypeVar("T")
 
@@ -121,16 +102,6 @@ def from_bytes(cls, raw: str) -> T:
     def from_file(cls, config_file: Path) -> T:
         return cls.from_bytes(config_file.read_text())
 
-    @classmethod
-    def _verify_version(cls, raw):
-        stored_version = raw.pop("version", None)
-        if stored_version != _CONFIG_VERSION:
-            msg = (
-                f"Unsupported config version: {stored_version}. "
-                f"UCX v{__version__} expects config version to be {_CONFIG_VERSION}"
-            )
-            raise ValueError(msg)
-
     def __post_init__(self):
         if self.connect is None:
             self.connect = ConnectConfig()
@@ -199,14 +170,19 @@ def to_account_client(self) -> AccountClient:
 @dataclass
 class WorkspaceConfig(_Config["WorkspaceConfig"]):
     inventory_database: str
-    groups: GroupsConfig
+
+    # Includes group names for migration. If not specified, all matching groups will be picked up
+    include_group_names: list[str] = None
+
+    renamed_group_prefix: str = "ucx-renamed-"
+
     instance_pool_id: str = None
     warehouse_id: str = None
     connect: ConnectConfig | None = None
     num_threads: int | None = 10
-    log_level: str | None = "INFO"
     database_to_catalog_mapping: dict[str, str] = None
     default_catalog: str = "ucx_default"
+    log_level: str = "INFO"
 
     # Starting path for notebooks and directories crawler
     workspace_start_path: str = "/"
@@ -215,24 +191,23 @@ class WorkspaceConfig(_Config["WorkspaceConfig"]):
 
     @classmethod
     def from_dict(cls, raw: dict):
-        cls._verify_version(raw)
-        return cls(
-            inventory_database=raw.get("inventory_database"),
-            groups=GroupsConfig.from_dict(raw.get("groups", {})),
-            connect=ConnectConfig.from_dict(raw.get("connect", {})),
-            instance_pool_id=raw.get("instance_pool_id", None),
-            warehouse_id=raw.get("warehouse_id", None),
-            num_threads=raw.get("num_threads", 10),
-            log_level=raw.get("log_level", "INFO"),
-            database_to_catalog_mapping=raw.get("database_to_catalog_mapping", None),
-            instance_profile=raw.get("instance_profile", None),
-            spark_conf=raw.get("spark_conf", None),
-            default_catalog=raw.get("default_catalog", "main"),
-            workspace_start_path=raw.get("workspace_start_path", "/"),
-        )
+        raw = cls._migrate_from_v1(raw)
+        return cls(**raw)
 
     def to_workspace_client(self) -> WorkspaceClient:
         return WorkspaceClient(config=self.to_databricks_config())
 
     def replace_inventory_variable(self, text: str) -> str:
         return text.replace("$inventory", f"hive_metastore.{self.inventory_database}")
+
+    @classmethod
+    def _migrate_from_v1(cls, raw: dict):
+        stored_version = raw.pop("version", None)
+        if stored_version == _CONFIG_VERSION:
+            return raw
+        if stored_version == 1:
+            raw["include_group_names"] = raw.get("groups", {"selected": []})["selected"]
+            raw["renamed_group_prefix"] = raw.get("groups", {"backup_group_prefix": "db-temp-"})["backup_group_prefix"]
+            return raw
+        msg = f"Unknown config version: {stored_version}"
+        raise ValueError(msg)

src/databricks/labs/ucx/install.py

@@ -30,7 +30,7 @@
     JobInfo,
     PipelineInfo,
 )
-from databricks.labs.ucx.config import GroupsConfig, WorkspaceConfig
+from databricks.labs.ucx.config import WorkspaceConfig
 from databricks.labs.ucx.framework.crawlers import (
     SchemaDeployer,
     SqlBackend,
@@ -379,10 +379,11 @@ def warehouse_type(_):
         groups_config_args = {
             "backup_group_prefix": backup_group_prefix,
         }
+
         if selected_groups != "<ALL>":
             groups_config_args["selected"] = [x.strip() for x in selected_groups.split(",")]
         else:
-            groups_config_args["auto"] = True
+            groups_config_args["selected"] = None
 
         # Checking for external HMS
         instance_profile = None
@@ -408,7 +409,8 @@ def warehouse_type(_):
 
         self._config = WorkspaceConfig(
             inventory_database=inventory_database,
-            groups=GroupsConfig(**groups_config_args),
+            include_group_names=groups_config_args["selected"],
+            renamed_group_prefix=groups_config_args["backup_group_prefix"],
             warehouse_id=warehouse_id,
             log_level=log_level,
             num_threads=num_threads,
@@ -718,7 +720,7 @@ def _job_clusters(self, names: set[str]):
                 num_workers=0,
             )
         )
-        if self._ws.config.is_aws:
+        if self._ws.config.is_aws and spec.aws_attributes is not None:
             aws_attributes = replace(spec.aws_attributes, instance_profile_arn=self._config.instance_profile)
             spec = replace(spec, aws_attributes=aws_attributes)
         if "main" in names:

src/databricks/labs/ucx/runtime.py

@@ -18,10 +18,7 @@
 from databricks.labs.ucx.hive_metastore.data_objects import ExternalLocationCrawler
 from databricks.labs.ucx.hive_metastore.mounts import Mounts
 from databricks.labs.ucx.workspace_access.generic import WorkspaceListing
-from databricks.labs.ucx.workspace_access.groups import (
-    GroupManager,
-    GroupMigrationState,
-)
+from databricks.labs.ucx.workspace_access.groups import GroupManager
 from databricks.labs.ucx.workspace_access.manager import PermissionManager
 
 logger = logging.getLogger(__name__)
@@ -194,10 +191,22 @@ def crawl_permissions(cfg: WorkspaceConfig):
     permission_manager.inventorize_permissions()
 
 
+@task("assessment")
+def crawl_groups(cfg: WorkspaceConfig):
+    """Scans all groups for the local group migration scope"""
+    sql_backend = RuntimeBackend()
+    ws = WorkspaceClient(config=cfg.to_databricks_config())
+    group_manager = GroupManager(
+        sql_backend, ws, cfg.inventory_database, cfg.include_group_names, cfg.renamed_group_prefix
+    )
+    group_manager.snapshot()
+
+
 @task(
     "assessment",
     depends_on=[
         crawl_grants,
+        crawl_groups,
         crawl_permissions,
         guess_external_locations,
         assess_jobs,
@@ -214,55 +223,29 @@ def assessment_report(_: WorkspaceConfig):
     dashboard _before_ all tasks have been completed, but then only already completed information is shown."""
 
 
-@task("002-apply-permissions-to-backup-groups", depends_on=[crawl_permissions], job_cluster="tacl")
-def apply_permissions_to_backup_groups(cfg: WorkspaceConfig):
-    """Second phase of the workspace-local group migration process. It does the following:
-      - Creates a backup of every workspace-local group, adding a prefix that can be set in the configuration
-      - Assigns the full set of permissions of the original group to the backup one
-
-    It covers local workspace-local permissions for all entities: Legacy Table ACLs, Entitlements,
-    AWS instance profiles, Clusters, Cluster policies, Instance Pools, Databricks SQL warehouses, Delta Live
-    Tables, Jobs, MLflow experiments, MLflow registry, SQL Dashboards & Queries, SQL Alerts, Token and Password usage
-    permissions, Secret Scopes, Notebooks, Directories, Repos, Files.
-
-    See [interactive tutorial here](https://app.getreprise.com/launch/myM3VNn/)."""
+@task("migrate-groups", depends_on=[crawl_groups])
+def rename_workspace_local_groups(cfg: WorkspaceConfig):
+    """Renames workspace local groups by adding `ucx-renamed-` prefix."""
+    sql_backend = RuntimeBackend()
     ws = WorkspaceClient(config=cfg.to_databricks_config())
-    group_manager = GroupManager(ws, cfg.groups)
-    group_manager.prepare_groups_in_environment()
-    if not group_manager.has_groups():
-        logger.info("Skipping group migration as no groups were found.")
-        return
-
-    backend = RuntimeBackend()
-    permission_manager = PermissionManager.factory(
-        ws,
-        backend,
-        cfg.inventory_database,
-        num_threads=cfg.num_threads,
-        workspace_start_path=cfg.workspace_start_path,
+    group_manager = GroupManager(
+        sql_backend, ws, cfg.inventory_database, cfg.include_group_names, cfg.renamed_group_prefix
     )
-    permission_manager.apply_group_permissions(group_manager.migration_state, destination="backup")
-    group_manager.migration_state.persist_migration_state(backend, cfg.inventory_database)
+    group_manager.rename_groups()
 
 
-@task("003-replace-workspace-local-with-account-groups", depends_on=[apply_permissions_to_backup_groups])
-def replace_workspace_groups_with_account_groups(cfg: WorkspaceConfig):
-    """Third phase of the workspace-local group migration process. It does the following:
-    - Creates an account-level group with the original name of the workspace-local one"""
+@task("migrate-groups", depends_on=[rename_workspace_local_groups])
+def reflect_account_groups_on_workspace(cfg: WorkspaceConfig):
+    """Adds matching account groups to this workspace."""
+    sql_backend = RuntimeBackend()
     ws = WorkspaceClient(config=cfg.to_databricks_config())
-    group_manager = GroupManager(ws, cfg.groups)
-    remote_state = GroupMigrationState().fetch_migration_state(RuntimeBackend(), cfg.inventory_database)
-    if len(remote_state.groups) == 0:
-        logger.info("Skipping group migration as no groups were found.")
-        return
-    group_manager.replace_workspace_groups_with_account_groups(remote_state)
+    group_manager = GroupManager(
+        sql_backend, ws, cfg.inventory_database, cfg.include_group_names, cfg.renamed_group_prefix
+    )
+    group_manager.reflect_account_groups_on_workspace()
 
 
-@task(
-    "004-apply-permissions-to-account-groups",
-    depends_on=[replace_workspace_groups_with_account_groups],
-    job_cluster="tacl",
-)
+@task("migrate-groups", depends_on=[reflect_account_groups_on_workspace], job_cluster="tacl")
 def apply_permissions_to_account_groups(cfg: WorkspaceConfig):
     """Fourth phase of the workspace-local group migration process. It does the following:
       - Assigns the full set of permissions of the original group to the account-level one
@@ -273,13 +256,11 @@ def apply_permissions_to_account_groups(cfg: WorkspaceConfig):
     permissions, Secret Scopes, Notebooks, Directories, Repos, Files.
 
     See [interactive tutorial here](https://app.getreprise.com/launch/myM3VNn/)."""
-    ws = WorkspaceClient(config=cfg.to_databricks_config())
     backend = RuntimeBackend()
+    ws = WorkspaceClient(config=cfg.to_databricks_config())
+    group_manager = GroupManager(backend, ws, cfg.inventory_database, cfg.include_group_names, cfg.renamed_group_prefix)
 
-    remote_state = GroupMigrationState().fetch_migration_state(backend, cfg.inventory_database)
-    migration_state = GroupManager.prepare_apply_permissions_to_account_groups(
-        ws, remote_state, cfg.groups.backup_group_prefix
-    )
+    migration_state = group_manager.get_migration_state()
     if len(migration_state.groups) == 0:
         logger.info("Skipping group migration as no groups were found.")
         return
@@ -291,17 +272,18 @@ def apply_permissions_to_account_groups(cfg: WorkspaceConfig):
         num_threads=cfg.num_threads,
         workspace_start_path=cfg.workspace_start_path,
     )
-    permission_manager.apply_group_permissions(migration_state, destination="account")
+    permission_manager.apply_group_permissions(migration_state)
 
 
-@task("005-remove-workspace-local-backup-groups", depends_on=[apply_permissions_to_account_groups])
+@task("remove-workspace-local-backup-groups", depends_on=[apply_permissions_to_account_groups])
 def delete_backup_groups(cfg: WorkspaceConfig):
     """Last step of the group migration process. Removes all workspace-level backup groups, along with their
     permissions. Execute this workflow only after you've confirmed that workspace-local migration worked
     successfully for all the groups involved."""
+    backend = RuntimeBackend()
     ws = WorkspaceClient(config=cfg.to_databricks_config())
-    group_manager = GroupManager(ws, cfg.groups)
-    group_manager.delete_backup_groups()
+    group_manager = GroupManager(backend, ws, cfg.inventory_database, cfg.include_group_names, cfg.renamed_group_prefix)
+    group_manager.delete_original_workspace_groups()
 
 
 @task("099-destroy-schema")

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

@@ -2,9 +2,8 @@
 from collections.abc import Callable, Iterator
 from dataclasses import dataclass
 from logging import Logger
-from typing import Literal
 
-from databricks.labs.ucx.workspace_access.groups import GroupMigrationState
+from databricks.labs.ucx.workspace_access.groups import MigrationState
 
 logger = Logger(__name__)
 
@@ -17,9 +16,6 @@ class Permissions:
     raw: str
 
 
-Destination = Literal["backup", "account"]
-
-
 class AclSupport:
     @abstractmethod
     def get_crawler_tasks(self) -> Iterator[Callable[..., Permissions | None]]:
@@ -29,9 +25,7 @@ def get_crawler_tasks(self) -> Iterator[Callable[..., Permissions | None]]:
         """
 
     @abstractmethod
-    def get_apply_task(
-        self, item: Permissions, migration_state: GroupMigrationState, destination: Destination
-    ) -> Callable[[], None] | None:
+    def get_apply_task(self, item: Permissions, migration_state: MigrationState) -> Callable[[], None] | None:
         """This method returns a Callable, that applies permissions to a destination group, based on
         the group migration state. The callable is required not to have any shared mutable state."""
 

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

@@ -15,12 +15,8 @@
 
 from databricks.labs.ucx.framework.crawlers import CrawlerBase, SqlBackend
 from databricks.labs.ucx.mixins.hardening import rate_limited
-from databricks.labs.ucx.workspace_access.base import (
-    AclSupport,
-    Destination,
-    Permissions,
-)
-from databricks.labs.ucx.workspace_access.groups import GroupMigrationState
+from databricks.labs.ucx.workspace_access.base import AclSupport, Permissions
+from databricks.labs.ucx.workspace_access.groups import MigrationState
 
 logger = logging.getLogger(__name__)
 
@@ -81,15 +77,15 @@ def object_types(self) -> set[str]:
                 all_object_types.add(object_type)
         return all_object_types
 
-    def get_apply_task(self, item: Permissions, migration_state: GroupMigrationState, destination: Destination):
+    def get_apply_task(self, item: Permissions, migration_state: MigrationState):
         if not self._is_item_relevant(item, migration_state):
             return None
         object_permissions = iam.ObjectPermissions.from_dict(json.loads(item.raw))
-        new_acl = self._prepare_new_acl(object_permissions, migration_state, destination)
+        new_acl = self._prepare_new_acl(object_permissions, migration_state)
         return partial(self._applier_task, item.object_type, item.object_id, new_acl)
 
     @staticmethod
-    def _is_item_relevant(item: Permissions, migration_state: GroupMigrationState) -> bool:
+    def _is_item_relevant(item: Permissions, migration_state: MigrationState) -> bool:
         # passwords and tokens are represented on the workspace-level
         if item.object_id in ("tokens", "passwords"):
             return True
@@ -239,7 +235,7 @@ def _safe_update_permissions(
                 raise RetryableError(message=msg) from e
 
     def _prepare_new_acl(
-        self, permissions: iam.ObjectPermissions, migration_state: GroupMigrationState, destination: Destination
+        self, permissions: iam.ObjectPermissions, migration_state: MigrationState
     ) -> list[iam.AccessControlRequest]:
         _acl = permissions.access_control_list
         acl_requests = []
@@ -248,7 +244,7 @@ def _prepare_new_acl(
             if not migration_state.is_in_scope(_item.group_name):
                 logger.debug(f"Skipping {_item} for {coord} because it is not in scope")
                 continue
-            new_group_name = migration_state.get_target_principal(_item.group_name, destination)
+            new_group_name = migration_state.get_target_principal(_item.group_name)
             if new_group_name is None:
                 logger.debug(f"Skipping {_item.group_name} for {coord} because it has no target principal")
                 continue