Skip to content

Commit 34321fc

Browse files
nfxnfx
authored
Added make_*_permissions fixtures (#159)
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) ```
1 parent 26c7246 commit 34321fc

File tree

3 files changed

+221
-4
lines changed

3 files changed

+221
-4
lines changed

‎src/databricks/labs/ucx/providers/mixins/fixtures.py‎

Lines changed: 209 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import pathlib
66
import string
77
import sys
8-
from typing import BinaryIO
8+
from typing import BinaryIO, Optional
99

1010
import pytest
1111
from databricks.sdk import AccountClient, WorkspaceClient
@@ -60,6 +60,214 @@ def acc() -> AccountClient:
6060
return AccountClient()
6161

6262

63+
def _permissions_mapping():
64+
from databricks.sdk.service.iam import PermissionLevel
65+
66+
def _simple(_, object_id):
67+
return object_id
68+
69+
def _path(ws, path):
70+
return ws.workspace.get_status(path).object_id
71+
72+
return [
73+
("cluster_policy", "cluster-policies", [PermissionLevel.CAN_USE], _simple),
74+
(
75+
"instance_pool",
76+
"instance-pools",
77+
[PermissionLevel.CAN_ATTACH_TO, PermissionLevel.CAN_MANAGE],
78+
_simple,
79+
),
80+
(
81+
"cluster",
82+
"clusters",
83+
[PermissionLevel.CAN_ATTACH_TO, PermissionLevel.CAN_RESTART, PermissionLevel.CAN_MANAGE],
84+
_simple,
85+
),
86+
(
87+
"pipeline",
88+
"pipelines",
89+
[PermissionLevel.CAN_VIEW, PermissionLevel.CAN_RUN, PermissionLevel.CAN_MANAGE, PermissionLevel.IS_OWNER],
90+
_simple,
91+
),
92+
(
93+
"job",
94+
"jobs",
95+
[
96+
PermissionLevel.CAN_VIEW,
97+
PermissionLevel.CAN_MANAGE_RUN,
98+
PermissionLevel.IS_OWNER,
99+
PermissionLevel.CAN_MANAGE,
100+
],
101+
_simple,
102+
),
103+
(
104+
"notebook",
105+
"notebooks",
106+
[PermissionLevel.CAN_READ, PermissionLevel.CAN_RUN, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_MANAGE],
107+
_path,
108+
),
109+
(
110+
"directory",
111+
"directories",
112+
[PermissionLevel.CAN_READ, PermissionLevel.CAN_RUN, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_MANAGE],
113+
_path,
114+
),
115+
(
116+
"workspace_file",
117+
"files",
118+
[PermissionLevel.CAN_READ, PermissionLevel.CAN_RUN, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_MANAGE],
119+
_simple,
120+
),
121+
(
122+
"workspace_file_path",
123+
"files",
124+
[PermissionLevel.CAN_READ, PermissionLevel.CAN_RUN, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_MANAGE],
125+
_path,
126+
),
127+
(
128+
"repo",
129+
"repos",
130+
[PermissionLevel.CAN_READ, PermissionLevel.CAN_RUN, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_MANAGE],
131+
_path,
132+
),
133+
("tokens_authorization", "authorization", [PermissionLevel.CAN_USE], _simple),
134+
("passwords_authorization", "authorization", [PermissionLevel.CAN_USE], _simple),
135+
(
136+
"warehouse",
137+
"sql/warehouses",
138+
[PermissionLevel.CAN_USE, PermissionLevel.CAN_MANAGE],
139+
_simple,
140+
),
141+
(
142+
"dashboard",
143+
"sql/dashboards",
144+
[PermissionLevel.CAN_EDIT, PermissionLevel.CAN_RUN, PermissionLevel.CAN_MANAGE, PermissionLevel.CAN_VIEW],
145+
_simple,
146+
),
147+
(
148+
"alert",
149+
"sql/alerts",
150+
[PermissionLevel.CAN_EDIT, PermissionLevel.CAN_RUN, PermissionLevel.CAN_MANAGE, PermissionLevel.CAN_VIEW],
151+
_simple,
152+
),
153+
(
154+
"query",
155+
"sql/queries",
156+
[PermissionLevel.CAN_EDIT, PermissionLevel.CAN_RUN, PermissionLevel.CAN_MANAGE, PermissionLevel.CAN_VIEW],
157+
_simple,
158+
),
159+
(
160+
"experiment",
161+
"experiments",
162+
[PermissionLevel.CAN_READ, PermissionLevel.CAN_EDIT, PermissionLevel.CAN_MANAGE],
163+
_simple,
164+
),
165+
(
166+
"registered_model",
167+
"registered-models",
168+
[
169+
PermissionLevel.CAN_READ,
170+
PermissionLevel.CAN_EDIT,
171+
PermissionLevel.CAN_MANAGE_STAGING_VERSIONS,
172+
PermissionLevel.CAN_MANAGE_PRODUCTION_VERSIONS,
173+
PermissionLevel.CAN_MANAGE,
174+
],
175+
_simple,
176+
),
177+
(
178+
"serving_endpoint",
179+
"serving-endpoints",
180+
[PermissionLevel.CAN_VIEW, PermissionLevel.CAN_MANAGE],
181+
_simple,
182+
),
183+
]
184+
185+
186+
class _PermissionsChange:
187+
def __init__(self, object_id: str, before: list[iam.AccessControlRequest], after: list[iam.AccessControlRequest]):
188+
self._object_id = object_id
189+
self._before = before
190+
self._after = after
191+
192+
@staticmethod
193+
def _principal(acr: iam.AccessControlRequest) -> str:
194+
if acr.user_name is not None:
195+
return f"user_name {acr.user_name}"
196+
elif acr.group_name is not None:
197+
return f"group_name {acr.group_name}"
198+
else:
199+
return f"service_principal_name {acr.service_principal_name}"
200+
201+
def _list(self, acl: list[iam.AccessControlRequest]):
202+
return ", ".join(f"{self._principal(_)} {_.permission_level.value}" for _ in acl)
203+
204+
def __repr__(self):
205+
return f"{self._object_id} [{self._list(self._before)}] -> [{self._list(self._after)}]"
206+
207+
208+
def _make_permissions_factory(name, resource_type, levels, id_retriever):
209+
def _non_inherited(x: iam.ObjectPermissions):
210+
return [
211+
iam.AccessControlRequest(
212+
permission_level=permission.permission_level,
213+
group_name=access_control.group_name,
214+
user_name=access_control.user_name,
215+
service_principal_name=access_control.service_principal_name,
216+
)
217+
for access_control in x.access_control_list
218+
for permission in access_control.all_permissions
219+
if not permission.inherited
220+
]
221+
222+
def _make_permissions(ws):
223+
def create(
224+
*,
225+
object_id: str,
226+
permission_level: iam.PermissionLevel | None = None,
227+
group_name: str | None = None,
228+
user_name: str | None = None,
229+
service_principal_name: str | None = None,
230+
access_control_list: Optional["list[iam.AccessControlRequest]"] = None,
231+
):
232+
nothing_specified = permission_level is None and access_control_list is None
233+
both_specified = permission_level is not None and access_control_list is not None
234+
if nothing_specified or both_specified:
235+
msg = "either permission_level or access_control_list has to be specified"
236+
raise ValueError(msg)
237+
238+
object_id = id_retriever(ws, object_id)
239+
initial = _non_inherited(ws.permissions.get(resource_type, object_id))
240+
if access_control_list is None:
241+
if permission_level not in levels:
242+
names = ", ".join(_.value for _ in levels)
243+
msg = f"invalid permission level: {permission_level.value}. Valid levels: {names}"
244+
raise ValueError(msg)
245+
access_control_list = [
246+
iam.AccessControlRequest(
247+
group_name=group_name,
248+
user_name=user_name,
249+
service_principal_name=service_principal_name,
250+
permission_level=permission_level,
251+
)
252+
]
253+
ws.permissions.set(resource_type, object_id, access_control_list=access_control_list)
254+
return _PermissionsChange(object_id, initial, access_control_list)
255+
256+
def remove(change: _PermissionsChange):
257+
ws.permissions.set(resource_type, change._object_id, access_control_list=change._before)
258+
259+
yield from factory(f"{name} permissions", create, remove)
260+
261+
return _make_permissions
262+
263+
264+
for name, resource_type, levels, id_retriever in _permissions_mapping():
265+
# wrap function factory, otherwise loop scope sticks the wrong way
266+
locals()[f"make_{name}_permissions"] = pytest.fixture(
267+
_make_permissions_factory(name, resource_type, levels, id_retriever)
268+
)
269+
270+
63271
@pytest.fixture
64272
def make_secret_scope(ws, make_random):
65273
def create(**kwargs):

‎tests/integration/conftest.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -580,7 +580,7 @@ def workspace_objects(ws: WorkspaceClient, env: EnvironmentInfo) -> WorkspaceObj
580580

581581
ws.permissions.set(
582582
request_object_type=RequestObjectType.DIRECTORIES,
583-
request_object_id=object_info.object_id,
583+
request_object_id=object_info._object_id,
584584
access_control_list=[
585585
AccessControlRequest(group_name=ws_group.display_name, permission_level=PermissionLevel.CAN_MANAGE)
586586
],
@@ -599,7 +599,7 @@ def workspace_objects(ws: WorkspaceClient, env: EnvironmentInfo) -> WorkspaceObj
599599
notebooks.append(_nb_obj)
600600
ws.permissions.set(
601601
request_object_type=RequestObjectType.NOTEBOOKS,
602-
request_object_id=_nb_obj.object_id,
602+
request_object_id=_nb_obj._object_id,
603603
access_control_list=[
604604
AccessControlRequest(group_name=random_group.display_name, permission_level=PermissionLevel.CAN_EDIT)
605605
],
@@ -609,7 +609,7 @@ def workspace_objects(ws: WorkspaceClient, env: EnvironmentInfo) -> WorkspaceObj
609609
root_dir=ObjectInfo(
610610
path=f"/{env.test_uid}",
611611
object_type=ObjectType.DIRECTORY,
612-
object_id=ws.workspace.get_status(f"/{env.test_uid}").object_id,
612+
object_id=ws.workspace.get_status(f"/{env.test_uid}")._object_id,
613613
),
614614
directories=base_dirs,
615615
notebooks=notebooks,

‎tests/integration/test_fixtures.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,15 @@ def test_notebook(make_notebook):
3131
logger.info(f"created {make_notebook()}")
3232

3333

34+
def test_notebook_permissions(make_notebook, make_notebook_permissions, make_group):
35+
group = make_group()
36+
notebook = make_notebook()
37+
acl = make_notebook_permissions(
38+
object_id=notebook, permission_level=iam.PermissionLevel.CAN_RUN, group_name=group.display_name # noqa: F405
39+
)
40+
logger.info(f"created {acl}")
41+
42+
3443
def test_directory(make_notebook, make_directory):
3544
logger.info(f'created {make_notebook(path=f"{make_directory()}/foo.py")}')
3645

0 commit comments

Comments
 (0)