From 11753d9b16492b5b3dd65fb7e592a3ebb22f8a72 Mon Sep 17 00:00:00 2001 From: Serge Smertin <259697+nfx@users.noreply.github.com> Date: Wed, 11 Oct 2023 16:47:14 +0200 Subject: [PATCH] Added recovery mode for workspace-local groups from temporary groups (#435) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This fix helps to recover from a state where - the temporary groups were created, - the account groups were not, - the original workspace-local groups were deleted. - The fix recreates the original state, with recreating all the workspace-local groups, which does have a temp pair but not an account level pair attached to the workspace. This helps to recreate the state of the migration with the already existing group_manager.migration_groups_provider --------- Co-authored-by: Fanni Jakó --- docs/local-group-migration.md | 10 +++++ .../labs/ucx/workspace_access/groups.py | 27 ++++++++++++ .../workspace_access/test_groups.py | 41 +++++++++++++++++++ 3 files changed, 78 insertions(+) diff --git a/docs/local-group-migration.md b/docs/local-group-migration.md index 50c46bd857..9843d013a4 100644 --- a/docs/local-group-migration.md +++ b/docs/local-group-migration.md @@ -238,4 +238,14 @@ ws_group_names = {_.display_name for _ in workspace_groups} ac_group_names = {_.display_name for _ in account_groups} group_names = list(ws_group_names.intersection(ac_group_names)) print(f"Found {len(group_names)} groups to migrate") +``` + +2. Recover workspace-local groups from backup groups from within a debug notebook: + +```python +from databricks.labs.ucx.workspace_access.groups import GroupManager +from databricks.labs.ucx.config import GroupsConfig + +group_manager = GroupManager(ws, GroupsConfig(auto=True)) +group_manager.ws_local_group_deletion_recovery() ``` \ No newline at end of file diff --git a/src/databricks/labs/ucx/workspace_access/groups.py b/src/databricks/labs/ucx/workspace_access/groups.py index 9eaec67495..7ca25cf0c9 100644 --- a/src/databricks/labs/ucx/workspace_access/groups.py +++ b/src/databricks/labs/ucx/workspace_access/groups.py @@ -272,3 +272,30 @@ def delete_backup_groups(self): self._delete_workspace_group(group) logger.info("Backup groups were successfully deleted") + + def ws_local_group_deletion_recovery(self): + workspace_groups = {_.display_name for _ in self._workspace_groups} + source_groups = [ + g + for g in self._workspace_groups + if g.display_name.removeprefix(self.config.backup_group_prefix) not in workspace_groups + ] + + logger.info( + f"Recovering from workspace-local group deletion. " + f"In total, {len(source_groups)} temporary groups found, which do not have corresponding workspace groups" + ) + + for source_group in source_groups: + ws_local_group = self._ws.groups.create( + display_name=source_group.display_name.removeprefix(self.config.backup_group_prefix), + meta=source_group.meta, + entitlements=source_group.entitlements, + roles=source_group.roles, + members=source_group.members, + ) + self._workspace_groups.append(ws_local_group) + self._migration_state.add(ws_local_group) + logger.info(f"Workspace-local group {ws_local_group} successfully recovered") + + logger.info("Workspace-local group deletion recovery completed") diff --git a/tests/integration/workspace_access/test_groups.py b/tests/integration/workspace_access/test_groups.py index 035c219fa3..10985065b7 100644 --- a/tests/integration/workspace_access/test_groups.py +++ b/tests/integration/workspace_access/test_groups.py @@ -52,3 +52,44 @@ def test_id_validity(ws: WorkspaceClient, make_ucx_group): manager = GroupManager(ws, GroupsConfig(selected=[ws_group.display_name])) assert ws_group.id == manager._get_group(ws_group.display_name, "workspace").id assert acc_group.id == manager._get_group(acc_group.display_name, "account").id + + +def test_recover_from_ws_local_deletion(ws, make_ucx_group): + ws_group, _ = make_ucx_group() + ws_group_two, _ = make_ucx_group() + + group_manager = GroupManager(ws, GroupsConfig(auto=True)) + group_manager.prepare_groups_in_environment() + + # simulate disaster + ws.groups.delete(ws_group.id) + ws.groups.delete(ws_group_two.id) + + # recovery run from a debug notebook + group_manager = GroupManager(ws, GroupsConfig(auto=True)) + group_manager.ws_local_group_deletion_recovery() + + # normal run after from a job + group_manager = GroupManager(ws, GroupsConfig(auto=True)) + group_manager.prepare_groups_in_environment() + + migration_state = group_manager.migration_groups_provider + + recovered_state = {} + for gi in migration_state.groups: + recovered_state[gi.workspace.display_name] = gi.workspace + + assert sorted([member.display for member in ws_group.members]) == sorted( + [member.display for member in recovered_state[ws_group.display_name].members] + ) + assert sorted([member.display for member in ws_group_two.members]) == sorted( + [member.display for member in recovered_state[ws_group_two.display_name].members] + ) + + assert sorted([member.value for member in ws_group.members]) == sorted( + [member.value for member in recovered_state[ws_group.display_name].members] + ) + + assert sorted([member.value for member in ws_group_two.members]) == sorted( + [member.value for member in recovered_state[ws_group_two.display_name].members] + )