test_client_rbac.py
1 from contextlib import contextmanager 2 3 import pytest 4 import requests 5 6 from mlflow import MlflowException 7 from mlflow.environment_variables import ( 8 MLFLOW_AUTH_CONFIG_PATH, 9 MLFLOW_FLASK_SERVER_SECRET_KEY, 10 MLFLOW_TRACKING_PASSWORD, 11 MLFLOW_TRACKING_USERNAME, 12 ) 13 from mlflow.protos.databricks_pb2 import ( 14 PERMISSION_DENIED, 15 RESOURCE_ALREADY_EXISTS, 16 RESOURCE_DOES_NOT_EXIST, 17 UNAUTHENTICATED, 18 ErrorCode, 19 ) 20 from mlflow.server.auth.client import AuthServiceClient 21 from mlflow.utils.os import is_windows 22 23 from tests.helper_functions import random_str 24 from tests.server.auth.auth_test_utils import ( 25 ADMIN_PASSWORD, 26 ADMIN_USERNAME, 27 User, 28 write_isolated_auth_config, 29 ) 30 from tests.tracking.integration_test_utils import _init_server 31 32 33 @pytest.fixture(autouse=True) 34 def clear_credentials(monkeypatch): 35 monkeypatch.delenv(MLFLOW_TRACKING_USERNAME.name, raising=False) 36 monkeypatch.delenv(MLFLOW_TRACKING_PASSWORD.name, raising=False) 37 38 39 @pytest.fixture 40 def client(tmp_path): 41 auth_config_path = write_isolated_auth_config(tmp_path) 42 path = tmp_path.joinpath("sqlalchemy.db").as_uri() 43 backend_uri = ("sqlite://" if is_windows() else "sqlite:////") + path[len("file://") :] 44 45 with _init_server( 46 backend_uri=backend_uri, 47 root_artifact_uri=tmp_path.joinpath("artifacts").as_uri(), 48 app="mlflow.server.auth:create_app", 49 extra_env={ 50 MLFLOW_FLASK_SERVER_SECRET_KEY.name: "my-secret-key", 51 MLFLOW_AUTH_CONFIG_PATH.name: str(auth_config_path), 52 }, 53 server_type="flask", 54 ) as url: 55 yield AuthServiceClient(url) 56 57 58 @contextmanager 59 def assert_unauthenticated(): 60 with pytest.raises(MlflowException, match=r"You are not authenticated.") as exception_context: 61 yield 62 assert exception_context.value.error_code == ErrorCode.Name(UNAUTHENTICATED) 63 64 65 @contextmanager 66 def assert_unauthorized(): 67 with pytest.raises(MlflowException, match=r"Permission denied.") as exception_context: 68 yield 69 assert exception_context.value.error_code == ErrorCode.Name(PERMISSION_DENIED) 70 71 72 def _create_admin_controlled_user(client, monkeypatch): 73 """Create a non-admin user via the admin account and return (username, password).""" 74 username = random_str() 75 password = random_str() 76 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 77 client.create_user(username, password) 78 return username, password 79 80 81 def _make_user_wp_admin(client, monkeypatch, username, workspace): 82 """Create a role with workspace-scope MANAGE in ``workspace`` and assign it to ``username``.""" 83 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 84 role = client.create_role(workspace=workspace, name=f"wp-admin-{random_str()}") 85 client.add_role_permission(role.id, "workspace", "*", "MANAGE") 86 client.assign_role(username, role.id) 87 return role 88 89 90 # ---- Role CRUD ---- 91 92 93 def test_create_role_requires_admin(client, monkeypatch): 94 username, password = _create_admin_controlled_user(client, monkeypatch) 95 96 with assert_unauthenticated(): 97 client.create_role(workspace="ws1", name="viewer") 98 99 with User(username, password, monkeypatch), assert_unauthorized(): 100 client.create_role(workspace="ws1", name="viewer") 101 102 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 103 role = client.create_role(workspace="ws1", name="viewer", description="read-only") 104 assert role.name == "viewer" 105 assert role.workspace == "ws1" 106 assert role.description == "read-only" 107 108 109 def test_create_role_duplicate(client, monkeypatch): 110 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 111 client.create_role(workspace="ws1", name="viewer") 112 with pytest.raises(MlflowException, match="already exists") as exc: 113 client.create_role(workspace="ws1", name="viewer") 114 assert exc.value.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS) 115 116 117 def test_get_role(client, monkeypatch): 118 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 119 created = client.create_role(workspace="ws1", name="viewer") 120 fetched = client.get_role(created.id) 121 assert fetched.id == created.id 122 assert fetched.name == "viewer" 123 124 125 def test_get_role_not_found(client, monkeypatch): 126 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 127 with pytest.raises(MlflowException, match="not found") as exc: 128 client.get_role(99999) 129 assert exc.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 130 131 132 def test_list_roles(client, monkeypatch): 133 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 134 client.create_role(workspace="ws1", name="viewer") 135 client.create_role(workspace="ws1", name="editor") 136 client.create_role(workspace="ws2", name="viewer") 137 ws1 = client.list_roles("ws1") 138 ws2 = client.list_roles("ws2") 139 assert {r.name for r in ws1} == {"viewer", "editor"} 140 assert {r.name for r in ws2} == {"viewer"} 141 142 143 def test_update_role(client, monkeypatch): 144 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 145 role = client.create_role(workspace="ws1", name="old") 146 updated = client.update_role(role.id, name="new", description="updated") 147 assert updated.name == "new" 148 assert updated.description == "updated" 149 150 151 def test_update_role_rejects_empty_update(client, monkeypatch): 152 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 153 role = client.create_role(workspace="ws1", name="viewer") 154 with pytest.raises(MlflowException, match="At least one of 'name' or 'description'"): 155 client.update_role(role.id) 156 157 158 def test_delete_role(client, monkeypatch): 159 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 160 role = client.create_role(workspace="ws1", name="viewer") 161 client.delete_role(role.id) 162 with pytest.raises(MlflowException, match="not found"): 163 client.get_role(role.id) 164 165 166 # ---- Role permission CRUD ---- 167 168 169 def test_role_permission_crud(client, monkeypatch): 170 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 171 role = client.create_role(workspace="ws1", name="viewer") 172 rp = client.add_role_permission(role.id, "experiment", "42", "READ") 173 assert rp.resource_type == "experiment" 174 assert rp.resource_pattern == "42" 175 assert rp.permission == "READ" 176 177 perms = client.list_role_permissions(role.id) 178 assert len(perms) == 1 179 180 updated = client.update_role_permission(rp.id, "EDIT") 181 assert updated.permission == "EDIT" 182 183 client.remove_role_permission(rp.id) 184 assert client.list_role_permissions(role.id) == [] 185 186 187 def test_add_role_permission_invalid_resource_type(client, monkeypatch): 188 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 189 role = client.create_role(workspace="ws1", name="viewer") 190 with pytest.raises(MlflowException, match="Invalid resource type"): 191 client.add_role_permission(role.id, "bogus", "42", "READ") 192 193 194 def test_add_workspace_resource_type(client, monkeypatch): 195 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 196 role = client.create_role(workspace="ws1", name="ws-admin") 197 rp = client.add_role_permission(role.id, "workspace", "*", "MANAGE") 198 assert rp.resource_type == "workspace" 199 assert rp.permission == "MANAGE" 200 201 202 # ---- User-role assignment ---- 203 204 205 def test_assign_unassign_role(client, monkeypatch): 206 username, _ = _create_admin_controlled_user(client, monkeypatch) 207 208 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 209 role = client.create_role(workspace="ws1", name="viewer") 210 assignment = client.assign_role(username, role.id) 211 assert assignment.role_id == role.id 212 213 user_roles = client.list_user_roles(username) 214 assert [r.id for r in user_roles] == [role.id] 215 216 role_users = client.list_role_users(role.id) 217 assert len(role_users) == 1 218 219 client.unassign_role(username, role.id) 220 assert client.list_user_roles(username) == [] 221 222 223 def test_assign_nonexistent_role(client, monkeypatch): 224 username, _ = _create_admin_controlled_user(client, monkeypatch) 225 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 226 with pytest.raises(MlflowException, match="not found"): 227 client.assign_role(username, 99999) 228 229 230 def test_assign_nonexistent_user(client, monkeypatch): 231 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 232 role = client.create_role(workspace="ws1", name="viewer") 233 with pytest.raises(MlflowException, match="not found"): 234 client.assign_role("nonexistent-user-" + random_str(), role.id) 235 236 237 # ---- Authorization: WP admin vs plain user vs super admin ---- 238 239 240 def test_wp_admin_can_manage_roles_in_own_workspace(client, monkeypatch): 241 wp_admin, wp_admin_pw = _create_admin_controlled_user(client, monkeypatch) 242 _make_user_wp_admin(client, monkeypatch, wp_admin, "ws1") 243 244 # WP admin can create and manage roles in their workspace. 245 with User(wp_admin, wp_admin_pw, monkeypatch): 246 role = client.create_role(workspace="ws1", name="editor") 247 client.add_role_permission(role.id, "experiment", "*", "EDIT") 248 assert len(client.list_roles("ws1")) >= 2 # the wp-admin role + this one 249 250 251 def test_wp_admin_cannot_manage_other_workspace(client, monkeypatch): 252 wp_admin, wp_admin_pw = _create_admin_controlled_user(client, monkeypatch) 253 _make_user_wp_admin(client, monkeypatch, wp_admin, "ws1") 254 255 # WP admin of ws1 cannot create roles in ws2. 256 with User(wp_admin, wp_admin_pw, monkeypatch), assert_unauthorized(): 257 client.create_role(workspace="ws2", name="editor") 258 259 260 def test_list_user_roles_self(client, monkeypatch): 261 username, password = _create_admin_controlled_user(client, monkeypatch) 262 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 263 role = client.create_role(workspace="ws1", name="viewer") 264 client.assign_role(username, role.id) 265 266 # Users can list their own role assignments. 267 with User(username, password, monkeypatch): 268 roles = client.list_user_roles(username) 269 assert [r.id for r in roles] == [role.id] 270 271 272 def test_list_user_roles_not_self_requires_admin(client, monkeypatch): 273 target, _ = _create_admin_controlled_user(client, monkeypatch) 274 requester, requester_pw = _create_admin_controlled_user(client, monkeypatch) 275 276 # A plain user cannot view another user's roles. 277 with User(requester, requester_pw, monkeypatch), assert_unauthorized(): 278 client.list_user_roles(target) 279 280 281 def test_list_user_roles_wp_admin_sees_only_own_workspace_scope(client, monkeypatch): 282 # Target has roles in two workspaces (ws1, ws2). Requester is WP admin of ws1 only. 283 # The response should be filtered to ws1 roles — ws2 membership must not leak. 284 target, _ = _create_admin_controlled_user(client, monkeypatch) 285 wp_admin, wp_admin_pw = _create_admin_controlled_user(client, monkeypatch) 286 _make_user_wp_admin(client, monkeypatch, wp_admin, "ws1") 287 288 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 289 ws1_role = client.create_role(workspace="ws1", name="ws1-viewer") 290 ws2_role = client.create_role(workspace="ws2", name="ws2-viewer") 291 client.assign_role(target, ws1_role.id) 292 client.assign_role(target, ws2_role.id) 293 294 with User(wp_admin, wp_admin_pw, monkeypatch): 295 visible = client.list_user_roles(target) 296 297 visible_names = {r.name for r in visible} 298 assert "ws1-viewer" in visible_names 299 assert "ws2-viewer" not in visible_names 300 301 302 def test_list_user_roles_super_admin_sees_all(client, monkeypatch): 303 target, _ = _create_admin_controlled_user(client, monkeypatch) 304 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 305 ws1_role = client.create_role(workspace="ws1", name="r1") 306 ws2_role = client.create_role(workspace="ws2", name="r2") 307 client.assign_role(target, ws1_role.id) 308 client.assign_role(target, ws2_role.id) 309 visible = client.list_user_roles(target) 310 311 assert {r.name for r in visible} == {"r1", "r2"} 312 313 314 # ---- Cross-workspace admin list ---- 315 316 317 def test_list_all_roles_admin_only(client, monkeypatch): 318 username, password = _create_admin_controlled_user(client, monkeypatch) 319 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 320 client.create_role(workspace="ws1", name="r1") 321 client.create_role(workspace="ws2", name="r2") 322 all_roles = client.list_all_roles() 323 names = {r.name for r in all_roles} 324 assert {"r1", "r2"}.issubset(names) 325 326 with User(username, password, monkeypatch), assert_unauthorized(): 327 client.list_all_roles() 328 329 330 # ---- /api/ and /ajax-api/ path parity ---- 331 332 333 @pytest.mark.parametrize("api_prefix", ["api", "ajax-api"]) 334 def test_rbac_endpoints_reachable_at_both_path_prefixes(client, monkeypatch, api_prefix): 335 # The MLflow frontend hits /ajax-api/ paths; the Python client hits /api/ paths. 336 # Every RBAC route must be reachable at both (see handlers._get_paths). 337 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 338 created = client.create_role(workspace="path-parity", name=f"r-{api_prefix}") 339 340 # GET list endpoint. 341 resp = requests.get( 342 f"{client.tracking_uri}/{api_prefix}/3.0/mlflow/roles/list", 343 params={"workspace": "path-parity"}, 344 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 345 ) 346 assert resp.status_code == 200, resp.text 347 names = {r["name"] for r in resp.json()["roles"]} 348 assert created.name in names 349 350 # GET single role by id. 351 resp = requests.get( 352 f"{client.tracking_uri}/{api_prefix}/3.0/mlflow/roles/get", 353 params={"role_id": str(created.id)}, 354 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 355 ) 356 assert resp.status_code == 200, resp.text 357 assert resp.json()["role"]["id"] == created.id