Split migrate-groups workflow into three different stages for reliability (#442)

Author: nfx

src/databricks/labs/ucx/mixins/fixtures.py

@@ -21,27 +21,27 @@
 
 from databricks.labs.ucx.framework.crawlers import StatementExecutionBackend
 
-_LOG = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
 
 
 def factory(name, create, remove):
     cleanup = []
 
     def inner(**kwargs):
         x = create(**kwargs)
-        _LOG.debug(f"added {name} fixture: {x}")
+        logger.debug(f"added {name} fixture: {x}")
         cleanup.append(x)
         return x
 
     yield inner
-    _LOG.debug(f"clearing {len(cleanup)} {name} fixtures")
+    logger.debug(f"clearing {len(cleanup)} {name} fixtures")
     for x in cleanup:
         try:
-            _LOG.debug(f"removing {name} fixture: {x}")
+            logger.debug(f"removing {name} fixture: {x}")
             remove(x)
         except DatabricksError as e:
             # TODO: fix on the databricks-labs-pytester level
-            _LOG.debug(f"ignoring error while {name} {x} teardown: {e}")
+            logger.debug(f"ignoring error while {name} {x} teardown: {e}")
 
 
 @pytest.fixture
@@ -122,7 +122,7 @@ def acc(product_info, debug_env) -> AccountClient:
     # Use variables from Unified Auth
     # See https://databricks-sdk-py.readthedocs.io/en/latest/authentication.html
     product_name, product_version = product_info
-    _LOG.debug(f"Running with {len(debug_env)} env variables")
+    logger.debug(f"Running with {len(debug_env)} env variables")
     return AccountClient(
         host=debug_env["DATABRICKS_HOST"],
         account_id=debug_env["DATABRICKS_ACCOUNT_ID"],
@@ -408,7 +408,7 @@ def _scim_values(ids: list[str]) -> list[iam.ComplexValue]:
     return [iam.ComplexValue(value=x) for x in ids]
 
 
-def _make_group(name, interface, make_random):
+def _make_group(name, cfg, interface, make_random):
     def create(
         *,
         members: list[str] | None = None,
@@ -425,19 +425,24 @@ def create(
         if entitlements is not None:
             kwargs["entitlements"] = _scim_values(entitlements)
         # TODO: REQUEST_LIMIT_EXCEEDED: GetUserPermissionsRequest RPC token bucket limit has been exceeded.
-        return interface.create(**kwargs)
+        group = interface.create(**kwargs)
+        if cfg.is_account_client:
+            logger.info(f"Account group {group.display_name}: {cfg.host}/users/groups/{group.id}/members")
+        else:
+            logger.info(f"Workspace group {group.display_name}: {cfg.host}#setting/accounts/groups/{group.id}")
+        return group
 
     yield from factory(name, create, lambda item: interface.delete(item.id))
 
 
 @pytest.fixture
 def make_group(ws, make_random):
-    yield from _make_group("workspace group", ws.groups, make_random)
+    yield from _make_group("workspace group", ws.config, ws.groups, make_random)
 
 
 @pytest.fixture
 def make_acc_group(acc, make_random):
-    yield from _make_group("account group", acc.groups, make_random)
+    yield from _make_group("account group", acc.config, acc.groups, make_random)
 
 
 @pytest.fixture
@@ -449,7 +454,11 @@ def create(*, name: str | None = None, **kwargs):
             kwargs["definition"] = json.dumps(
                 {"spark_conf.spark.databricks.delta.preview.enabled": {"type": "fixed", "value": "true"}}
             )
-        return ws.cluster_policies.create(name, **kwargs)
+        cluster_policy = ws.cluster_policies.create(name, **kwargs)
+        logger.info(
+            f"Cluster policy: {ws.config.host}#setting/clusters/cluster-policies/view/{cluster_policy.policy_id}"
+        )
+        return cluster_policy
 
     yield from factory("cluster policy", create, lambda item: ws.cluster_policies.delete(item.policy_id))
 
@@ -565,7 +574,9 @@ def create(**kwargs):
                         timeout_seconds=0,
                     )
                 ]
-        return ws.jobs.create(**kwargs)
+        job = ws.jobs.create(**kwargs)
+        logger.info(f"Job: {ws.config.host}#job/{job.job_id}")
+        return job
 
     yield from factory("job", create, lambda item: ws.jobs.delete(item.job_id))
 
@@ -726,13 +737,18 @@ def create() -> CatalogInfo:
 
 
 @pytest.fixture
-def make_schema(sql_backend, make_random) -> Callable[..., SchemaInfo]:
+def make_schema(ws, sql_backend, make_random) -> Callable[..., SchemaInfo]:
     def create(*, catalog_name: str = "hive_metastore", name: str | None = None) -> SchemaInfo:
         if name is None:
             name = f"ucx_S{make_random(4)}"
         full_name = f"{catalog_name}.{name}".lower()
         sql_backend.execute(f"CREATE SCHEMA {full_name}")
-        return SchemaInfo(catalog_name=catalog_name, name=name, full_name=full_name)
+        schema_info = SchemaInfo(catalog_name=catalog_name, name=name, full_name=full_name)
+        logger.info(
+            f"Schema {schema_info.full_name}: "
+            f"{ws.config.host}/explore/data/{schema_info.catalog_name}/{schema_info.name}"
+        )
+        return schema_info
 
     yield from factory(
         "schema",
@@ -742,7 +758,7 @@ def create(*, catalog_name: str = "hive_metastore", name: str | None = None) ->
 
 
 @pytest.fixture
-def make_table(sql_backend, make_schema, make_random) -> Callable[..., TableInfo]:
+def make_table(ws, sql_backend, make_schema, make_random) -> Callable[..., TableInfo]:
     def create(
         *,
         catalog_name="hive_metastore",
@@ -781,7 +797,12 @@ def create(
             ddl = f"{ddl} TBLPROPERTIES ({tbl_properties})"
 
         sql_backend.execute(ddl)
-        return TableInfo(catalog_name=catalog_name, schema_name=schema_name, name=name, full_name=full_name)
+        table_info = TableInfo(catalog_name=catalog_name, schema_name=schema_name, name=name, full_name=full_name)
+        logger.info(
+            f"Table {table_info.full_name}: "
+            f"{ws.config.host}/explore/data/{table_info.catalog_name}/{table_info.schema_name}/{table_info.name}"
+        )
+        return table_info
 
     def remove(table_info: TableInfo):
         try:

src/databricks/labs/ucx/runtime.py

@@ -203,13 +203,11 @@ def assessment_report(_: WorkspaceConfig):
     dashboard _before_ all tasks have been completed, but then only already completed information is shown."""
 
 
-@task("migrate-groups", depends_on=[crawl_permissions], job_cluster="tacl")
-def migrate_permissions(cfg: WorkspaceConfig):
-    """Main phase of the group migration process. It does the following:
+@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
-      - Creates an account-level group with the original name of the workspace-local one
-      - Assigns the full set of permissions of the original group to the account-level 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
@@ -231,13 +229,55 @@ def migrate_permissions(cfg: WorkspaceConfig):
         num_threads=cfg.num_threads,
         workspace_start_path=cfg.workspace_start_path,
     )
-
     permission_manager.apply_group_permissions(group_manager.migration_state, destination="backup")
+
+
+@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"""
+    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
     group_manager.replace_workspace_groups_with_account_groups()
+
+
+@task(
+    "004-apply-permissions-to-account-groups",
+    depends_on=[replace_workspace_groups_with_account_groups],
+    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
+
+    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/)."""
+    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
+
+    permission_manager = PermissionManager.factory(
+        ws,
+        RuntimeBackend(),
+        cfg.inventory_database,
+        num_threads=cfg.num_threads,
+        workspace_start_path=cfg.workspace_start_path,
+    )
     permission_manager.apply_group_permissions(group_manager.migration_state, destination="account")
 
 
-@task("migrate-groups-cleanup", depends_on=[migrate_permissions])
+@task("005-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
@@ -247,7 +287,7 @@ def delete_backup_groups(cfg: WorkspaceConfig):
     group_manager.delete_backup_groups()
 
 
-@task("destroy-schema")
+@task("099-destroy-schema")
 def destroy_schema(cfg: WorkspaceConfig):
     """This _clean-up_ workflow allows to removes the `$inventory` database, with all the inventory tables created by
     the previous workflow runs. Use this to reset the entire state and start with the assessment step again."""

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

@@ -1,3 +1,4 @@
+import collections
 import json
 import logging
 import typing
@@ -110,6 +111,28 @@ def delete_backup_groups(self):
             self._delete_workspace_group(group)
         logger.info("Backup groups were successfully deleted")
 
+    def get_workspace_membership(self, resource_type: str = "WorkspaceGroup"):
+        membership = collections.defaultdict(set)
+        for g in self._ws.groups.list(attributes=self._SCIM_ATTRIBUTES):
+            if g.display_name in self._SYSTEM_GROUPS:
+                continue
+            if g.meta.resource_type != resource_type:
+                continue
+            if g.members is None:
+                continue
+            for m in g.members:
+                membership[g.display_name].add(m.display)
+        return membership
+
+    def get_account_membership(self):
+        membership = collections.defaultdict(set)
+        for g in self._account_groups:
+            if g.members is None:
+                continue
+            for m in g.members:
+                membership[g.display_name].add(m.display)
+        return membership
+
     def _list_workspace_groups(self) -> list[iam.Group]:
         logger.info("Listing workspace groups...")
         workspace_groups = [
@@ -159,7 +182,7 @@ def _get_or_create_backup_group(self, source_group_name: str, source_group: iam.
             entitlements=source_group.entitlements,
             roles=source_group.roles,
             members=source_group.members,
-        )
+        )  # TODO: there still could be a corner case, where we get `Group with name db-temp-XXX already exists.`
         self._workspace_groups.append(backup_group)
         logger.info(f"Backup group {backup_group_name} successfully created")