Check if permissions are up to date (#421)

Attempt to fix #36

Goal of this PR is to add a validation in each class that implements the
Crawler and Applier class.
This validation step will verify that a permissions that has been
applied to an object is correct, if it's not the case, it attemps to
re-apply this permission multiple times

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

@@ -1,9 +1,11 @@
 import datetime
 import json
 import logging
+import time
 from collections.abc import Callable, Iterator
 from dataclasses import dataclass
 from functools import partial
+from typing import Optional
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.core import DatabricksError
@@ -83,10 +85,38 @@ def _is_item_relevant(item: Permissions, migration_state: GroupMigrationState) -
                 return True
         return False
 
+    @staticmethod
+    def _response_to_request(
+        acls: Optional["list[iam.AccessControlResponse]"] = None,
+    ) -> list[iam.AccessControlRequest]:
+        results = []
+        for acl in acls:
+            for permission in acl.all_permissions:
+                results.append(
+                    iam.AccessControlRequest(
+                        acl.group_name, permission.permission_level, acl.service_principal_name, acl.user_name
+                    )
+                )
+        return results
+
     @rate_limited(max_requests=30)
     def _applier_task(self, object_type: str, object_id: str, acl: list[iam.AccessControlRequest]):
-        self._ws.permissions.update(object_type, object_id, access_control_list=acl)
-        return True
+        for _i in range(0, 3):
+            self._ws.permissions.update(object_type, object_id, access_control_list=acl)
+
+            remote_permission = self._safe_get_permissions(object_type, object_id)
+            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
+
+            logger.warning(
+                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}
+            """
+            )
+            time.sleep(1 + _i)
+        return False
 
     @rate_limited(max_requests=100)
     def _crawler_task(self, object_type: str, object_id: str) -> Permissions | None:

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

@@ -1,5 +1,6 @@
 import dataclasses
 import json
+import time
 from collections.abc import Callable
 from dataclasses import dataclass
 from functools import partial
@@ -104,8 +105,21 @@ 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, access_control_list=acl)
-        return True
+        for _i in range(0, 3):
+            self._ws.dbsql_permissions.set(object_type=object_type, object_id=object_id, access_control_list=acl)
+
+            remote_permission = self._safe_get_dbsql_permissions(object_type, object_id)
+            if all(elem in remote_permission.access_control_list for elem in acl):
+                return True
+
+            logger.warning(
+                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}
+            """
+            )
+            time.sleep(1 + _i)
+        return False
 
     def _prepare_new_acl(
         self, acl: list[sql.AccessControl], migration_state: GroupMigrationState, destination: Destination

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

@@ -1,5 +1,7 @@
 import json
+import time
 from functools import partial
+from logging import Logger
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.core import DatabricksError
@@ -14,6 +16,8 @@
 )
 from databricks.labs.ucx.workspace_access.groups import GroupMigrationState
 
+logger = Logger(__name__)
+
 
 class ScimSupport(AclSupport):
     def __init__(self, ws: WorkspaceClient):
@@ -58,7 +62,24 @@ def _crawler_task(group: iam.Group, property_name: str):
 
     @rate_limited(max_requests=10)
     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=[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)
