Added default owner group selection to the installer (#3370)

related to #3111

---------

Co-authored-by: Serge Smertin <[email protected]>

Author: FastLee

src/databricks/labs/ucx/cli.py

@@ -20,6 +20,7 @@
 from databricks.labs.ucx.hive_metastore.tables import What
 from databricks.labs.ucx.install import AccountInstaller
 from databricks.labs.ucx.source_code.linters.files import LocalCodeLinter
+from databricks.labs.ucx.workspace_access.groups import AccountGroupLookup
 
 ucx = App(__file__)
 logger = get_logger(__file__)
@@ -657,7 +658,7 @@ def assign_owner_group(
     else:
         workspace_contexts = _get_workspace_contexts(w, a, run_as_collection)
 
-    owner_group = workspace_contexts[0].group_manager.pick_owner_group(prompts)
+    owner_group = AccountGroupLookup(workspace_contexts[0].workspace_client).pick_owner_group(prompts)
     if not owner_group:
         return
     for workspace_context in workspace_contexts:

src/databricks/labs/ucx/hive_metastore/ownership.py

@@ -115,7 +115,7 @@ def _application_principal(self) -> str | None:
     @cached_property
     def _static_owner(self) -> str | None:
         # If the default owner group is not valid, fall back to the application principal
-        if self._default_owner_group and self._group_manager.validate_owner_group(self._default_owner_group):
+        if self._default_owner_group and self._group_manager.current_user_in_owner_group(self._default_owner_group):
             logger.warning("Default owner group is not valid, falling back to administrator ownership.")
             return self._default_owner_group
         return self._application_principal

src/databricks/labs/ucx/install.py

@@ -20,28 +20,23 @@
 from databricks.labs.blueprint.parallel import ManyError, Threads
 from databricks.labs.blueprint.tui import Prompts
 from databricks.labs.blueprint.upgrades import Upgrades
-from databricks.labs.blueprint.wheels import (
-    ProductInfo,
-    Version,
-    find_project_root,
-)
+from databricks.labs.blueprint.wheels import ProductInfo, Version, find_project_root
 from databricks.labs.lsql.backends import SqlBackend, StatementExecutionBackend
 from databricks.labs.lsql.dashboards import DashboardMetadata, Dashboards
 from databricks.labs.lsql.deployment import SchemaDeployer
-from databricks.sdk import WorkspaceClient, AccountClient
-from databricks.sdk.useragent import with_extra
+from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.errors import (
     AlreadyExists,
     BadRequest,
     DeadlineExceeded,
     InternalError,
     InvalidParameterValue,
     NotFound,
+    OperationFailed,
     PermissionDenied,
     ResourceAlreadyExists,
     ResourceDoesNotExist,
     Unauthenticated,
-    OperationFailed,
 )
 from databricks.sdk.retries import retried
 from databricks.sdk.service.dashboards import LifecycleState
@@ -51,7 +46,7 @@
     EndpointInfoWarehouseType,
     SpotInstancePolicy,
 )
-
+from databricks.sdk.useragent import with_extra
 from databricks.labs.ucx.__about__ import __version__
 from databricks.labs.ucx.assessment.azure import AzureServicePrincipalInfo
 from databricks.labs.ucx.assessment.clusters import ClusterInfo, PolicyInfo
@@ -81,7 +76,7 @@
 from databricks.labs.ucx.source_code.queries import QueryProblem
 from databricks.labs.ucx.workspace_access.base import Permissions
 from databricks.labs.ucx.workspace_access.generic import WorkspaceObjectInfo
-from databricks.labs.ucx.workspace_access.groups import ConfigureGroups, MigratedGroup
+from databricks.labs.ucx.workspace_access.groups import AccountGroupLookup, ConfigureGroups, MigratedGroup
 
 TAG_STEP = "step"
 WAREHOUSE_PREFIX = "Unity Catalog Migration"
@@ -245,6 +240,12 @@ def _prompt_for_new_installation(self) -> WorkspaceConfig:
         configure_groups = ConfigureGroups(self.prompts)
         configure_groups.run()
         include_databases = self._select_databases()
+
+        # Checking if the user wants to define a default owner group.
+        default_owner_group = None
+        if self.prompts.confirm("Do you want to define a default owner group for all tables and schemas? "):
+            default_owner_group = AccountGroupLookup(self.workspace_client).pick_owner_group(self.prompts)
+
         upload_dependencies = self.prompts.confirm(
             f"Does given workspace {self.workspace_client.get_workspace_id()} block Internet access?"
         )
@@ -267,6 +268,7 @@ def _prompt_for_new_installation(self) -> WorkspaceConfig:
             trigger_job=trigger_job,
             recon_tolerance_percent=recon_tolerance_percent,
             upload_dependencies=upload_dependencies,
+            default_owner_group=default_owner_group,
         )
 
     def _compare_remote_local_versions(self):

