Add a command to create account level groups if they do not exist

labs.yml

@@ -125,4 +125,8 @@ commands:
       Requires a working setup of AWS CLI.
     flags:
       - name: aws-profile
-        description: AWS Profile to use for authentication
+        description: AWS Profile to use for authentication
+
+  - name: create-account-groups
+    is_account_level: true
+    description: Creates account level groups for all groups in all workspaces.

src/databricks/labs/ucx/account.py

@@ -6,6 +6,7 @@
 from databricks.labs.blueprint.tui import Prompts
 from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.errors import NotFound
+from databricks.sdk.service.iam import ComplexValue, Group, Patch, PatchOp, PatchSchema
 from databricks.sdk.service.provisioning import Workspace
 
 from databricks.labs.ucx.__about__ import __version__
@@ -69,6 +70,84 @@
             for installation in Installation.existing(ws, "ucx"):
                 installation.save(workspaces, filename=self.SYNC_FILE_NAME)
 
+    def create_account_level_groups(self):
+        acc_groups = self._get_account_groups()
+        all_valid_workspace_groups = self._get_valid_workspaces_groups()
+
+        for group_name, valid_group in all_valid_workspace_groups.items():
+            if group_name in acc_groups:
+                logger.info(f"Group {group_name} already exist in the account, ignoring")
+                continue
+
+            acc_group = self._ac.groups.create(display_name=group_name)
+
+            if len(acc_group.members) > 0:
+                self._add_members_to_acc_group(acc_group.id, group_name, valid_group)
+
+            logger.info(f"Group {group_name} created in the account")
+
+    def _add_members_to_acc_group(self, acc_group_id: str, group_name: str, valid_group: Group):
+        for chunk in self._chunks(valid_group.members, 20):
+            logger.debug(f"Adding 20 members to acc group {group_name}")
+            self._ac.groups.patch(
+                acc_group_id,
+                operations=[Patch(op=PatchOp.ADD, path="members", value=[x.as_dict() for x in chunk])],
+                schemas=[PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP],
+            )
+
+    def _chunks(self, lst, n):
+        """Yield successive n-sized chunks from lst."""
+        for i in range(0, len(lst), n):
+            yield lst[i : i + n]
+
+    def _get_valid_workspaces_groups(self) -> dict[str, Group]:
+        all_workspaces_groups: dict[str, Group] = {}
+
+        for workspace in self._workspaces():
+            client = self.client_for(workspace)
+            ws_group_ids = client.groups.list(attributes="id")
+            for group_id in ws_group_ids:
+                if not group_id.id:
+                    continue
+
+                full_workspace_group = client.groups.get(group_id.id)
+                group_name = full_workspace_group.display_name
+
+                if group_name in all_workspaces_groups:
+                    if self._has_same_members(all_workspaces_groups[group_name], full_workspace_group):
+                        logger.info(f"Workspace group {group_name} already found, ignoring")
+                        continue
+
+                    logger.warning(
+                        f"Group {group_name} does not have the same amount of members "
+                        f"in workspace {client.config.host}, it will be created with account "
+                        f"name {workspace.workspace_name}_{group_name}"
+                    )
+                    all_workspaces_groups[f"{workspace.workspace_name}_{group_name}"] = full_workspace_group
+                    continue
+
+                if not group_name:
+                    continue
+
+                logger.info(f"Found new group {group_name}")
+                all_workspaces_groups[group_name] = full_workspace_group
+
+        return all_workspaces_groups
+
+    def _has_same_members(self, group_1: Group, group_2: Group) -> bool:
+        ws_members_set_1 = set([m.display for m in group_1.members] if group_1.members else [])
+        ws_members_set_2 = set([m.display for m in group_2.members] if group_2.members else [])
+        return not bool((ws_members_set_1 - ws_members_set_2).union(ws_members_set_2 - ws_members_set_1))
+
+    def _get_account_groups(self) -> dict[str | None, list[ComplexValue] | None]:
+        acc_groups = {}
+        for acc_grp_id in self._ac.groups.list(attributes="id"):
+            if not acc_grp_id.id:
+                continue
+            full_account_group = self._ac.groups.get(acc_grp_id.id)
+            acc_groups[full_account_group.display_name] = full_account_group.members
+        return acc_groups
+
 
 class WorkspaceInfo:
     def __init__(self, installation: Installation, ws: WorkspaceClient):

src/databricks/labs/ucx/cli.py

@@ -89,6 +89,23 @@ def sync_workspace_info(a: AccountClient):
     workspaces.sync_workspace_info()
 
 
