Skip to content

Commit

Permalink
Added make_*_permissions fixtures (#159)
Browse files Browse the repository at this point in the history
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)
```
  • Loading branch information
nfx authored Sep 5, 2023
1 parent 26c7246 commit 34321fc
Show file tree
Hide file tree
Showing 3 changed files with 221 additions and 4 deletions.
210 changes: 209 additions & 1 deletion src/databricks/labs/ucx/providers/mixins/fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down
6 changes: 3 additions & 3 deletions tests/integration/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
],
Expand All @@ -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)
],
Expand All @@ -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,
Expand Down
9 changes: 9 additions & 0 deletions tests/integration/test_fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")}')

Expand Down

0 comments on commit 34321fc

Please sign in to comment.