♻️ Rename MigrationGroupsProvider to GroupMigrationState (#81)

This PR adds more clarity to what this container data structure is
responsible for.

src/uc_migration_toolkit/managers/group.py

@@ -7,8 +7,8 @@
 from uc_migration_toolkit.generic import StrEnum
 from uc_migration_toolkit.providers.client import ImprovedWorkspaceClient
 from uc_migration_toolkit.providers.groups_info import (
+    GroupMigrationState,
     MigrationGroupInfo,
-    MigrationGroupsProvider,
 )
 from uc_migration_toolkit.providers.logger import logger
 from uc_migration_toolkit.utils import ThreadedExecution
@@ -25,7 +25,7 @@ class GroupManager:
     def __init__(self, ws: ImprovedWorkspaceClient, groups: GroupsConfig):
         self._ws = ws
         self.config = groups
-        self._migration_groups_provider: MigrationGroupsProvider = MigrationGroupsProvider()
+        self._migration_state: GroupMigrationState = GroupMigrationState()
 
     # please keep the internal methods below this line
 
@@ -81,10 +81,10 @@ def get_group_info(name: str):
         executables = [partial(get_group_info, group_name) for group_name in groups_names]
 
         collected_groups = ThreadedExecution[MigrationGroupInfo](executables).run()
+        for g in collected_groups:
+            self._migration_state.add(g)
 
-        self._migration_groups_provider.groups = collected_groups
-
-        logger.info(f"Prepared {len(self._migration_groups_provider.groups)} groups for migration")
+        logger.info(f"Prepared {len(collected_groups)} groups for migration")
 
     def _replace_group(self, migration_info: MigrationGroupInfo):
         ws_group = migration_info.workspace
@@ -118,9 +118,9 @@ def prepare_groups_in_environment(self):
         logger.info("Environment prepared successfully")
 
     @property
-    def migration_groups_provider(self) -> MigrationGroupsProvider:
-        assert len(self._migration_groups_provider.groups) > 0, "Migration groups were not loaded or initialized"
-        return self._migration_groups_provider
+    def migration_groups_provider(self) -> GroupMigrationState:
+        assert len(self._migration_state.groups) > 0, "Migration groups were not loaded or initialized"
+        return self._migration_state
 
     def replace_workspace_groups_with_account_groups(self):
         logger.info("Replacing the workspace groups with account-level groups")

src/uc_migration_toolkit/managers/inventory/inventorizer.py

@@ -23,7 +23,7 @@
     RequestObjectType,
 )
 from uc_migration_toolkit.providers.client import ImprovedWorkspaceClient
-from uc_migration_toolkit.providers.groups_info import MigrationGroupsProvider
+from uc_migration_toolkit.providers.groups_info import GroupMigrationState
 from uc_migration_toolkit.providers.logger import logger
 from uc_migration_toolkit.utils import ProgressReporter, ThreadedExecution
 
@@ -305,15 +305,17 @@ class RolesAndEntitlementsInventorizer(BaseInventorizer[InventoryObject]):
     def logical_object_types(self) -> list[LogicalObjectType]:
         return [LogicalObjectType.ROLES, LogicalObjectType.ENTITLEMENTS]
 
-    def __init__(self, ws: ImprovedWorkspaceClient, migration_provider: MigrationGroupsProvider):
+    def __init__(self, ws: ImprovedWorkspaceClient, migration_state: GroupMigrationState):
         self._ws = ws
-        self._migration_provider = migration_provider
+        self._migration_state = migration_state
         self._group_info: list[Group] = []
 
     def preload(self):
         logger.info("Please note that group roles and entitlements will be ONLY inventorized for migration groups")
         self._group_info: list[Group] = [
-            self._ws.groups.get(id=g.workspace.id) for g in self._migration_provider.groups
+            # TODO: why do we load group twice from platform? this really looks unnecessary
+            self._ws.groups.get(id=g.workspace.id)
+            for g in self._migration_state.groups
         ]
         logger.info("Group roles and entitlements preload completed")
 
@@ -360,9 +362,9 @@ def inner() -> Iterator[ModelDatabricks]:
 
 class Inventorizers:
     @staticmethod
