Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove exception and added proper logging for groups in the list that… #357

Merged
merged 2 commits into from
Oct 3, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 24 additions & 8 deletions src/databricks/labs/ucx/workspace_access/groups.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,40 +161,56 @@ def prepare_groups_in_environment(self):
"exist and are of the correct type. If some temporary groups are missing, they'll be created"
)
group_names = self.config.selected
valid_group_names = []
if group_names:
logger.info("Using the provided group listing")

for g in group_names:
assert g not in self.SYSTEM_GROUPS, f"Cannot migrate system group {g}"
assert self._get_group(g, "workspace"), f"Group {g} not found on the workspace level"
assert self._get_group(g, "account"), f"Group {g} not found on the account level"

if not group_names:
if g in self.SYSTEM_GROUPS:
logger.info(f"Cannot migrate system group {g}. {g} will be skipped.")
continue
if not self._get_group(g, "workspace"):
logger.info(f"Group {g} not found on the workspace level. {g} will be skipped.")
continue
if not self._get_group(g, "account"):
logger.info(
f"Group {g} not found on the account level. {g} will be skipped. You can add {g} "
f"to the account and rerun the job."
)
continue
valid_group_names.append(g)
else:
logger.info(
"No group listing provided, all available workspace-level groups that have an account-level "
"group with the same name will be used"
)
ws_group_names = {_.display_name for _ in self._workspace_groups}
ac_group_names = {_.display_name for _ in self._account_groups}
group_names = list(ws_group_names.intersection(ac_group_names))
valid_group_names = list(ws_group_names.intersection(ac_group_names))

self._set_migration_groups(group_names)
self._set_migration_groups(valid_group_names)
logger.info("Environment prepared successfully")

@property
def migration_groups_provider(self) -> GroupMigrationState:
assert len(self._migration_state.groups) > 0, "Migration groups were not loaded or initialized"
if len(self._migration_state.groups) == 0:
logger.info("No groups were loaded or initialized, nothing to do")
return self._migration_state

def replace_workspace_groups_with_account_groups(self):
logger.info("Replacing the workspace groups with account-level groups")
if len(self._migration_state.groups) == 0:
logger.info("No groups were loaded or initialized, nothing to do")
return
ThreadedExecution.gather(
"groups: workspace -> account",
[partial(self._replace_group, migration_info) for migration_info in self.migration_groups_provider.groups],
)
logger.info("Workspace groups were successfully replaced with account-level groups")

