Improved validating groups membership cli command (#816)

Author: mwojtyczka

labs.yml

@@ -107,8 +107,8 @@ commands:
   - name: validate-groups-membership
     description: Validate the groups to see if the groups at account level and workspace level have different membership
     table_template: |-
-      Workflow Group Name\tAccount Group Name
-      {{range .}}{{.name_in_workspace}}\t{{.name_in_group}}
+      Workspace Group Name\tMembers Count\tAccount Group Name\tMembers Count
+      {{range .}}{{.wf_group_name}}\t{{.wf_group_members_count}}\t{{.acc_group_name}}\t{{.acc_group_members_count}}
       {{end}}
 
   - name: save-aws-iam-profiles

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

@@ -409,19 +409,24 @@ def validate_group_membership(self) -> list[dict]:
         strategy = self._get_strategy(workspace_groups_in_workspace, account_groups_in_account)
         migrated_groups = strategy.generate_migrated_groups()
         mismatch_group = []
-        for groups in migrated_groups:
-            ws_members_set = set([m.get("display") for m in json.loads(groups.members)] if groups.members else [])
-            account_group = account_groups_in_account[groups.name_in_account]
-            ac_members_set = set(
-                [a.as_dict().get("display") for a in account_group.members] if account_group.members else []
-            )
-            set_diff = (ws_members_set - ac_members_set).union(ac_members_set - ws_members_set)
+        retry_on_internal_error = retried(on=[InternalError], timeout=self._verify_timeout)
+        get_account_group = retry_on_internal_error(self._get_account_group)
+        for ws_group in migrated_groups:
+            # Users with the same display name but different email will be deduplicated!
+            ws_members_set = {m.get("display") for m in json.loads(ws_group.members)} if ws_group.members else set()
+            acc_group = get_account_group(account_groups_in_account[ws_group.name_in_account].id)
+            if not acc_group:
+                continue  # group not present anymore
+            acc_members_set = {a.as_dict().get("display") for a in acc_group.members} if acc_group.members else set()
+            set_diff = (ws_members_set - acc_members_set).union(acc_members_set - ws_members_set)
             if not set_diff:
                 continue
             mismatch_group.append(
                 {
-                    "wf_group_name": groups.name_in_workspace,
-                    "ac_group_name": groups.name_in_account,
+                    "wf_group_name": ws_group.name_in_workspace,
+                    "wf_group_members_count": len(ws_members_set),
+                    "acc_group_name": ws_group.name_in_account,
+                    "acc_group_members_count": len(acc_members_set),
                 }
             )
         if not mismatch_group:
@@ -500,6 +505,16 @@ def _get_group(self, group_id: str) -> iam.Group | None:
             # which will cause timeout errors because of groups no longer there.
             return None
 
+    @rate_limited(max_requests=20)
+    def _get_account_group(self, group_id: str) -> Group | None:
+        try:
+            raw = self._ws.api_client.do("GET", f"/api/2.0/account/scim/v2/Groups/{group_id}")
+            return iam.Group.from_dict(raw)  # type: ignore[arg-type]
+        except NotFound:
+            # the given group has been removed from the account after getting the group and before running this method
+            logger.warning("Group with ID: %s does not exist anymore in the Databricks account.", group_id)
+            return None
+
     def _list_account_groups(self, scim_attributes: str) -> list[iam.Group]:
         # TODO: we should avoid using this method, as it's not documented
         # get account-level groups even if they're not (yet) assigned to a workspace

tests/integration/workspace_access/test_groups.py

@@ -238,6 +238,36 @@ def test_group_matching_names_with_diff_users(
     assert len(t) > 0
 
 
+@retried(on=[NotFound], timeout=timedelta(minutes=2))
+def test_group_matching_names_with_same_users(
+    ws, sql_backend, inventory_schema, make_random, make_user, make_group, make_acc_group
+):
+    rand_elem = make_random(4)
+    workspace_group_name = f"test_group_{rand_elem}"
+    account_group_name = f"same_group_[{rand_elem}]"
+    user1 = make_user()
+    members1 = [user1.id]
+    members2 = [user1.id]
+    ws_group = make_group(display_name=workspace_group_name, members=members1, entitlements=["allow-cluster-create"])
+    acc_group = make_acc_group(display_name=account_group_name, members=members2)
+
+    logger.info(
+        f"Attempting Mapping From Workspace Group {ws_group.display_name} to Account Group {acc_group.display_name}"
+    )
+    group_manager = GroupManager(
+        sql_backend,
+        ws,
+        inventory_schema,
+        [ws_group.display_name],
+        "ucx-temp-",
+        workspace_group_regex=r"([0-9a-zA-Z]*)$",
+        account_group_regex=r"\[([0-9a-zA-Z]*)\]",
+    )
+
+    t = group_manager.validate_group_membership()
+    assert len(t) == 0
+
+
 # average runtime is 100 seconds
 @retried(on=[NotFound], timeout=timedelta(minutes=3))
 def test_replace_workspace_groups_with_account_groups(

tests/unit/workspace_access/test_groups.py

@@ -6,7 +6,7 @@
 from databricks.labs.blueprint.parallel import ManyError
 from databricks.labs.blueprint.tui import MockPrompts
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.errors import DatabricksError, ResourceDoesNotExist
+from databricks.sdk.errors import DatabricksError, NotFound, ResourceDoesNotExist
 from databricks.sdk.service import iam
 from databricks.sdk.service.iam import ComplexValue, Group, ResourceMeta
 
@@ -836,23 +836,121 @@ def test_validate_group_diff_membership():
         external_id="1234",
         display_name="test_(1234)",
         meta=ResourceMeta(resource_type="WorkspaceGroup"),
-        members=[ComplexValue(display="test-user-1", value="20"), ComplexValue(display="test-user-2", value="21")],
+        members=[ComplexValue(display="test-user-1", value="1"), ComplexValue(display="test-user-2", value="2")],
         roles=[
             ComplexValue(value="arn:aws:iam::123456789098:instance-profile/ip1"),
             ComplexValue(value="arn:aws:iam::123456789098:instance-profile/ip2"),
         ],
         entitlements=[ComplexValue(value="allow-cluster-create"), ComplexValue(value="allow-instance-pool-create")],
     )
     wsclient.groups.list.return_value = [group]
-    wsclient.groups.get.return_value = group
-    account_admins_group = Group(id="1234", external_id="1234", display_name="ac_test_1234")
-    wsclient.api_client.do.return_value = {
-        "Resources": [g.as_dict() for g in [account_admins_group]],
-    }
+    account_admins_group = Group(
+        id="1234",
+        external_id="1234",
+        display_name="ac_test_1234",
+        members=[ComplexValue(display="test-user-3", value="3")],
+    )
+
+    def do_api_side_effect(*args, **kwargs):
+        if args[0] == "GET":
+            if args[1] == "/api/2.0/account/scim/v2/Groups":
+                return {"Resources": [g.as_dict() for g in [account_admins_group]]}
+            else:
+                return account_admins_group.as_dict()
+        else:
+            raise RuntimeError()
+
+    wsclient.api_client.do.side_effect = do_api_side_effect
+    wsclient.groups.get.side_effect = lambda group_id: group if group_id == "1" else account_admins_group
+    grp_membership = GroupManager(
+        backend, wsclient, inventory_database="inv", workspace_group_regex=r"\(([1-9]+)\)", account_group_regex="[1-9]+"
+    ).validate_group_membership()
+    assert grp_membership == [
+        {
+            "wf_group_name": "test_(1234)",
+            "wf_group_members_count": 2,
+            "acc_group_name": "ac_test_1234",
+            "acc_group_members_count": 1,
+        }
+    ]
+
+
+def test_validate_group_diff_membership_no_members():
+    backend = create_autospec(SqlBackend)
+    wsclient = create_autospec(WorkspaceClient)
+    group = Group(
+        id="1",
+        external_id="1234",
+        display_name="test_(1234)",
+        meta=ResourceMeta(resource_type="WorkspaceGroup"),
+        members=None,
+        roles=[
+            ComplexValue(value="arn:aws:iam::123456789098:instance-profile/ip1"),
+        ],
+        entitlements=[ComplexValue(value="allow-cluster-create"), ComplexValue(value="allow-instance-pool-create")],
+    )
+    wsclient.groups.list.return_value = [group]
+    account_admins_group = Group(
+        id="1234",
+        external_id="1234",
+        display_name="ac_test_1234",
+        members=None,
+    )
+
+    def do_api_side_effect(*args, **kwargs):
+        if args[0] == "GET":
+            if args[1] == "/api/2.0/account/scim/v2/Groups":
+                return {"Resources": [g.as_dict() for g in [account_admins_group]]}
+            else:
+                return account_admins_group.as_dict()
+        else:
+            raise RuntimeError()
+
+    wsclient.api_client.do.side_effect = do_api_side_effect
+    wsclient.groups.get.side_effect = lambda group_id: group if group_id == "1" else account_admins_group
     grp_membership = GroupManager(
         backend, wsclient, inventory_database="inv", workspace_group_regex=r"\(([1-9]+)\)", account_group_regex="[1-9]+"
     ).validate_group_membership()
-    assert grp_membership == [{"wf_group_name": "test_(1234)", "ac_group_name": "ac_test_1234"}]
+    assert grp_membership == []
+
+
+def test_validate_group_diff_membership_no_account_group_found():
+    backend = create_autospec(SqlBackend)
+    wsclient = create_autospec(WorkspaceClient)
+    group = Group(
+        id="1",
+        external_id="1234",
+        display_name="test_(1234)",
+        meta=ResourceMeta(resource_type="WorkspaceGroup"),
+        members=None,
+        roles=[
+            ComplexValue(value="arn:aws:iam::123456789098:instance-profile/ip1"),
+        ],
+        entitlements=[ComplexValue(value="allow-cluster-create"), ComplexValue(value="allow-instance-pool-create")],
+    )
+    wsclient.groups.list.return_value = [group]
+    account_admins_group = Group(
+        id="1234",
+        external_id="1234",
+        display_name="ac_test_1234",
+        members=None,
+    )
+
+    def do_api_side_effect(*args, **kwargs):
+        if args[0] == "GET":
+            if args[1] == "/api/2.0/account/scim/v2/Groups":
+                return {"Resources": [g.as_dict() for g in [account_admins_group]]}
+            else:
+                return account_admins_group.as_dict()
+        else:
+            raise RuntimeError()
+
+    wsclient.api_client.do.side_effect = do_api_side_effect
+    wsclient.groups.get.side_effect = lambda group_id: group if group_id == "1" else None
+    grp_membership = GroupManager(
+        backend, wsclient, inventory_database="inv", workspace_group_regex=r"\(([1-9]+)\)", account_group_regex="[1-9]+"
+    ).validate_group_membership()
+    assert grp_membership == []
 
 
 def test_validate_group_same_membership():
@@ -878,9 +976,57 @@ def test_validate_group_same_membership():
         display_name="ac_test_1234",
         members=[ComplexValue(display="test-user-1", value="01"), ComplexValue(display="test-user-2", value="02")],
     )
-    wsclient.api_client.do.return_value = {
-        "Resources": [g.as_dict() for g in [account_admins_group]],
-    }
+
+    def do_api_side_effect(*args, **kwargs):
+        if args[0] == "GET":
+            if args[1] == "/api/2.0/account/scim/v2/Groups":
+                return {"Resources": [g.as_dict() for g in [account_admins_group]]}
+            else:
+                return account_admins_group.as_dict()
+        else:
+            raise RuntimeError()
+
+    wsclient.api_client.do.side_effect = do_api_side_effect
+    grp_membership = GroupManager(
+        backend, wsclient, inventory_database="inv", workspace_group_regex=r"\(([1-9]+)\)", account_group_regex="[1-9]+"
+    ).validate_group_membership()
+    assert grp_membership == []
+
+
+def test_validate_acc_group_removed_after_listing():
+    backend = MockBackend()
+    wsclient = MagicMock()
+    group = Group(
+        id="1",
+        external_id="1234",
+        display_name="test_(1234)",
+        meta=ResourceMeta(resource_type="WorkspaceGroup"),
+        members=[ComplexValue(display="test-user-1", value="01"), ComplexValue(display="test-user-2", value="02")],
+        roles=[
+            ComplexValue(value="arn:aws:iam::123456789098:instance-profile/test_ip1"),
+            ComplexValue(value="arn:aws:iam::123456789098:instance-profile/test_ip2"),
+        ],
+        entitlements=[ComplexValue(value="allow-cluster-create"), ComplexValue(value="allow-instance-pool-create")],
+    )
+    wsclient.groups.list.return_value = [group]
+    wsclient.groups.get.return_value = group
+    account_admins_group = Group(
+        id="1234",
+        external_id="1234",
+        display_name="ac_test_1234",
+        members=[ComplexValue(display="test-user-1", value="01"), ComplexValue(display="test-user-2", value="02")],
+    )
+
+    def do_api_side_effect(*args, **kwargs):
+        if args[0] == "GET":
+            if args[1] == "/api/2.0/account/scim/v2/Groups":
+                return {"Resources": [g.as_dict() for g in [account_admins_group]]}
+            else:
+                raise NotFound()
+        else:
+            raise RuntimeError()
+
+    wsclient.api_client.do.side_effect = do_api_side_effect
     grp_membership = GroupManager(
         backend, wsclient, inventory_database="inv", workspace_group_regex=r"\(([1-9]+)\)", account_group_regex="[1-9]+"
     ).validate_group_membership()