src/databricks/labs/ucx/workspace_access/groups.py

@@ -6,7 +6,6 @@
 from collections.abc import Iterable, Collection
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import ClassVar
 
 from databricks.labs.blueprint.limiter import rate_limited
 from databricks.labs.blueprint.parallel import ManyError, Threads
@@ -22,13 +21,15 @@
 )
 from databricks.sdk.retries import retried
 from databricks.sdk.service import iam
-from databricks.sdk.service.iam import Group
+from databricks.sdk.service.iam import Group, User
 
 from databricks.labs.ucx.framework.crawlers import CrawlerBase
 from databricks.labs.ucx.framework.utils import escape_sql_identifier
 
 logger = logging.getLogger(__name__)
 
+SYSTEM_GROUPS: list[str] = ["users", "admins", "account users"]
+
 
 @dataclass
 class MigratedGroup:
@@ -402,7 +403,6 @@ def __init__(self, group_id: str, old_name: str, new_name: str) -> None:
 
 
 class GroupManager(CrawlerBase[MigratedGroup]):
-    _SYSTEM_GROUPS: ClassVar[list[str]] = ["users", "admins", "account users"]
 
     def __init__(  # pylint: disable=too-many-arguments
         self,
@@ -430,6 +430,7 @@ def __init__(  # pylint: disable=too-many-arguments
         self._account_group_regex = account_group_regex
         self._external_id_match = external_id_match
         self._verify_timeout = verify_timeout
+        self._account_groups_lookup = AccountGroupLookup(ws)
 
     def rename_groups(self):
         account_groups_in_workspace = self._account_groups_in_workspace()
@@ -484,6 +485,9 @@ def rename_groups(self):
         # Step 3: Wait for enumeration to also reflect the updated information.
         self._wait_for_renamed_groups(renamed_groups)
 
+    def current_user_in_owner_group(self, group_name: str) -> bool:
+        return self._account_groups_lookup.user_in_group(group_name, self._ws.current_user.me())
+
     def _rename_group(self, group_id: str, old_group_name: str, new_group_name: str) -> None:
         logger.debug(f"Renaming group: {old_group_name} (id={group_id}) -> {new_group_name}")
         self._rate_limited_rename_group_with_retry(group_id, new_group_name)
@@ -555,7 +559,7 @@ def _check_for_renamed_groups(self, expected_groups: Collection[tuple[str, str]]
 
     def reflect_account_groups_on_workspace(self):
         tasks = []
-        account_groups_in_account = self._account_groups_in_account()
+        account_groups_in_account = self._account_groups_lookup.get_mapping()
         account_groups_in_workspace = self._account_groups_in_workspace()
         workspace_groups_in_workspace = self._workspace_groups_in_workspace()
         groups_to_migrate = self.get_migration_state().groups
@@ -624,50 +628,6 @@ 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
-        user_id = self._ws.current_user.me().id
-        if not user_id:
-            logger.error("Couldn't find the user id of the current user.")
-            return None
-        groups = self._user_account_groups(user_id)
-        if not groups:
-            logger.warning("No account groups found for the current user.")
-            return None
-        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 validate_owner_group(self, group_name: str) -> bool:
-        # This method is used to validate that the current owner is a member of the group
-        user_id = self._ws.current_user.me().id
-        if not user_id:
-            logger.warning("No user found for the current session.")
-            return False
-        groups = self._user_account_groups(user_id)
-        if not groups:
-            logger.warning("No account groups found for the current user.")
-            return False
-        group_names = [group.display_name for group in groups]
-        return group_name in group_names
-
-    def _user_account_groups(self, user_id: str) -> list[Group]:
-        # This method is used to find all the account groups that a user is a member of.
-        groups: list[Group] = []
-        account_groups = self._list_account_groups("id,displayName,externalId,members")
-        if not account_groups:
-            return groups
-        for group in account_groups:
-            if not group.members:
-                continue
-            for member in group.members:
-                if member.value == user_id:
-                    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)}"):
@@ -690,13 +650,13 @@ def _try_fetch(self) -> Iterable[MigratedGroup]:
 
     def _crawl(self) -> Iterable[MigratedGroup]:
         workspace_groups_in_workspace = self._workspace_groups_in_workspace()
-        account_groups_in_account = self._account_groups_in_account()
+        account_groups_in_account = self._account_groups_lookup.get_mapping()
         strategy = self._get_strategy(workspace_groups_in_workspace, account_groups_in_account)
         yield from strategy.generate_migrated_groups()
 
     def validate_group_membership(self) -> list[dict]:
         workspace_groups_in_workspace = self._workspace_groups_in_workspace()
-        account_groups_in_account = self._account_groups_in_account()
+        account_groups_in_account = self._account_groups_lookup.get_mapping()
         strategy = self._get_strategy(workspace_groups_in_workspace, account_groups_in_account)
         migrated_groups = list(strategy.generate_migrated_groups())
         mismatch_group = []
@@ -762,17 +722,9 @@ def _account_groups_in_workspace(self) -> dict[str, Group]:
             groups[group.display_name] = group
         return groups
 
-    def _account_groups_in_account(self) -> dict[str, Group]:
-        groups = {}
-        for group in self._list_account_groups("id,displayName,externalId"):
-            if not group.display_name:
-                logger.debug(f"Ignoring account group in without name: {group.id}")
-                continue
-            groups[group.display_name] = group
-        return groups
-
-    def _is_group_out_of_scope(self, group: iam.Group, resource_type: str) -> bool:
-        if group.display_name in self._SYSTEM_GROUPS:
+    @staticmethod
+    def _is_group_out_of_scope(group: iam.Group, resource_type: str) -> bool:
+        if group.display_name in SYSTEM_GROUPS:
             return True
         meta = group.meta
         if not meta:
@@ -826,23 +778,6 @@ def _get_account_group(self, group_id: str) -> Group | None:
             logger.warning(f"Group with ID {group_id} does not exist anymore in the Databricks account.")
             return None
 
-    def _list_account_groups(self, scim_attributes: str) -> list[iam.Group]:
-        # TODO: we should avoid using this method, as it's not documented
-        # get account-level groups even if they're not (yet) assigned to a workspace
-        logger.info(f"Listing account groups with {scim_attributes}...")
-        account_groups = []
-        raw = self._ws.api_client.do("GET", "/api/2.0/account/scim/v2/Groups", query={"attributes": scim_attributes})
-        for resource in raw.get("Resources", []):  # type: ignore[union-attr]
-            group = iam.Group.from_dict(resource)
-            if group.display_name in self._SYSTEM_GROUPS:
-                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 if _.display_name else ""
-        )  # 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:
         logger.debug(f"Deleting workspace group: {display_name} (id={group_id})")
         self._delete_workspace_group(group_id, display_name)
