Added databricks labs ucx validate-groups-membership command to validate groups to see if they have same membership across acount and workspace level (#772)

Author: prajin-29

labs.yml

@@ -86,3 +86,10 @@ commands:
       - name: subscription-id
         description: Subscription to scan storage account in
 
+
+  - 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}}
+      {{end}}

src/databricks/labs/ucx/cli.py

@@ -18,6 +18,7 @@
 from databricks.labs.ucx.hive_metastore.table_migrate import TableMove, TablesMigrate
 from databricks.labs.ucx.install import WorkspaceInstaller
 from databricks.labs.ucx.installer import InstallationManager
+from databricks.labs.ucx.workspace_access.groups import GroupManager
 
 ucx = App(__file__)
 logger = get_logger(__file__)
@@ -142,6 +143,37 @@ def repair_run(w: WorkspaceClient, step):
     installer.repair_run(step)
 
 
+@ucx.command
+def validate_groups_membership(w: WorkspaceClient):
+    """Validate the groups to see if the groups at account level and workspace level has different membership"""
+    installation_manager = InstallationManager(w)
+    installation = installation_manager.for_user(w.current_user.me())
+    if not installation:
+        logger.error(CANT_FIND_UCX_MSG)
+        return None
+    warehouse_id = installation.config.warehouse_id
+    inventory_database = installation.config.inventory_database
+    renamed_group_prefix = installation.config.renamed_group_prefix
+    workspace_group_regex = installation.config.workspace_group_regex
+    workspace_group_replace = installation.config.workspace_group_replace
+    account_group_regex = installation.config.account_group_regex
+    include_group_names = installation.config.include_group_names
+    sql_backend = StatementExecutionBackend(w, warehouse_id)
+    logger.info("Validating Groups which are having different memberships between account and workspace")
+    group_manager = GroupManager(
+        sql_backend=sql_backend,
+        ws=w,
+        inventory_database=inventory_database,
+        include_group_names=include_group_names,
+        renamed_group_prefix=renamed_group_prefix,
+        workspace_group_regex=workspace_group_regex,
+        workspace_group_replace=workspace_group_replace,
+        account_group_regex=account_group_regex,
+    )
+    mismatch_groups = group_manager.validate_group_membership()
+    print(json.dumps(mismatch_groups))
+
+
 @ucx.command
 def revert_migrated_tables(w: WorkspaceClient, schema: str, table: str, *, delete_managed: bool = False):
     """remove notation on a migrated table for re-migration"""

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

@@ -406,6 +406,33 @@ def _crawler(self) -> Iterable[MigratedGroup]:
         strategy = self._get_strategy(workspace_groups_in_workspace, account_groups_in_account)
         yield from strategy.generate_migrated_groups()
 
+    def validate_group_membership(self) -> list[dict]:
+        workspace_groups_in_workspace = self._workspace_groups_in_workspace()
+        account_groups_in_account = self._account_groups_in_account()
+        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)
+            if not set_diff:
+                continue
+            mismatch_group.append(
+                {
+                    "wf_group_name": groups.name_in_workspace,
+                    "ac_group_name": groups.name_in_account,
+                }
+            )
+        if not mismatch_group:
+            logger.info("There are no groups with different membership between account and workspace")
+        else:
+            logger.info("There are groups with different membership between account and workspace")
+        return mismatch_group
+
     def _workspace_groups_in_workspace(self) -> dict[str, Group]:
         attributes = "id,displayName,meta,externalId,members,roles,entitlements"
         groups = {}

tests/integration/workspace_access/test_groups.py