-        return True
+        for _i in range(0, 3):
+            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)
+
+            group = self._ws.groups.get(group_id)
+            if property_name == "roles" and group.roles:
+                if all(elem in group.roles for elem in value):
+                    return True
+            elif property_name == "entitlements" and group.entitlements:
+                if all(elem in group.entitlements for elem in value):
+                    return True
+
+            logger.warning(
+                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()}
+                    """
+            )
+            time.sleep(1 + _i)
+        return False

tests/unit/workspace_access/test_generic.py

@@ -1,4 +1,5 @@
 import json
+import unittest.mock
 from unittest.mock import MagicMock
 
 from databricks.sdk.core import DatabricksError
@@ -55,6 +56,11 @@ def test_crawler():
 
 def test_apply(migration_state):
     ws = MagicMock()
+
+    acl1 = iam.AccessControlResponse(
+        all_permissions=[iam.Permission(permission_level=iam.PermissionLevel.CAN_USE)], group_name="db-temp-test"
+    )
+    ws.permissions.get.return_value = iam.ObjectPermissions(access_control_list=[acl1])
     sup = GenericPermissionsSupport(ws=ws, listings=[])  # no listings since only apply is tested
 
     item = Permissions(
@@ -148,9 +154,17 @@ def test_passwords_tokens_crawler(migration_state):
         )
     ]
 
+    temp_acl = [
+        iam.AccessControlResponse(
+            group_name="db-temp-test",
+            all_permissions=[iam.Permission(inherited=False, permission_level=iam.PermissionLevel.CAN_USE)],
+        )
+    ]
     ws.permissions.get.side_effect = [
         iam.ObjectPermissions(object_id="passwords", object_type="authorization", access_control_list=basic_acl),
         iam.ObjectPermissions(object_id="tokens", object_type="authorization", access_control_list=basic_acl),
+        iam.ObjectPermissions(object_id="passwords", object_type="authorization", access_control_list=temp_acl),
+        iam.ObjectPermissions(object_id="tokens", object_type="authorization", access_control_list=temp_acl),
     ]
 
     sup = GenericPermissionsSupport(ws=ws, listings=[Listing(tokens_and_passwords, "object_id", "authorization")])
@@ -205,3 +219,156 @@ def test_experiment_listing():
     for res in results:
         assert res.request_type == "experiments"
         assert res.object_id in ["test", "test2"]
+
+
+def test_response_to_request_mapping():
+    permissions1 = [
+        iam.Permission(permission_level=iam.PermissionLevel.CAN_BIND),
+        iam.Permission(permission_level=iam.PermissionLevel.CAN_MANAGE),
+    ]
+    response1 = iam.AccessControlResponse(all_permissions=permissions1, user_name="test1212")
+
+    permissions2 = [iam.Permission(permission_level=iam.PermissionLevel.CAN_ATTACH_TO)]
+    response2 = iam.AccessControlResponse(all_permissions=permissions2, group_name="data-engineers")
+
+    permissions3 = [iam.Permission(permission_level=iam.PermissionLevel.CAN_MANAGE_PRODUCTION_VERSIONS)]
+    response3 = iam.AccessControlResponse(all_permissions=permissions3, service_principal_name="sp1")
+
+    object_permissions = iam.ObjectPermissions(access_control_list=[response1, response2, response3])
+
+    sup = GenericPermissionsSupport(ws=MagicMock(), listings=[])
+    results = sup._response_to_request(object_permissions.access_control_list)
+
+    assert results == [
+        iam.AccessControlRequest(permission_level=iam.PermissionLevel.CAN_BIND, user_name="test1212"),
+        iam.AccessControlRequest(permission_level=iam.PermissionLevel.CAN_MANAGE, user_name="test1212"),
+        iam.AccessControlRequest(permission_level=iam.PermissionLevel.CAN_ATTACH_TO, group_name="data-engineers"),
+        iam.AccessControlRequest(
+            permission_level=iam.PermissionLevel.CAN_MANAGE_PRODUCTION_VERSIONS, service_principal_name="sp1"
+        ),
+    ]
+
+
+def test_applier_task_should_return_true_if_permission_is_up_to_date():
+    ws = MagicMock()
+    acl1 = iam.AccessControlResponse(
+        all_permissions=[iam.Permission(permission_level=iam.PermissionLevel.CAN_USE)], group_name="group"
+    )
+    acl2 = iam.AccessControlResponse(
+        all_permissions=[iam.Permission(permission_level=iam.PermissionLevel.CAN_RUN)], group_name="group2"
+    )
+    ws.permissions.get.return_value = iam.ObjectPermissions(access_control_list=[acl1, acl2])
+
+    sup = GenericPermissionsSupport(ws=ws, listings=[])
+    result = sup._applier_task(
+        object_type="clusters",
+        object_id="cluster_id",
+        acl=[iam.AccessControlRequest(group_name="group", permission_level=iam.PermissionLevel.CAN_USE)],
+    )
+    assert result
+
+
+def test_applier_task_should_return_true_if_permission_is_up_to_date_with_multiple_permissions():
+    ws = MagicMock()
+    acl = iam.AccessControlResponse(
+        all_permissions=[
+            iam.Permission(permission_level=iam.PermissionLevel.CAN_USE),
+            iam.Permission(permission_level=iam.PermissionLevel.CAN_ATTACH_TO),
+            iam.Permission(permission_level=iam.PermissionLevel.CAN_RUN),
+        ],
+        group_name="group",
+    )
+
+    ws.permissions.get.return_value = iam.ObjectPermissions(access_control_list=[acl])
+    sup = GenericPermissionsSupport(ws=ws, listings=[])
+
+    result = sup._applier_task(
+        object_type="clusters",
+        object_id="cluster_id",
+        acl=[
+            iam.AccessControlRequest(group_name="group", permission_level=iam.PermissionLevel.CAN_USE),
+            iam.AccessControlRequest(group_name="group", permission_level=iam.PermissionLevel.CAN_ATTACH_TO),
+        ],
+    )
+    assert result
+
+
+def test_applier_task_should_return_false_if_permission_couldnt_be_applied():
+    ws = MagicMock()
+    acl = iam.AccessControlResponse(all_permissions=[], group_name="group")
+
+    ws.permissions.update.return_value = iam.ObjectPermissions(access_control_list=[acl])
+    ws.permissions.get.return_value = iam.ObjectPermissions(access_control_list=[acl])
+    sup = GenericPermissionsSupport(ws=ws, listings=[])
+
+    result = sup._applier_task(
+        object_type="clusters",
+        object_id="cluster_id",
+        acl=[iam.AccessControlRequest(group_name="group", permission_level=iam.PermissionLevel.CAN_USE)],
+    )
+    assert not result
+
+
+def test_applier_task_should_return_false_if_all_permission_couldnt_be_applied():
+    ws = MagicMock()
+    group_1_acl = iam.AccessControlResponse(
+        all_permissions=[
+            iam.Permission(permission_level=iam.PermissionLevel.CAN_USE),
+            iam.Permission(permission_level=iam.PermissionLevel.CAN_ATTACH_TO),
+            iam.Permission(permission_level=iam.PermissionLevel.CAN_RUN),
+        ],
+        group_name="group_1",
+    )
+    group_2_acl = iam.AccessControlResponse(
+        all_permissions=[
+            iam.Permission(permission_level=iam.PermissionLevel.CAN_USE),
+            iam.Permission(permission_level=iam.PermissionLevel.CAN_MANAGE),
+        ],
+        group_name="group_2",
+    )
+    ws.permissions.update.return_value = iam.ObjectPermissions(access_control_list=[group_1_acl, group_2_acl])
+    ws.permissions.get.return_value = iam.ObjectPermissions(access_control_list=[group_1_acl, group_2_acl])
+    sup = GenericPermissionsSupport(ws=ws, listings=[])
+
+    result = sup._applier_task(
+        object_type="clusters",
+        object_id="cluster_id",
+        acl=[
+            iam.AccessControlRequest(group_name="group_1", permission_level=iam.PermissionLevel.CAN_USE),
+            iam.AccessControlRequest(group_name="group_1", permission_level=iam.PermissionLevel.CAN_MANAGE),
+        ],
+    )
+    assert not result
+
+
+def test_applier_task_should_be_called_three_times_if_permission_couldnt_be_applied():
+    ws = MagicMock()
+    acl = iam.AccessControlResponse(all_permissions=[], group_name="group")
+
+    ws.permissions.update.return_value = iam.ObjectPermissions(access_control_list=[acl])
+    ws.permissions.get.return_value = iam.ObjectPermissions(access_control_list=[acl])
+    sup = GenericPermissionsSupport(ws=ws, listings=[])
+
+    input_acl = [iam.AccessControlRequest(group_name="group", permission_level=iam.PermissionLevel.CAN_USE)]
+    sup._applier_task(
+        object_type="clusters",
+        object_id="cluster_id",
+        acl=input_acl,
+    )
+
+    assert len(ws.permissions.update.mock_calls) == 3
+    assert ws.permissions.update.has_calls(
+        [
+            unittest.mock.call(object_type="clusters", object_id="cluster_id", acl=input_acl),
+            unittest.mock.call(object_type="clusters", object_id="cluster_id", acl=input_acl),
+            unittest.mock.call(object_type="clusters", object_id="cluster_id", acl=input_acl),
+        ]
+    )
+    assert len(ws.permissions.get.mock_calls) == 3
+    assert ws.permissions.get.has_calls(
+        [
+            unittest.mock.call(object_type="clusters", object_id="cluster_id"),
+            unittest.mock.call(object_type="clusters", object_id="cluster_id"),
+            unittest.mock.call(object_type="clusters", object_id="cluster_id"),
+        ]
+    )

tests/unit/workspace_access/test_groups.py

@@ -47,8 +47,9 @@ def test_scim_crawler():
 
 def test_scim_apply(migration_state):
     ws = MagicMock()
-    sup = ScimSupport(ws=ws)
     sample_permissions = [iam.ComplexValue(value="role1"), iam.ComplexValue(value="role2")]
+    ws.groups.get.return_value = Group(id="test-backup", roles=sample_permissions)
+    sup = ScimSupport(ws=ws)
     item = Permissions(
         object_id="test-ws",
         object_type="roles",