Validate toolkit notebook test (#183)

Toolkit notebook test validated. Switched order of crawling TACL and
group migration in notebook toolkit.py.

---------

Co-authored-by: Sara Ivanyos <[email protected]>

Author: saraivdbx

notebooks/toolkit.py

@@ -94,6 +94,16 @@
 
 # COMMAND ----------
 
+# MAGIC %md
+# MAGIC
+# MAGIC ## Inventorize Table ACL's
+
+# COMMAND ----------
+
+tacltoolkit.grants_snapshot()
+
+# COMMAND ----------
+
 # MAGIC %md
 # MAGIC
 # MAGIC ## Inventorize the permissions
@@ -148,15 +158,6 @@
 
 toolkit.delete_backup_groups()
 
-# COMMAND ----------
-
-# MAGIC %md
-# MAGIC
-# MAGIC ## Inventorize Table ACL's
-
-# COMMAND ----------
-
-tacltoolkit.grants_snapshot()
 
 # COMMAND ----------
 

src/databricks/labs/ucx/toolkits/group_migration.py

@@ -5,7 +5,9 @@
 from databricks.labs.ucx.config import MigrationConfig
 from databricks.labs.ucx.inventory.inventorizer import Inventorizers
 from databricks.labs.ucx.inventory.permissions import PermissionManager
-from databricks.labs.ucx.inventory.permissions_inventory import PermissionsInventoryTable
+from databricks.labs.ucx.inventory.permissions_inventory import (
+    PermissionsInventoryTable,
+)
 from databricks.labs.ucx.managers.group import GroupManager
 
 

tests/integration/test_installation.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import random
 import shutil
 import subprocess
 import sys
@@ -8,9 +9,13 @@
 
 import pytest
 from databricks.sdk.service import compute, jobs
+from databricks.sdk.service.iam import PermissionLevel
 from databricks.sdk.service.workspace import ImportFormat
 
+from databricks.labs.ucx.inventory.types import RequestObjectType
 from databricks.labs.ucx.providers.mixins.compute import CommandExecutor
+from databricks.labs.ucx.tacl.grants import Grant
+from databricks.labs.ucx.tacl.tables import Table
 
 logger = logging.getLogger(__name__)
 
@@ -96,63 +101,55 @@ def test_sql_backend_works(ws, wsfs_wheel):
 def test_toolkit_notebook(
     ws,
     sql_exec,
+    sql_fetch_all,
     wsfs_wheel,
-    make_cluster,
     make_cluster_policy,
-    make_directory,
+    make_cluster_policy_permissions,
     make_ucx_group,
-    make_instance_pool,
     make_job,
-    make_notebook,
-    make_pipeline,
+    make_job_permissions,
     make_random,
-    make_repo,
-    make_secret_scope,
     make_schema,
     make_table,
-    make_user,
 ):
     logger.info("setting up fixtures")
 
-    user_a = make_user()
-    user_b = make_user()
-    user_c = make_user()
-
-    logger.info(f"user_a={user_a}, user_b={user_b}, user_c={user_c}, ")
-
-    # TODO add users to groups
     ws_group_a, acc_group_a = make_ucx_group()
+    members_src_a = sorted([_.display for _ in ws.groups.get(id=ws_group_a.id).members])
     ws_group_b, acc_group_b = make_ucx_group()
+    members_src_b = sorted([_.display for _ in ws.groups.get(id=ws_group_b.id).members])
     ws_group_c, acc_group_c = make_ucx_group()
+    members_src_c = sorted([_.display for _ in ws.groups.get(id=ws_group_c.id).members])
 
     selected_groups = ",".join([ws_group_a.display_name, ws_group_b.display_name, ws_group_c.display_name])
 
     logger.info(f"group_a={ws_group_a}, group_b={ws_group_b}, group_c={ws_group_c}, ")
 
-    cluster = make_cluster(instance_pool_id=os.environ["TEST_INSTANCE_POOL_ID"], single_node=True)
     cluster_policy = make_cluster_policy()
-    directory = make_directory()
-    instance_pool = make_instance_pool()
+    make_cluster_policy_permissions(
+        object_id=cluster_policy.policy_id,
+        permission_level=random.choice([PermissionLevel.CAN_USE]),
+        group_name=ws_group_a.display_name,
+    )
+    cpp_src = ws.permissions.get(RequestObjectType.CLUSTER_POLICIES, cluster_policy.policy_id)
+    cluster_policy_src_permissions = sorted(
+        [_ for _ in cpp_src.access_control_list if _.group_name == ws_group_a.display_name],
+        key=lambda p: p.group_name,
+    )
     job = make_job()
-    notebook = make_notebook()
-    pipeline = make_pipeline()
-    repo = make_repo()
-    secret_scope = make_secret_scope()
-
-    logger.info(
-        f"cluster={cluster}, "
-        f"cluster_policy={cluster_policy}, "
-        f"directory={directory}, "
-        f"instance_pool={instance_pool}, "
-        f"job={job}, "
-        f"notebook={notebook}, "
-        f"pipeline={pipeline}"
-        f"repo={repo}, "
-        f"secret_scope={secret_scope}, "
+    make_job_permissions(
+        object_id=job.job_id,
+        permission_level=random.choice(
+            [PermissionLevel.CAN_VIEW, PermissionLevel.CAN_MANAGE_RUN, PermissionLevel.CAN_MANAGE]
+        ),
+        group_name=ws_group_b.display_name,
     )
-
-    # TODO create fixtures for DBSQL assets
-    # TODO set permissions
+    jp_src = ws.permissions.get(RequestObjectType.JOBS, job.job_id)
+    job_src_permissions = sorted(
+        [_ for _ in jp_src.access_control_list if _.group_name == ws_group_b.display_name],
+        key=lambda p: p.group_name,
+    )
+    logger.info(f"cluster_policy={cluster_policy}, job={job}, ")
 
     schema_a = make_schema()
     schema_b = make_schema()
@@ -222,7 +219,91 @@ def test_toolkit_notebook(
 
     try:
         ws.jobs.run_now(created_job.job_id).result()
-        # TODO Validate migration, tacl
+
+        logger.info("validating group ids")
+
+        dst_ws_group_a = ws.groups.list(filter=f"displayName eq {ws_group_a.display_name}")[0]
+        assert (
+            ws_group_a.id != dst_ws_group_a.id
+        ), f"Group id for target group {ws_group_a.display_name} should differ from group id of source group"
+
+        dst_ws_group_b = ws.groups.list(filter=f"displayName eq {ws_group_b.display_name}")[0]
+        assert (
+            ws_group_b.id != dst_ws_group_b.id
+        ), f"Group id for target group {ws_group_b.display_name} should differ from group id of source group"
+
+        dst_ws_group_c = ws.groups.list(filter=f"displayName eq {ws_group_c.display_name}")[0]
+        assert (
+            ws_group_c.id != dst_ws_group_c.id
+        ), f"Group id for target group {ws_group_c.display_name} should differ from group id of source group"
+
+        logger.info("validating group members")
+
+        members_dst_a = sorted([_.display for _ in ws.groups.get(id=dst_ws_group_a.id).members])
+        assert members_dst_a == members_src_a, f"Members from {ws_group_a.display_name} were not migrated correctly"
+
+        members_dst_b = sorted([_.display for _ in ws.groups.get(id=dst_ws_group_b.id).members])
+        assert members_dst_b == members_src_b, f"Members in {ws_group_b.display_name} were not migrated correctly"
+
+        members_dst_c = sorted([_.display for _ in ws.groups.get(id=dst_ws_group_c.id).members])
+        assert members_dst_c == members_src_c, f"Members in {ws_group_c.display_name} were not migrated correctly"
+
+        logger.info("validating permissions")
+
+        cp_dst = ws.permissions.get(RequestObjectType.CLUSTER_POLICIES, cluster_policy.policy_id)
+        cluster_policy_dst_permissions = sorted(
+            [_ for _ in cp_dst.access_control_list if _.group_name == ws_group_a.display_name],
+            key=lambda p: p.group_name,
+        )
+        assert len(cluster_policy_dst_permissions) == len(
+            cluster_policy_src_permissions
+        ), "Target permissions were not applied correctly for cluster policies"
+        assert [t.all_permissions for t in cluster_policy_dst_permissions] == [
+            s.all_permissions for s in cluster_policy_src_permissions
+        ], "Target permissions were not applied correctly for cluster policies"
+
+        jp_dst = ws.permissions.get(RequestObjectType.JOBS, job.job_id)
+        job_dst_permissions = sorted(
+            [_ for _ in jp_dst.access_control_list if _.group_name == ws_group_b.display_name],
+            key=lambda p: p.group_name,
+        )
+        assert len(job_dst_permissions) == len(
+            job_src_permissions
+        ), f"Target permissions were not applied correctly for {RequestObjectType.JOBS}/{job.job_id}"
+        assert [t.all_permissions for t in job_dst_permissions] == [
+            s.all_permissions for s in job_src_permissions
+        ], f"Target permissions were not applied correctly for {RequestObjectType.JOBS}/{job.job_id}"
+
+        logger.info("validating tacl")
+
+        tables = sql_fetch_all(f"SELECT * FROM hive_metastore.{inventory_database}.tables")
+        print(list(sql_fetch_all(f"SELECT * FROM hive_metastore.{inventory_database}.tables")))
+
+        all_tables = {}
+        for t_row in tables:
+            table = Table(*t_row)
+            all_tables[table.key] = table
+
+        assert len(all_tables) >= 2, "must have at least two tables"
+
+        logger.debug(f"all tables={all_tables}, ")
+
+        grants = sql_fetch_all(f"SELECT * FROM hive_metastore.{inventory_database}.grants")
+        all_grants = {}
+        for g_row in grants:
+            grant = Grant(*g_row)
+            if grant.table:
+                all_grants[f"{grant.principal}.{grant.catalog}.{grant.database}.{grant.table}"] = grant.action_type
+            else:
+                all_grants[f"{grant.principal}.{grant.catalog}.{grant.database}"] = grant.action_type
+
+        logger.debug(f"all grants={all_grants}, ")
+
+        assert len(all_grants) >= 3, "must have at least three grants"
+        assert all_grants[f"{ws_group_a.display_name}.{table_a}"] == "SELECT"
+        assert all_grants[f"{ws_group_b.display_name}.{table_b}"] == "SELECT"
+        assert all_grants[f"{ws_group_b.display_name}.{schema_b}"] == "MODIFY"
+
     finally:
         logger.info("deleting workbook")