fix and cleanup e2e after refactoring (#185)

After running the e2e tests, I've located some small bugs appeared after
the refactoring.

I've provided both fixes and additional tests to verify that everything
works smoothly.

Author: renardeinside

src/databricks/labs/ucx/inventory/permissions_inventory.py

@@ -1,6 +1,5 @@
 import logging
 
-import pandas as pd
 from databricks.sdk import WorkspaceClient
 from pyspark.sql import DataFrame
 from pyspark.sql.types import StringType, StructField, StructType
@@ -39,7 +38,7 @@ def cleanup(self):
     def save(self, items: list[PermissionsInventoryItem]):
         # TODO: update instead of append
         logger.info(f"Saving {len(items)} items to inventory table {self._table}")
-        serialized_items = pd.DataFrame([item.as_dict() for item in items])
+        serialized_items = [item.as_dict() for item in items]
         df = self.spark.createDataFrame(serialized_items, schema=self._table_schema)
         df.write.mode("append").format("delta").saveAsTable(self._table)
         logger.info("Successfully saved the items to inventory table")

src/databricks/labs/ucx/inventory/verification.py

@@ -1,14 +1,15 @@
 from typing import Literal
 
 from databricks.sdk import WorkspaceClient
-from databricks.sdk.service import workspace
 
 from databricks.labs.ucx.providers.groups_info import GroupMigrationState
+from databricks.labs.ucx.support.secrets import SecretScopesSupport
 
 
 class VerificationManager:
-    def __init__(self, ws: WorkspaceClient):
+    def __init__(self, ws: WorkspaceClient, secrets_support: SecretScopesSupport):
         self._ws = ws
+        self._secrets_support = secrets_support
 
     def verify(
         self, migration_state: GroupMigrationState, target: Literal["backup", "account"], tuples: list[tuple[str, str]]
@@ -51,8 +52,8 @@ def verify_applied_scope_acls(
         for mi in migration_state.groups:
             src_name = getattr(mi, base_attr).display_name
             dst_name = getattr(mi, target).display_name
-            src_permission = self._secret_scope_permission(scope_name, src_name)
-            dst_permission = self._secret_scope_permission(scope_name, dst_name)
+            src_permission = self._secrets_support.secret_scope_permission(scope_name, src_name)
+            dst_permission = self._secrets_support.secret_scope_permission(scope_name, dst_name)
             assert src_permission == dst_permission, "Scope ACLs were not applied correctly"
 
     def verify_roles_and_entitlements(self, migration_state: GroupMigrationState, target: Literal["backup", "account"]):
@@ -65,9 +66,3 @@ def verify_roles_and_entitlements(self, migration_state: GroupMigrationState, ta
 
             assert base_group_info.roles == target_group_info.roles
             assert base_group_info.entitlements == target_group_info.entitlements
-
-    def _secret_scope_permission(self, scope_name: str, group_name: str) -> workspace.AclPermission | None:
-        for acl in self._ws.secrets.list_acls(scope=scope_name):
-            if acl.principal == group_name:
-                return acl.permission
-        return None

src/databricks/labs/ucx/support/group_level.py

@@ -20,7 +20,7 @@ def _crawler_task(self, group: iam.Group, property_name: str):
     @sleep_and_retry
     @limits(calls=10, period=1)
     def _applier_task(self, group_id: str, value: list[iam.ComplexValue], property_name: str):
-        operations = [iam.Patch(op=iam.PatchOp.ADD, path=property_name, value=value)]
+        operations = [iam.Patch(op=iam.PatchOp.ADD, path=property_name, value=[e.as_dict() for e in value])]
         schemas = [iam.PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP]
         self._ws.groups.patch(id=group_id, operations=operations, schemas=schemas)
 

src/databricks/labs/ucx/support/impl.py

@@ -30,11 +30,11 @@ def __init__(self, ws: WorkspaceClient, num_threads: int, workspace_start_path:
             ws=ws,
             listings=[
                 listing_wrapper(ws.clusters.list, "cluster_id", RequestObjectType.CLUSTERS),
-                listing_wrapper(ws.cluster_policies.list, "cluster_policy_id", RequestObjectType.CLUSTER_POLICIES),
+                listing_wrapper(ws.cluster_policies.list, "policy_id", RequestObjectType.CLUSTER_POLICIES),
                 listing_wrapper(ws.instance_pools.list, "instance_pool_id", RequestObjectType.INSTANCE_POOLS),
                 listing_wrapper(ws.warehouses.list, "id", RequestObjectType.SQL_WAREHOUSES),
                 listing_wrapper(ws.jobs.list, "job_id", RequestObjectType.JOBS),
-                listing_wrapper(ws.pipelines.list, "pipeline_id", RequestObjectType.PIPELINES),
+                listing_wrapper(ws.pipelines.list_pipelines, "pipeline_id", RequestObjectType.PIPELINES),
                 listing_wrapper(experiments_listing(ws), "experiment_id", RequestObjectType.EXPERIMENTS),
                 listing_wrapper(models_listing(ws), "id", RequestObjectType.REGISTERED_MODELS),
                 workspace_listing(ws, num_threads=num_threads, start_path=workspace_start_path),
@@ -64,13 +64,13 @@ def supports(self) -> dict[str, BaseSupport]:
             "roles": self._scim_support,
             # generic API
             "clusters": self._generic_support,
-            "cluster_policies": self._generic_support,
-            "instance_pools": self._generic_support,
-            "sql_warehouses": self._generic_support,
+            "cluster-policies": self._generic_support,
+            "instance-pools": self._generic_support,
+            "sql/warehouses": self._generic_support,
             "jobs": self._generic_support,
             "pipelines": self._generic_support,
             "experiments": self._generic_support,
-            "registered_models": self._generic_support,
+            "registered-models": self._generic_support,
             "tokens": self._generic_support,
             "passwords": self._generic_support,
             # workspace objects

src/databricks/labs/ucx/support/permissions.py

@@ -39,7 +39,12 @@ def _safe_get_permissions(
             permissions = ws.permissions.get(request_object_type, object_id)
             return permissions
         except DatabricksError as e:
-            if e.error_code in ["RESOURCE_DOES_NOT_EXIST", "RESOURCE_NOT_FOUND", "PERMISSION_DENIED"]:
+            if e.error_code in [
+                "RESOURCE_DOES_NOT_EXIST",
+                "RESOURCE_NOT_FOUND",
+                "PERMISSION_DENIED",
+                "FEATURE_DISABLED",
+            ]:
                 logger.warning(f"Could not get permissions for {request_object_type} {object_id} due to {e.error_code}")
                 return None
             else:
@@ -79,7 +84,7 @@ def _prepare_new_acl(
     def _applier_task(
         self, ws: WorkspaceClient, object_id: str, acl: list[iam.AccessControlRequest], request_type: RequestObjectType
     ):
-        ws.permissions.update(request_type, object_id, acl)
+        ws.permissions.update(request_object_type=request_type, request_object_id=object_id, access_control_list=acl)
 
     @sleep_and_retry
     @limits(calls=100, period=1)
@@ -90,10 +95,13 @@ def _crawler_task(
         request_type: RequestObjectType,
     ) -> PermissionsInventoryItem | None:
         permissions = self._safe_get_permissions(ws, request_type, object_id)
+
+        support = object_id if request_type == RequestObjectType.AUTHORIZATION else request_type.value
+
         if permissions:
             return PermissionsInventoryItem(
                 object_id=object_id,
-                support=request_type.value,
+                support=support,
                 raw_object_permissions=json.dumps(permissions.as_dict()),
             )
 
@@ -103,10 +111,17 @@ def _get_apply_task(
         new_acl = self._prepare_new_acl(
             iam.ObjectPermissions.from_dict(json.loads(item.raw_object_permissions)), migration_state, destination
         )
+
+        request_type = (
+            RequestObjectType.AUTHORIZATION
+            if item.support in ("passwords", "tokens")
+            else RequestObjectType(item.support)
+        )
+
         return partial(
             self._applier_task,
             ws=self._ws,
-            request_type=RequestObjectType(item.support),
+            request_type=request_type,
             acl=new_acl,
             object_id=item.object_id,
         )

src/databricks/labs/ucx/support/secrets.py

@@ -1,4 +1,6 @@
 import json
+import random
+import time
 from functools import partial
 
 from databricks.sdk import WorkspaceClient
@@ -33,10 +35,42 @@ def is_item_relevant(self, item: PermissionsInventoryItem, migration_state: Grou
         mentioned_groups = [acl.principal for acl in acls]
         return any(g in mentioned_groups for g in [info.workspace.display_name for info in migration_state.groups])
 
+    def secret_scope_permission(self, scope_name: str, group_name: str) -> workspace.AclPermission | None:
+        for acl in self._ws.secrets.list_acls(scope=scope_name):
+            if acl.principal == group_name:
+                return acl.permission
+        return None
+
+    def _inflight_check(
+        self, scope_name: str, group_name: str, expected_permission: workspace.AclPermission, num_retries: int = 5
+    ):
+        # in-flight check for the applied permissions
+        # the api might be inconsistent, therefore we need to check that the permissions were applied
+        # TODO: add mixin to SDK
+        retries_left = num_retries
+        while retries_left > 0:
+            time.sleep(random.random() * 2)
+            applied_permission = self.secret_scope_permission(scope_name=scope_name, group_name=group_name)
+            if applied_permission:
+                if applied_permission == expected_permission:
+                    return
+                else:
+                    msg = (
+                        f"Applied permission {applied_permission} is not "
+                        f"equal to expected permission {expected_permission}"
+                    )
+                    raise ValueError(msg)
+
+            retries_left -= 1
+
+        msg = f"Failed to apply permissions for {group_name} on scope {scope_name} in {num_retries} retries"
+        raise ValueError(msg)
+
     @sleep_and_retry
     @limits(calls=30, period=1)
     def _rate_limited_put_acl(self, object_id: str, principal: str, permission: workspace.AclPermission):
         self._ws.secrets.put_acl(object_id, principal, permission)
+        self._inflight_check(scope_name=object_id, group_name=principal, expected_permission=permission)
 
     def _get_apply_task(
         self, item: PermissionsInventoryItem, migration_state: GroupMigrationState, destination: Destination

src/databricks/labs/ucx/support/sql.py

@@ -64,7 +64,7 @@ def _applier_task(self, object_type: sql.ObjectTypePlural, object_id: str, acl:
         Please note that we only have SET option (DBSQL Permissions API doesn't support UPDATE operation).
         This affects the way how we prepare the new ACL request.
         """
-        self._ws.dbsql_permissions.set(object_type=object_type, object_id=object_id, acl=acl)
+        self._ws.dbsql_permissions.set(object_type=object_type, object_id=object_id, access_control_list=acl)
 
     def get_crawler_tasks(self):
         for listing in self._listings:

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

@@ -28,12 +28,11 @@ def __init__(self, config: MigrationConfig):
 
         self._group_manager = GroupManager(self._ws, config.groups)
         self._permissions_inventory = PermissionsInventoryTable(config.inventory_database, self._ws)
+        self._supports_provider = SupportsProvider(self._ws, self._num_threads, self._workspace_start_path)
         self._permissions_manager = PermissionManager(
-            self._ws,
-            self._permissions_inventory,
-            supports_provider=SupportsProvider(self._ws, self._num_threads, self._workspace_start_path),
+            self._ws, self._permissions_inventory, supports_provider=self._supports_provider
         )
-        self._verification_manager = VerificationManager(self._ws)
+        self._verification_manager = VerificationManager(self._ws, self._supports_provider.supports["secrets"])
 
     @staticmethod
     def _verify_ws_client(w: WorkspaceClient):

tests/unit/support/test_group_level.py

@@ -57,7 +57,7 @@ def test_scim_apply(migration_state):
     task()
     ws.groups.patch.assert_called_once_with(
         id="test-backup",
-        operations=[iam.Patch(op=iam.PatchOp.ADD, path="roles", value=sample_permissions)],
+        operations=[iam.Patch(op=iam.PatchOp.ADD, path="roles", value=[p.as_dict() for p in sample_permissions])],
         schemas=[iam.PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP],
     )
 

tests/unit/support/test_impl.py

@@ -9,13 +9,13 @@ def test_supports_provider():
         "entitlements",
         "roles",
         "clusters",
-        "cluster_policies",
-        "instance_pools",
-        "sql_warehouses",
+        "cluster-policies",
+        "instance-pools",
+        "sql/warehouses",
         "jobs",
         "pipelines",
         "experiments",
-        "registered_models",
+        "registered-models",
         "tokens",
         "passwords",
         "notebooks",