Added integration test and converted from username to userid for memb…

…ership lookup

src/databricks/labs/ucx/cli.py

@@ -652,11 +652,11 @@ def assign_owner_group(
     else:
         workspace_contexts = _get_workspace_contexts(w, a, run_as_collection)
 
-    username = w.current_user.me().user_name
-    if not username:
-        logger.error("Couldn't find the username of the current user.")
+    user_id = w.current_user.me().id
+    if not user_id:
+        logger.error("Couldn't find the user id of the current user.")
         return
-    owner_group = workspace_contexts[0].group_manager.pick_owner_group(prompts, username)
+    owner_group = workspace_contexts[0].group_manager.pick_owner_group(prompts, user_id)
     if not owner_group:
         return
     for workspace_context in workspace_contexts:

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

@@ -624,10 +624,10 @@ 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, username: str) -> str | None:
+    def pick_owner_group(self, prompt: Prompts, user_id: str) -> 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(username)
+        groups = self._user_account_groups(user_id)
         if not groups:
             logger.warning("No account groups found for the current user.")
             return None
@@ -638,18 +638,18 @@ def pick_owner_group(self, prompt: Prompts, username: str) -> str | None:
 
     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
-        username = self._ws.current_user.me().user_name
-        if not username:
+        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(username)
+        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, username: str) -> list[Group]:
+    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")
@@ -659,7 +659,7 @@ def _user_account_groups(self, username: str) -> list[Group]:
             if not group.members:
                 continue
             for member in group.members:
-                if member.display == username:
+                if member.value == user_id:
                     groups.append(group)
                     break
         return groups

tests/integration/hive_metastore/test_grants.py

@@ -1,5 +1,6 @@
 import logging
 from collections import defaultdict
+from dataclasses import replace
 from datetime import timedelta
 
 from databricks.sdk.errors import NotFound
