Skip to content

Commit

Permalink
Added CLI Option
Browse files Browse the repository at this point in the history
  • Loading branch information
FastLee committed Nov 6, 2024
1 parent 91d6c24 commit a651c3b
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 2 deletions.
3 changes: 2 additions & 1 deletion labs.yml
Original file line number Diff line number Diff line change
Expand Up @@ -358,4 +358,5 @@ commands:
- name: enable-hms-federation
description: (EXPERIMENTAL) Enable HMS federation based migration flow. When this is enabled, UCX will create a federated HMS catalog which syncs from the workspace HMS.


- name: assign-owner-group
description: Assign owner group to the workspace. This group will be assigned as an owner to all migrated tables and views.
27 changes: 27 additions & 0 deletions src/databricks/labs/ucx/cli.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from dataclasses import replace
from io import BytesIO
import json
import webbrowser
Expand Down Expand Up @@ -634,6 +635,32 @@ def create_ucx_catalog(w: WorkspaceClient, prompts: Prompts, ctx: WorkspaceConte
workspace_context.verify_progress_tracking.verify()


@ucx.command
def assign_owner_group(
w: WorkspaceClient,
prompts: Prompts,
*,
ctx: WorkspaceContext | None = None,
run_as_collection: bool = False,
a: AccountClient | None = None,
) -> None:
"""
Pick owner group. This group will be assigned as owner to all the migrated tables.
"""
if ctx:
workspace_contexts = [ctx]
else:
workspace_contexts = _get_workspace_contexts(w, a, run_as_collection)

owner_group = workspace_contexts[0].group_manager.pick_owner_group(prompts)
if not owner_group:
return
for workspace_context in workspace_contexts:
config = workspace_context.installation.load(WorkspaceConfig)
config = replace(config, default_owner_group=owner_group)
workspace_context.installation.save(config)


@ucx.command
def migrate_tables(
w: WorkspaceClient,
Expand Down
25 changes: 24 additions & 1 deletion src/databricks/labs/ucx/workspace_access/groups.py
Original file line number Diff line number Diff line change
Expand Up @@ -624,6 +624,28 @@ def delete_original_workspace_groups(self):
# Step 3: Confirm that enumeration no longer returns the deleted groups.
self._wait_for_deleted_workspace_groups(deleted_groups)

def pick_owner_group(self, prompt: Prompts) -> str | None:
# This method is used to select the group that will be used as the owner group.
# The owner group will be assigned by default to all migrated tables/schemas
groups = self._user_account_groups(self._ws.current_user.me().user_name)
if not groups:
logger.warning("No account groups found for the current user.")
return
if len(groups) == 1:
return groups[0].display_name
group_names = [group.display_name for group in groups]
return prompt.choice("Select the group to be used as the owner group", group_names, max_attempts=3)

def _user_account_groups(self, username: str) -> list[Group]:
# This method is used to find all the account groups that a user is a member of.
groups = []
for group in self._list_account_groups("id,displayName,externalId,members"):
for member in group.members:
if member.display == username:
groups.append(group)
break
return groups

def _try_fetch(self) -> Iterable[MigratedGroup]:
state = []
for row in self._sql_backend.fetch(f"SELECT * FROM {escape_sql_identifier(self.full_name)}"):
Expand Down Expand Up @@ -794,7 +816,8 @@ def _list_account_groups(self, scim_attributes: str) -> list[iam.Group]:
continue
account_groups.append(group)
logger.info(f"Found {len(account_groups)} account groups")
sorted_groups: list[iam.Group] = sorted(account_groups, key=lambda _: _.display_name) # type: ignore[arg-type,return-value]
sorted_groups: list[iam.Group] = sorted(account_groups,
key=lambda _: _.display_name) # type: ignore[arg-type,return-value]
return sorted_groups

def _delete_workspace_group_and_wait_for_deletion(self, group_id: str, display_name: str) -> str:
Expand Down

0 comments on commit a651c3b

Please sign in to comment.