def delete_backup_groups(self):
if len(self._migration_state.groups) == 0:
return
logger.info(
f"Deleting the workspace-level backup groups. "
f"In total, {len(self.migration_groups_provider.groups)} group(s) to be deleted"
Expand Down
3 changes: 3 additions & 0 deletions src/databricks/labs/ucx/workspace_access/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,9 @@ def inventorize_permissions(self):

def apply_group_permissions(self, migration_state: GroupMigrationState, destination: Literal["backup", "account"]):
# list shall be sorted prior to using group by
if len(migration_state.groups) == 0:
logger.info("No valid groups selected, nothing to do.")
return
items = sorted(self._load_all(), key=lambda i: i.object_type)
logger.info(
f"Applying the permissions to {destination} groups. "
Expand Down
45 changes: 41 additions & 4 deletions tests/unit/workspace_access/test_groups.py
Original file line number Diff line number Diff line change
Expand Up @@ -216,7 +216,7 @@ def test_prepare_groups_in_environment_with_multiple_groups_in_conf_should_retur
assert compare(manager._migration_state.groups, [ds_group_info, de_group_info])


def test_prepare_groups_in_environment_should_throw_when_account_group_doesnt_exist():
def test_prepare_groups_in_environment_should_not_throw_when_account_group_doesnt_exist():
de_group = Group(display_name="de", meta=ResourceMeta(resource_type="WorkspaceGroup"))

client = Mock()
Expand All @@ -227,9 +227,8 @@ def test_prepare_groups_in_environment_should_throw_when_account_group_doesnt_ex
group_conf = GroupsConfig(selected=["de"], backup_group_prefix="dbr_backup_")
manager = GroupManager(client, group_conf)

with pytest.raises(AssertionError) as e_info:
manager.prepare_groups_in_environment()
assert str(e_info.value) == "Group de not found on the account level"
manager.prepare_groups_in_environment()
assert len(manager.migration_groups_provider.groups) == 0


def test_prepare_groups_in_environment_should_throw_when_workspace_group_doesnt_exist():
Expand Down Expand Up @@ -319,6 +318,44 @@ def test_replace_workspace_groups_with_account_groups_should_call_delete_and_do(
)


def test_system_groups():
client = Mock()
test_ws_group = Group(display_name="admins", meta=ResourceMeta(resource_type="WorkspaceGroup"))
test_acc_group = Group(display_name="admins", meta=ResourceMeta(resource_type="Group"))
backup_group_id = "100"
client.groups.list.return_value = [test_ws_group]
client.groups.create.return_value = Group(
display_name="dbr_backup_de", meta=ResourceMeta(resource_type="WorkspaceGroup"), id=backup_group_id
)
client.api_client.do.return_value = {
"Resources": [g.as_dict() for g in [test_acc_group]],
}

group_conf = GroupsConfig(backup_group_prefix="dbr_backup_", selected=["admins"])
manager = GroupManager(client, group_conf)
manager.prepare_groups_in_environment()
assert len(manager._migration_state.groups) == 0


def test_workspace_only_groups():
client = Mock()
test_ws_group = Group(display_name="ws_group", meta=ResourceMeta(resource_type="WorkspaceGroup"))
test_acc_group = Group(display_name="acc_group", meta=ResourceMeta(resource_type="Group"))
backup_group_id = "100"
client.groups.list.return_value = [test_ws_group, test_acc_group]
client.groups.create.return_value = Group(
display_name="dbr_backup_de", meta=ResourceMeta(resource_type="WorkspaceGroup"), id=backup_group_id
)
client.api_client.do.return_value = {
"Resources": [g.as_dict() for g in [test_acc_group]],
}

group_conf = GroupsConfig(backup_group_prefix="dbr_backup_", selected=["ws_group"])
manager = GroupManager(client, group_conf)
manager.prepare_groups_in_environment()
assert len(manager._migration_state.groups) == 0


def test_delete_backup_groups():
client = Mock()

Expand Down
19 changes: 16 additions & 3 deletions tests/unit/workspace_access/test_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,13 @@

import pytest
from databricks.sdk.service import iam
from databricks.sdk.service.iam import Group, ResourceMeta

from databricks.labs.ucx.mixins.sql import Row
from databricks.labs.ucx.workspace_access.groups import (
GroupMigrationState,
MigrationGroupInfo,
)
from databricks.labs.ucx.workspace_access.manager import PermissionManager, Permissions

from ..framework.mocks import MockBackend
Expand Down Expand Up @@ -141,7 +146,16 @@ def test_manager_apply(mocker):
"cluster-policies": mock_applier,
},
)
pm.apply_group_permissions(MagicMock(), "backup")
group_migration_state: GroupMigrationState = MagicMock()
group_migration_state.groups = [
MigrationGroupInfo(
Group(display_name="group", meta=ResourceMeta(resource_type="WorkspaceGroup")),
Group(display_name="group_backup", meta=ResourceMeta(resource_type="WorkspaceGroup")),
Group(display_name="group", meta=ResourceMeta(resource_type="AccountGroup")),
)
]

pm.apply_group_permissions(group_migration_state, "backup")

assert {"test2 test2 backup", "test test backup"} == applied_items

Expand All @@ -155,5 +169,4 @@ def test_unregistered_support():
}
)
pm = PermissionManager(b, "test", [], {})
with pytest.raises(ValueError):
pm.apply_group_permissions(migration_state=MagicMock(), destination="backup")
pm.apply_group_permissions(migration_state=MagicMock(), destination="backup")