+@ucx.command(is_account=True)
+def create_account_groups(a: AccountClient):
+    """
+    Crawl all workspaces, and create account level groups if a WS local group is not present in the account.
+    The feature is not configurable, meaning that it fetches all workspaces groups and all account groups.
+
+    The following scenarios are supported, if a group X:
+    - Exist in workspaces A,B,C and it has same members in there, it will be created in the account
+    - Exist in workspaces A,B but not in C, it will be created in the account
+    - Exist in workspaces A,B,C and it has same members in A,B, but not in C, then, X and C_X will be created in the
+    account
+    """
+    logger.info(f"Account ID: {a.config.account_id}")
+    workspaces = AccountWorkspaces(a)
+    workspaces.create_account_level_groups()
+
+
 @ucx.command
 def manual_workspace_info(w: WorkspaceClient):
     """only supposed to be run if cannot get admins to run `databricks labs ucx sync-workspace-info`"""

src/databricks/labs/ucx/framework/dashboards.py

@@ -45,7 +45,7 @@ def viz_args(self) -> dict:
 class VizColumn:  # pylint: disable=too-many-instance-attributes)
     name: str
     title: str
-    type: str = "string"
+    type: str = "string"  # noqa: A003
     imageUrlTemplate: str = "{{ @ }}"  # noqa: N815
     imageTitleTemplate: str = "{{ @ }}"  # noqa: N815
     linkUrlTemplate: str = "{{ @ }}"  # noqa: N815

tests/unit/test_account.py

@@ -6,7 +6,8 @@
 from databricks.labs.blueprint.tui import MockPrompts
 from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.config import Config
-from databricks.sdk.service.iam import User
+from databricks.sdk.service import iam
+from databricks.sdk.service.iam import ComplexValue, Group, User
 from databricks.sdk.service.provisioning import Workspace
 
 from databricks.labs.ucx.account import AccountWorkspaces, WorkspaceInfo
@@ -68,3 +69,199 @@ def test_manual_workspace_info(mocker):
         wir.manual_workspace_info(prompts)
 
     ws.workspace.upload.assert_called()