@@ -975,6 +910,81 @@ def _get_strategy(
         )
 
 
+class AccountGroupLookup:
+    def __init__(self, ws: WorkspaceClient):
+        self._ws = ws
+
+    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
+        user_id = self._ws.current_user.me().id
+        if not user_id:
+            logger.error("Couldn't find the user id of the current user.")
+            return None
+        groups = self._user_account_groups(user_id)
+        if not groups:
+            logger.warning("No account groups found for the current user.")
+            return None
+        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_in_group(self, group_name: str, user: User) -> bool:
+        # This method is used to validate that the current user is a member of the group
+        user_id = user.id
+        if not user_id:
+            logger.warning("No user found for the current session.")
+            return False
+        groups = self._user_account_groups(user_id)
+        if not groups:
+            logger.warning("No account groups found for the current user.")
+            return False
+        group_names = [group.display_name for group in groups]
+        return group_name in group_names
+
+    def _user_account_groups(self, user_id: str) -> list[Group]:
+        # This method is used to find all the account groups that a user is a member of.
+        groups: list[Group] = []
+        account_groups = self._list_account_groups("id,displayName,externalId,members")
+        if not account_groups:
+            return groups
+        for group in account_groups:
+            if not group.members:
+                continue
+            for member in group.members:
+                if member.value == user_id:
+                    groups.append(group)
+                    break
+        return groups
+
+    def _list_account_groups(self, scim_attributes: str) -> list[iam.Group]:
+        # TODO: we should avoid using this method, as it's not documented
+        # get account-level groups even if they're not (yet) assigned to a workspace
+        logger.info(f"Listing account groups with {scim_attributes}...")
+        account_groups = []
+        raw = self._ws.api_client.do("GET", "/api/2.0/account/scim/v2/Groups", query={"attributes": scim_attributes})
+        for resource in raw.get("Resources", []):  # type: ignore[union-attr]
+            group = iam.Group.from_dict(resource)
+            if group.display_name in SYSTEM_GROUPS:
+                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 if _.display_name else ""
+        )  # type: ignore[arg-type,return-value]
+        return sorted_groups
+
+    def get_mapping(self) -> dict[str, Group]:
+        groups: dict[str, Group] = {}
+        for group in self._list_account_groups("id,displayName,externalId"):
+            if not group.display_name:
+                logger.debug(f"Ignoring account group in without name: {group.id}")
+                continue
+            groups[group.display_name] = group
+        return groups
+
+
 class ConfigureGroups:
     renamed_group_prefix = "db-temp-"
     workspace_group_regex = None

