Group migration: improve robustness when renaming groups (#2263)

## Changes

This PR updates the way we check for the completion of group renames;
similar to deletion we now double-check before assuming the rename has
taken effect.

### Linked issues

Follows on from #2247.

### Functionality

- modified existing workflows:

  - `migrate-groups`
  - `migrate-groups-experimental`

### Tests

- updated unit tests
- existing integration tests

Author: asnare

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

@@ -437,11 +437,20 @@ def has_groups(self) -> bool:
         return len(self.snapshot()) > 0
 
     def rename_groups(self):
-        tasks = []
         account_groups_in_workspace = self._account_groups_in_workspace()
         workspace_groups_in_workspace = self._workspace_groups_in_workspace()
+        # Renaming a group is eventually consistent, and not monotonically consistent, with a rather long time to
+        # converge: internally the Databricks API caches some things for up to 60s. To avoid excessive wait times when
+        # large numbers of groups need to be deleted (some deployments have >10K groups) we use the following steps:
+        #  1. Rename all the groups.
+        #  2. Confirm for each group that direct GETs yield the new name.
+        #  3. Confirm that group enumeration no longer includes the deleted groups.
+        # This caution is necessary because otherwise downstream tasks (like reflect_account_groups_on_workspace()) may
+        # skip a renamed group because it doesn't appear to be present.
         groups_to_migrate = self.get_migration_state().groups
-
+        rename_tasks = []
+        waiting_tasks = []
+        renamed_groups = []
         logger.info(f"Starting to rename {len(groups_to_migrate)} groups for migration...")
         for migrated_group in groups_to_migrate:
             if migrated_group.name_in_account in account_groups_in_workspace:
@@ -450,51 +459,82 @@ def rename_groups(self):
             if migrated_group.temporary_name in workspace_groups_in_workspace:
                 logger.info(f"Skipping {migrated_group.name_in_workspace}: already renamed")
                 continue
-            logger.info(f"Renaming: {migrated_group.name_in_workspace} -> {migrated_group.temporary_name}")
-            tasks.append(
+            rename_tasks.append(
                 functools.partial(
-                    self._rename_group_and_wait_for_rename,
+                    self._rename_group,
                     migrated_group.id_in_workspace,
                     migrated_group.name_in_workspace,
                     migrated_group.temporary_name,
                 )
             )
-        renamed_groups = Threads.strict("rename groups in the workspace", tasks)
-        # Renaming is eventually consistent, and the tasks above have each polled to verify their rename completed.
-        # Here we also check that enumeration yields the updated names; this is necessary because otherwise downstream
-        # tasks (like reflect_account_groups_on_workspace()) may skip a renamed group because it doesn't appear to be
-        # present.
+            waiting_tasks.append(
+                functools.partial(
+                    self._wait_for_group_rename,
+                    migrated_group.id_in_workspace,
+                    migrated_group.name_in_workspace,
+                    migrated_group.temporary_name,
+                )
+            )
+            renamed_groups.append((migrated_group.id_in_workspace, migrated_group.temporary_name))
+        # Step 1: Rename all the groups.
+        _, errors = Threads.gather("rename groups in the workspace", rename_tasks)
+        if errors:
+            logger.error(f"During renaming of workspace groups {len(errors)} errors occurred. See debug logs.")
+            raise ManyError(errors)
+        # Step 2: Confirm that direct GETs yield the updated information.
+        _, errors = Threads.gather("waiting for renamed groups in the workspace", waiting_tasks)
+        if errors:
+            logger.error(f"While waiting for renamed workspace groups {len(errors)} errors occurred. See debug logs.")
+            raise ManyError(errors)
+        # Step 3: Wait for enumeration to also reflect the updated information.
         self._wait_for_renamed_groups(renamed_groups)
 
-    def _rename_group_and_wait_for_rename(self, group_id: str, old_group_name, new_group_name: str) -> tuple[str, str]:
-        logger.debug(f"Renaming group {group_id}: {old_group_name} -> {new_group_name}")
-        self._rename_group(group_id, new_group_name)
-        logger.debug(f"Waiting for group {group_id} rename to take effect: {old_group_name} -> {new_group_name}")
-        self._wait_for_rename(group_id, old_group_name, new_group_name)
-        return group_id, new_group_name
+    def _rename_group(self, group_id: str, old_group_name: str, new_group_name: str) -> None:
+        logger.debug(f"Renaming group: {old_group_name} (id={group_id}) -> {new_group_name}")
+        self._rate_limited_rename_group_with_retry(group_id, new_group_name)
 
     @retried(on=[InternalError, ResourceConflict, DeadlineExceeded])
     @rate_limited(max_requests=10, burst_period_seconds=60)
-    def _rename_group(self, group_id: str, new_group_name: str) -> None:
+    def _rate_limited_rename_group_with_retry(self, group_id: str, new_group_name: str) -> None:
         ops = [iam.Patch(iam.PatchOp.REPLACE, "displayName", new_group_name)]
         self._ws.groups.patch(group_id, operations=ops)
 
-    @retried(on=[GroupRenameIncompleteError], timeout=timedelta(minutes=2))
-    def _wait_for_rename(self, group_id: str, old_group_name: str, new_group_name: str) -> None:
+    @retried(on=[GroupRenameIncompleteError], timeout=timedelta(seconds=90))
+    def _wait_for_group_rename(self, group_id: str, old_group_name: str, new_group_name: str) -> None:
+        # The groups API is eventually consistent, but not monotonically consistent. Here we verify that the effects of
+        # the rename have taken effect, and try to compensate for the lack of monotonic consistency by requiring two
+        # subsequent calls to confirm the rename. REST API internals cache things for up to 60s, and we see times close
+        # to this during testing. The retry timeout reflects this: if it's taking much longer then something else is
+        # wrong.
+        self._check_group_rename(group_id, old_group_name, new_group_name, logging.DEBUG)
+        self._check_group_rename(group_id, old_group_name, new_group_name, logging.WARNING)
+        logger.debug(f"Group rename is assumed complete: {old_group_name} (id={group_id}) -> {new_group_name}")
+
+    def _check_group_rename(self, group_id, old_group_name: str, new_group_name: str, pending_log_level: int) -> None:
         group = self._ws.groups.get(group_id)
         if group.display_name == old_group_name:
-            logger.debug(
-                f"Group {group_id} still has old name; still waiting for rename to take effect: {old_group_name} -> {new_group_name}"
+            logger.log(
+                pending_log_level,
+                f"Group still has old name; still waiting for rename to take effect: {old_group_name} (id={group_id}) -> {new_group_name}",
             )
             raise GroupRenameIncompleteError(group_id, old_group_name, new_group_name)
         if group.display_name != new_group_name:
             # Group has an entirely unexpected name; something else is interfering.
-            msg = f"While waiting for group {group_id} rename ({old_group_name} -> {new_group_name}) an unexpected name was observed: {group.display_name}"
+            msg = f"While waiting for group rename ({old_group_name} (id={group_id}) -> {new_group_name}) an unexpected name was observed: {group.display_name}"
             raise RuntimeError(msg)
-        logger.debug(f"Group {group_id} rename has taken effect: {old_group_name} -> {new_group_name}")
+        logger.debug(f"Group rename has possibly taken effect: {old_group_name} (id={group_id}) -> {new_group_name}")
 
     @retried(on=[ManyError], timeout=timedelta(minutes=2))
     def _wait_for_renamed_groups(self, expected_groups: Collection[tuple[str, str]]) -> None:
+        # The groups API is eventually consistent, but not monotonically consistent. Here we verify that the group
+        # has been deleted, and try to compensate for the lack of monotonic consistency by requiring two subsequent
+        # calls to confirm deletion. REST API internals cache things for up to 60s, and we see times close to this
+        # during testing. The retry timeout reflects this: if it's taking much longer then something else is wrong.
+        self._check_for_renamed_groups(expected_groups, logging.DEBUG)
+        self._check_for_renamed_groups(expected_groups, logging.WARNING)
+        logger.debug(f"Group enumeration showed all {len(expected_groups)} renamed groups; assuming complete.")
+
+    def _check_for_renamed_groups(self, expected_groups: Collection[tuple[str, str]], pending_log_level: int) -> None:
         attributes = "id,displayName"
         found_groups = {
             group.id: group.display_name
@@ -508,8 +548,9 @@ def _wait_for_renamed_groups(self, expected_groups: Collection[tuple[str, str]])
                 logger.warning(f"Group enumeration omits renamed group: {group_id} (renamed to {expected_name})")
                 pending_renames.append(RuntimeError(f"Missing group with id: {group_id} (renamed to {expected_name})"))
             elif found_name != expected_name:
-                logger.debug(
-                    f"Group enumeration does not yet reflect rename: {group_id} (renamed to {expected_name} but currently {found_name})"
+                logger.log(
+                    pending_log_level,
+                    f"Group enumeration does not yet reflect rename: {group_id} (renamed to {expected_name} but currently {found_name})",
                 )
                 pending_renames.append(GroupRenameIncompleteError(group_id, found_name, expected_name))
             else:
@@ -544,11 +585,11 @@ def delete_original_workspace_groups(self):
         migrated_groups = self.snapshot()
         logger.info(f"Starting to remove {len(migrated_groups)} migrated workspace groups...")
         # Group deletion is eventually consistent, and not monotonically consistent, with a rather long time to
-        # converge: internally API caches some things for up to 60s. To avoid excessive wait times when large numbers of
-        # groups need to be deleted (some deployments have >10K groups) we use the following steps:
-        #  1. Delete the groups.
-        #  2. Confirm that direct GETs no longer see the group.
-        #  3. Confirm that account enumeration no longer includes the deleted groups.
+        # converge: internally the Databricks API caches some things for up to 60s. To avoid excessive wait times when
+        # large numbers of groups need to be deleted (some deployments have >10K groups) we use the following steps:
+        #  1. Delete all the groups.
+        #  2. Confirm for each group that direct GETs no longer see the group.
+        #  3. Confirm that group enumeration no longer includes the deleted groups.
         deletion_tasks = []
         waiting_tasks = []
         deleted_groups = []
@@ -573,12 +614,12 @@ def delete_original_workspace_groups(self):
             deleted_groups.append(migrated_group)
         # Step 1: Delete the groups.
         _, errors = Threads.gather("removing original workspace groups", deletion_tasks)
-        if len(errors) > 0:
+        if errors:
             logger.error(f"During deletion of workspace groups got {len(errors)} errors. See debug logs.")
             raise ManyError(errors)
-        # Step 2: Confirm that direct gets no longer return the deleted group.
+        # Step 2: Confirm that direct GETs no longer return the deleted group.
         _, errors = Threads.gather("waiting for removal of original workspace groups", waiting_tasks)
-        if len(errors) > 0:
+        if errors:
             logger.error(f"Waiting for deletion of workspace groups got {len(errors)} errors. See debug logs.")
             raise ManyError(errors)
         # Step 3: Confirm that enumeration no longer returns the deleted groups.
@@ -786,11 +827,11 @@ def _wait_for_workspace_group_deletion(self, group_id: str, display_name: str) -
         self._check_workspace_group_deletion(group_id, display_name, logging.WARNING)
         logger.debug(f"Workspace group is assumed deleted: {display_name} (id={group_id})")
 
-    def _check_workspace_group_deletion(self, group_id: str, display_name: str, still_present_level_level: int) -> None:
+    def _check_workspace_group_deletion(self, group_id: str, display_name: str, pending_log_level: int) -> None:
         try:
             _ = self._ws.groups.get(id=group_id)
             logger.log(
-                still_present_level_level,
+                pending_log_level,
                 f"Deleted group is still present; still waiting for deletion to take effect: {display_name} (id={group_id})",
             )
             # Deletion is still pending.

tests/unit/workspace_access/test_groups.py

@@ -269,24 +269,30 @@ def test_snapshot_should_rename_groups_defined_in_conf():
     ]
 
 
-def test_rename_groups_should_patch_eligible_groups():
+def test_rename_groups_should_patch_eligible_groups(fake_sleep: Mock) -> None:
     backend = MockBackend()
     wsclient = create_autospec(WorkspaceClient)
     group1 = Group(id="1", display_name="de", meta=ResourceMeta(resource_type="WorkspaceGroup"))
     updated_group1 = dataclasses.replace(group1, display_name="test-group-de")
-    wsclient.groups.list.side_effect = (
+    list_side_effects = (
         # Preparing for the rename.
         *[[group1]] * 3,
-        # Checking the rename completed
+        # Checking (twice) that enumeration reflects the rename.
+        [updated_group1],
         [updated_group1],
     )
-    wsclient.groups.get.side_effect = (
+    wsclient.groups.list.side_effect = list_side_effects
+    get_side_effects = (
         # Preparing for the rename.
         *[group1] * 2,
-        # Checking the rename completed.
+        # Checking that the rename completed: emulate "not yet".
+        updated_group1,
+        group1,
+        # Checking (twice) that the rename completed.
+        updated_group1,
         updated_group1,
     )
-    wsclient.groups.get.return_value = group1
+    wsclient.groups.get.side_effect = get_side_effects
     account_admins_group_1 = Group(id="11", display_name="de")
     wsclient.api_client.do.return_value = {
         "Resources": [g.as_dict() for g in (account_admins_group_1,)],
@@ -296,6 +302,9 @@ def test_rename_groups_should_patch_eligible_groups():
         "1",
         operations=[iam.Patch(iam.PatchOp.REPLACE, "displayName", "test-group-de")],
     )
+    assert wsclient.groups.list.call_count == len(list_side_effects)
+    assert wsclient.groups.get.call_count == len(get_side_effects)
+    fake_sleep.assert_called()
 
 
 def test_rename_groups_should_filter_account_groups_in_workspace():
@@ -337,10 +346,14 @@ def test_rename_groups_should_wait_for_renames_to_complete(fake_sleep: Mock) ->
         *[[original_group]] * 3,
         # Checking the rename completed; simulate:
         # 1. Rename incomplete: original group still visible.
-        # 2. Still incomplete: group missing entirely.
-        # 3. Rename has finished: updated group is now visible.
         [original_group],
+        # 2. Still incomplete: group missing entirely.
         [],
+        # 3. Maybe complete: update visible, but not twice in a row.
+        [updated_group],
+        [original_group],
+        # 4. Rename has finished: updated group is now visible (twice).
+        [updated_group],
         [updated_group],
     )
     wsclient.groups.list.side_effect = list_side_effect_values
@@ -355,9 +368,11 @@ def test_rename_groups_should_wait_for_renames_to_complete(fake_sleep: Mock) ->
     get_return_values = (
         # Enumerating the workspace groups and loading the migration state.
         *[original_group] * 2,
-        # First call after rename: simulate the rename not yet being visible.
+        # First calls after the rename: simulate the rename not yet being visible.
+        updated_group,
         original_group,
-        # Second call, simulate things now being complete.
+        # Final calls after the rename: simulate things now being complete.
+        updated_group,
         updated_group,
     )
     wsclient.groups.get.side_effect = get_return_values
@@ -380,7 +395,8 @@ def test_rename_groups_should_retry_on_internal_error(fake_sleep: Mock) -> None:
     wsclient.groups.list.side_effect = (
         # Preparing for the rename.
         *[[original_group]] * 3,
-        # Checking the rename completed.
+        # Checking (twice) that the rename completed.
+        [updated_group],
         [updated_group],
     )
     matching_account_admin_group = dataclasses.replace(
@@ -392,7 +408,8 @@ def test_rename_groups_should_retry_on_internal_error(fake_sleep: Mock) -> None:
     wsclient.groups.get.side_effect = (
         # Preparing for the rename.
         *[original_group] * 2,
-        # Checking the rename completed.
+        # Checking (twice) that the rename completed.
+        updated_group,
         updated_group,
     )
     # The response to the PATCH call is ignored; set things up to fail on the first call and succeed on the second.