-    def provide(ws: ImprovedWorkspaceClient, migration_provider: MigrationGroupsProvider, num_threads: int):
+    def provide(ws: ImprovedWorkspaceClient, migration_state: GroupMigrationState, num_threads: int):
         return [
-            RolesAndEntitlementsInventorizer(ws, migration_provider),
+            RolesAndEntitlementsInventorizer(ws, migration_state),
             TokensAndPasswordsInventorizer(ws),
             StandardInventorizer(
                 ws,

src/uc_migration_toolkit/managers/inventory/permissions.py

@@ -19,7 +19,7 @@
     RolesAndEntitlements,
 )
 from uc_migration_toolkit.providers.client import ImprovedWorkspaceClient
-from uc_migration_toolkit.providers.groups_info import MigrationGroupsProvider
+from uc_migration_toolkit.providers.groups_info import GroupMigrationState
 from uc_migration_toolkit.providers.logger import logger
 from uc_migration_toolkit.utils import ThreadedExecution, safe_get_acls
 
@@ -47,6 +47,7 @@ class RolesAndEntitlementsRequestPayload:
 AnyRequestPayload = PermissionRequestPayload | SecretsPermissionRequestPayload | RolesAndEntitlementsRequestPayload
 
 
+# TODO: this class has too many @staticmethod and they must not be such. write a unit test for this logic.
 class PermissionManager:
     def __init__(self, ws: ImprovedWorkspaceClient, inventory_table_manager: InventoryTableManager):
         self._ws = ws
@@ -75,16 +76,19 @@ def inventorize_permissions(self):
     @staticmethod
     def __prepare_request_for_permissions_api(
         item: PermissionsInventoryItem,
-        migration_groups_provider: MigrationGroupsProvider,
+        migration_state: GroupMigrationState,
         destination: Literal["backup", "account"],
     ) -> PermissionRequestPayload:
         _existing_permissions: ObjectPermissions = item.typed_object_permissions
         _acl = _existing_permissions.access_control_list
         acl_requests = []
 
         for _item in _acl:
-            if _item.group_name in [g.workspace.display_name for g in migration_groups_provider.groups]:
-                migration_info = migration_groups_provider.get_by_workspace_group_name(_item.group_name)
+            # TODO: we have a double iteration over migration_state.groups
+            #  (also by migration_state.get_by_workspace_group_name).
+            #  Has to be be fixed by iterating just on .groups
+            if _item.group_name in [g.workspace.display_name for g in migration_state.groups]:
+                migration_info = migration_state.get_by_workspace_group_name(_item.group_name)
                 assert migration_info is not None, f"Group {_item.group_name} is not in the migration groups provider"
                 destination_group: Group = getattr(migration_info, destination)
                 _item.group_name = destination_group.display_name
@@ -110,7 +114,7 @@ def __prepare_request_for_permissions_api(
     @staticmethod
     def _prepare_permission_request_for_secrets_api(
         item: PermissionsInventoryItem,
-        migration_groups_provider: MigrationGroupsProvider,
+        migration_state: GroupMigrationState,
         destination: Literal["backup", "account"],
     ) -> SecretsPermissionRequestPayload:
         _existing_acl_container: AclItemsContainer = item.typed_object_permissions
@@ -121,8 +125,8 @@ def _prepare_permission_request_for_secrets_api(
         for _existing_acl in _existing_acl_container.acls:
             _new_acl = deepcopy(_existing_acl)
 
-            if _existing_acl.principal in [g.workspace.display_name for g in migration_groups_provider.groups]:
-                migration_info = migration_groups_provider.get_by_workspace_group_name(_existing_acl.principal)
+            if _existing_acl.principal in [g.workspace.display_name for g in migration_state.groups]:
+                migration_info = migration_state.get_by_workspace_group_name(_existing_acl.principal)
                 assert (
                     migration_info is not None
                 ), f"Group {_existing_acl.principal} is not in the migration groups provider"
@@ -139,27 +143,28 @@ def _prepare_permission_request_for_secrets_api(
 
     @staticmethod
     def __prepare_request_for_roles_and_entitlements(
-        item: PermissionsInventoryItem, migration_groups_provider: MigrationGroupsProvider, destination
+        item: PermissionsInventoryItem, migration_state: GroupMigrationState, destination
     ) -> RolesAndEntitlementsRequestPayload:
-        migration_info = migration_groups_provider.get_by_workspace_group_name(item.object_id)
+        # TODO: potential BUG - why does item.object_id hold a group name and not ID?
+        migration_info = migration_state.get_by_workspace_group_name(item.object_id)
         assert migration_info is not None, f"Group {item.object_id} is not in the migration groups provider"
         destination_group: Group = getattr(migration_info, destination)
         return RolesAndEntitlementsRequestPayload(payload=item.typed_object_permissions, group_id=destination_group.id)
 
     def _prepare_new_permission_request(
         self,
         item: PermissionsInventoryItem,
-        migration_groups_provider: MigrationGroupsProvider,
+        migration_state: GroupMigrationState,
         destination: Literal["backup", "account"],
     ) -> AnyRequestPayload:
         if isinstance(item.request_object_type, RequestObjectType) and isinstance(
             item.typed_object_permissions, ObjectPermissions
         ):
-            return self.__prepare_request_for_permissions_api(item, migration_groups_provider, destination)
+            return self.__prepare_request_for_permissions_api(item, migration_state, destination)
         elif item.logical_object_type == LogicalObjectType.SECRET_SCOPE:
-            return self._prepare_permission_request_for_secrets_api(item, migration_groups_provider, destination)
+            return self._prepare_permission_request_for_secrets_api(item, migration_state, destination)
         elif item.logical_object_type in [LogicalObjectType.ROLES, LogicalObjectType.ENTITLEMENTS]:
-            return self.__prepare_request_for_roles_and_entitlements(item, migration_groups_provider, destination)
+            return self.__prepare_request_for_roles_and_entitlements(item, migration_state, destination)
         else:
             logger.warning(
                 f"Unsupported permissions payload for object {item.object_id} "
@@ -217,17 +222,15 @@ def _apply_permissions_in_parallel(
         execution = ThreadedExecution[None](executables)
         execution.run()
 
-    def apply_group_permissions(
-        self, migration_groups_provider: MigrationGroupsProvider, destination: Literal["backup", "account"]
-    ):
+    def apply_group_permissions(self, migration_state: GroupMigrationState, destination: Literal["backup", "account"]):
         logger.info(f"Applying the permissions to {destination} groups")
-        logger.info(f"Total groups to apply permissions: {len(migration_groups_provider.groups)}")
+        logger.info(f"Total groups to apply permissions: {len(migration_state.groups)}")
 
         permissions_on_source = self.inventory_table_manager.load_for_groups(
-            groups=[g.workspace.display_name for g in migration_groups_provider.groups]
+            groups=[g.workspace.display_name for g in migration_state.groups]
         )
         permission_payloads: list[AnyRequestPayload] = [
-            self._prepare_new_permission_request(item, migration_groups_provider, destination=destination)
+            self._prepare_new_permission_request(item, migration_state, destination=destination)
             for item in permissions_on_source
         ]
         logger.info(f"Applying {len(permission_payloads)} permissions")

src/uc_migration_toolkit/providers/groups_info.py

@@ -10,8 +10,9 @@ class MigrationGroupInfo:
     account: Group
 
 
-# TODO: this is not a provider, this is a container, e.g. State or Cache
-class MigrationGroupsProvider:
+class GroupMigrationState:
+    """Holds migration state of workspace-to-account groups"""
+
     def __init__(self):
         self.groups: list[MigrationGroupInfo] = []
 

tests/integration/test_e2e.py

@@ -19,7 +19,7 @@
 )
 from uc_migration_toolkit.managers.inventory.types import RequestObjectType
 from uc_migration_toolkit.providers.client import ImprovedWorkspaceClient
-from uc_migration_toolkit.providers.groups_info import MigrationGroupsProvider
+from uc_migration_toolkit.providers.groups_info import GroupMigrationState
 from uc_migration_toolkit.providers.logger import logger
 from uc_migration_toolkit.toolkits.group_migration import GroupMigrationToolkit
 from uc_migration_toolkit.utils import safe_get_acls
@@ -146,11 +146,11 @@ def _verify_group_permissions(
 
 
 def _verify_roles_and_entitlements(
-    migration_provider: MigrationGroupsProvider,
+    migration_state: GroupMigrationState,
     ws: ImprovedWorkspaceClient,
     target: Literal["backup", "account"],
 ):
-    for el in migration_provider.groups:
+    for el in migration_state.groups:
         comparison_base = getattr(el, "workspace" if target == "backup" else "backup")
         comparison_target = getattr(el, target)
 

tests/unit/test_secrets_api.py

@@ -6,8 +6,8 @@
     PermissionsInventoryItem,
 )
 from uc_migration_toolkit.providers.groups_info import (
+    GroupMigrationState,
     MigrationGroupInfo,
-    MigrationGroupsProvider,
 )
 
 
@@ -23,21 +23,21 @@ def test_secrets_api():
         ]}""",
     )
 
-    groups_provider = MigrationGroupsProvider()
-    groups_provider.groups = [
+    migration_state = GroupMigrationState()
+    migration_state.groups = [
         MigrationGroupInfo(
             account=Group(display_name="g1"),
             workspace=Group(display_name="g1"),
             backup=Group(display_name="some-prefix-g1"),
         )
     ]
 
-    apply_backup = PermissionManager._prepare_permission_request_for_secrets_api(item, groups_provider, "backup")
+    apply_backup = PermissionManager._prepare_permission_request_for_secrets_api(item, migration_state, "backup")
 
     assert len(apply_backup.access_control_list) == 1
     assert apply_backup.access_control_list[0].principal == "some-prefix-g1"
 
-    apply_account = PermissionManager._prepare_permission_request_for_secrets_api(item, groups_provider, "account")
+    apply_account = PermissionManager._prepare_permission_request_for_secrets_api(item, migration_state, "account")
 
     assert len(apply_account.access_control_list) == 1
     assert apply_account.access_control_list[0].principal == "g1"