+
+
+def test_create_acc_groups_should_create_acc_group_if_no_group_found_in_account(mocker):
+    acc_client = create_autospec(AccountClient)
+    acc_client.config = Config(host="https://accounts.cloud.databricks.com", account_id="123", token="123")
+
+    acc_client.workspaces.list.return_value = [
+        Workspace(workspace_name="foo", workspace_id=123, workspace_status_message="Running", deployment_name="abc")
+    ]
+
+    ws = create_autospec(WorkspaceClient)
+
+    def workspace_client(**kwargs) -> WorkspaceClient:
+        return ws
+
+    account_workspaces = AccountWorkspaces(acc_client, workspace_client)
+
+    group = Group(
+        id="12",
+        display_name="de",
+        members=[ComplexValue(display="test-user-1", value="20"), ComplexValue(display="test-user-2", value="21")],
+    )
+
+    ws.groups.list.return_value = [group]
+    ws.groups.get.return_value = group
+    acc_client.groups.create.return_value = group
+
+    account_workspaces.create_account_level_groups()
+
+    acc_client.groups.create.assert_called_with(
+        display_name="de",
+    )
+    acc_client.groups.patch.assert_called_with(
+        "12",
+        operations=[
+            iam.Patch(
+                op=iam.PatchOp.ADD,
+                path='members',
+                value=[{'display': 'test-user-1', 'value': '20'}, {'display': 'test-user-2', 'value': '21'}],
+            )
+        ],
+        schemas=[iam.PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP],
+    )
+
+
+def test_create_acc_groups_should_not_create_group_if_exists_in_account(mocker):
+    acc_client = create_autospec(AccountClient)
+    acc_client.config = Config(host="https://accounts.cloud.databricks.com", account_id="123", token="123")
+
+    group = Group(
+        id="12",
+        display_name="de",
+        members=[ComplexValue(display="test-user-1", value="20"), ComplexValue(display="test-user-2", value="21")],
+    )
+    acc_client.groups.list.return_value = [group]
+    acc_client.groups.get.return_value = group
+    acc_client.workspaces.list.return_value = [
+        Workspace(workspace_name="foo", workspace_id=123, workspace_status_message="Running", deployment_name="abc")
+    ]
+
+    ws = create_autospec(WorkspaceClient)
+
+    def workspace_client(**kwargs) -> WorkspaceClient:
+        return ws
+
+    ws.groups.list.return_value = [group]
+    ws.groups.get.return_value = group
+
+    account_workspaces = AccountWorkspaces(acc_client, workspace_client)
+    account_workspaces.create_account_level_groups()
+
+    acc_client.groups.create.assert_not_called()
+
+
+def test_create_acc_groups_should_create_groups_accross_workspaces(mocker):
+    acc_client = create_autospec(AccountClient)
+    acc_client.config = Config(host="https://accounts.cloud.databricks.com", account_id="123", token="123")
+
+    acc_client.workspaces.list.return_value = [
+        Workspace(workspace_name="foo", workspace_id=123, workspace_status_message="Running", deployment_name="abc"),
+        Workspace(workspace_name="bar", workspace_id=456, workspace_status_message="Running", deployment_name="def"),
+    ]
+
+    ws1 = create_autospec(WorkspaceClient)
+    ws2 = create_autospec(WorkspaceClient)
+
+    def workspace_client(host, product, **kwargs) -> WorkspaceClient:
+        if host == "https://abc.cloud.databricks.com":
+            return ws1
+        else:
+            return ws2
+
+    group = Group(id="12", display_name="de")
+    group2 = Group(id="12", display_name="security_grp")
+
+    ws1.groups.list.return_value = [group]
+    ws1.groups.get.return_value = group
+
+    ws2.groups.list.return_value = [group2]
+    ws2.groups.get.return_value = group2
+
+    account_workspaces = AccountWorkspaces(acc_client, workspace_client)
+    account_workspaces.create_account_level_groups()
+
+    acc_client.groups.create.assert_any_call(display_name="de")
+    acc_client.groups.create.assert_any_call(display_name="security_grp")
+
+
+def test_create_acc_groups_should_filter_groups_accross_workspaces(mocker):
+    acc_client = create_autospec(AccountClient)
+    acc_client.config = Config(host="https://accounts.cloud.databricks.com", account_id="123", token="123")
+
+    acc_client.workspaces.list.return_value = [
+        Workspace(workspace_name="foo", workspace_id=123, workspace_status_message="Running", deployment_name="abc"),
+        Workspace(workspace_name="bar", workspace_id=456, workspace_status_message="Running", deployment_name="def"),
+    ]
+
+    ws1 = create_autospec(WorkspaceClient)
+    ws2 = create_autospec(WorkspaceClient)
+
+    def workspace_client(host, product, **kwargs) -> WorkspaceClient:
+        if host == "https://abc.cloud.databricks.com":
+            return ws1
+        else:
+            return ws2
+
+    group = Group(
+        id="12",
+        display_name="de",
+        members=[ComplexValue(display="test-user-1", value="20"), ComplexValue(display="test-user-2", value="21")],
+    )
+
+    ws1.groups.list.return_value = [group]
+    ws1.groups.get.return_value = group
+
+    ws2.groups.list.return_value = [group]
+    ws2.groups.get.return_value = group
+    acc_client.groups.create.return_value = group
+
+    account_workspaces = AccountWorkspaces(acc_client, workspace_client)
+    account_workspaces.create_account_level_groups()
+
+    acc_client.groups.create.assert_called_once_with(display_name="de")
+    acc_client.groups.patch.assert_called_once_with(
+        "12",
+        operations=[
+            iam.Patch(
+                op=iam.PatchOp.ADD,
+                path='members',
+                value=[{'display': 'test-user-1', 'value': '20'}, {'display': 'test-user-2', 'value': '21'}],
+            )
+        ],
+        schemas=[iam.PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP],
+    )
+
+
+def test_create_acc_groups_should_create_acc_group_if_exist_in_other_workspaces_but_not_same_members(mocker):
+    acc_client = create_autospec(AccountClient)
+    acc_client.config = Config(host="https://accounts.cloud.databricks.com", account_id="123", token="123")
+
+    acc_client.workspaces.list.return_value = [
+        Workspace(workspace_name="ws1", workspace_id=123, workspace_status_message="Running", deployment_name="abc"),
+        Workspace(workspace_name="ws2", workspace_id=456, workspace_status_message="Running", deployment_name="def"),
+    ]
+
+    ws1 = create_autospec(WorkspaceClient)
+    ws2 = create_autospec(WorkspaceClient)
+
+    def workspace_client(host, product, **kwargs) -> WorkspaceClient:
+        if host == "https://abc.cloud.databricks.com":
+            return ws1
+        else:
+            return ws2
+
+    group = Group(
+        id="12",
+        display_name="de",
+        members=[ComplexValue(display="test-user-1", value="20"), ComplexValue(display="test-user-2", value="21")],
+    )
+    group_2 = Group(
+        id="12",
+        display_name="de",
+        members=[ComplexValue(display="test-user-1", value="20")],
+    )
+
+    ws1.groups.list.return_value = [group]
+    ws1.groups.get.return_value = group
+
+    ws2.groups.list.return_value = [group_2]
+    ws2.groups.get.return_value = group_2
+
+    account_workspaces = AccountWorkspaces(acc_client, workspace_client)
+    account_workspaces.create_account_level_groups()
+
+    acc_client.groups.create.assert_any_call(display_name="de")
+    acc_client.groups.create.assert_any_call(display_name="ws2_de")

tests/unit/test_cli.py

@@ -11,6 +11,7 @@
 
 from databricks.labs.ucx.cli import (
     alias,
+    create_account_groups,
     create_table_mapping,
     ensure_assessment_run,
     installations,
@@ -128,6 +129,12 @@ def test_sync_workspace_info():
         s.assert_called_once()
 
 
+def test_create_account_groups():
+    a = create_autospec(AccountClient)
+    create_account_groups(a)
+    a.groups.list.assert_called_with(attributes="id")
+
+
 def test_manual_workspace_info(ws):
     with patch("databricks.labs.ucx.account.WorkspaceInfo.manual_workspace_info", return_value=None) as m:
         manual_workspace_info(ws)