Skip to content

Commit

Permalink
Use all account-level groups with matching names to workspace-level g…
Browse files Browse the repository at this point in the history
…roups in case no explicit configuration (#277)

Fixes #272
Fixes #236
  • Loading branch information
nfx authored Sep 25, 2023
1 parent dfa350a commit 20d1b42
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 9 deletions.
4 changes: 3 additions & 1 deletion src/databricks/labs/ucx/install.py
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,9 @@ def _configure(self):
warehouse_id = new_warehouse.id

selected_groups = self._question(
"Comma-separated list of workspace group names to migrate (empty means all)", default="<ALL>"
"Comma-separated list of workspace group names to migrate. If not specified, we'll wse all "
"account-level groups with matching names to workspace-level groups.",
default="<ALL>",
)
backup_group_prefix = self._question("Backup prefix", default="db-temp-")
log_level = self._question("Log level", default="INFO").upper()
Expand Down
22 changes: 14 additions & 8 deletions src/databricks/labs/ucx/workspace_access/groups.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ def _list_workspace_groups(self) -> list[iam.Group]:

def _list_account_groups(self) -> list[iam.Group]:
# TODO: we should avoid using this method, as it's not documented
# unfortunately, there's no other way to consistently get the list of account groups
# get account-level groups even if they're not (yet) assigned to a workspace
logger.debug("Listing account groups...")
account_groups = [
iam.Group.from_dict(r)
Expand Down Expand Up @@ -160,19 +160,25 @@ def prepare_groups_in_environment(self):
"Preparing groups in the current environment. At this step we'll verify that all groups "
"exist and are of the correct type. If some temporary groups are missing, they'll be created"
)
if self.config.selected:
group_names = self.config.selected
if group_names:
logger.info("Using the provided group listing")

for g in self.config.selected:
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"

self._set_migration_groups(self.config.selected)
else:
logger.info("No group listing provided, all available workspace-level groups will be used")
available_group_names = [g.display_name for g in self._workspace_groups]
self._set_migration_groups(groups_names=available_group_names)
if not group_names:
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))

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

@property
Expand Down
21 changes: 21 additions & 0 deletions tests/integration/workspace_access/test_groups.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,27 @@ def test_prepare_environment(ws, make_ucx_group):
assert _ws_members == _backup_members


def test_prepare_environment_no_groups_selected(ws, make_ucx_group, make_group, make_acc_group):
make_group()
make_acc_group()
for_test = [make_ucx_group(), make_ucx_group()]

group_manager = GroupManager(ws, GroupsConfig(auto=True))
group_manager.prepare_groups_in_environment()

group_migration_state = group_manager.migration_groups_provider
for _info in group_migration_state.groups:
_ws = ws.groups.get(id=_info.workspace.id)
_backup = ws.groups.get(id=_info.backup.id)
# https://github.com/databricks/databricks-sdk-py/pull/361 may fix the NPE gotcha with empty members
_ws_members = sorted([m.value for m in _ws.members])
_backup_members = sorted([m.value for m in _backup.members])
assert _ws_members == _backup_members

for g, _ in for_test:
assert group_migration_state.get_by_workspace_group_name(g.display_name) is not None


def test_group_listing(ws: WorkspaceClient, make_ucx_group):
ws_group, acc_group = make_ucx_group()
manager = GroupManager(ws, GroupsConfig(selected=[ws_group.display_name]))
Expand Down

0 comments on commit 20d1b42

Please sign in to comment.