@@ -115,15 +116,18 @@ def test_all_grants_for_other_objects(
     assert {"DENIED_SELECT"} == found_anonymous_function_grants[group_d.display_name]
 
 
-def test_grant_ownership(ws, runtime_ctx, inventory_schema, sql_backend) -> None:
-    """Verify the ownership can be determined for crawled grants."""
-    # This currently isn't very useful: we can't locate specific owners for grants.
+def test_grant_ownership(ws, runtime_ctx, inventory_schema, sql_backend, make_random, make_acc_group) -> None:
+    """Verify the ownership can be determined for crawled grants.
+    This currently isn't very useful: we can't locate specific owners for grants"""
 
     schema = runtime_ctx.make_schema()
     this_user = ws.current_user.me()
     sql_backend.execute(f"GRANT SELECT ON SCHEMA {escape_sql_identifier(schema.full_name)} TO `{this_user.user_name}`")
     table_crawler = TablesCrawler(sql_backend, schema=inventory_schema, include_databases=[schema.name])
     udf_crawler = UdfsCrawler(sql_backend, schema=inventory_schema, include_databases=[schema.name])
+    current_user = ws.current_user.me()
+    admin_group_name = f"admin_group_{make_random()}"
+    make_acc_group(display_name=admin_group_name, members=[current_user.id], wait_for_provisioning=True)
 
     # Produce the crawled records.
     crawler = GrantsCrawler(table_crawler, udf_crawler, include_databases=[schema.name])
@@ -135,3 +139,8 @@ def test_grant_ownership(ws, runtime_ctx, inventory_schema, sql_backend) -> None
     # Verify ownership can be made.
     ownership = GrantOwnership(runtime_ctx.administrator_locator)
     assert ownership.owner_of(grant_record) == runtime_ctx.administrator_locator.get_workspace_administrator()
+
+    # Test default group ownership
+    test_config = replace(runtime_ctx.config, default_owner_group=admin_group_name)
+    default_owner_group = runtime_ctx.replace(config=test_config).default_owner_group
+    assert default_owner_group == admin_group_name

tests/unit/__init__.py

@@ -163,7 +163,7 @@ def mock_workspace_client(
 ):
     ws = create_autospec(WorkspaceClient)
     ws.current_user.me.side_effect = lambda: iam.User(
-        user_name="[email protected]", groups=[iam.ComplexValue(display="admins")]
+        id="666", user_name="[email protected]", groups=[iam.ComplexValue(display="admins")]
     )
     ws.api_client.do.return_value = {}
     ws.permissions.get.return_value = {}

tests/unit/hive_metastore/test_grants.py

@@ -8,15 +8,13 @@
 import yaml
 from databricks.labs.lsql.backends import MockBackend
 from databricks.labs.lsql.core import Row
-from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.iam import ComplexValue, Group
 
 from databricks.labs.ucx.__about__ import __version__ as ucx_version
 from databricks.labs.ucx.contexts.workspace_cli import WorkspaceContext
 from databricks.labs.ucx.framework.owners import AdministratorLocator
 from databricks.labs.ucx.hive_metastore.catalog_schema import Catalog, Schema
 from databricks.labs.ucx.hive_metastore.grants import Grant, GrantsCrawler, MigrateGrants, GrantOwnership
-from databricks.labs.ucx.hive_metastore.ownership import DefaultSecurableOwnership
 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
@@ -223,7 +221,7 @@ def test_crawler_crawl() -> None:
             "SHOW GRANTS ON CATALOG `hive_metastore`": SHOW_GRANTS[("princ1", "USE", "CATALOG$", "hive_metastore"),],
             "SHOW GRANTS ON DATABASE `hive_metastore`\\.`database_one`": SHOW_GRANTS[
                 ("princ2", "OWN", "DATABASE", "database_one"),
-                    # Enumerating database grants can include some grants for the catalog.
+                # Enumerating database grants can include some grants for the catalog.
                 ("princ1", "SELECT", "CATALOG$", None),
             ],
             "SHOW GRANTS ON VIEW `hive_metastore`\\.`database_one`\\.`table_one`": SHOW_GRANTS[
@@ -959,14 +957,9 @@ def test_grant_supports_history(mock_backend, grant_record: Grant, history_recor
     assert rows == [history_record]
 
 
-@pytest.mark.parametrize(
-    "member, expected",
-    [
-        ("[email protected]", "owners"),
-        ("[email protected]", None)
-    ]
-)
-def test_default_owner(member, expected) -> None:
+# Testing the validation in retrival of the default owner group. 666 is the current_user user_id.
+@pytest.mark.parametrize("user_id, expected", [("666", "owners"), ("777", None)])
+def test_default_owner(user_id, expected) -> None:
     ws = mock_workspace_client()
     download_yaml = {
         'config.yml': yaml.dump(
@@ -984,9 +977,9 @@ def test_default_owner(member, expected) -> None:
     }
 
     ws.workspace.download.side_effect = lambda file_name: io.StringIO(download_yaml[os.path.basename(file_name)])
-    account_admins_group = Group(id="1234",
-                                 display_name="owners",
-                                 members=[ComplexValue(display=member, value="1")])
+    account_admins_group = Group(
+        id="1234", display_name="owners", members=[ComplexValue(display="User Name", value=user_id)]
+    )
     ws.api_client.do.return_value = {
         "Resources": [account_admins_group.as_dict()],
     }

tests/unit/test_cli.py

@@ -137,7 +137,9 @@ def mock_download(path: str, **_) -> io.StringIO | io.BytesIO:
     workspace_client = create_autospec(WorkspaceClient)
     workspace_client.get_workspace_id.return_value = workspace_id
     workspace_client.config.host = 'https://localhost'
-    workspace_client.current_user.me.return_value = User(user_name="foo", groups=[ComplexValue(display="admins")])
+    workspace_client.current_user.me.return_value = User(
+        id="666", user_name="foo", groups=[ComplexValue(display="admins")]
+    )
     workspace_client.workspace.download.side_effect = mock_download
     workspace_client.statement_execution.execute_statement.return_value = sql.StatementResponse(
         status=sql.StatementStatus(state=sql.StatementState.SUCCEEDED),
@@ -1284,7 +1286,7 @@ def test_assign_owner_group(tmp_path, workspace_clients, acc_client, run_as_coll
     prompts = MockPrompts({"Please provide the group name to assign as owner": "test_group"})
 
     for workspace_client in workspace_clients:
-        group1 = Group(id="1", display_name="test_group", members=[ComplexValue(display="foo", value="1")])
+        group1 = Group(id="1", display_name="test_group", members=[ComplexValue(display="foo", value="666")])
         workspace_client.api_client.do.return_value = {"Resources": [group1.as_dict()]}
 
     assign_owner_group(workspace_clients[0], prompts, run_as_collection=run_as_collection, a=acc_client)