From 173c353f64b1c97ab0880b1827203a967d039428 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Fri, 22 Sep 2023 19:46:23 +0200 Subject: [PATCH 1/2] Use all account-level groups with matching names to workspace-level groups in case no explicit configuration Fixes #272 Fixes #236 --- src/databricks/labs/ucx/install.py | 3 ++- .../labs/ucx/workspace_access/groups.py | 22 ++++++++++++------- .../workspace_access/test_groups.py | 21 ++++++++++++++++++ 3 files changed, 37 insertions(+), 9 deletions(-) diff --git a/src/databricks/labs/ucx/install.py b/src/databricks/labs/ucx/install.py index a208af22df..39793268f6 100644 --- a/src/databricks/labs/ucx/install.py +++ b/src/databricks/labs/ucx/install.py @@ -187,7 +187,8 @@ 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="" + "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="" ) backup_group_prefix = self._question("Backup prefix", default="db-temp-") log_level = self._question("Log level", default="INFO").upper() diff --git a/src/databricks/labs/ucx/workspace_access/groups.py b/src/databricks/labs/ucx/workspace_access/groups.py index de4143a942..3ecaad7c8d 100644 --- a/src/databricks/labs/ucx/workspace_access/groups.py +++ b/src/databricks/labs/ucx/workspace_access/groups.py @@ -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) @@ -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 diff --git a/tests/integration/workspace_access/test_groups.py b/tests/integration/workspace_access/test_groups.py index d6c987ac49..310e62121e 100644 --- a/tests/integration/workspace_access/test_groups.py +++ b/tests/integration/workspace_access/test_groups.py @@ -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])) From dc82b69b24b7669450f495f0ff4f4fadc566b3b1 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Mon, 25 Sep 2023 14:54:58 +0200 Subject: [PATCH 2/2] fmt --- src/databricks/labs/ucx/install.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/databricks/labs/ucx/install.py b/src/databricks/labs/ucx/install.py index 39793268f6..c1e28dd956 100644 --- a/src/databricks/labs/ucx/install.py +++ b/src/databricks/labs/ucx/install.py @@ -188,7 +188,8 @@ def _configure(self): selected_groups = self._question( "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="" + "account-level groups with matching names to workspace-level groups.", + default="", ) backup_group_prefix = self._question("Backup prefix", default="db-temp-") log_level = self._question("Log level", default="INFO").upper()