..

nfx · nfx · commit ad02beb882fd · 2023-10-11T22:58:50.000+02:00
diff --git a/src/databricks/labs/ucx/hive_metastore/grants.py b/src/databricks/labs/ucx/hive_metastore/grants.py
@@ -5,6 +5,7 @@
 from functools import partial
 
 from databricks.sdk.service.catalog import SchemaInfo, TableInfo
+
 from databricks.labs.ucx.framework.crawlers import CrawlerBase
 from databricks.labs.ucx.framework.parallel import Threads
 from databricks.labs.ucx.hive_metastore.tables import TablesCrawler
diff --git a/src/databricks/labs/ucx/workspace_access/groups.py b/src/databricks/labs/ucx/workspace_access/groups.py
@@ -270,16 +270,23 @@ def ws_local_group_deletion_recovery(self):
             f"In total, {len(source_groups)} temporary groups found, which do not have corresponding workspace groups"
         )
 
-        for source_group in source_groups:
+        for backup_group in source_groups:
             ws_local_group = self._ws.groups.create(
-                display_name=source_group.display_name.removeprefix(self._backup_group_prefix),
-                meta=source_group.meta,
-                entitlements=source_group.entitlements,
-                roles=source_group.roles,
-                members=source_group.members,
+                display_name=backup_group.display_name.removeprefix(self._backup_group_prefix),
+                meta=backup_group.meta,
+                entitlements=backup_group.entitlements,
+                roles=backup_group.roles,
+                members=backup_group.members,
             )
             self._workspace_groups.append(ws_local_group)
-            self._migration_state.add(ws_local_group)
+
+            account_groups = [_ for _ in self._account_groups if _.display_name == ws_local_group.display_name]
+            if len(account_groups) == 0:
+                logger.error(f"Cannot find matching group on account level: {ws_local_group.display_name}")
+                continue
+
+            account_group = account_groups[0]
+            self._migration_state.add(ws_local_group, backup_group, account_group)
             logger.info(f"Workspace-local group {ws_local_group} successfully recovered")
 
-        logger.info("Workspace-local group deletion recovery completed")
+        logger.info("Workspace-local group recovery completed")
diff --git a/tests/integration/workspace_access/test_groups.py b/tests/integration/workspace_access/test_groups.py
@@ -6,8 +6,7 @@
 from databricks.labs.ucx.config import GroupsConfig
 from databricks.labs.ucx.hive_metastore import GrantsCrawler, TablesCrawler
 from databricks.labs.ucx.workspace_access.generic import (
-    GenericPermissionsSupport,
-    listing_wrapper,
+    GenericPermissionsSupport, Listing,
 )
 from databricks.labs.ucx.workspace_access.groups import GroupManager
 from databricks.labs.ucx.workspace_access.manager import PermissionManager
@@ -88,11 +87,9 @@ def test_replace_workspace_groups_with_account_groups(
     group_info = group_manager.migration_state.get_by_workspace_group_name(ws_group.display_name)
 
     generic_permissions = GenericPermissionsSupport(
-        ws, [listing_wrapper(ws.cluster_policies.list, "policy_id", "cluster-policies")]
-    )
-    permission_manager = PermissionManager(
-        sql_backend, inventory_schema, [generic_permissions], {"cluster-policies": generic_permissions}
+        ws, [Listing(ws.cluster_policies.list, "policy_id", "cluster-policies")]
     )
+    permission_manager = PermissionManager(sql_backend, inventory_schema, [generic_permissions])
     tables = TablesCrawler(sql_backend, inventory_schema)
     grants = GrantsCrawler(tables)
 
@@ -167,10 +164,8 @@ def test_recover_from_ws_local_deletion(ws, make_ucx_group):
     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:
+    for gi in group_manager.migration_state.groups:
         recovered_state[gi.workspace.display_name] = gi.workspace
 
     assert sorted([member.display for member in ws_group.members]) == sorted(
diff --git a/tests/unit/workspace_access/test_base.py b/tests/unit/workspace_access/test_base.py
@@ -10,7 +10,6 @@
 )
 from databricks.labs.ucx.workspace_access.groups import (
     GroupMigrationState,
-    MigrationGroupInfo,
 )
 
 
diff --git a/tests/unit/workspace_access/test_tacl.py b/tests/unit/workspace_access/test_tacl.py
@@ -51,13 +51,7 @@ def test_tacl_applier(mocker):
         ),
     )
     migration_state = GroupMigrationState()
-    migration_state.add(
-        MigrationGroupInfo(
-            workspace=Group(display_name="abc"),
-            backup=Group(display_name="tmp-backup-abc"),
-            account=Group(display_name="account-abc"),
-        )
-    )
+    migration_state.add(Group(display_name="abc"), Group(display_name="tmp-backup-abc"), Group(display_name="account-abc"))
     task = table_acl_support.get_apply_task(permissions, migration_state, "backup")
     task()
 
@@ -83,12 +77,9 @@ def test_tacl_applier_no_target_principal(mocker):
     )
     migration_state = GroupMigrationState()
     migration_state.add(
-        MigrationGroupInfo(
-            workspace=Group(display_name="abc"),
-            backup=Group(display_name="tmp-backup-abc"),
-            account=Group(display_name="account-abc"),
-        )
-    )
+        Group(display_name="abc"),
+        Group(display_name="tmp-backup-abc"),
+        Group(display_name="account-abc"))
     task = table_acl_support.get_apply_task(permissions, migration_state, "backup")
     assert task is None
 

Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,6 @@`
`10`	`10`	`)`
`11`	`11`	`from databricks.labs.ucx.workspace_access.groups import (`
`12`	`12`	`GroupMigrationState,`
`13`		`- MigrationGroupInfo,`
`14`	`13`	`)`
`15`	`14`
`16`	`15`