From a651c3b7255b18a159725a0f8e2cbcaa5a1ee7f2 Mon Sep 17 00:00:00 2001
From: Liran Bareket <liran.bareket@databricks.com>
Date: Mon, 4 Nov 2024 15:25:40 -0500
Subject: [PATCH] Added CLI Option

---
 labs.yml                                      |  3 ++-
 src/databricks/labs/ucx/cli.py                | 27 +++++++++++++++++++
 .../labs/ucx/workspace_access/groups.py       | 25 ++++++++++++++++-
 3 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/labs.yml b/labs.yml
index 1ab75fa6e9..58fd85800e 100644
--- a/labs.yml
+++ b/labs.yml
@@ -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.
diff --git a/src/databricks/labs/ucx/cli.py b/src/databricks/labs/ucx/cli.py
index 2215f79fb6..02b4fb782a 100644
--- a/src/databricks/labs/ucx/cli.py
+++ b/src/databricks/labs/ucx/cli.py
@@ -1,3 +1,4 @@
+from dataclasses import replace
 from io import BytesIO
 import json
 import webbrowser
@@ -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,
diff --git a/src/databricks/labs/ucx/workspace_access/groups.py b/src/databricks/labs/ucx/workspace_access/groups.py
index b361ed773d..4c33d17f1c 100644
--- a/src/databricks/labs/ucx/workspace_access/groups.py
+++ b/src/databricks/labs/ucx/workspace_access/groups.py
@@ -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)}"):
@@ -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: