Added verification of group permissions (#841)

## Changes
Added `verify` method to all `AclSupport` subclasses. Existing
`_inflight_check` and similar names were renamed to _verify for
consistency. Expose a method to validate groups permissions and added
cli command.

### Functionality 

- [x] added new CLI command (`validate-groups-permissions`) to run
verification of groups permissions

### Tests

- [x] manually tested
- [x] added unit tests
- [x] added integration tests

---------

Co-authored-by: mwojtyczka <[email protected]>
Co-authored-by: Marcin Wojtyczka <[email protected]>

Author: andrascsillag-db

labs.yml

@@ -103,7 +103,6 @@ commands:
       - name: subscription-id
         description: Subscription to scan storage account in
 
-
   - name: validate-groups-membership
     description: Validate the groups to see if the groups at account level and workspace level have different membership
     table_template: |-

src/databricks/labs/ucx/runtime.py

@@ -310,6 +310,20 @@ def apply_permissions_to_account_groups(cfg: WorkspaceConfig, ws: WorkspaceClien
     permission_manager.apply_group_permissions(migration_state)
 
 
+@task("validate-groups-permissions")
+def validate_groups_permissions(cfg: WorkspaceConfig, ws: WorkspaceClient, sql_backend: SqlBackend):
+    """Validate that all the crawled permissions are applied correctly to the destination groups."""
+    logger.info("Running validation of permissions applied to destination groups.")
+    permission_manager = PermissionManager.factory(
+        ws,
+        sql_backend,
+        cfg.inventory_database,
+        num_threads=cfg.num_threads,
+        workspace_start_path=cfg.workspace_start_path,
+    )
+    permission_manager.verify_group_permissions()
+
+
 @task("remove-workspace-local-backup-groups", depends_on=[apply_permissions_to_account_groups])
 def delete_backup_groups(cfg: WorkspaceConfig, ws: WorkspaceClient, sql_backend: SqlBackend):
     """Last step of the group migration process. Removes all workspace-level backup groups, along with their

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

@@ -29,6 +29,10 @@ def get_apply_task(self, item: Permissions, migration_state: MigrationState) ->
         """This method returns a Callable, that applies permissions to a destination group, based on
         the group migration state. The callable is required not to have any shared mutable state."""
 
+    def get_verify_task(self, item: Permissions) -> Callable[[], bool] | None:
+        """This method returns a Callable that verifies that all the crawled permissions are applied correctly to the
+        destination group."""
+
     @abstractmethod
     def object_types(self) -> set[str]:
         """This method returns a set of strings, that represent object types that are applicable by this instance."""

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

@@ -119,21 +119,32 @@ def _response_to_request(
                 )
         return results
 
-    def _inflight_check(self, object_type: str, object_id: str, acl: list[iam.AccessControlRequest]):
+    @rate_limited(max_requests=100)
+    def _verify(self, object_type: str, object_id: str, acl: list[iam.AccessControlRequest]):
         # in-flight check for the applied permissions
         # the api might be inconsistent, therefore we need to check that the permissions were applied
         remote_permission = self._safe_get_permissions(object_type, object_id)
         if remote_permission:
             remote_permission_as_request = self._response_to_request(remote_permission.access_control_list)
             if all(elem in remote_permission_as_request for elem in acl):
                 return True
-            msg = f"""Couldn't apply appropriate permission for object type {object_type} with id {object_id}
-                acl to be applied={acl}
-                acl found in the object={remote_permission_as_request}
-                """
+            msg = (
+                f"Couldn't find permission for object type {object_type} with id {object_id}\n"
+                f"acl to be applied={acl}\n"
+                f"acl found in the object={remote_permission_as_request}\n"
+            )
             raise ValueError(msg)
         return False
 
+    def get_verify_task(self, item: Permissions) -> Callable[[], bool]:
+        acl = iam.ObjectPermissions.from_dict(json.loads(item.raw))
+        if not acl.access_control_list:
+            raise ValueError(
+                f"Access control list not present for object type " f"{item.object_type} and object id {item.object_id}"
+            )
+        permissions_as_request = self._response_to_request(acl.access_control_list)
+        return partial(self._verify, item.object_type, item.object_id, permissions_as_request)
+
     @rate_limited(max_requests=30)
     def _applier_task(self, object_type: str, object_id: str, acl: list[iam.AccessControlRequest]):
         retryable_exceptions = [InternalError, NotFound, ResourceConflict, TemporarilyUnavailable, DeadlineExceeded]
@@ -145,7 +156,7 @@ def _applier_task(self, object_type: str, object_id: str, acl: list[iam.AccessCo
         update_retried_check(object_type, object_id, acl)
 
         retry_on_value_error = retried(on=[*retryable_exceptions, ValueError], timeout=self._verify_timeout)
-        retried_check = retry_on_value_error(self._inflight_check)
+        retried_check = retry_on_value_error(self._verify)
         return retried_check(object_type, object_id, acl)
 
     @rate_limited(max_requests=100)

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

@@ -91,7 +91,7 @@ def inventorize_permissions(self):
         self._save(items)
         logger.info(f"Saved {len(items)} to {self._full_name}")
 
-    def apply_group_permissions(self, migration_state: MigrationState):
+    def apply_group_permissions(self, migration_state: MigrationState) -> bool:
         # list shall be sorted prior to using group by
         if len(migration_state) == 0:
             logger.info("No valid groups selected, nothing to do.")
@@ -140,6 +140,34 @@ def apply_group_permissions(self, migration_state: MigrationState):
         logger.info("Permissions were applied")
         return True
 
+    def verify_group_permissions(self) -> bool:
+        items = sorted(self.load_all(), key=lambda i: i.object_type)
+        logger.info(f"Total permissions found: {len(items)}")
+        verifier_tasks: list[Callable[..., bool]] = []
+        appliers = self._appliers()
+
+        for object_type, items_subset in groupby(items, key=lambda i: i.object_type):
+            if object_type not in appliers:
+                msg = f"Could not find support for {object_type}. Please check the inventory table."
+                raise ValueError(msg)
+
+            relevant_support = appliers[object_type]
+            tasks_for_support: list[Callable[..., bool]] = []
+            for item in items_subset:
+                task = relevant_support.get_verify_task(item)
+                if not task:
+                    continue
+                tasks_for_support.append(task)
+
+            logger.info(f"Total tasks for {object_type}: {len(tasks_for_support)}")
+            verifier_tasks.extend(tasks_for_support)
+
+        logger.info(f"Starting to verify permissions. Total tasks: {len(verifier_tasks)}")
+        Threads.strict("verify group permissions", verifier_tasks)
+        logger.info("All permissions validated successfully. No issues found.")
+
+        return True
+
     def _appliers(self) -> dict[str, AclSupport]:
         appliers: dict[str, AclSupport] = {}
         for support in self._acl_support:

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

@@ -1,3 +1,4 @@
+# pylint: disable=duplicate-code
 import dataclasses
 import json
 import logging
@@ -141,22 +142,31 @@ def _crawler_task(self, object_id: str, object_type: sql.ObjectTypePlural) -> Pe
             )
         return None
 
-    def _inflight_check(self, object_type: sql.ObjectTypePlural, object_id: str, acl: list[sql.AccessControl]):
+    @rate_limited(max_requests=100)
+    def _verify(self, object_type: sql.ObjectTypePlural, object_id: str, acl: list[sql.AccessControl]):
         # in-flight check for the applied permissions
         # the api might be inconsistent, therefore we need to check that the permissions were applied
         remote_permission = self._safe_get_dbsql_permissions(object_type, object_id)
         if remote_permission:
             assert remote_permission.access_control_list is not None
             if all(elem in remote_permission.access_control_list for elem in acl):
                 return True
-            msg = f"""
-            Couldn't apply appropriate permission for object type {object_type} with id {object_id}
-            acl to be applied={acl}
-            acl found in the object={remote_permission}
-            """
+            msg = (
+                f"Couldn't find permission for object type {object_type} with id {object_id}\n"
+                f"acl to be applied={acl}\n"
+                f"acl found in the object={remote_permission}\n"
+            )
             raise ValueError(msg)
         return False
 
+    def get_verify_task(self, item: Permissions) -> Callable[[], bool]:
+        acl = sql.GetResponse.from_dict(json.loads(item.raw))
+        if not acl.access_control_list:
+            raise ValueError(
+                f"Access control list not present for object type " f"{item.object_type} and object id {item.object_id}"
+            )
+        return partial(self._verify, sql.ObjectTypePlural(item.object_type), item.object_id, acl.access_control_list)
+
     @rate_limited(max_requests=30)
     def _applier_task(self, object_type: sql.ObjectTypePlural, object_id: str, acl: list[sql.AccessControl]):
         """
@@ -169,7 +179,7 @@ def _applier_task(self, object_type: sql.ObjectTypePlural, object_id: str, acl:
         set_retried_check(object_type, object_id, acl)
 
         retry_on_value_error = retried(on=[InternalError, ValueError], timeout=self._verify_timeout)
-        retried_check = retry_on_value_error(self._inflight_check)
+        retried_check = retry_on_value_error(self._verify)
         return retried_check(object_type, object_id, acl)
 
     def _prepare_new_acl(

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

@@ -1,5 +1,6 @@
 import json
 import logging
+from collections.abc import Callable
 from datetime import timedelta
 from functools import partial
 
@@ -94,7 +95,8 @@ def _crawler_task(group: iam.Group, property_name: str):
             raw=json.dumps([e.as_dict() for e in getattr(group, property_name)]),
         )
 
-    def _inflight_check(self, group_id: str, value: list[iam.ComplexValue], property_name: str):
+    @rate_limited(max_requests=255, burst_period_seconds=60)
+    def _verify(self, group_id: str, value: list[iam.ComplexValue], property_name: str):
         # in-flight check for the applied permissions
         # the api might be inconsistent, therefore we need to check that the permissions were applied
         group = self._safe_get_group(group_id)
@@ -105,13 +107,18 @@ def _inflight_check(self, group_id: str, value: list[iam.ComplexValue], property
             if property_name == "entitlements" and group.entitlements:
                 if all(elem in group.entitlements for elem in value):
                     return True
-            msg = f"""Couldn't apply appropriate role for group {group_id}
-                            acl to be applied={[e.as_dict() for e in value]}
-                            acl found in the object={group.as_dict()}
-                            """
+            msg = (
+                f"Couldn't find role for group {group_id}\n"
+                f"acl to be applied={[e.as_dict() for e in value]}\n"
+                f"acl found in the object={group.as_dict()}\n"
+            )
             raise ValueError(msg)
         return False
 
+    def get_verify_task(self, item: Permissions) -> Callable[[], bool]:
+        value = [iam.ComplexValue.from_dict(e) for e in json.loads(item.raw)]
+        return partial(self._verify, item.object_id, value, item.object_type)
+
     @rate_limited(max_requests=10, burst_period_seconds=60)
     def _applier_task(self, group_id: str, value: list[iam.ComplexValue], property_name: str):
         retryable_errors = [ResourceConflict, Aborted, DeadlineExceeded, InternalError]
@@ -126,7 +133,7 @@ def _applier_task(self, group_id: str, value: list[iam.ComplexValue], property_n
         patch_retried_check(group_id, operations, schemas)
 
         retry_on_value_error = retried(on=[*retryable_errors, ValueError], timeout=self._verify_timeout)
-        retried_check = retry_on_value_error(self._inflight_check)
+        retried_check = retry_on_value_error(self._verify)
         return retried_check(group_id, value, property_name)
 
     def _safe_patch_group(

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

@@ -1,12 +1,14 @@
 import json
 import logging
+from collections.abc import Callable, Iterable
 from datetime import timedelta
 from functools import partial
 
 from databricks.labs.blueprint.limiter import rate_limited
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.retries import retried
 from databricks.sdk.service import workspace
+from databricks.sdk.service.workspace import AclItem
 
 from databricks.labs.ucx.workspace_access.base import AclSupport, Permissions
 from databricks.labs.ucx.workspace_access.groups import MigrationState
@@ -58,7 +60,7 @@ def get_apply_task(self, item: Permissions, migration_state: MigrationState):
 
         def apply_acls():
             for acl in new_acls:
-                self._rate_limited_put_acl(item.object_id, acl.principal, acl.permission)
+                self._applier_task(item.object_id, acl.principal, acl.permission)
             return True
 
         return partial(apply_acls)
@@ -78,21 +80,42 @@ def secret_scope_permission(self, scope_name: str, group_name: str) -> workspace
     def _reapply_on_failure(self, scope_name: str, group_name: str, expected_permission: workspace.AclPermission):
         # in-flight check for the applied permissions
         # the api might be inconsistent, therefore we need to check that the permissions were applied
-        applied_permission = self.secret_scope_permission(scope_name, group_name)
-        if applied_permission != expected_permission:
-            # try to apply again if the permissions are not equal: sometimes the list_acls api is inconsistent
+        try:
+            self._verify(scope_name, group_name, expected_permission)
+        except ValueError:
             logger.info(f"Applying permissions again {expected_permission} to {group_name} for {scope_name}")
             self._ws.secrets.put_acl(scope_name, group_name, expected_permission)
+            raise
+        return True
+
+    @rate_limited(max_requests=1100, burst_period_seconds=60)
+    def _verify(self, scope_name: str, group_name: str, expected_permission: workspace.AclPermission):
+        # in-flight check for the applied permissions
+        # the api might be inconsistent, therefore we need to check that the permissions were applied
+        applied_permission = self.secret_scope_permission(scope_name, group_name)
+        if applied_permission != expected_permission:
             msg = (
-                f"Applied permission {applied_permission} is not equal to expected "
-                f"permission {expected_permission} for {scope_name} and {group_name}!"
+                f"Couldn't find permission for scope {scope_name} and group {group_name}\n"
+                f"acl to be applied={expected_permission}\n"
+                f"acl found in the object={applied_permission}\n"
             )
             raise ValueError(msg)
-        logger.info(f"Permissions matched for {scope_name}, {group_name} and {expected_permission}!")
         return True
 
+    def get_verify_task(self, item: Permissions) -> Callable[[], bool]:
+        acls = [workspace.AclItem.from_dict(acl) for acl in json.loads(item.raw)]
+
+        def _verify_acls(scope_name: str, acls: Iterable[AclItem]):
+            for acl in acls:
+                assert acl.permission is not None
+                assert acl.principal is not None
+                self._verify(scope_name, acl.principal, acl.permission)
+            return True
+
+        return partial(_verify_acls, item.object_id, acls)
+
     @rate_limited(max_requests=1100, burst_period_seconds=60)
-    def _rate_limited_put_acl(self, object_id: str, principal: str, permission: workspace.AclPermission):
+    def _applier_task(self, object_id: str, principal: str, permission: workspace.AclPermission):
         self._ws.secrets.put_acl(object_id, principal, permission)
         retry_on_value_error = retried(on=[ValueError], timeout=self._verify_timeout)
         retried_check = retry_on_value_error(self._reapply_on_failure)

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

@@ -3,8 +3,11 @@
 import functools
 import json
 from collections.abc import Callable, Iterator
+from datetime import timedelta
 from functools import partial
 
+from databricks.sdk.retries import retried
+
 from databricks.labs.ucx.framework.crawlers import SqlBackend
 from databricks.labs.ucx.hive_metastore import GrantsCrawler
 from databricks.labs.ucx.hive_metastore.grants import Grant
@@ -13,9 +16,15 @@
 
 
 class TableAclSupport(AclSupport):
-    def __init__(self, grants_crawler: GrantsCrawler, sql_backend: SqlBackend):
+    def __init__(
+        self,
+        grants_crawler: GrantsCrawler,
+        sql_backend: SqlBackend,
+        verify_timeout: timedelta | None = timedelta(minutes=1),
+    ):
         self._grants_crawler = grants_crawler
         self._sql_backend = sql_backend
+        self._verify_timeout = verify_timeout
 
     def get_crawler_tasks(self) -> Iterator[Callable[..., Permissions | None]]:
         # Table ACL permissions (grant/revoke and ownership) are not atomic. When granting the permissions,
@@ -97,4 +106,33 @@ def _apply_grant(self, grant: Grant):
         """
         for sql in grant.hive_grant_sql():
             self._sql_backend.execute(sql)
-        return True
+
+        object_type, object_id = grant.this_type_and_key()
+        retry_on_value_error = retried(on=[ValueError], timeout=self._verify_timeout)
+        retried_check = retry_on_value_error(self._verify)
+        return retried_check(object_type, object_id, grant)
+
+    def _verify(self, object_type: str, object_id: str, acl: Grant) -> bool:
+        grant_dict = dataclasses.asdict(acl)
+        del grant_dict["action_type"]
+        del grant_dict["principal"]
+        grants_on_object = self._grants_crawler._grants(**grant_dict)
+
+        if grants_on_object:
+            action_types_for_current_principal = [
+                grant.action_type for grant in grants_on_object if grant.principal == acl.principal
+            ]
+            acl_action_types = acl.action_type.split(", ")
+            if all(action_type in action_types_for_current_principal for action_type in acl_action_types):
+                return True
+            msg = (
+                f"Couldn't find permission for object type {object_type}, id {object_id} and principal {acl.principal}\n"
+                f"acl to be applied={acl_action_types}\n"
+                f"acl found in the object={action_types_for_current_principal}\n"
+            )
+            raise ValueError(msg)
+        return False
+
+    def get_verify_task(self, item: Permissions) -> Callable[[], bool]:
+        grant = Grant(**json.loads(item.raw))
+        return partial(self._verify, item.object_type, item.object_id, grant)