From 8b2751b705d5fc5562aa832625954ddbaf0f290b Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Tue, 5 Sep 2023 14:39:25 -0700 Subject: [PATCH] Added `make_*_permissions` fixtures This PR introduces new permissions fixtures. Example usage: ```python def test_notebook_permissions(make_notebook, make_notebook_permissions, make_group): group = make_group() notebook = make_notebook() make_notebook_permissions(object_id=notebook, permission_level=PermissionLevel.CAN_RUN, group_name=group.display_name) ``` --- .../labs/ucx/providers/mixins/fixtures.py | 210 +++++++++++++++++- tests/integration/conftest.py | 6 +- tests/integration/test_fixtures.py | 9 + 3 files changed, 221 insertions(+), 4 deletions(-) diff --git a/src/databricks/labs/ucx/providers/mixins/fixtures.py b/src/databricks/labs/ucx/providers/mixins/fixtures.py index 693d07f1f5..74e5785be7 100644 --- a/src/databricks/labs/ucx/providers/mixins/fixtures.py +++ b/src/databricks/labs/ucx/providers/mixins/fixtures.py @@ -5,7 +5,7 @@ import pathlib import string import sys -from typing import BinaryIO +from typing import BinaryIO, Optional import pytest from databricks.sdk import AccountClient, WorkspaceClient @@ -60,6 +60,214 @@ def acc() -> AccountClient: return AccountClient() +def _permissions_mapping(): + from databricks.sdk.service.iam import PermissionLevel + + def _simple(_, object_id): + return object_id + + def _path(ws, path): + return ws.workspace.get_status(path).object_id + + return [ + ("cluster_policy", "cluster-policies", [PermissionLevel.CAN_USE], _simple), + ( + "instance_pool", + "instance-pools", + [PermissionLevel.CAN_ATTACH_TO, PermissionLevel.CAN_MANAGE], + _simple, + ), + ( + "cluster", + "clusters", + [PermissionLevel.CAN_ATTACH_TO, PermissionLevel.CAN_RESTART, PermissionLevel.CAN_MANAGE], + _simple, + ), + ( + "pipeline", + "pipelines", + [PermissionLevel.CAN_VIEW, PermissionLevel.CAN_RUN, PermissionLevel.CAN_MANAGE, PermissionLevel.IS_OWNER], + _simple, + ), + ( + "job", + "jobs", + [ + PermissionLevel.CAN_VIEW, + PermissionLevel.CAN_MANAGE_RUN, + PermissionLevel.IS_OWNER, + PermissionLevel.CAN_MANAGE, + ], + _simple, + ), + ( + "notebook", + "notebooks", + [PermissionLevel.CAN_READ, PermissionLevel.CAN_RUN, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_MANAGE], + _path, + ), + ( + "directory", + "directories", + [PermissionLevel.CAN_READ, PermissionLevel.CAN_RUN, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_MANAGE], + _path, + ), + ( + "workspace_file", + "files", + [PermissionLevel.CAN_READ, PermissionLevel.CAN_RUN, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_MANAGE], + _simple, + ), + ( + "workspace_file_path", + "files", + [PermissionLevel.CAN_READ, PermissionLevel.CAN_RUN, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_MANAGE], + _path, + ), + ( + "repo", + "repos", + [PermissionLevel.CAN_READ, PermissionLevel.CAN_RUN, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_MANAGE], + _path, + ), + ("tokens_authorization", "authorization", [PermissionLevel.CAN_USE], _simple), + ("passwords_authorization", "authorization", [PermissionLevel.CAN_USE], _simple), + ( + "warehouse", + "sql/warehouses", + [PermissionLevel.CAN_USE, PermissionLevel.CAN_MANAGE], + _simple, + ), + ( + "dashboard", + "sql/dashboards", + [PermissionLevel.CAN_EDIT, PermissionLevel.CAN_RUN, PermissionLevel.CAN_MANAGE, PermissionLevel.CAN_VIEW], + _simple, + ), + ( + "alert", + "sql/alerts", + [PermissionLevel.CAN_EDIT, PermissionLevel.CAN_RUN, PermissionLevel.CAN_MANAGE, PermissionLevel.CAN_VIEW], + _simple, + ), + ( + "query", + "sql/queries", + [PermissionLevel.CAN_EDIT, PermissionLevel.CAN_RUN, PermissionLevel.CAN_MANAGE, PermissionLevel.CAN_VIEW], + _simple, + ), + ( + "experiment", + "experiments", + [PermissionLevel.CAN_READ, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_MANAGE], + _simple, + ), + ( + "registered_model", + "registered-models", + [ + PermissionLevel.CAN_READ, + PermissionLevel.CAN_EDIT, + PermissionLevel.CAN_MANAGE_STAGING_VERSIONS, + PermissionLevel.CAN_MANAGE_PRODUCTION_VERSIONS, + PermissionLevel.CAN_MANAGE, + ], + _simple, + ), + ( + "serving_endpoint", + "serving-endpoints", + [PermissionLevel.CAN_VIEW, PermissionLevel.CAN_MANAGE], + _simple, + ), + ] + + +class _PermissionsChange: + def __init__(self, object_id: str, before: list[iam.AccessControlRequest], after: list[iam.AccessControlRequest]): + self._object_id = object_id + self._before = before + self._after = after + + @staticmethod + def _principal(acr: iam.AccessControlRequest) -> str: + if acr.user_name is not None: + return f"user_name {acr.user_name}" + elif acr.group_name is not None: + return f"group_name {acr.group_name}" + else: + return f"service_principal_name {acr.service_principal_name}" + + def _list(self, acl: list[iam.AccessControlRequest]): + return ", ".join(f"{self._principal(_)} {_.permission_level.value}" for _ in acl) + + def __repr__(self): + return f"{self._object_id} [{self._list(self._before)}] -> [{self._list(self._after)}]" + + +def _make_permissions_factory(name, resource_type, levels, id_retriever): + def _non_inherited(x: iam.ObjectPermissions): + return [ + iam.AccessControlRequest( + permission_level=permission.permission_level, + group_name=access_control.group_name, + user_name=access_control.user_name, + service_principal_name=access_control.service_principal_name, + ) + for access_control in x.access_control_list + for permission in access_control.all_permissions + if not permission.inherited + ] + + def _make_permissions(ws): + def create( + *, + object_id: str, + permission_level: iam.PermissionLevel | None = None, + group_name: str | None = None, + user_name: str | None = None, + service_principal_name: str | None = None, + access_control_list: Optional["list[iam.AccessControlRequest]"] = None, + ): + nothing_specified = permission_level is None and access_control_list is None + both_specified = permission_level is not None and access_control_list is not None + if nothing_specified or both_specified: + msg = "either permission_level or access_control_list has to be specified" + raise ValueError(msg) + + object_id = id_retriever(ws, object_id) + initial = _non_inherited(ws.permissions.get(resource_type, object_id)) + if access_control_list is None: + if permission_level not in levels: + names = ", ".join(_.value for _ in levels) + msg = f"invalid permission level: {permission_level.value}. Valid levels: {names}" + raise ValueError(msg) + access_control_list = [ + iam.AccessControlRequest( + group_name=group_name, + user_name=user_name, + service_principal_name=service_principal_name, + permission_level=permission_level, + ) + ] + ws.permissions.set(resource_type, object_id, access_control_list=access_control_list) + return _PermissionsChange(object_id, initial, access_control_list) + + def remove(change: _PermissionsChange): + ws.permissions.set(resource_type, change._object_id, access_control_list=change._before) + + yield from factory(f"{name} permissions", create, remove) + + return _make_permissions + + +for name, resource_type, levels, id_retriever in _permissions_mapping(): + # wrap function factory, otherwise loop scope sticks the wrong way + locals()[f"make_{name}_permissions"] = pytest.fixture( + _make_permissions_factory(name, resource_type, levels, id_retriever) + ) + + @pytest.fixture def make_secret_scope(ws, make_random): def create(**kwargs): diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 436d1bd2ba..c986d98621 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -580,7 +580,7 @@ def workspace_objects(ws: WorkspaceClient, env: EnvironmentInfo) -> WorkspaceObj ws.permissions.set( request_object_type=RequestObjectType.DIRECTORIES, - request_object_id=object_info.object_id, + request_object_id=object_info._object_id, access_control_list=[ AccessControlRequest(group_name=ws_group.display_name, permission_level=PermissionLevel.CAN_MANAGE) ], @@ -599,7 +599,7 @@ def workspace_objects(ws: WorkspaceClient, env: EnvironmentInfo) -> WorkspaceObj notebooks.append(_nb_obj) ws.permissions.set( request_object_type=RequestObjectType.NOTEBOOKS, - request_object_id=_nb_obj.object_id, + request_object_id=_nb_obj._object_id, access_control_list=[ AccessControlRequest(group_name=random_group.display_name, permission_level=PermissionLevel.CAN_EDIT) ], @@ -609,7 +609,7 @@ def workspace_objects(ws: WorkspaceClient, env: EnvironmentInfo) -> WorkspaceObj root_dir=ObjectInfo( path=f"/{env.test_uid}", object_type=ObjectType.DIRECTORY, - object_id=ws.workspace.get_status(f"/{env.test_uid}").object_id, + object_id=ws.workspace.get_status(f"/{env.test_uid}")._object_id, ), directories=base_dirs, notebooks=notebooks, diff --git a/tests/integration/test_fixtures.py b/tests/integration/test_fixtures.py index 303115f695..167c54c55b 100644 --- a/tests/integration/test_fixtures.py +++ b/tests/integration/test_fixtures.py @@ -31,6 +31,15 @@ def test_notebook(make_notebook): logger.info(f"created {make_notebook()}") +def test_notebook_permissions(make_notebook, make_notebook_permissions, make_group): + group = make_group() + notebook = make_notebook() + acl = make_notebook_permissions( + object_id=notebook, permission_level=iam.PermissionLevel.CAN_RUN, group_name=group.display_name # noqa: F405 + ) + logger.info(f"created {acl}") + + def test_directory(make_notebook, make_directory): logger.info(f'created {make_notebook(path=f"{make_directory()}/foo.py")}')