tests/unit/hive_metastore/test_grants.py

@@ -14,7 +14,7 @@
 from databricks.labs.ucx.hive_metastore.tables import Table, TablesCrawler
 from databricks.labs.ucx.hive_metastore.udfs import UdfsCrawler
 from databricks.labs.ucx.progress.history import ProgressEncoder
-from databricks.labs.ucx.workspace_access.groups import GroupManager
+from databricks.labs.ucx.workspace_access.groups import GroupManager, AccountGroupLookup
 from tests.unit import mock_workspace_client
 
 
@@ -956,7 +956,6 @@ def test_grant_supports_history(mock_backend, grant_record: Grant, history_recor
 # Testing the validation in retrival of the default owner group. 666 is the current_user user_id.
 @pytest.mark.parametrize("user_id, expected", [("666", True), ("777", False)])
 def test_default_owner(user_id, expected) -> None:
-    sql_backend = MockBackend()
     ws = mock_workspace_client()
 
     account_admins_group = Group(
@@ -966,5 +965,5 @@ def test_default_owner(user_id, expected) -> None:
         "Resources": [account_admins_group.as_dict()],
     }
 
-    group_manager = GroupManager(sql_backend, ws, "ucx")
-    assert group_manager.validate_owner_group("owners") == expected
+    account_group_lookup = AccountGroupLookup(ws)
+    assert account_group_lookup.user_in_group("owners", ws.current_user.me()) == expected

tests/unit/hive_metastore/test_tables.py

@@ -836,7 +836,7 @@ def test_default_securable_ownership(
         Table("main", "foo", "baz", "VIEW", "UNKNOWN", None, "select * from bar"),
     ]
     group_manager = create_autospec(GroupManager)
-    group_manager.validate_owner_group.return_value = valid_admin
+    group_manager.current_user_in_owner_group.return_value = valid_admin
 
     ownership = DefaultSecurableOwnership(
         admin_locator, table_crawler, group_manager, default_owner_group, lambda: cli_user

tests/unit/install/conftest.py

@@ -77,7 +77,7 @@ def download(path: str) -> io.StringIO | io.BytesIO:
     workspace_client = create_autospec(WorkspaceClient)
 
     workspace_client.current_user.me = lambda: iam.User(
-        user_name="[email protected]", groups=[iam.ComplexValue(display="admins")]
+        user_name="[email protected]", id="666", groups=[iam.ComplexValue(display="admins")]
     )
     workspace_client.config.host = "https://foo"
     workspace_client.config.is_aws = True