@@ -206,6 +206,38 @@ def test_group_matching_names(ws, sql_backend, inventory_schema, make_ucx_group,
     validate_migrate_groups(group_manager, ws_group, accnt_group)
 
 
+@retried(on=[NotFound], timeout=timedelta(minutes=2))
+def test_group_matching_names_with_diff_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()
+    user2 = make_user()
+    members1 = [user1.id]
+    members2 = [user2.id]
+    ws_group = make_group(display_name=workspace_group_name, members=members1, entitlements=["allow-cluster-create"])
+    accnt_group = make_acc_group(display_name=account_group_name, members=members2)
+
+    logger.info(
+        f"Attempting Mapping From Workspace Group {ws_group.display_name} to "
+        f"Account Group {accnt_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/test_cli.py

@@ -19,8 +19,11 @@
     skip,
     sync_workspace_info,
     validate_external_locations,
+    validate_groups_membership,
     workflows,
 )
+from databricks.labs.ucx.config import WorkspaceConfig
+from databricks.labs.ucx.installer import Installation
 
 
 def test_workflow(caplog):
@@ -320,3 +323,41 @@ def test_save_azure_storage_accounts(mocker, caplog):
     ) as m:
         save_azure_storage_accounts(w, "test")
         m.assert_called_once()
+
+
+def test_validate_groups_membership(mocker, caplog):
+    w = create_autospec(WorkspaceClient)
+    inst_data = Installation(
+        config=WorkspaceConfig(
+            inventory_database="test_database",
+            workspace_group_regex=None,
+            workspace_group_replace=None,
+            account_group_regex=None,
+            group_match_by_external_id=False,
+            include_group_names=None,
+            renamed_group_prefix="db-temp-",
+            warehouse_id="test_id",
+            database_to_catalog_mapping=None,
+            default_catalog="ucx_default",
+        ),
+        user="test_user",
+        path="/Users/[email protected]/.ucx",
+    )
+    w.current_user.me = lambda: iam.User(user_name="[email protected]", groups=[iam.ComplexValue(display="admins")])
+    mocker.patch("databricks.labs.ucx.installer.InstallationManager.__init__", return_value=None)
+    mocker.patch("databricks.labs.ucx.installer.InstallationManager.for_user", return_value=inst_data)
+    mocker.patch("databricks.labs.ucx.workspace_access.groups.GroupManager.__init__", return_value=None)
+    mocker.patch(
+        "databricks.labs.ucx.workspace_access.groups.GroupManager.validate_group_membership",
+        return_value={"wf_group_name": "test_group", "ac_group_name": "test_group"},
+    )
+    validate_groups_membership(w)
+    assert caplog.messages == ["Validating Groups which are having different memberships between account and workspace"]
+
+
+def test_validate_group_no_ucx(mocker, caplog):
+    w = create_autospec(WorkspaceClient)
+    w.current_user.me = lambda: iam.User(user_name="test_user", groups=[iam.ComplexValue(display="admins")])
+    mocker.patch("databricks.labs.ucx.installer.InstallationManager.for_user", return_value=None)
+    validate_groups_membership(w)
+    assert "Couldn't find UCX configuration" in caplog.messages[0]

tests/unit/workspace_access/test_groups.py

@@ -1,14 +1,16 @@
 import json
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, create_autospec
 
 import pytest
 from _pytest.outcomes import fail
 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.service import iam
 from databricks.sdk.service.iam import ComplexValue, Group, ResourceMeta
 
+from databricks.labs.ucx.framework.crawlers import SqlBackend
 from databricks.labs.ucx.workspace_access.groups import (
     ConfigureGroups,
     GroupManager,
@@ -824,3 +826,62 @@ def test_state():
     assert not state.get_target_principal("invalid_group_name")
     assert not state.get_temp_principal("invalid_group_name")
     assert not state.is_in_scope("invalid_group_name")
+
+
+def test_validate_group_diff_membership():
+    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=[ComplexValue(display="test-user-1", value="20"), ComplexValue(display="test-user-2", value="21")],
+        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]],
+    }
+    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"}]
+
+
+def test_validate_group_same_membership():
+    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")],
+    )
+    wsclient.api_client.do.return_value = {
+        "Resources": [g.as_dict() for g in [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 == []