test_auth.py
1 """ 2 Integration test which starts a local Tracking Server on an ephemeral port, 3 and ensures authentication is working. 4 """ 5 6 import base64 7 import re 8 import subprocess 9 import sys 10 import time 11 from pathlib import Path 12 from unittest import mock 13 14 import jwt 15 import pytest 16 import requests 17 from cachetools import TTLCache 18 from cryptography.fernet import Fernet 19 20 import mlflow 21 from mlflow import MlflowClient 22 from mlflow.entities.logged_model_status import LoggedModelStatus 23 from mlflow.environment_variables import ( 24 _MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN, 25 MLFLOW_FLASK_SERVER_SECRET_KEY, 26 MLFLOW_TRACKING_PASSWORD, 27 MLFLOW_TRACKING_USERNAME, 28 ) 29 from mlflow.exceptions import MlflowException 30 from mlflow.protos.databricks_pb2 import ( 31 RESOURCE_DOES_NOT_EXIST, 32 UNAUTHENTICATED, 33 ErrorCode, 34 ) 35 from mlflow.server import auth as auth_module 36 from mlflow.server.auth import _authenticate_fastapi_request, _re_compile_path 37 from mlflow.server.auth.routes import ( 38 AJAX_LIST_USERS, 39 CREATE_REGISTERED_MODEL_PERMISSION, 40 GET_REGISTERED_MODEL_PERMISSION, 41 GET_SCORER_PERMISSION, 42 LIST_USERS, 43 ) 44 from mlflow.server.handlers import STATIC_PREFIX_ENV_VAR, _get_ajax_path 45 from mlflow.utils.os import is_windows 46 from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME 47 48 from tests.helper_functions import kill_process_tree, random_str 49 from tests.server.auth.auth_test_utils import ADMIN_PASSWORD, ADMIN_USERNAME, User, create_user 50 from tests.tracking.integration_test_utils import ( 51 _init_server, 52 _send_rest_tracking_post_request, 53 get_safe_port, 54 ) 55 56 57 @pytest.fixture 58 def client(request, tmp_path): 59 path = tmp_path.joinpath("sqlalchemy.db").as_uri() 60 backend_uri = ("sqlite://" if is_windows() else "sqlite:////") + path[len("file://") :] 61 extra_env = getattr(request, "param", {}) 62 extra_env[MLFLOW_FLASK_SERVER_SECRET_KEY.name] = "my-secret-key" 63 64 with _init_server( 65 backend_uri=backend_uri, 66 root_artifact_uri=tmp_path.joinpath("artifacts").as_uri(), 67 extra_env=extra_env, 68 app="mlflow.server.auth:create_app", 69 server_type="flask", 70 ) as url: 71 yield MlflowClient(url) 72 73 74 @pytest.fixture 75 def fastapi_client(request, tmp_path): 76 """FastAPI client fixture for testing FastAPI-specific middleware (e.g., gateway routes).""" 77 path = tmp_path.joinpath("sqlalchemy.db").as_uri() 78 backend_uri = ("sqlite://" if is_windows() else "sqlite:////") + path[len("file://") :] 79 extra_env = getattr(request, "param", {}) 80 extra_env[MLFLOW_FLASK_SERVER_SECRET_KEY.name] = "my-secret-key" 81 # Set _MLFLOW_SGI_NAME to "uvicorn" so auth module returns FastAPI app 82 extra_env["_MLFLOW_SGI_NAME"] = "uvicorn" 83 84 with _init_server( 85 backend_uri=backend_uri, 86 root_artifact_uri=tmp_path.joinpath("artifacts").as_uri(), 87 extra_env=extra_env, 88 app="mlflow.server.auth:create_app", 89 server_type="fastapi", 90 ) as url: 91 yield MlflowClient(url) 92 93 94 def test_authenticate(client, monkeypatch): 95 # unauthenticated 96 monkeypatch.delenv(MLFLOW_TRACKING_USERNAME.name, raising=False) 97 monkeypatch.delenv(MLFLOW_TRACKING_PASSWORD.name, raising=False) 98 with pytest.raises(MlflowException, match=r"You are not authenticated.") as exception_context: 99 client.search_experiments() 100 assert exception_context.value.error_code == ErrorCode.Name(UNAUTHENTICATED) 101 102 # authenticated 103 username, password = create_user(client.tracking_uri) 104 with User(username, password, monkeypatch): 105 client.search_experiments() 106 107 108 @pytest.mark.parametrize( 109 ("username", "password"), 110 [ 111 ("", "password"), 112 ("username", ""), 113 ("", ""), 114 ], 115 ) 116 def test_validate_username_and_password(client, username, password): 117 with pytest.raises(requests.exceptions.HTTPError, match=r"BAD REQUEST"): 118 create_user(client.tracking_uri, username=username, password=password) 119 120 121 def test_proxy_artifact_path_detection(): 122 assert auth_module._is_proxy_artifact_path("/api/2.0/mlflow-artifacts/artifacts/foo") 123 assert auth_module._is_proxy_artifact_path("/ajax-api/2.0/mlflow-artifacts/artifacts/foo") 124 125 126 def test_proxy_artifact_mpu_path_detection(): 127 # MPU create/complete/abort paths should be recognized as proxy artifact paths 128 for action in ("create", "complete", "abort"): 129 assert auth_module._is_proxy_artifact_path( 130 f"/api/2.0/mlflow-artifacts/mpu/{action}/1/run-id/artifacts/model" 131 ) 132 assert auth_module._is_proxy_artifact_path( 133 f"/ajax-api/2.0/mlflow-artifacts/mpu/{action}/1/run-id/artifacts/model" 134 ) 135 136 # Non-artifact paths should not match 137 assert not auth_module._is_proxy_artifact_path("/api/2.0/mlflow/experiments/get") 138 139 140 def test_proxy_artifact_mpu_validator_returns_update_for_post(): 141 validator = auth_module._get_proxy_artifact_validator( 142 "POST", {"artifact_path": "1/run-id/artifacts/model"} 143 ) 144 assert validator is auth_module.validate_can_update_experiment_artifact_proxy 145 146 147 def test_proxy_artifact_authorization_required(client, monkeypatch): 148 username1, password1 = create_user(client.tracking_uri) 149 username2, password2 = create_user(client.tracking_uri) 150 151 with User(username1, password1, monkeypatch): 152 experiment_id = client.create_experiment("proxy-artifact-authz-test") 153 154 response = requests.put( 155 url=( 156 client.tracking_uri 157 + f"/ajax-api/2.0/mlflow-artifacts/artifacts/{experiment_id}/test.txt" 158 ), 159 data=b"forbidden", 160 auth=(username2, password2), 161 ) 162 assert response.status_code == 403 163 164 165 @pytest.mark.parametrize( 166 "client", 167 [{"MLFLOW_AUTH_CONFIG_PATH": "tests/server/auth/fixtures/no_permission_auth.ini"}], 168 indirect=True, 169 ) 170 def test_proxy_artifact_list_query_param_uses_experiment_permission(client, monkeypatch): 171 # Regression test for https://github.com/mlflow/mlflow/issues/21201: 172 # When default_permission is NO_PERMISSIONS, a user with explicit experiment permission 173 # should be able to list artifacts via query parameter path (GET ?path=<experiment_id>/...). 174 username1, password1 = create_user(client.tracking_uri) 175 username2, password2 = create_user(client.tracking_uri) 176 177 with User(username1, password1, monkeypatch): 178 experiment_id = client.create_experiment("proxy-artifact-list-query-param-test") 179 180 # user1 has MANAGE on experiment — list via query param path should be allowed (HTTP 200) 181 response = requests.get( 182 url=client.tracking_uri + "/api/2.0/mlflow-artifacts/artifacts", 183 params={"path": f"{experiment_id}/models/m-abc123/artifacts"}, 184 auth=(username1, password1), 185 ) 186 assert response.status_code != 403 187 188 # user2 has no permission on the experiment — expect 403 189 response = requests.get( 190 url=client.tracking_uri + "/api/2.0/mlflow-artifacts/artifacts", 191 params={"path": f"{experiment_id}/models/m-abc123/artifacts"}, 192 auth=(username2, password2), 193 ) 194 assert response.status_code == 403 195 196 197 @pytest.mark.parametrize("mpu_action", ["create", "complete", "abort"]) 198 def test_mpu_authorization_required(client, monkeypatch, mpu_action): 199 username1, password1 = create_user(client.tracking_uri) 200 username2, password2 = create_user(client.tracking_uri) 201 202 with User(username1, password1, monkeypatch): 203 experiment_id = client.create_experiment(f"mpu-authz-test-{mpu_action}") 204 205 # user2 has no permission on user1's experiment — expect 403 206 response = requests.post( 207 url=( 208 client.tracking_uri 209 + f"/api/2.0/mlflow-artifacts/mpu/{mpu_action}/{experiment_id}/artifacts/model" 210 ), 211 json={"path": "python_model.pkl", "num_parts": 1}, 212 auth=(username2, password2), 213 ) 214 assert response.status_code == 403 215 216 217 def _mlflow_search_experiments_rest(base_uri, headers): 218 response = requests.post( 219 f"{base_uri}/api/2.0/mlflow/experiments/search", 220 headers=headers, 221 json={ 222 "max_results": 100, 223 }, 224 ) 225 response.raise_for_status() 226 return response 227 228 229 def _mlflow_create_user_rest(base_uri, headers): 230 username = random_str() 231 password = random_str() 232 response = requests.post( 233 f"{base_uri}/api/2.0/mlflow/users/create", 234 headers=headers, 235 json={ 236 "username": username, 237 "password": password, 238 }, 239 ) 240 response.raise_for_status() 241 return username, password 242 243 244 @pytest.mark.parametrize( 245 "client", 246 [ 247 { 248 "MLFLOW_AUTH_CONFIG_PATH": "tests/server/auth/fixtures/jwt_auth.ini", 249 "PYTHONPATH": str(Path.cwd() / "examples" / "jwt_auth"), 250 } 251 ], 252 indirect=True, 253 ) 254 def test_authenticate_jwt(client): 255 # unauthenticated 256 with pytest.raises(requests.HTTPError, match=r"401 Client Error: UNAUTHORIZED") as e: 257 _mlflow_search_experiments_rest(client.tracking_uri, {}) 258 assert e.value.response.status_code == 401 # Unauthorized 259 260 # authenticated 261 # we need to use jwt to authenticate as admin so that we can create a new user 262 bearer_token = jwt.encode({"username": ADMIN_USERNAME}, "secret", algorithm="HS256") 263 headers = {"Authorization": f"Bearer {bearer_token}"} 264 username, password = _mlflow_create_user_rest(client.tracking_uri, headers) 265 266 # authenticate with the newly created user 267 headers = { 268 "Authorization": f"Bearer {jwt.encode({'username': username}, 'secret', algorithm='HS256')}" 269 } 270 _mlflow_search_experiments_rest(client.tracking_uri, headers) 271 272 # invalid token 273 bearer_token = jwt.encode({"username": username}, "invalid", algorithm="HS256") 274 headers = {"Authorization": f"Bearer {bearer_token}"} 275 with pytest.raises(requests.HTTPError, match=r"401 Client Error: UNAUTHORIZED") as e: 276 _mlflow_search_experiments_rest(client.tracking_uri, headers) 277 assert e.value.response.status_code == 401 # Unauthorized 278 279 280 def test_search_experiments(client, monkeypatch): 281 """ 282 Use user1 to create 10 experiments, 283 grant READ permission to user2 on experiments [0, 3, 4, 5, 6, 8]. 284 Test whether user2 can search only and all the readable experiments, 285 both paged and un-paged. 286 """ 287 username1, password1 = create_user(client.tracking_uri) 288 username2, password2 = create_user(client.tracking_uri) 289 290 readable = [0, 3, 4, 5, 6, 8] 291 292 with User(username1, password1, monkeypatch): 293 for i in range(10): 294 experiment_id = client.create_experiment(f"exp{i}") 295 _send_rest_tracking_post_request( 296 client.tracking_uri, 297 "/api/2.0/mlflow/experiments/permissions/create", 298 json_payload={ 299 "experiment_id": experiment_id, 300 "username": username2, 301 "permission": "READ" if i in readable else "NO_PERMISSIONS", 302 }, 303 auth=(username1, password1), 304 ) 305 306 # test un-paged search 307 with User(username1, password1, monkeypatch): 308 experiments = client.search_experiments( 309 max_results=100, 310 filter_string="name LIKE 'exp%'", 311 order_by=["name ASC"], 312 ) 313 names = sorted([exp.name for exp in experiments]) 314 assert names == [f"exp{i}" for i in range(10)] 315 316 with User(username2, password2, monkeypatch): 317 experiments = client.search_experiments( 318 max_results=100, 319 filter_string="name LIKE 'exp%'", 320 order_by=["name ASC"], 321 ) 322 names = sorted([exp.name for exp in experiments]) 323 assert names == [f"exp{i}" for i in readable] 324 325 # test paged search 326 with User(username1, password1, monkeypatch): 327 page_token = "" 328 experiments = [] 329 while True: 330 res = client.search_experiments( 331 max_results=4, 332 filter_string="name LIKE 'exp%'", 333 order_by=["name ASC"], 334 page_token=page_token, 335 ) 336 experiments.extend(res) 337 page_token = res.token 338 if not page_token: 339 break 340 341 names = sorted([exp.name for exp in experiments]) 342 assert names == [f"exp{i}" for i in range(10)] 343 344 with User(username2, password2, monkeypatch): 345 page_token = "" 346 experiments = [] 347 while True: 348 res = client.search_experiments( 349 max_results=4, 350 filter_string="name LIKE 'exp%'", 351 order_by=["name ASC"], 352 page_token=page_token, 353 ) 354 experiments.extend(res) 355 page_token = res.token 356 if not page_token: 357 break 358 359 names = sorted([exp.name for exp in experiments]) 360 assert names == [f"exp{i}" for i in readable] 361 362 363 def test_search_registered_models(client, monkeypatch): 364 """ 365 Use user1 to create 10 registered_models, 366 grant READ permission to user2 on registered_models [0, 3, 4, 5, 6, 8]. 367 Test whether user2 can search only and all the readable registered_models, 368 both paged and un-paged. 369 """ 370 username1, password1 = create_user(client.tracking_uri) 371 username2, password2 = create_user(client.tracking_uri) 372 373 readable = [0, 3, 4, 5, 6, 8] 374 375 with User(username1, password1, monkeypatch): 376 for i in range(10): 377 rm = client.create_registered_model(f"rm{i}") 378 _send_rest_tracking_post_request( 379 client.tracking_uri, 380 "/api/2.0/mlflow/registered-models/permissions/create", 381 json_payload={ 382 "name": rm.name, 383 "username": username2, 384 "permission": "READ" if i in readable else "NO_PERMISSIONS", 385 }, 386 auth=(username1, password1), 387 ) 388 389 # test un-paged search 390 with User(username1, password1, monkeypatch): 391 registered_models = client.search_registered_models( 392 max_results=100, 393 filter_string="name LIKE 'rm%'", 394 order_by=["name ASC"], 395 ) 396 names = sorted([rm.name for rm in registered_models]) 397 assert names == [f"rm{i}" for i in range(10)] 398 399 with User(username2, password2, monkeypatch): 400 registered_models = client.search_registered_models( 401 max_results=100, 402 filter_string="name LIKE 'rm%'", 403 order_by=["name ASC"], 404 ) 405 names = sorted([rm.name for rm in registered_models]) 406 assert names == [f"rm{i}" for i in readable] 407 408 # test paged search 409 with User(username1, password1, monkeypatch): 410 page_token = "" 411 registered_models = [] 412 while True: 413 res = client.search_registered_models( 414 max_results=4, 415 filter_string="name LIKE 'rm%'", 416 order_by=["name ASC"], 417 page_token=page_token, 418 ) 419 registered_models.extend(res) 420 page_token = res.token 421 if not page_token: 422 break 423 424 names = sorted([rm.name for rm in registered_models]) 425 assert names == [f"rm{i}" for i in range(10)] 426 427 with User(username2, password2, monkeypatch): 428 page_token = "" 429 registered_models = [] 430 while True: 431 res = client.search_registered_models( 432 max_results=4, 433 filter_string="name LIKE 'rm%'", 434 order_by=["name ASC"], 435 page_token=page_token, 436 ) 437 registered_models.extend(res) 438 page_token = res.token 439 if not page_token: 440 break 441 442 names = sorted([rm.name for rm in registered_models]) 443 assert names == [f"rm{i}" for i in readable] 444 445 446 def test_search_model_versions(client, monkeypatch): 447 username1, password1 = create_user(client.tracking_uri) 448 username2, password2 = create_user(client.tracking_uri) 449 450 readable = [0, 2, 4] 451 452 with User(username1, password1, monkeypatch): 453 experiment_id = client.create_experiment("mv_test_exp") 454 run = client.create_run(experiment_id) 455 run_id = run.info.run_id 456 for i in range(5): 457 rm = client.create_registered_model(f"mv_model{i}") 458 client.create_model_version(rm.name, f"runs:/{run_id}/model", run_id=run_id) 459 _send_rest_tracking_post_request( 460 client.tracking_uri, 461 "/api/2.0/mlflow/registered-models/permissions/create", 462 json_payload={ 463 "name": rm.name, 464 "username": username2, 465 "permission": "READ" if i in readable else "NO_PERMISSIONS", 466 }, 467 auth=(username1, password1), 468 ) 469 470 # user1 (owner) sees all model versions 471 with User(username1, password1, monkeypatch): 472 versions = client.search_model_versions(filter_string="name LIKE 'mv_model%'") 473 names = sorted({mv.name for mv in versions}) 474 assert names == [f"mv_model{i}" for i in range(5)] 475 476 # user2 only sees model versions for readable models 477 with User(username2, password2, monkeypatch): 478 versions = client.search_model_versions(filter_string="name LIKE 'mv_model%'") 479 names = sorted({mv.name for mv in versions}) 480 assert names == [f"mv_model{i}" for i in readable] 481 482 483 def test_graphql_search_model_versions(client, monkeypatch): 484 username1, password1 = create_user(client.tracking_uri) 485 username2, password2 = create_user(client.tracking_uri) 486 487 readable = [0, 2, 4] 488 489 with User(username1, password1, monkeypatch): 490 experiment_id = client.create_experiment("gql_mv_test_exp") 491 run = client.create_run(experiment_id) 492 run_id = run.info.run_id 493 for i in range(5): 494 rm = client.create_registered_model(f"gql_mv_model{i}") 495 client.create_model_version(rm.name, f"runs:/{run_id}/model", run_id=run_id) 496 _send_rest_tracking_post_request( 497 client.tracking_uri, 498 "/api/2.0/mlflow/registered-models/permissions/create", 499 json_payload={ 500 "name": rm.name, 501 "username": username2, 502 "permission": "READ" if i in readable else "NO_PERMISSIONS", 503 }, 504 auth=(username1, password1), 505 ) 506 507 query = """ 508 query SearchModelVersions($input: MlflowSearchModelVersionsInput){ 509 mlflowSearchModelVersions(input: $input){ 510 modelVersions { name version } 511 } 512 } 513 """ 514 variables = {"input": {"filter": "name LIKE 'gql_mv_model%'"}} 515 516 # user1 (owner) sees all via GraphQL 517 resp = requests.post( 518 f"{client.tracking_uri}/graphql", 519 json={"query": query, "variables": variables}, 520 auth=(username1, password1), 521 ) 522 resp.raise_for_status() 523 payload = resp.json() 524 assert payload.get("errors") in (None, []) 525 names = sorted({ 526 mv["name"] for mv in payload["data"]["mlflowSearchModelVersions"]["modelVersions"] 527 }) 528 assert names == [f"gql_mv_model{i}" for i in range(5)] 529 530 # user2 only sees versions for readable models via GraphQL 531 resp = requests.post( 532 f"{client.tracking_uri}/graphql", 533 json={"query": query, "variables": variables}, 534 auth=(username2, password2), 535 ) 536 resp.raise_for_status() 537 payload = resp.json() 538 assert payload.get("errors") in (None, []) 539 names = sorted({ 540 mv["name"] for mv in payload["data"]["mlflowSearchModelVersions"]["modelVersions"] 541 }) 542 assert names == [f"gql_mv_model{i}" for i in readable] 543 544 545 def test_create_and_delete_registered_model(client, monkeypatch): 546 username1, password1 = create_user(client.tracking_uri) 547 548 # create a registered model 549 with User(username1, password1, monkeypatch): 550 rm = client.create_registered_model("test_model") 551 552 # confirm the permission is set correctly 553 with User(username1, password1, monkeypatch): 554 response = requests.get( 555 url=client.tracking_uri + GET_REGISTERED_MODEL_PERMISSION, 556 params={"name": rm.name, "username": username1}, 557 auth=(username1, password1), 558 ) 559 560 permission = response.json()["registered_model_permission"] 561 assert permission["name"] == rm.name 562 assert permission["permission"] == "MANAGE" 563 assert permission["workspace"] == DEFAULT_WORKSPACE_NAME 564 565 # trying to create a model with the same name should fail 566 with User(username1, password1, monkeypatch): 567 with pytest.raises(MlflowException, match=r"RESOURCE_ALREADY_EXISTS"): 568 client.create_registered_model("test_model") 569 570 # delete the registered model 571 with User(username1, password1, monkeypatch): 572 client.delete_registered_model(rm.name) 573 574 # confirm the registered model permission is also deleted 575 with User(username1, password1, monkeypatch): 576 response = requests.get( 577 url=client.tracking_uri + GET_REGISTERED_MODEL_PERMISSION, 578 params={"name": rm.name, "username": username1}, 579 # Check with admin because the user permission is deleted 580 auth=("admin", "password1234"), 581 ) 582 583 assert response.status_code == 404 584 assert response.json()["error_code"] == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 585 expected_message = ( 586 "Registered model permission with " 587 f"workspace={DEFAULT_WORKSPACE_NAME}, name={rm.name} " 588 f"and username={username1} not found" 589 ) 590 assert response.json()["message"] == expected_message 591 592 # now we should be able to create a model with the same name 593 with User(username1, password1, monkeypatch): 594 rm = client.create_registered_model("test_model") 595 assert rm.name == "test_model" 596 597 598 def test_delete_registered_model_clears_all_permissions(client, monkeypatch): 599 username1, password1 = create_user(client.tracking_uri) 600 username2, _password2 = create_user(client.tracking_uri) 601 602 # create a registered model and grant user2 READ 603 with User(username1, password1, monkeypatch): 604 rm = client.create_registered_model("test_model_permissions") 605 _send_rest_tracking_post_request( 606 client.tracking_uri, 607 CREATE_REGISTERED_MODEL_PERMISSION, 608 json_payload={"name": rm.name, "username": username2, "permission": "READ"}, 609 auth=(username1, password1), 610 ) 611 612 # confirm permission exists for user2 613 with User(username1, password1, monkeypatch): 614 response = requests.get( 615 url=client.tracking_uri + GET_REGISTERED_MODEL_PERMISSION, 616 params={"name": rm.name, "username": username2}, 617 auth=(username1, password1), 618 ) 619 assert response.ok 620 assert response.json()["registered_model_permission"]["permission"] == "READ" 621 622 # delete the registered model 623 with User(username1, password1, monkeypatch): 624 client.delete_registered_model(rm.name) 625 626 # confirm permissions are deleted for *all* users (check as admin) 627 for username in (username1, username2): 628 response = requests.get( 629 url=client.tracking_uri + GET_REGISTERED_MODEL_PERMISSION, 630 params={"name": rm.name, "username": username}, 631 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 632 ) 633 assert response.status_code == 404 634 assert response.json()["error_code"] == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 635 636 # recreate model with the same name; user2 should not regain access implicitly 637 with User(username1, password1, monkeypatch): 638 rm2 = client.create_registered_model(rm.name) 639 assert rm2.name == rm.name 640 641 response = requests.get( 642 url=client.tracking_uri + GET_REGISTERED_MODEL_PERMISSION, 643 params={"name": rm.name, "username": username2}, 644 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 645 ) 646 assert response.status_code == 404 647 648 649 def _wait(url: str, timeout: int = 10) -> None: 650 t = time.time() 651 while time.time() - t < timeout: 652 try: 653 if requests.get(f"{url}/health").ok: 654 return 655 except requests.exceptions.ConnectionError: 656 pass 657 time.sleep(0.5) 658 659 pytest.fail("Server did not start") 660 661 662 def test_proxy_log_artifacts(monkeypatch, tmp_path): 663 backend_uri = f"sqlite:///{tmp_path / 'sqlalchemy.db'}" 664 port = get_safe_port() 665 host = "localhost" 666 with subprocess.Popen( 667 [ 668 sys.executable, 669 "-m", 670 "mlflow", 671 "server", 672 "--app-name", 673 "basic-auth", 674 "--backend-store-uri", 675 backend_uri, 676 "--host", 677 host, 678 "--port", 679 str(port), 680 "--workers", 681 "1", 682 "--gunicorn-opts", 683 "--log-level debug", 684 ], 685 env={MLFLOW_FLASK_SERVER_SECRET_KEY.name: "my-secret-key"}, 686 ) as prc: 687 try: 688 url = f"http://{host}:{port}" 689 _wait(url) 690 691 mlflow.set_tracking_uri(url) 692 client = MlflowClient(url) 693 tmp_file = tmp_path / "test.txt" 694 tmp_file.touch() 695 username1, password1 = create_user(url) 696 with User(username1, password1, monkeypatch): 697 exp_id = client.create_experiment("exp") 698 run = client.create_run(exp_id) 699 client.log_artifact(run.info.run_id, tmp_file) 700 client.list_artifacts(run.info.run_id) 701 702 username2, password2 = create_user(url) 703 with User(username2, password2, monkeypatch): 704 client.list_artifacts(run.info.run_id) 705 with pytest.raises(requests.HTTPError, match="Permission denied"): 706 client.log_artifact(run.info.run_id, tmp_file) 707 708 # Ensure that the regular expression captures an experiment ID correctly 709 tmp_file_with_numbers = tmp_path / "123456.txt" 710 tmp_file_with_numbers.touch() 711 with pytest.raises(requests.HTTPError, match="Permission denied"): 712 client.log_artifact(run.info.run_id, tmp_file_with_numbers) 713 finally: 714 # Kill the server process to prevent `prc.wait()` (called when exiting the context 715 # manager) from waiting forever. 716 kill_process_tree(prc.pid) 717 718 719 def test_create_user_from_ui_fails_without_csrf_token(client): 720 response = requests.post( 721 client.tracking_uri + "/api/2.0/mlflow/users/create-ui", 722 json={"username": "test", "password": "test"}, 723 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 724 headers={"Content-Type": "application/x-www-form-urlencoded"}, 725 ) 726 727 assert "The CSRF token is missing" in response.text 728 729 730 def test_create_user_ui(client): 731 # needs to be a session as the CSRF protection will set some 732 # cookies that need to be present for server side validation 733 with requests.Session() as session: 734 page = session.get(client.tracking_uri + "/signup", auth=(ADMIN_USERNAME, ADMIN_PASSWORD)) 735 736 csrf_regex = re.compile(r"name=\"csrf_token\" value=\"([\S]+)\"") 737 match = csrf_regex.search(page.text) 738 739 # assert that the CSRF token is sent in the form 740 assert match is not None 741 742 csrf_token = match.group(1) 743 744 response = session.post( 745 client.tracking_uri + "/api/2.0/mlflow/users/create-ui", 746 data={"username": random_str(), "password": random_str(), "csrf_token": csrf_token}, 747 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 748 headers={"Content-Type": "application/x-www-form-urlencoded"}, 749 ) 750 751 assert "Successfully signed up user" in response.text 752 753 754 def test_logged_model(client: MlflowClient, monkeypatch: pytest.MonkeyPatch): 755 username1, password1 = create_user(client.tracking_uri) 756 username2, password2 = create_user(client.tracking_uri) 757 758 class Model(mlflow.pyfunc.PythonModel): 759 def predict(self, context, model_input): 760 return model_input 761 762 with User(username1, password1, monkeypatch): 763 exp_id = client.create_experiment("exp") 764 model = client.create_logged_model(experiment_id=exp_id) 765 client.finalize_logged_model(model_id=model.model_id, status=LoggedModelStatus.READY) 766 client.set_logged_model_tags(model_id=model.model_id, tags={"key": "value"}) 767 client.delete_logged_model_tag(model_id=model.model_id, key="key") 768 models = client.search_logged_models(experiment_ids=[exp_id]) 769 assert len(models) == 1 770 771 with User(username2, password2, monkeypatch): 772 loaded_model = client.get_logged_model(model.model_id) 773 assert loaded_model.model_id == model.model_id 774 775 models = client.search_logged_models(experiment_ids=[exp_id]) 776 assert len(models) == 1 777 778 with pytest.raises(MlflowException, match="Permission denied"): 779 client.finalize_logged_model(model_id=model.model_id, status=LoggedModelStatus.READY) 780 with pytest.raises(MlflowException, match="Permission denied"): 781 client.set_logged_model_tags(model_id=model.model_id, tags={"key": "value"}) 782 with pytest.raises(MlflowException, match="Permission denied"): 783 client.delete_logged_model_tag(model_id=model.model_id, key="key") 784 with pytest.raises(MlflowException, match="Permission denied"): 785 client.delete_logged_model(model_id=model.model_id) 786 787 788 @pytest.mark.parametrize( 789 "client", 790 [{"MLFLOW_AUTH_CONFIG_PATH": "tests/server/auth/fixtures/no_permission_auth.ini"}], 791 indirect=True, 792 ) 793 def test_logged_model_artifact_authorization(client: MlflowClient, monkeypatch: pytest.MonkeyPatch): 794 username1, password1 = create_user(client.tracking_uri) 795 username2, password2 = create_user(client.tracking_uri) 796 797 with User(username1, password1, monkeypatch): 798 exp_id = client.create_experiment("logged-model-artifact-authz-test") 799 model = client.create_logged_model(experiment_id=exp_id) 800 801 # user1 (owner) should be able to access the artifact endpoint (404 since no artifact 802 # exists, but should NOT be 403) 803 response = requests.get( 804 url=( 805 client.tracking_uri 806 + f"/ajax-api/2.0/mlflow/logged-models/{model.model_id}/artifacts/files" 807 ), 808 params={"artifact_file_path": "test.txt"}, 809 auth=(username1, password1), 810 ) 811 assert response.status_code != 403 812 813 # user2 has no permission on the experiment — expect 403 814 response = requests.get( 815 url=( 816 client.tracking_uri 817 + f"/ajax-api/2.0/mlflow/logged-models/{model.model_id}/artifacts/files" 818 ), 819 params={"artifact_file_path": "test.txt"}, 820 auth=(username2, password2), 821 ) 822 assert response.status_code == 403 823 824 # Also verify the list-artifacts (directories) endpoint 825 # user1 (owner) should be able to list artifacts 826 response = requests.get( 827 url=( 828 client.tracking_uri 829 + f"/api/2.0/mlflow/logged-models/{model.model_id}/artifacts/directories" 830 ), 831 auth=(username1, password1), 832 ) 833 assert response.status_code != 403 834 835 # user2 has no permission — expect 403 836 response = requests.get( 837 url=( 838 client.tracking_uri 839 + f"/api/2.0/mlflow/logged-models/{model.model_id}/artifacts/directories" 840 ), 841 auth=(username2, password2), 842 ) 843 assert response.status_code == 403 844 845 846 def test_logged_model_artifact_validator_respects_static_prefix( 847 monkeypatch: pytest.MonkeyPatch, 848 ): 849 base = "/mlflow/logged-models/<model_id>/artifacts/files" 850 851 # Without prefix — should match the bare path 852 pat_no_prefix = _re_compile_path(_get_ajax_path(base)) 853 assert pat_no_prefix.fullmatch("/ajax-api/2.0/mlflow/logged-models/abc123/artifacts/files") 854 855 # With prefix — should match the prefixed path 856 monkeypatch.setenv(STATIC_PREFIX_ENV_VAR, "/custom-prefix") 857 _re_compile_path.cache_clear() 858 pat_with_prefix = _re_compile_path(_get_ajax_path(base)) 859 assert pat_with_prefix.fullmatch( 860 "/custom-prefix/ajax-api/2.0/mlflow/logged-models/abc123/artifacts/files" 861 ) 862 # bare path should NOT match the prefixed pattern 863 assert not pat_with_prefix.fullmatch( 864 "/ajax-api/2.0/mlflow/logged-models/abc123/artifacts/files" 865 ) 866 867 _re_compile_path.cache_clear() 868 869 870 def test_search_logged_models(client: MlflowClient, monkeypatch: pytest.MonkeyPatch): 871 username1, password1 = create_user(client.tracking_uri) 872 username2, password2 = create_user(client.tracking_uri) 873 readable = [0, 3, 4, 5, 6, 8] 874 with User(username1, password1, monkeypatch): 875 experiment_ids: list[str] = [] 876 for i in range(10): 877 experiment_id = client.create_experiment(f"exp-{i}") 878 experiment_ids.append(experiment_id) 879 _send_rest_tracking_post_request( 880 client.tracking_uri, 881 "/api/2.0/mlflow/experiments/permissions/create", 882 json_payload={ 883 "experiment_id": experiment_id, 884 "username": username2, 885 "permission": "READ" if (i in readable) else "NO_PERMISSIONS", 886 }, 887 auth=(username1, password1), 888 ) 889 client.create_logged_model(experiment_id=experiment_id) 890 891 models = client.search_logged_models(experiment_ids=experiment_ids) 892 assert len(models) == 10 893 894 # Pagination 895 models = client.search_logged_models(experiment_ids=experiment_ids, max_results=2) 896 assert len(models) == 2 897 assert models.token is not None 898 899 models = client.search_logged_models( 900 experiment_ids=experiment_ids, max_results=2, page_token=models.token 901 ) 902 assert len(models) == 2 903 assert models.token is not None 904 905 models = client.search_logged_models(experiment_ids=experiment_ids, page_token=models.token) 906 assert len(models) == 6 907 assert models.token is None 908 909 with User(username2, password2, monkeypatch): 910 models = client.search_logged_models(experiment_ids=experiment_ids) 911 assert len(models) == len(readable) 912 913 # Pagination 914 models = client.search_logged_models(experiment_ids=experiment_ids, max_results=2) 915 assert len(models) == 2 916 assert models.token is not None 917 918 models = client.search_logged_models( 919 experiment_ids=experiment_ids, max_results=2, page_token=models.token 920 ) 921 assert len(models) == 2 922 assert models.token is not None 923 924 models = client.search_logged_models(experiment_ids=experiment_ids, page_token=models.token) 925 assert len(models) == 2 926 assert models.token is None 927 928 929 def test_search_runs(client: MlflowClient, monkeypatch: pytest.MonkeyPatch): 930 username1, password1 = create_user(client.tracking_uri) 931 username2, password2 = create_user(client.tracking_uri) 932 933 readable = [0, 2] 934 935 with User(username1, password1, monkeypatch): 936 experiment_ids: list[str] = [] 937 run_counts = [8, 10, 7] 938 all_runs = {} 939 940 for i in range(3): 941 experiment_id = client.create_experiment(f"exp-{i}") 942 experiment_ids.append(experiment_id) 943 _send_rest_tracking_post_request( 944 client.tracking_uri, 945 "/api/2.0/mlflow/experiments/permissions/create", 946 json_payload={ 947 "experiment_id": experiment_id, 948 "username": username2, 949 "permission": "READ" if i in readable else "NO_PERMISSIONS", 950 }, 951 auth=(username1, password1), 952 ) 953 954 all_runs[experiment_id] = [] 955 for _ in range(run_counts[i]): 956 run = client.create_run(experiment_id) 957 all_runs[experiment_id].append(run.info.run_id) 958 959 expected_readable_runs = set(all_runs[experiment_ids[0]] + all_runs[experiment_ids[2]]) 960 961 with User(username1, password1, monkeypatch): 962 runs = client.search_runs(experiment_ids=experiment_ids) 963 assert len(runs) == sum(run_counts) 964 965 with User(username2, password2, monkeypatch): 966 runs = client.search_runs(experiment_ids=experiment_ids) 967 returned_run_ids = {run.info.run_id for run in runs} 968 assert returned_run_ids == expected_readable_runs 969 assert len(runs) == len(expected_readable_runs) 970 971 page_token = None 972 all_paginated_runs = [] 973 while True: 974 runs = client.search_runs( 975 experiment_ids=experiment_ids, 976 max_results=3, 977 page_token=page_token, 978 ) 979 all_paginated_runs.extend([run.info.run_id for run in runs]) 980 page_token = runs.token 981 if not page_token: 982 break 983 984 assert len(all_paginated_runs) == len(set(all_paginated_runs)) 985 assert set(all_paginated_runs) == expected_readable_runs 986 987 inaccessible_runs = set(all_runs[experiment_ids[1]]) 988 returned_inaccessible = set(all_paginated_runs) & inaccessible_runs 989 assert len(returned_inaccessible) == 0 990 991 992 def test_register_and_delete_scorer(client, monkeypatch): 993 username1, password1 = create_user(client.tracking_uri) 994 995 with User(username1, password1, monkeypatch): 996 experiment_id = client.create_experiment("test_experiment") 997 998 scorer_json = '{"name": "test_scorer", "type": "pyfunc"}' 999 1000 with User(username1, password1, monkeypatch): 1001 response = _send_rest_tracking_post_request( 1002 client.tracking_uri, 1003 "/api/3.0/mlflow/scorers/register", 1004 json_payload={ 1005 "experiment_id": experiment_id, 1006 "name": "test_scorer", 1007 "serialized_scorer": scorer_json, 1008 }, 1009 auth=(username1, password1), 1010 ) 1011 1012 scorer_name = response.json()["name"] 1013 assert scorer_name == "test_scorer" 1014 1015 with User(username1, password1, monkeypatch): 1016 response = requests.get( 1017 url=client.tracking_uri + GET_SCORER_PERMISSION, 1018 params={ 1019 "experiment_id": experiment_id, 1020 "scorer_name": scorer_name, 1021 "username": username1, 1022 }, 1023 auth=(username1, password1), 1024 ) 1025 1026 permission = response.json()["scorer_permission"] 1027 assert permission["experiment_id"] == experiment_id 1028 assert permission["scorer_name"] == scorer_name 1029 assert permission["permission"] == "MANAGE" 1030 1031 with User(username1, password1, monkeypatch): 1032 requests.delete( 1033 url=client.tracking_uri + "/api/3.0/mlflow/scorers/delete", 1034 json={ 1035 "experiment_id": experiment_id, 1036 "name": scorer_name, 1037 }, 1038 auth=(username1, password1), 1039 ) 1040 1041 with User(username1, password1, monkeypatch): 1042 response = requests.get( 1043 url=client.tracking_uri + GET_SCORER_PERMISSION, 1044 params={ 1045 "experiment_id": experiment_id, 1046 "scorer_name": scorer_name, 1047 "username": username1, 1048 }, 1049 auth=("admin", "password1234"), 1050 ) 1051 1052 assert response.status_code == 404 1053 assert response.json()["error_code"] == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 1054 1055 1056 def test_reregister_scorer_does_not_raise(client, monkeypatch): 1057 username1, password1 = create_user(client.tracking_uri) 1058 1059 with User(username1, password1, monkeypatch): 1060 experiment_id = client.create_experiment("test_experiment") 1061 1062 scorer_json = '{"name": "test_scorer", "type": "pyfunc"}' 1063 1064 # First registration 1065 with User(username1, password1, monkeypatch): 1066 response = _send_rest_tracking_post_request( 1067 client.tracking_uri, 1068 "/api/3.0/mlflow/scorers/register", 1069 json_payload={ 1070 "experiment_id": experiment_id, 1071 "name": "test_scorer", 1072 "serialized_scorer": scorer_json, 1073 }, 1074 auth=(username1, password1), 1075 ) 1076 assert response.status_code == 200 1077 assert response.json()["version"] == 1 1078 1079 # Re-registration with the same name should succeed (not raise RESOURCE_ALREADY_EXISTS) 1080 updated_scorer_json = '{"name": "test_scorer", "type": "pyfunc", "updated": true}' 1081 with User(username1, password1, monkeypatch): 1082 response = _send_rest_tracking_post_request( 1083 client.tracking_uri, 1084 "/api/3.0/mlflow/scorers/register", 1085 json_payload={ 1086 "experiment_id": experiment_id, 1087 "name": "test_scorer", 1088 "serialized_scorer": updated_scorer_json, 1089 }, 1090 auth=(username1, password1), 1091 ) 1092 assert response.status_code == 200 1093 assert response.json()["version"] == 2 1094 1095 1096 def test_scorer_permission_denial(client, monkeypatch): 1097 username1, password1 = create_user(client.tracking_uri) 1098 username2, password2 = create_user(client.tracking_uri) 1099 1100 with User(username1, password1, monkeypatch): 1101 experiment_id = client.create_experiment("test_experiment") 1102 1103 scorer_json = '{"name": "test_scorer", "type": "pyfunc"}' 1104 1105 with User(username1, password1, monkeypatch): 1106 response = _send_rest_tracking_post_request( 1107 client.tracking_uri, 1108 "/api/3.0/mlflow/scorers/register", 1109 json_payload={ 1110 "experiment_id": experiment_id, 1111 "name": "test_scorer", 1112 "serialized_scorer": scorer_json, 1113 }, 1114 auth=(username1, password1), 1115 ) 1116 1117 scorer_name = response.json()["name"] 1118 1119 with User(username2, password2, monkeypatch): 1120 # user2 has default READ permission, so they CAN read the scorer 1121 response = requests.get( 1122 url=client.tracking_uri + "/api/3.0/mlflow/scorers/get", 1123 params={ 1124 "experiment_id": experiment_id, 1125 "name": scorer_name, 1126 }, 1127 auth=(username2, password2), 1128 ) 1129 response.raise_for_status() 1130 assert response.json()["scorer"]["scorer_name"] == scorer_name 1131 1132 # But they CANNOT delete it (READ permission doesn't allow delete) 1133 response = requests.delete( 1134 url=client.tracking_uri + "/api/3.0/mlflow/scorers/delete", 1135 json={ 1136 "experiment_id": experiment_id, 1137 "name": scorer_name, 1138 }, 1139 auth=(username2, password2), 1140 ) 1141 with pytest.raises(requests.HTTPError, match="403"): 1142 response.raise_for_status() 1143 1144 1145 def test_scorer_read_permission(client, monkeypatch): 1146 username1, password1 = create_user(client.tracking_uri) 1147 username2, password2 = create_user(client.tracking_uri) 1148 1149 with User(username1, password1, monkeypatch): 1150 experiment_id = client.create_experiment("test_experiment") 1151 1152 scorer_json = '{"name": "test_scorer", "type": "pyfunc"}' 1153 1154 with User(username1, password1, monkeypatch): 1155 response = _send_rest_tracking_post_request( 1156 client.tracking_uri, 1157 "/api/3.0/mlflow/scorers/register", 1158 json_payload={ 1159 "experiment_id": experiment_id, 1160 "name": "test_scorer", 1161 "serialized_scorer": scorer_json, 1162 }, 1163 auth=(username1, password1), 1164 ) 1165 1166 scorer_name = response.json()["name"] 1167 1168 _send_rest_tracking_post_request( 1169 client.tracking_uri, 1170 "/api/3.0/mlflow/scorers/permissions/create", 1171 json_payload={ 1172 "experiment_id": experiment_id, 1173 "scorer_name": scorer_name, 1174 "username": username2, 1175 "permission": "READ", 1176 }, 1177 auth=(username1, password1), 1178 ) 1179 1180 with User(username2, password2, monkeypatch): 1181 response = requests.get( 1182 url=client.tracking_uri + "/api/3.0/mlflow/scorers/get", 1183 params={ 1184 "experiment_id": experiment_id, 1185 "name": scorer_name, 1186 }, 1187 auth=(username2, password2), 1188 ) 1189 response.raise_for_status() 1190 assert response.json()["scorer"]["scorer_name"] == scorer_name 1191 1192 with User(username2, password2, monkeypatch): 1193 response = requests.delete( 1194 url=client.tracking_uri + "/api/3.0/mlflow/scorers/delete", 1195 json={ 1196 "experiment_id": experiment_id, 1197 "name": scorer_name, 1198 }, 1199 auth=(username2, password2), 1200 ) 1201 with pytest.raises(requests.HTTPError, match="403"): 1202 response.raise_for_status() 1203 1204 1205 def _graphql_query(tracking_uri, query, variables=None, auth=None): 1206 return requests.post( 1207 f"{tracking_uri}/graphql", 1208 json={"query": query, "variables": variables or {}}, 1209 auth=auth, 1210 ) 1211 1212 1213 def test_graphql_requires_authentication(client, monkeypatch): 1214 monkeypatch.delenv(MLFLOW_TRACKING_USERNAME.name, raising=False) 1215 monkeypatch.delenv(MLFLOW_TRACKING_PASSWORD.name, raising=False) 1216 1217 query = """ 1218 query { 1219 mlflowGetExperiment(input: {experimentId: "0"}) { 1220 experiment { 1221 experimentId 1222 name 1223 } 1224 } 1225 } 1226 """ 1227 response = _graphql_query(client.tracking_uri, query) 1228 assert response.status_code == 401 1229 1230 1231 def test_graphql_get_experiment_authorization(client, monkeypatch): 1232 username1, password1 = create_user(client.tracking_uri) 1233 username2, password2 = create_user(client.tracking_uri) 1234 1235 with User(username1, password1, monkeypatch): 1236 experiment_id = client.create_experiment("graphql_test_exp") 1237 _send_rest_tracking_post_request( 1238 client.tracking_uri, 1239 "/api/2.0/mlflow/experiments/permissions/create", 1240 json_payload={ 1241 "experiment_id": experiment_id, 1242 "username": username2, 1243 "permission": "NO_PERMISSIONS", 1244 }, 1245 auth=(username1, password1), 1246 ) 1247 1248 query = """ 1249 query($expId: String!) { 1250 mlflowGetExperiment(input: {experimentId: $expId}) { 1251 experiment { 1252 experimentId 1253 name 1254 } 1255 } 1256 } 1257 """ 1258 1259 # user1 (creator) should be able to read the experiment 1260 response = _graphql_query( 1261 client.tracking_uri, 1262 query, 1263 variables={"expId": experiment_id}, 1264 auth=(username1, password1), 1265 ) 1266 assert response.status_code == 200 1267 data = response.json() 1268 experiment_data = data["data"]["mlflowGetExperiment"]["experiment"] 1269 assert experiment_data["experimentId"] == experiment_id 1270 assert experiment_data["name"] == "graphql_test_exp" 1271 1272 # user2 (NO_PERMISSIONS) should NOT be able to read the experiment 1273 response = _graphql_query( 1274 client.tracking_uri, 1275 query, 1276 variables={"expId": experiment_id}, 1277 auth=(username2, password2), 1278 ) 1279 assert response.status_code == 200 1280 data = response.json() 1281 # With authorization denied, the result should be null 1282 assert data.get("data", {}).get("mlflowGetExperiment") is None 1283 1284 1285 def test_graphql_get_run_authorization(client, monkeypatch): 1286 username1, password1 = create_user(client.tracking_uri) 1287 username2, password2 = create_user(client.tracking_uri) 1288 1289 with User(username1, password1, monkeypatch): 1290 experiment_id = client.create_experiment("graphql_run_test_exp") 1291 run = client.create_run(experiment_id) 1292 run_id = run.info.run_id 1293 client.set_terminated(run_id) 1294 1295 _send_rest_tracking_post_request( 1296 client.tracking_uri, 1297 "/api/2.0/mlflow/experiments/permissions/create", 1298 json_payload={ 1299 "experiment_id": experiment_id, 1300 "username": username2, 1301 "permission": "NO_PERMISSIONS", 1302 }, 1303 auth=(username1, password1), 1304 ) 1305 1306 query = """ 1307 query($runId: String!) { 1308 mlflowGetRun(input: {runId: $runId}) { 1309 run { 1310 info { 1311 runId 1312 experimentId 1313 } 1314 } 1315 } 1316 } 1317 """ 1318 1319 # user1 (creator) should be able to read the run 1320 response = _graphql_query( 1321 client.tracking_uri, 1322 query, 1323 variables={"runId": run_id}, 1324 auth=(username1, password1), 1325 ) 1326 assert response.status_code == 200 1327 data = response.json() 1328 run_data = data["data"]["mlflowGetRun"]["run"] 1329 assert run_data["info"]["runId"] == run_id 1330 assert run_data["info"]["experimentId"] == experiment_id 1331 1332 # user2 (NO_PERMISSIONS) should NOT be able to read the run 1333 response = _graphql_query( 1334 client.tracking_uri, 1335 query, 1336 variables={"runId": run_id}, 1337 auth=(username2, password2), 1338 ) 1339 assert response.status_code == 200 1340 data = response.json() 1341 assert data.get("data", {}).get("mlflowGetRun") is None 1342 1343 1344 def test_graphql_search_runs_authorization(client, monkeypatch): 1345 username1, password1 = create_user(client.tracking_uri) 1346 username2, password2 = create_user(client.tracking_uri) 1347 1348 with User(username1, password1, monkeypatch): 1349 exp1_id = client.create_experiment("graphql_search_exp1") 1350 exp2_id = client.create_experiment("graphql_search_exp2") 1351 1352 run1 = client.create_run(exp1_id) 1353 client.set_terminated(run1.info.run_id) 1354 1355 run2 = client.create_run(exp2_id) 1356 client.set_terminated(run2.info.run_id) 1357 1358 # Grant READ on exp1 to user2, NO_PERMISSIONS on exp2 1359 _send_rest_tracking_post_request( 1360 client.tracking_uri, 1361 "/api/2.0/mlflow/experiments/permissions/create", 1362 json_payload={ 1363 "experiment_id": exp1_id, 1364 "username": username2, 1365 "permission": "READ", 1366 }, 1367 auth=(username1, password1), 1368 ) 1369 _send_rest_tracking_post_request( 1370 client.tracking_uri, 1371 "/api/2.0/mlflow/experiments/permissions/create", 1372 json_payload={ 1373 "experiment_id": exp2_id, 1374 "username": username2, 1375 "permission": "NO_PERMISSIONS", 1376 }, 1377 auth=(username1, password1), 1378 ) 1379 1380 query = """ 1381 query($expIds: [String]!) { 1382 mlflowSearchRuns(input: {experimentIds: $expIds}) { 1383 runs { 1384 info { 1385 runId 1386 experimentId 1387 } 1388 } 1389 } 1390 } 1391 """ 1392 1393 # user1 should see both runs 1394 response = _graphql_query( 1395 client.tracking_uri, 1396 query, 1397 variables={"expIds": [exp1_id, exp2_id]}, 1398 auth=(username1, password1), 1399 ) 1400 assert response.status_code == 200 1401 data = response.json() 1402 runs = data.get("data", {}).get("mlflowSearchRuns", {}).get("runs", []) 1403 assert len(runs) == 2 1404 1405 # user2 should only see run from exp1 (exp2 is filtered out) 1406 response = _graphql_query( 1407 client.tracking_uri, 1408 query, 1409 variables={"expIds": [exp1_id, exp2_id]}, 1410 auth=(username2, password2), 1411 ) 1412 assert response.status_code == 200 1413 data = response.json() 1414 runs = data.get("data", {}).get("mlflowSearchRuns", {}).get("runs", []) 1415 assert len(runs) == 1 1416 assert runs[0]["info"]["experimentId"] == exp1_id 1417 1418 1419 def test_graphql_list_artifacts_authorization(client, monkeypatch): 1420 username1, password1 = create_user(client.tracking_uri) 1421 username2, password2 = create_user(client.tracking_uri) 1422 1423 with User(username1, password1, monkeypatch): 1424 experiment_id = client.create_experiment("graphql_artifacts_test_exp") 1425 run = client.create_run(experiment_id) 1426 run_id = run.info.run_id 1427 client.set_terminated(run_id) 1428 1429 _send_rest_tracking_post_request( 1430 client.tracking_uri, 1431 "/api/2.0/mlflow/experiments/permissions/create", 1432 json_payload={ 1433 "experiment_id": experiment_id, 1434 "username": username2, 1435 "permission": "NO_PERMISSIONS", 1436 }, 1437 auth=(username1, password1), 1438 ) 1439 1440 query = """ 1441 query($runId: String!) { 1442 mlflowListArtifacts(input: {runId: $runId}) { 1443 rootUri 1444 files { 1445 path 1446 } 1447 } 1448 } 1449 """ 1450 1451 # user1 (creator) should be able to list artifacts 1452 response = _graphql_query( 1453 client.tracking_uri, 1454 query, 1455 variables={"runId": run_id}, 1456 auth=(username1, password1), 1457 ) 1458 assert response.status_code == 200 1459 data = response.json() 1460 assert data.get("data", {}).get("mlflowListArtifacts") is not None 1461 1462 # user2 (NO_PERMISSIONS) should NOT be able to list artifacts 1463 response = _graphql_query( 1464 client.tracking_uri, 1465 query, 1466 variables={"runId": run_id}, 1467 auth=(username2, password2), 1468 ) 1469 assert response.status_code == 200 1470 data = response.json() 1471 assert data.get("data", {}).get("mlflowListArtifacts") is None 1472 1473 1474 def test_graphql_nonexistent_experiment(client, monkeypatch): 1475 username, password = create_user(client.tracking_uri) 1476 1477 query = """ 1478 query($expId: String!) { 1479 mlflowGetExperiment(input: {experimentId: $expId}) { 1480 experiment { 1481 experimentId 1482 name 1483 } 1484 } 1485 } 1486 """ 1487 1488 response = _graphql_query( 1489 client.tracking_uri, 1490 query, 1491 variables={"expId": "999999999"}, 1492 auth=(username, password), 1493 ) 1494 assert response.status_code == 200 1495 data = response.json() 1496 assert data.get("data", {}).get("mlflowGetExperiment") is None 1497 1498 1499 def test_graphql_nonexistent_run(client, monkeypatch): 1500 username, password = create_user(client.tracking_uri) 1501 1502 query = """ 1503 query($runId: String!) { 1504 mlflowGetRun(input: {runId: $runId}) { 1505 run { 1506 info { 1507 runId 1508 experimentId 1509 } 1510 } 1511 } 1512 } 1513 """ 1514 1515 response = _graphql_query( 1516 client.tracking_uri, 1517 query, 1518 variables={"runId": "00000000000000000000000000000000"}, 1519 auth=(username, password), 1520 ) 1521 assert response.status_code == 200 1522 data = response.json() 1523 assert data.get("data", {}).get("mlflowGetRun") is None 1524 1525 1526 def test_get_metric_history_bulk_interval_auth(client: MlflowClient, monkeypatch): 1527 username1, password1 = create_user(client.tracking_uri) 1528 username2, password2 = create_user(client.tracking_uri) 1529 1530 with User(username1, password1, monkeypatch): 1531 experiment_id = client.create_experiment("test_metric_history_experiment") 1532 run = client.create_run(experiment_id) 1533 run_id = run.info.run_id 1534 client.log_metric(run_id, "test_metric", 1.0, step=0) 1535 client.log_metric(run_id, "test_metric", 2.0, step=1) 1536 1537 _send_rest_tracking_post_request( 1538 client.tracking_uri, 1539 "/api/2.0/mlflow/experiments/permissions/create", 1540 json_payload={ 1541 "experiment_id": experiment_id, 1542 "username": username2, 1543 "permission": "READ", 1544 }, 1545 auth=(username1, password1), 1546 ) 1547 1548 with User(username2, password2, monkeypatch): 1549 response = requests.get( 1550 url=client.tracking_uri + "/ajax-api/2.0/mlflow/metrics/get-history-bulk-interval", 1551 params={ 1552 "run_ids": run_id, 1553 "metric_key": "test_metric", 1554 "max_results": 100, 1555 }, 1556 auth=(username2, password2), 1557 ) 1558 response.raise_for_status() 1559 data = response.json() 1560 assert "metrics" in data 1561 assert len(data["metrics"]) == 2 1562 1563 1564 def test_gateway_secrets_permissions(client, monkeypatch): 1565 user1, password1 = create_user(client.tracking_uri) 1566 user2, password2 = create_user(client.tracking_uri) 1567 1568 with User(user1, password1, monkeypatch): 1569 response = requests.post( 1570 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/create", 1571 json={ 1572 "secret_name": "user1_secret", 1573 "secret_value": {"api_key": "test-key"}, 1574 "provider": "openai", 1575 }, 1576 auth=(user1, password1), 1577 ) 1578 response.raise_for_status() 1579 user1_secret_id = response.json()["secret"]["secret_id"] 1580 1581 with User(user1, password1, monkeypatch): 1582 response = requests.get( 1583 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/get", 1584 params={"secret_id": user1_secret_id}, 1585 auth=(user1, password1), 1586 ) 1587 response.raise_for_status() 1588 1589 with User(user1, password1, monkeypatch): 1590 response = requests.post( 1591 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/update", 1592 json={ 1593 "secret_id": user1_secret_id, 1594 "secret_value": {"api_key": "updated-key"}, 1595 }, 1596 auth=(user1, password1), 1597 ) 1598 response.raise_for_status() 1599 1600 # User2 can read secrets by default (READ permission is default) 1601 with User(user2, password2, monkeypatch): 1602 response = requests.get( 1603 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/get", 1604 params={"secret_id": user1_secret_id}, 1605 auth=(user2, password2), 1606 ) 1607 response.raise_for_status() 1608 1609 # User2 cannot update secrets without explicit permission 1610 with User(user2, password2, monkeypatch): 1611 response = requests.post( 1612 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/update", 1613 json={ 1614 "secret_id": user1_secret_id, 1615 "secret_value": {"api_key": "hacked-key"}, 1616 }, 1617 auth=(user2, password2), 1618 ) 1619 assert response.status_code == 403 1620 1621 # User2 cannot delete secrets without explicit permission 1622 with User(user2, password2, monkeypatch): 1623 response = requests.delete( 1624 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/delete", 1625 json={"secret_id": user1_secret_id}, 1626 auth=(user2, password2), 1627 ) 1628 assert response.status_code == 403 1629 1630 with User(user1, password1, monkeypatch): 1631 response = requests.get( 1632 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/list", 1633 auth=(user1, password1), 1634 ) 1635 response.raise_for_status() 1636 1637 with User(user1, password1, monkeypatch): 1638 response = requests.get( 1639 url=client.tracking_uri + "/ajax-api/3.0/mlflow/gateway/secrets/config", 1640 auth=(user1, password1), 1641 ) 1642 response.raise_for_status() 1643 assert "secrets_available" in response.json() 1644 1645 with User(user1, password1, monkeypatch): 1646 response = requests.delete( 1647 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/delete", 1648 json={"secret_id": user1_secret_id}, 1649 auth=(user1, password1), 1650 ) 1651 response.raise_for_status() 1652 1653 1654 def test_gateway_endpoints_permissions(client, monkeypatch): 1655 user1, password1 = create_user(client.tracking_uri) 1656 user2, password2 = create_user(client.tracking_uri) 1657 1658 with User(user1, password1, monkeypatch): 1659 response = requests.post( 1660 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/create", 1661 json={ 1662 "secret_name": "user1_secret_for_endpoint", 1663 "secret_value": {"api_key": "test-key"}, 1664 "provider": "openai", 1665 }, 1666 auth=(user1, password1), 1667 ) 1668 response.raise_for_status() 1669 secret_id = response.json()["secret"]["secret_id"] 1670 1671 with User(user1, password1, monkeypatch): 1672 response = requests.post( 1673 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/create", 1674 json={ 1675 "name": "user1_model_def", 1676 "secret_id": secret_id, 1677 "provider": "openai", 1678 "model_name": "gpt-4", 1679 }, 1680 auth=(user1, password1), 1681 ) 1682 response.raise_for_status() 1683 model_definition_id = response.json()["model_definition"]["model_definition_id"] 1684 1685 with User(user1, password1, monkeypatch): 1686 response = requests.post( 1687 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/create", 1688 json={ 1689 "name": "user1_endpoint", 1690 "model_configs": [ 1691 { 1692 "model_definition_id": model_definition_id, 1693 "linkage_type": "PRIMARY", 1694 } 1695 ], 1696 }, 1697 auth=(user1, password1), 1698 ) 1699 response.raise_for_status() 1700 endpoint_id = response.json()["endpoint"]["endpoint_id"] 1701 1702 with User(user1, password1, monkeypatch): 1703 response = requests.get( 1704 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/list", 1705 auth=(user1, password1), 1706 ) 1707 response.raise_for_status() 1708 1709 with User(user1, password1, monkeypatch): 1710 response = requests.get( 1711 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/get", 1712 params={"endpoint_id": endpoint_id}, 1713 auth=(user1, password1), 1714 ) 1715 response.raise_for_status() 1716 1717 with User(user1, password1, monkeypatch): 1718 response = requests.post( 1719 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/update", 1720 json={ 1721 "endpoint_id": endpoint_id, 1722 "name": "updated_endpoint", 1723 }, 1724 auth=(user1, password1), 1725 ) 1726 response.raise_for_status() 1727 1728 # User2 can read endpoints by default (READ permission is default) 1729 with User(user2, password2, monkeypatch): 1730 response = requests.get( 1731 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/get", 1732 params={"endpoint_id": endpoint_id}, 1733 auth=(user2, password2), 1734 ) 1735 response.raise_for_status() 1736 1737 # User2 cannot update endpoints without explicit permission 1738 with User(user2, password2, monkeypatch): 1739 response = requests.post( 1740 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/update", 1741 json={ 1742 "endpoint_id": endpoint_id, 1743 "name": "hacked_endpoint", 1744 }, 1745 auth=(user2, password2), 1746 ) 1747 assert response.status_code == 403 1748 1749 # User2 cannot delete endpoints without explicit permission 1750 with User(user2, password2, monkeypatch): 1751 response = requests.delete( 1752 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/delete", 1753 json={"endpoint_id": endpoint_id}, 1754 auth=(user2, password2), 1755 ) 1756 assert response.status_code == 403 1757 1758 with User(user1, password1, monkeypatch): 1759 response = requests.delete( 1760 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/delete", 1761 json={"endpoint_id": endpoint_id}, 1762 auth=(user1, password1), 1763 ) 1764 response.raise_for_status() 1765 1766 with User(user1, password1, monkeypatch): 1767 response = requests.delete( 1768 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/delete", 1769 json={"model_definition_id": model_definition_id}, 1770 auth=(user1, password1), 1771 ) 1772 response.raise_for_status() 1773 1774 with User(user1, password1, monkeypatch): 1775 response = requests.delete( 1776 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/delete", 1777 json={"secret_id": secret_id}, 1778 auth=(user1, password1), 1779 ) 1780 response.raise_for_status() 1781 1782 1783 def test_gateway_model_definitions_permissions(client, monkeypatch): 1784 user1, password1 = create_user(client.tracking_uri) 1785 user2, password2 = create_user(client.tracking_uri) 1786 1787 with User(user1, password1, monkeypatch): 1788 response = requests.post( 1789 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/create", 1790 json={ 1791 "secret_name": "user1_secret_for_model_def", 1792 "secret_value": {"api_key": "test-key"}, 1793 "provider": "openai", 1794 }, 1795 auth=(user1, password1), 1796 ) 1797 response.raise_for_status() 1798 secret_id = response.json()["secret"]["secret_id"] 1799 1800 with User(user1, password1, monkeypatch): 1801 response = requests.post( 1802 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/create", 1803 json={ 1804 "name": "user1_model_def", 1805 "secret_id": secret_id, 1806 "provider": "openai", 1807 "model_name": "gpt-4", 1808 }, 1809 auth=(user1, password1), 1810 ) 1811 response.raise_for_status() 1812 model_definition_id = response.json()["model_definition"]["model_definition_id"] 1813 1814 with User(user1, password1, monkeypatch): 1815 response = requests.get( 1816 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/list", 1817 auth=(user1, password1), 1818 ) 1819 response.raise_for_status() 1820 1821 with User(user1, password1, monkeypatch): 1822 response = requests.get( 1823 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/get", 1824 params={"model_definition_id": model_definition_id}, 1825 auth=(user1, password1), 1826 ) 1827 response.raise_for_status() 1828 1829 with User(user1, password1, monkeypatch): 1830 response = requests.post( 1831 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/update", 1832 json={ 1833 "model_definition_id": model_definition_id, 1834 "name": "updated_model_def", 1835 }, 1836 auth=(user1, password1), 1837 ) 1838 response.raise_for_status() 1839 1840 # User2 can read model definitions by default (READ permission is default) 1841 with User(user2, password2, monkeypatch): 1842 response = requests.get( 1843 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/get", 1844 params={"model_definition_id": model_definition_id}, 1845 auth=(user2, password2), 1846 ) 1847 response.raise_for_status() 1848 1849 # User2 cannot update model definitions without explicit permission 1850 with User(user2, password2, monkeypatch): 1851 response = requests.post( 1852 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/update", 1853 json={ 1854 "model_definition_id": model_definition_id, 1855 "name": "hacked_model_def", 1856 }, 1857 auth=(user2, password2), 1858 ) 1859 assert response.status_code == 403 1860 1861 # User2 cannot delete model definitions without explicit permission 1862 with User(user2, password2, monkeypatch): 1863 response = requests.delete( 1864 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/delete", 1865 json={"model_definition_id": model_definition_id}, 1866 auth=(user2, password2), 1867 ) 1868 assert response.status_code == 403 1869 1870 with User(user1, password1, monkeypatch): 1871 response = requests.delete( 1872 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/delete", 1873 json={"model_definition_id": model_definition_id}, 1874 auth=(user1, password1), 1875 ) 1876 response.raise_for_status() 1877 1878 with User(user1, password1, monkeypatch): 1879 response = requests.delete( 1880 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/delete", 1881 json={"secret_id": secret_id}, 1882 auth=(user1, password1), 1883 ) 1884 response.raise_for_status() 1885 1886 1887 def test_gateway_budget_policy_admin_only(client, monkeypatch): 1888 user1, password1 = create_user(client.tracking_uri) 1889 1890 # Admin creates a budget policy 1891 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 1892 response = requests.post( 1893 url=client.tracking_uri + "/api/3.0/mlflow/gateway/budgets/create", 1894 json={ 1895 "budget_unit": "USD", 1896 "budget_amount": 100.0, 1897 "duration": {"unit": "DAYS", "value": 30}, 1898 "target_scope": "GLOBAL", 1899 "budget_action": "ALERT", 1900 }, 1901 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 1902 ) 1903 response.raise_for_status() 1904 budget_policy_id = response.json()["budget_policy"]["budget_policy_id"] 1905 1906 # Non-admin can list budget policies 1907 with User(user1, password1, monkeypatch): 1908 response = requests.get( 1909 url=client.tracking_uri + "/api/3.0/mlflow/gateway/budgets/list", 1910 auth=(user1, password1), 1911 ) 1912 response.raise_for_status() 1913 1914 # Non-admin can get a budget policy 1915 with User(user1, password1, monkeypatch): 1916 response = requests.get( 1917 url=client.tracking_uri + "/api/3.0/mlflow/gateway/budgets/get", 1918 params={"budget_policy_id": budget_policy_id}, 1919 auth=(user1, password1), 1920 ) 1921 response.raise_for_status() 1922 1923 # Non-admin cannot create a budget policy 1924 with User(user1, password1, monkeypatch): 1925 response = requests.post( 1926 url=client.tracking_uri + "/api/3.0/mlflow/gateway/budgets/create", 1927 json={ 1928 "budget_unit": "USD", 1929 "budget_amount": 50.0, 1930 "duration": {"unit": "DAYS", "value": 7}, 1931 "target_scope": "GLOBAL", 1932 "budget_action": "REJECT", 1933 }, 1934 auth=(user1, password1), 1935 ) 1936 assert response.status_code == 403 1937 1938 # Non-admin cannot update a budget policy 1939 with User(user1, password1, monkeypatch): 1940 response = requests.post( 1941 url=client.tracking_uri + "/api/3.0/mlflow/gateway/budgets/update", 1942 json={ 1943 "budget_policy_id": budget_policy_id, 1944 "budget_amount": 200.0, 1945 }, 1946 auth=(user1, password1), 1947 ) 1948 assert response.status_code == 403 1949 1950 # Non-admin cannot delete a budget policy 1951 with User(user1, password1, monkeypatch): 1952 response = requests.delete( 1953 url=client.tracking_uri + "/api/3.0/mlflow/gateway/budgets/delete", 1954 json={"budget_policy_id": budget_policy_id}, 1955 auth=(user1, password1), 1956 ) 1957 assert response.status_code == 403 1958 1959 # Admin can delete the budget policy 1960 with User(ADMIN_USERNAME, ADMIN_PASSWORD, monkeypatch): 1961 response = requests.delete( 1962 url=client.tracking_uri + "/api/3.0/mlflow/gateway/budgets/delete", 1963 json={"budget_policy_id": budget_policy_id}, 1964 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 1965 ) 1966 response.raise_for_status() 1967 1968 1969 def test_gateway_ajax_routes_permissions(client, monkeypatch): 1970 username, password = create_user(client.tracking_uri) 1971 1972 with User(username, password, monkeypatch): 1973 response = requests.get( 1974 url=client.tracking_uri + "/ajax-api/3.0/mlflow/gateway/supported-providers", 1975 auth=(username, password), 1976 ) 1977 response.raise_for_status() 1978 assert "providers" in response.json() 1979 1980 with User(username, password, monkeypatch): 1981 response = requests.get( 1982 url=client.tracking_uri + "/ajax-api/3.0/mlflow/gateway/supported-models", 1983 auth=(username, password), 1984 ) 1985 response.raise_for_status() 1986 1987 with User(username, password, monkeypatch): 1988 response = requests.get( 1989 url=client.tracking_uri + "/ajax-api/3.0/mlflow/gateway/provider-config", 1990 params={"provider": "openai"}, 1991 auth=(username, password), 1992 ) 1993 response.raise_for_status() 1994 1995 with User(username, password, monkeypatch): 1996 response = requests.get( 1997 url=client.tracking_uri + "/ajax-api/3.0/mlflow/gateway/secrets/config", 1998 auth=(username, password), 1999 ) 2000 response.raise_for_status() 2001 assert "secrets_available" in response.json() 2002 2003 2004 def test_gateway_unauthenticated_access_denied(client, monkeypatch): 2005 monkeypatch.delenv(MLFLOW_TRACKING_USERNAME.name, raising=False) 2006 monkeypatch.delenv(MLFLOW_TRACKING_PASSWORD.name, raising=False) 2007 2008 response = requests.get( 2009 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/list", 2010 ) 2011 assert response.status_code == 401 2012 2013 response = requests.get( 2014 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/list", 2015 ) 2016 assert response.status_code == 401 2017 2018 response = requests.get( 2019 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/list", 2020 ) 2021 assert response.status_code == 401 2022 2023 response = requests.get( 2024 url=client.tracking_uri + "/ajax-api/3.0/mlflow/gateway/supported-providers", 2025 ) 2026 assert response.status_code == 401 2027 2028 2029 def test_gateway_endpoint_use_permission(fastapi_client, monkeypatch): 2030 user1, password1 = create_user(fastapi_client.tracking_uri) 2031 user2, password2 = create_user(fastapi_client.tracking_uri) 2032 2033 # User1 creates a secret, model definition, and endpoint 2034 with User(user1, password1, monkeypatch): 2035 response = requests.post( 2036 url=fastapi_client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/create", 2037 json={ 2038 "secret_name": "test_secret", 2039 "secret_value": {"api_key": "test-key"}, 2040 "provider": "openai", 2041 }, 2042 auth=(user1, password1), 2043 ) 2044 response.raise_for_status() 2045 secret_id = response.json()["secret"]["secret_id"] 2046 2047 with User(user1, password1, monkeypatch): 2048 response = requests.post( 2049 url=fastapi_client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/create", 2050 json={ 2051 "name": "test_model_def", 2052 "secret_id": secret_id, 2053 "provider": "openai", 2054 "model_name": "gpt-4", 2055 }, 2056 auth=(user1, password1), 2057 ) 2058 response.raise_for_status() 2059 model_definition_id = response.json()["model_definition"]["model_definition_id"] 2060 2061 with User(user1, password1, monkeypatch): 2062 response = requests.post( 2063 url=fastapi_client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/create", 2064 json={ 2065 "name": "test_endpoint", 2066 "model_configs": [ 2067 { 2068 "model_definition_id": model_definition_id, 2069 "linkage_type": "PRIMARY", 2070 } 2071 ], 2072 }, 2073 auth=(user1, password1), 2074 ) 2075 response.raise_for_status() 2076 endpoint_id = response.json()["endpoint"]["endpoint_id"] 2077 endpoint_name = response.json()["endpoint"]["name"] 2078 2079 # User2 without permission cannot invoke the endpoint 2080 with User(user2, password2, monkeypatch): 2081 response = requests.post( 2082 url=fastapi_client.tracking_uri + f"/gateway/{endpoint_name}/mlflow/invocations", 2083 json={"messages": [{"role": "user", "content": "test"}]}, 2084 auth=(user2, password2), 2085 ) 2086 assert response.status_code == 403 2087 2088 # Grant USE permission to user2 2089 with User(user1, password1, monkeypatch): 2090 _send_rest_tracking_post_request( 2091 fastapi_client.tracking_uri, 2092 "/api/3.0/mlflow/gateway/endpoints/permissions/create", 2093 json_payload={ 2094 "endpoint_id": endpoint_id, 2095 "username": user2, 2096 "permission": "USE", 2097 }, 2098 auth=(user1, password1), 2099 ) 2100 2101 # User2 with USE permission can invoke 2102 with User(user2, password2, monkeypatch): 2103 response = requests.post( 2104 url=fastapi_client.tracking_uri + f"/gateway/{endpoint_name}/mlflow/invocations", 2105 json={"messages": [{"role": "user", "content": "test"}]}, 2106 auth=(user2, password2), 2107 ) 2108 # Will fail because we don't have real LLM credentials, but should pass auth (not 403) 2109 assert response.status_code != 403 2110 2111 # Cleanup 2112 with User(user1, password1, monkeypatch): 2113 requests.delete( 2114 url=fastapi_client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/delete", 2115 json={"endpoint_id": endpoint_id}, 2116 auth=(user1, password1), 2117 ).raise_for_status() 2118 requests.delete( 2119 url=fastapi_client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/delete", 2120 json={"model_definition_id": model_definition_id}, 2121 auth=(user1, password1), 2122 ).raise_for_status() 2123 requests.delete( 2124 url=fastapi_client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/delete", 2125 json={"secret_id": secret_id}, 2126 auth=(user1, password1), 2127 ).raise_for_status() 2128 2129 2130 def test_gateway_model_definition_requires_secret_use_permission(client, monkeypatch): 2131 user1, password1 = create_user(client.tracking_uri) 2132 user2, password2 = create_user(client.tracking_uri) 2133 2134 # User1 creates a secret 2135 with User(user1, password1, monkeypatch): 2136 response = requests.post( 2137 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/create", 2138 json={ 2139 "secret_name": "user1_secret", 2140 "secret_value": {"api_key": "test-key"}, 2141 "provider": "openai", 2142 }, 2143 auth=(user1, password1), 2144 ) 2145 response.raise_for_status() 2146 secret_id = response.json()["secret"]["secret_id"] 2147 2148 # User2 cannot create a model definition using user1's secret (no permission) 2149 with User(user2, password2, monkeypatch): 2150 response = requests.post( 2151 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/create", 2152 json={ 2153 "name": "model_def_1", 2154 "secret_id": secret_id, 2155 "provider": "openai", 2156 "model_name": "gpt-4", 2157 }, 2158 auth=(user2, password2), 2159 ) 2160 assert response.status_code == 403 2161 2162 # Grant USE permission to user2 2163 with User(user1, password1, monkeypatch): 2164 _send_rest_tracking_post_request( 2165 client.tracking_uri, 2166 "/api/3.0/mlflow/gateway/secrets/permissions/create", 2167 json_payload={"secret_id": secret_id, "username": user2, "permission": "USE"}, 2168 auth=(user1, password1), 2169 ) 2170 2171 # User2 can now create a model definition using user1's secret 2172 with User(user2, password2, monkeypatch): 2173 response = requests.post( 2174 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/create", 2175 json={ 2176 "name": "model_def_1", 2177 "secret_id": secret_id, 2178 "provider": "openai", 2179 "model_name": "gpt-4", 2180 }, 2181 auth=(user2, password2), 2182 ) 2183 response.raise_for_status() 2184 model_def_id = response.json()["model_definition"]["model_definition_id"] 2185 2186 # User1 creates another secret 2187 with User(user1, password1, monkeypatch): 2188 response = requests.post( 2189 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/create", 2190 json={ 2191 "secret_name": "user1_secret_2", 2192 "secret_value": {"api_key": "test-key-2"}, 2193 "provider": "anthropic", 2194 }, 2195 auth=(user1, password1), 2196 ) 2197 response.raise_for_status() 2198 secret_id_2 = response.json()["secret"]["secret_id"] 2199 2200 # User2 cannot update the model definition to use secret_id_2 (no permission on that secret) 2201 with User(user2, password2, monkeypatch): 2202 response = requests.post( 2203 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/update", 2204 json={ 2205 "model_definition_id": model_def_id, 2206 "secret_id": secret_id_2, 2207 "provider": "anthropic", 2208 }, 2209 auth=(user2, password2), 2210 ) 2211 assert response.status_code == 403 2212 2213 # Grant USE permission to user2 on secret_id_2 2214 with User(user1, password1, monkeypatch): 2215 _send_rest_tracking_post_request( 2216 client.tracking_uri, 2217 "/api/3.0/mlflow/gateway/secrets/permissions/create", 2218 json_payload={"secret_id": secret_id_2, "username": user2, "permission": "USE"}, 2219 auth=(user1, password1), 2220 ) 2221 2222 # User2 can now update the model definition to use secret_id_2 2223 with User(user2, password2, monkeypatch): 2224 response = requests.post( 2225 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/update", 2226 json={ 2227 "model_definition_id": model_def_id, 2228 "secret_id": secret_id_2, 2229 "provider": "anthropic", 2230 }, 2231 auth=(user2, password2), 2232 ) 2233 response.raise_for_status() 2234 2235 # Cleanup 2236 with User(user2, password2, monkeypatch): 2237 requests.delete( 2238 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/delete", 2239 json={"model_definition_id": model_def_id}, 2240 auth=(user2, password2), 2241 ).raise_for_status() 2242 2243 with User(user1, password1, monkeypatch): 2244 requests.delete( 2245 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/delete", 2246 json={"secret_id": secret_id}, 2247 auth=(user1, password1), 2248 ).raise_for_status() 2249 requests.delete( 2250 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/delete", 2251 json={"secret_id": secret_id_2}, 2252 auth=(user1, password1), 2253 ).raise_for_status() 2254 2255 2256 def test_gateway_endpoint_requires_model_definition_use_permission(client, monkeypatch): 2257 user1, password1 = create_user(client.tracking_uri) 2258 user2, password2 = create_user(client.tracking_uri) 2259 2260 # User1 creates a secret and model definition 2261 with User(user1, password1, monkeypatch): 2262 response = requests.post( 2263 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/create", 2264 json={ 2265 "secret_name": "user1_secret", 2266 "secret_value": {"api_key": "test-key"}, 2267 "provider": "openai", 2268 }, 2269 auth=(user1, password1), 2270 ) 2271 response.raise_for_status() 2272 secret_id = response.json()["secret"]["secret_id"] 2273 2274 response = requests.post( 2275 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/create", 2276 json={ 2277 "name": "model_def_1", 2278 "secret_id": secret_id, 2279 "provider": "openai", 2280 "model_name": "gpt-4", 2281 }, 2282 auth=(user1, password1), 2283 ) 2284 response.raise_for_status() 2285 model_def_id = response.json()["model_definition"]["model_definition_id"] 2286 2287 # User2 cannot create an endpoint using user1's model definition (no permission) 2288 with User(user2, password2, monkeypatch): 2289 response = requests.post( 2290 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/create", 2291 json={ 2292 "name": "endpoint_1", 2293 "model_configs": [ 2294 { 2295 "model_definition_id": model_def_id, 2296 "linkage_type": "PRIMARY", 2297 } 2298 ], 2299 }, 2300 auth=(user2, password2), 2301 ) 2302 assert response.status_code == 403 2303 2304 # Grant USE permission to user2 on the model definition 2305 with User(user1, password1, monkeypatch): 2306 _send_rest_tracking_post_request( 2307 client.tracking_uri, 2308 "/api/3.0/mlflow/gateway/model-definitions/permissions/create", 2309 json_payload={ 2310 "model_definition_id": model_def_id, 2311 "username": user2, 2312 "permission": "USE", 2313 }, 2314 auth=(user1, password1), 2315 ) 2316 2317 # User2 can now create an endpoint using user1's model definition 2318 with User(user2, password2, monkeypatch): 2319 response = requests.post( 2320 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/create", 2321 json={ 2322 "name": "endpoint_1", 2323 "model_configs": [ 2324 { 2325 "model_definition_id": model_def_id, 2326 "linkage_type": "PRIMARY", 2327 } 2328 ], 2329 }, 2330 auth=(user2, password2), 2331 ) 2332 response.raise_for_status() 2333 endpoint_id = response.json()["endpoint"]["endpoint_id"] 2334 2335 # User1 creates another model definition 2336 with User(user1, password1, monkeypatch): 2337 response = requests.post( 2338 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/create", 2339 json={ 2340 "name": "model_def_2", 2341 "secret_id": secret_id, 2342 "provider": "openai", 2343 "model_name": "gpt-3.5-turbo", 2344 }, 2345 auth=(user1, password1), 2346 ) 2347 response.raise_for_status() 2348 model_def_id_2 = response.json()["model_definition"]["model_definition_id"] 2349 2350 # User2 cannot update the endpoint to use model_def_id_2 (no permission on that model def) 2351 with User(user2, password2, monkeypatch): 2352 response = requests.post( 2353 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/update", 2354 json={ 2355 "endpoint_id": endpoint_id, 2356 "model_configs": [ 2357 { 2358 "model_definition_id": model_def_id_2, 2359 "linkage_type": "PRIMARY", 2360 } 2361 ], 2362 }, 2363 auth=(user2, password2), 2364 ) 2365 assert response.status_code == 403 2366 2367 # Grant USE permission to user2 on model_def_id_2 2368 with User(user1, password1, monkeypatch): 2369 _send_rest_tracking_post_request( 2370 client.tracking_uri, 2371 "/api/3.0/mlflow/gateway/model-definitions/permissions/create", 2372 json_payload={ 2373 "model_definition_id": model_def_id_2, 2374 "username": user2, 2375 "permission": "USE", 2376 }, 2377 auth=(user1, password1), 2378 ) 2379 2380 # User2 can now update the endpoint to use model_def_id_2 2381 with User(user2, password2, monkeypatch): 2382 response = requests.post( 2383 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/update", 2384 json={ 2385 "endpoint_id": endpoint_id, 2386 "model_configs": [ 2387 { 2388 "model_definition_id": model_def_id_2, 2389 "linkage_type": "PRIMARY", 2390 } 2391 ], 2392 }, 2393 auth=(user2, password2), 2394 ) 2395 response.raise_for_status() 2396 2397 # Cleanup 2398 with User(user2, password2, monkeypatch): 2399 requests.delete( 2400 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/delete", 2401 json={"endpoint_id": endpoint_id}, 2402 auth=(user2, password2), 2403 ).raise_for_status() 2404 2405 with User(user1, password1, monkeypatch): 2406 requests.delete( 2407 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/delete", 2408 json={"model_definition_id": model_def_id}, 2409 auth=(user1, password1), 2410 ).raise_for_status() 2411 requests.delete( 2412 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/delete", 2413 json={"model_definition_id": model_def_id_2}, 2414 auth=(user1, password1), 2415 ).raise_for_status() 2416 requests.delete( 2417 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/delete", 2418 json={"secret_id": secret_id}, 2419 auth=(user1, password1), 2420 ).raise_for_status() 2421 2422 2423 def test_gateway_endpoint_requires_fallback_model_definition_use_permission(client, monkeypatch): 2424 user1, password1 = create_user(client.tracking_uri) 2425 user2, password2 = create_user(client.tracking_uri) 2426 2427 # User1 creates secrets and model definitions 2428 with User(user1, password1, monkeypatch): 2429 response = requests.post( 2430 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/create", 2431 json={ 2432 "secret_name": "user1_secret", 2433 "secret_value": {"api_key": "test-key"}, 2434 "provider": "openai", 2435 }, 2436 auth=(user1, password1), 2437 ) 2438 response.raise_for_status() 2439 secret_id = response.json()["secret"]["secret_id"] 2440 2441 # Create primary model definition 2442 response = requests.post( 2443 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/create", 2444 json={ 2445 "name": "primary_model", 2446 "secret_id": secret_id, 2447 "provider": "openai", 2448 "model_name": "gpt-4", 2449 }, 2450 auth=(user1, password1), 2451 ) 2452 response.raise_for_status() 2453 primary_model_def_id = response.json()["model_definition"]["model_definition_id"] 2454 2455 # Create fallback model definition 2456 response = requests.post( 2457 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/create", 2458 json={ 2459 "name": "fallback_model", 2460 "secret_id": secret_id, 2461 "provider": "openai", 2462 "model_name": "gpt-3.5-turbo", 2463 }, 2464 auth=(user1, password1), 2465 ) 2466 response.raise_for_status() 2467 fallback_model_def_id = response.json()["model_definition"]["model_definition_id"] 2468 2469 # Grant USE permission to user2 on primary model but not fallback 2470 with User(user1, password1, monkeypatch): 2471 _send_rest_tracking_post_request( 2472 client.tracking_uri, 2473 "/api/3.0/mlflow/gateway/model-definitions/permissions/create", 2474 json_payload={ 2475 "model_definition_id": primary_model_def_id, 2476 "username": user2, 2477 "permission": "USE", 2478 }, 2479 auth=(user1, password1), 2480 ) 2481 2482 # User2 cannot create an endpoint with fallback model (no permission on fallback) 2483 with User(user2, password2, monkeypatch): 2484 response = requests.post( 2485 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/create", 2486 json={ 2487 "name": "endpoint_with_fallback", 2488 "model_configs": [ 2489 { 2490 "model_definition_id": primary_model_def_id, 2491 "linkage_type": "PRIMARY", 2492 }, 2493 { 2494 "model_definition_id": fallback_model_def_id, 2495 "linkage_type": "FALLBACK", 2496 "fallback_order": 1, 2497 }, 2498 ], 2499 }, 2500 auth=(user2, password2), 2501 ) 2502 assert response.status_code == 403 2503 2504 # Grant USE permission to user2 on fallback model 2505 with User(user1, password1, monkeypatch): 2506 _send_rest_tracking_post_request( 2507 client.tracking_uri, 2508 "/api/3.0/mlflow/gateway/model-definitions/permissions/create", 2509 json_payload={ 2510 "model_definition_id": fallback_model_def_id, 2511 "username": user2, 2512 "permission": "USE", 2513 }, 2514 auth=(user1, password1), 2515 ) 2516 2517 # User2 can now create an endpoint with both primary and fallback models 2518 with User(user2, password2, monkeypatch): 2519 response = requests.post( 2520 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/create", 2521 json={ 2522 "name": "endpoint_with_fallback", 2523 "model_configs": [ 2524 { 2525 "model_definition_id": primary_model_def_id, 2526 "linkage_type": "PRIMARY", 2527 }, 2528 { 2529 "model_definition_id": fallback_model_def_id, 2530 "linkage_type": "FALLBACK", 2531 "fallback_order": 1, 2532 }, 2533 ], 2534 }, 2535 auth=(user2, password2), 2536 ) 2537 response.raise_for_status() 2538 endpoint_id = response.json()["endpoint"]["endpoint_id"] 2539 2540 # Cleanup 2541 with User(user2, password2, monkeypatch): 2542 requests.delete( 2543 url=client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/delete", 2544 json={"endpoint_id": endpoint_id}, 2545 auth=(user2, password2), 2546 ).raise_for_status() 2547 2548 with User(user1, password1, monkeypatch): 2549 requests.delete( 2550 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/delete", 2551 json={"model_definition_id": primary_model_def_id}, 2552 auth=(user1, password1), 2553 ).raise_for_status() 2554 requests.delete( 2555 url=client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/delete", 2556 json={"model_definition_id": fallback_model_def_id}, 2557 auth=(user1, password1), 2558 ).raise_for_status() 2559 requests.delete( 2560 url=client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/delete", 2561 json={"secret_id": secret_id}, 2562 auth=(user1, password1), 2563 ).raise_for_status() 2564 2565 2566 def test_prompt_optimization_job_search_permissions(client, monkeypatch): 2567 user1, password1 = create_user(client.tracking_uri) 2568 user2, password2 = create_user(client.tracking_uri) 2569 2570 # user1 creates an experiment 2571 with User(user1, password1, monkeypatch): 2572 experiment_id = client.create_experiment("prompt_optimization_search_test") 2573 2574 # Grant NO_PERMISSIONS to user2 on the experiment 2575 _send_rest_tracking_post_request( 2576 client.tracking_uri, 2577 "/api/2.0/mlflow/experiments/permissions/create", 2578 json_payload={ 2579 "experiment_id": experiment_id, 2580 "username": user2, 2581 "permission": "NO_PERMISSIONS", 2582 }, 2583 auth=(user1, password1), 2584 ) 2585 2586 # user1 can search jobs in the experiment 2587 response = requests.post( 2588 url=client.tracking_uri + "/api/3.0/mlflow/prompt-optimization/jobs/search", 2589 json={"experiment_id": experiment_id}, 2590 auth=(user1, password1), 2591 ) 2592 assert response.status_code != 403 2593 2594 # user2 cannot search jobs in the experiment (NO_PERMISSIONS) 2595 response = requests.post( 2596 url=client.tracking_uri + "/api/3.0/mlflow/prompt-optimization/jobs/search", 2597 json={"experiment_id": experiment_id}, 2598 auth=(user2, password2), 2599 ) 2600 assert response.status_code == 403 2601 2602 # Grant READ permission to user2 2603 response = requests.patch( 2604 url=client.tracking_uri + "/api/2.0/mlflow/experiments/permissions/update", 2605 json={ 2606 "experiment_id": experiment_id, 2607 "username": user2, 2608 "permission": "READ", 2609 }, 2610 auth=(user1, password1), 2611 ) 2612 assert response.status_code == 200 2613 2614 # user2 can now search jobs (READ grants can_read) 2615 response = requests.post( 2616 url=client.tracking_uri + "/api/3.0/mlflow/prompt-optimization/jobs/search", 2617 json={"experiment_id": experiment_id}, 2618 auth=(user2, password2), 2619 ) 2620 assert response.status_code != 403 2621 2622 2623 def test_prompt_optimization_job_create_permissions(client, monkeypatch): 2624 user1, password1 = create_user(client.tracking_uri) 2625 user2, password2 = create_user(client.tracking_uri) 2626 2627 # user1 creates an experiment 2628 with User(user1, password1, monkeypatch): 2629 experiment_id = client.create_experiment("prompt_optimization_create_test") 2630 2631 # Grant READ permission to user2 (not enough for create) 2632 _send_rest_tracking_post_request( 2633 client.tracking_uri, 2634 "/api/2.0/mlflow/experiments/permissions/create", 2635 json_payload={ 2636 "experiment_id": experiment_id, 2637 "username": user2, 2638 "permission": "READ", 2639 }, 2640 auth=(user1, password1), 2641 ) 2642 2643 # user2 cannot create jobs (READ doesn't grant update) 2644 response = requests.post( 2645 url=client.tracking_uri + "/api/3.0/mlflow/prompt-optimization/jobs", 2646 json={ 2647 "experiment_id": experiment_id, 2648 "source_prompt_uri": "prompts:/test/1", 2649 "config": { 2650 "optimizer_type": 1, # GEPA 2651 "dataset_id": "test-dataset", 2652 "scorers": ["Correctness"], 2653 }, 2654 }, 2655 auth=(user2, password2), 2656 ) 2657 assert response.status_code == 403 2658 2659 # Grant EDIT permission to user2 2660 response = requests.patch( 2661 url=client.tracking_uri + "/api/2.0/mlflow/experiments/permissions/update", 2662 json={ 2663 "experiment_id": experiment_id, 2664 "username": user2, 2665 "permission": "EDIT", 2666 }, 2667 auth=(user1, password1), 2668 ) 2669 assert response.status_code == 200 2670 2671 # user2 can now create jobs (EDIT grants can_update) 2672 # The request will fail for other reasons (missing prompt, dataset, etc.) 2673 # but should pass the permission check 2674 response = requests.post( 2675 url=client.tracking_uri + "/api/3.0/mlflow/prompt-optimization/jobs", 2676 json={ 2677 "experiment_id": experiment_id, 2678 "source_prompt_uri": "prompts:/test/1", 2679 "config": { 2680 "optimizer_type": 1, # GEPA 2681 "dataset_id": "test-dataset", 2682 "scorers": ["Correctness"], 2683 }, 2684 }, 2685 auth=(user2, password2), 2686 ) 2687 # Should not be 403 (permission denied) 2688 assert response.status_code != 403 2689 2690 2691 def test_gateway_endpoint_invocation_requires_use_permission(fastapi_client, monkeypatch): 2692 user1, password1 = create_user(fastapi_client.tracking_uri) 2693 user2, password2 = create_user(fastapi_client.tracking_uri) 2694 2695 # User1 creates a secret, model definition, and endpoint 2696 with User(user1, password1, monkeypatch): 2697 response = requests.post( 2698 url=fastapi_client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/create", 2699 json={ 2700 "secret_name": "user1_secret", 2701 "secret_value": {"api_key": "test-key"}, 2702 "provider": "openai", 2703 }, 2704 auth=(user1, password1), 2705 ) 2706 response.raise_for_status() 2707 secret_id = response.json()["secret"]["secret_id"] 2708 2709 response = requests.post( 2710 url=fastapi_client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/create", 2711 json={ 2712 "name": "test_model_def", 2713 "secret_id": secret_id, 2714 "provider": "openai", 2715 "model_name": "gpt-4", 2716 }, 2717 auth=(user1, password1), 2718 ) 2719 response.raise_for_status() 2720 model_def_id = response.json()["model_definition"]["model_definition_id"] 2721 2722 response = requests.post( 2723 url=fastapi_client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/create", 2724 json={ 2725 "name": "test_endpoint", 2726 "model_configs": [ 2727 { 2728 "model_definition_id": model_def_id, 2729 "linkage_type": "PRIMARY", 2730 } 2731 ], 2732 }, 2733 auth=(user1, password1), 2734 ) 2735 response.raise_for_status() 2736 endpoint_id = response.json()["endpoint"]["endpoint_id"] 2737 2738 # User2 cannot invoke the endpoint (no permission) 2739 with User(user2, password2, monkeypatch): 2740 response = requests.post( 2741 url=fastapi_client.tracking_uri + "/gateway/test_endpoint/mlflow/invocations", 2742 json={"messages": [{"role": "user", "content": "Hello"}]}, 2743 auth=(user2, password2), 2744 ) 2745 assert response.status_code == 403 2746 2747 # Grant READ permission to user2 (not enough for invocation) 2748 with User(user1, password1, monkeypatch): 2749 _send_rest_tracking_post_request( 2750 fastapi_client.tracking_uri, 2751 "/api/3.0/mlflow/gateway/endpoints/permissions/create", 2752 json_payload={ 2753 "endpoint_id": endpoint_id, 2754 "username": user2, 2755 "permission": "READ", 2756 }, 2757 auth=(user1, password1), 2758 ) 2759 2760 # User2 still cannot invoke (READ is not sufficient) 2761 with User(user2, password2, monkeypatch): 2762 response = requests.post( 2763 url=fastapi_client.tracking_uri + "/gateway/test_endpoint/mlflow/invocations", 2764 json={"messages": [{"role": "user", "content": "Hello"}]}, 2765 auth=(user2, password2), 2766 ) 2767 assert response.status_code == 403 2768 2769 # Upgrade to USE permission 2770 with User(user1, password1, monkeypatch): 2771 response = requests.patch( 2772 url=fastapi_client.tracking_uri 2773 + "/api/3.0/mlflow/gateway/endpoints/permissions/update", 2774 json={ 2775 "endpoint_id": endpoint_id, 2776 "username": user2, 2777 "permission": "USE", 2778 }, 2779 auth=(user1, password1), 2780 ) 2781 response.raise_for_status() 2782 2783 # User2 can now invoke the endpoint (though it will fail due to invalid API key) 2784 # We just check that we get past the permission check (403) to a different error 2785 with User(user2, password2, monkeypatch): 2786 response = requests.post( 2787 url=fastapi_client.tracking_uri + "/gateway/test_endpoint/mlflow/invocations", 2788 json={"messages": [{"role": "user", "content": "Hello"}]}, 2789 auth=(user2, password2), 2790 ) 2791 # Should not be 403 anymore (permission granted) 2792 # Will likely be 400 or 500 due to invalid API key, but that's fine 2793 assert response.status_code != 403 2794 2795 # Cleanup 2796 with User(user1, password1, monkeypatch): 2797 requests.delete( 2798 url=fastapi_client.tracking_uri + "/api/3.0/mlflow/gateway/endpoints/delete", 2799 json={"endpoint_id": endpoint_id}, 2800 auth=(user1, password1), 2801 ).raise_for_status() 2802 requests.delete( 2803 url=fastapi_client.tracking_uri + "/api/3.0/mlflow/gateway/model-definitions/delete", 2804 json={"model_definition_id": model_def_id}, 2805 auth=(user1, password1), 2806 ).raise_for_status() 2807 requests.delete( 2808 url=fastapi_client.tracking_uri + "/api/3.0/mlflow/gateway/secrets/delete", 2809 json={"secret_id": secret_id}, 2810 auth=(user1, password1), 2811 ).raise_for_status() 2812 2813 2814 def test_otel_unauthenticated_access_denied(fastapi_client, monkeypatch): 2815 monkeypatch.delenv(MLFLOW_TRACKING_USERNAME.name, raising=False) 2816 monkeypatch.delenv(MLFLOW_TRACKING_PASSWORD.name, raising=False) 2817 2818 response = requests.post( 2819 url=fastapi_client.tracking_uri + "/v1/traces", 2820 headers={"Content-Type": "application/x-protobuf", "X-Mlflow-Experiment-Id": "1"}, 2821 data=b"", 2822 ) 2823 assert response.status_code == 401 2824 2825 2826 def test_otel_experiment_permission(fastapi_client, monkeypatch): 2827 user1, password1 = create_user(fastapi_client.tracking_uri) 2828 user2, password2 = create_user(fastapi_client.tracking_uri) 2829 2830 # user1 creates an experiment 2831 with User(user1, password1, monkeypatch): 2832 experiment_id = fastapi_client.create_experiment("otel_permission_test") 2833 2834 # Grant READ permission to user2 (not enough for writing traces) 2835 _send_rest_tracking_post_request( 2836 fastapi_client.tracking_uri, 2837 "/api/2.0/mlflow/experiments/permissions/create", 2838 json_payload={ 2839 "experiment_id": experiment_id, 2840 "username": user2, 2841 "permission": "READ", 2842 }, 2843 auth=(user1, password1), 2844 ) 2845 2846 # user2 cannot write traces (READ doesn't grant can_update) 2847 response = requests.post( 2848 url=fastapi_client.tracking_uri + "/v1/traces", 2849 headers={ 2850 "Content-Type": "application/x-protobuf", 2851 "X-Mlflow-Experiment-Id": experiment_id, 2852 }, 2853 data=b"", 2854 auth=(user2, password2), 2855 ) 2856 assert response.status_code == 403 2857 2858 # Grant EDIT permission to user2 2859 requests.patch( 2860 url=fastapi_client.tracking_uri + "/api/2.0/mlflow/experiments/permissions/update", 2861 json={ 2862 "experiment_id": experiment_id, 2863 "username": user2, 2864 "permission": "EDIT", 2865 }, 2866 auth=(user1, password1), 2867 ) 2868 2869 # user2 can now write traces (EDIT grants can_update) 2870 # The request may fail for other reasons (invalid protobuf) but should pass permission check 2871 response = requests.post( 2872 url=fastapi_client.tracking_uri + "/v1/traces", 2873 headers={ 2874 "Content-Type": "application/x-protobuf", 2875 "X-Mlflow-Experiment-Id": experiment_id, 2876 }, 2877 data=b"", 2878 auth=(user2, password2), 2879 ) 2880 assert response.status_code != 403 2881 2882 2883 def test_job_api_unauthenticated_access_denied(fastapi_client, monkeypatch): 2884 monkeypatch.delenv(MLFLOW_TRACKING_USERNAME.name, raising=False) 2885 monkeypatch.delenv(MLFLOW_TRACKING_PASSWORD.name, raising=False) 2886 2887 response = requests.post( 2888 url=fastapi_client.tracking_uri + "/ajax-api/3.0/jobs/search", 2889 json={}, 2890 ) 2891 assert response.status_code == 401 2892 2893 2894 def test_assistant_unauthenticated_access_denied(fastapi_client, monkeypatch): 2895 monkeypatch.delenv(MLFLOW_TRACKING_USERNAME.name, raising=False) 2896 monkeypatch.delenv(MLFLOW_TRACKING_PASSWORD.name, raising=False) 2897 2898 response = requests.post( 2899 url=fastapi_client.tracking_uri + "/ajax-api/3.0/mlflow/assistant/chat", 2900 json={"messages": []}, 2901 ) 2902 assert response.status_code == 401 2903 2904 2905 def test_get_online_scoring_configs_with_auth(client, monkeypatch): 2906 username, password = create_user(client.tracking_uri) 2907 2908 with User(username, password, monkeypatch): 2909 experiment_id = client.create_experiment("test_experiment") 2910 2911 # Register a scorer 2912 scorer_json = '{"name": "test_scorer", "type": "pyfunc"}' 2913 response = _send_rest_tracking_post_request( 2914 client.tracking_uri, 2915 "/api/3.0/mlflow/scorers/register", 2916 json_payload={ 2917 "experiment_id": experiment_id, 2918 "name": "test_scorer", 2919 "serialized_scorer": scorer_json, 2920 }, 2921 auth=(username, password), 2922 ) 2923 scorer_id = response.json()["scorer_id"] 2924 2925 # Test the online scoring configs endpoint (GET) 2926 # This should not raise a TypeError as it did before when the endpoint 2927 # was incorrectly included in AFTER_REQUEST_HANDLERS 2928 response = requests.get( 2929 url=client.tracking_uri + "/ajax-api/3.0/mlflow/scorers/online-configs", 2930 params={"scorer_ids": scorer_id}, 2931 auth=(username, password), 2932 ) 2933 2934 # Should return 200 (not 500 with TypeError) 2935 assert response.status_code == 200 2936 data = response.json() 2937 assert "configs" in data 2938 assert isinstance(data["configs"], list) 2939 2940 2941 def test_list_users(client): 2942 username1, password1 = create_user(client.tracking_uri) 2943 username2, _password2 = create_user(client.tracking_uri) 2944 2945 # Admin can list all users 2946 response = requests.get( 2947 url=client.tracking_uri + LIST_USERS, 2948 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 2949 ) 2950 assert response.status_code == 200 2951 data = response.json() 2952 assert "users" in data 2953 usernames = [u["username"] for u in data["users"]] 2954 assert ADMIN_USERNAME in usernames 2955 assert username1 in usernames 2956 assert username2 in usernames 2957 for user in data["users"]: 2958 assert "id" in user 2959 assert "username" in user 2960 assert "password" not in user 2961 assert "password_hash" not in user 2962 2963 # Unauthenticated request should fail 2964 response = requests.get(url=client.tracking_uri + LIST_USERS) 2965 assert response.status_code == 401 2966 2967 # Non-admin user should not be able to list all users 2968 response = requests.get( 2969 url=client.tracking_uri + LIST_USERS, 2970 auth=(username1, password1), 2971 ) 2972 assert response.status_code == 403 2973 2974 # Ajax API path should also work for admin 2975 response = requests.get( 2976 url=client.tracking_uri + AJAX_LIST_USERS, 2977 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 2978 ) 2979 assert response.status_code == 200 2980 data = response.json() 2981 assert "users" in data 2982 assert len(data["users"]) >= 3 2983 2984 2985 @pytest.mark.parametrize( 2986 "client", 2987 [{"MLFLOW_WEBHOOK_SECRET_ENCRYPTION_KEY": Fernet.generate_key().decode("utf-8")}], 2988 indirect=True, 2989 ) 2990 def test_webhook_admin_only_permissions(client, monkeypatch): 2991 user1, password1 = create_user(client.tracking_uri) 2992 2993 # Non-admin: create webhook should be forbidden 2994 with User(user1, password1, monkeypatch): 2995 response = requests.post( 2996 url=client.tracking_uri + "/api/2.0/mlflow/webhooks", 2997 json={ 2998 "name": "test-webhook", 2999 "url": "https://example.com/webhook", 3000 "events": [{"entity": "MODEL_VERSION", "action": "CREATED"}], 3001 }, 3002 auth=(user1, password1), 3003 ) 3004 assert response.status_code == 403 3005 3006 # Non-admin: list webhooks should be forbidden 3007 with User(user1, password1, monkeypatch): 3008 response = requests.get( 3009 url=client.tracking_uri + "/api/2.0/mlflow/webhooks", 3010 auth=(user1, password1), 3011 ) 3012 assert response.status_code == 403 3013 3014 # Admin: create webhook should succeed 3015 response = requests.post( 3016 url=client.tracking_uri + "/api/2.0/mlflow/webhooks", 3017 json={ 3018 "name": "admin-webhook", 3019 "url": "https://example.com/webhook", 3020 "events": [{"entity": "MODEL_VERSION", "action": "CREATED"}], 3021 }, 3022 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 3023 ) 3024 response.raise_for_status() 3025 webhook_id = response.json()["webhook"]["webhook_id"] 3026 3027 # Admin: list webhooks should succeed 3028 response = requests.get( 3029 url=client.tracking_uri + "/api/2.0/mlflow/webhooks", 3030 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 3031 ) 3032 response.raise_for_status() 3033 3034 # Non-admin: get webhook should be forbidden 3035 with User(user1, password1, monkeypatch): 3036 response = requests.get( 3037 url=client.tracking_uri + f"/api/2.0/mlflow/webhooks/{webhook_id}", 3038 auth=(user1, password1), 3039 ) 3040 assert response.status_code == 403 3041 3042 # Admin: get webhook should succeed 3043 response = requests.get( 3044 url=client.tracking_uri + f"/api/2.0/mlflow/webhooks/{webhook_id}", 3045 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 3046 ) 3047 response.raise_for_status() 3048 3049 # Non-admin: update webhook should be forbidden 3050 with User(user1, password1, monkeypatch): 3051 response = requests.patch( 3052 url=client.tracking_uri + f"/api/2.0/mlflow/webhooks/{webhook_id}", 3053 json={"name": "updated-name"}, 3054 auth=(user1, password1), 3055 ) 3056 assert response.status_code == 403 3057 3058 # Admin: update webhook should succeed 3059 response = requests.patch( 3060 url=client.tracking_uri + f"/api/2.0/mlflow/webhooks/{webhook_id}", 3061 json={"name": "updated-name"}, 3062 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 3063 ) 3064 response.raise_for_status() 3065 3066 # Non-admin: test webhook should be forbidden 3067 with User(user1, password1, monkeypatch): 3068 response = requests.post( 3069 url=client.tracking_uri + f"/api/2.0/mlflow/webhooks/{webhook_id}/test", 3070 json={}, 3071 auth=(user1, password1), 3072 ) 3073 assert response.status_code == 403 3074 3075 # Admin: test webhook should succeed 3076 response = requests.post( 3077 url=client.tracking_uri + f"/api/2.0/mlflow/webhooks/{webhook_id}/test", 3078 json={}, 3079 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 3080 ) 3081 response.raise_for_status() 3082 3083 # Non-admin: delete webhook should be forbidden 3084 with User(user1, password1, monkeypatch): 3085 response = requests.delete( 3086 url=client.tracking_uri + f"/api/2.0/mlflow/webhooks/{webhook_id}", 3087 auth=(user1, password1), 3088 ) 3089 assert response.status_code == 403 3090 3091 # Admin: delete webhook should succeed 3092 response = requests.delete( 3093 url=client.tracking_uri + f"/api/2.0/mlflow/webhooks/{webhook_id}", 3094 auth=(ADMIN_USERNAME, ADMIN_PASSWORD), 3095 ) 3096 response.raise_for_status() 3097 3098 3099 # -- Unit tests for _authenticate_fastapi_request -- 3100 3101 3102 @pytest.fixture 3103 def mock_auth_store(): 3104 if auth_module._USER_AUTH_CACHE is not None: 3105 with auth_module._USER_AUTH_CACHE_LOCK: 3106 auth_module._USER_AUTH_CACHE.clear() 3107 with mock.patch("mlflow.server.auth.store") as mock_store: 3108 mock_store.get_user.side_effect = lambda username: mock.Mock(username=username) 3109 mock_store.authenticate_user.return_value = True 3110 yield mock_store 3111 if auth_module._USER_AUTH_CACHE is not None: 3112 with auth_module._USER_AUTH_CACHE_LOCK: 3113 auth_module._USER_AUTH_CACHE.clear() 3114 3115 3116 @pytest.fixture 3117 def mock_auth_config(): 3118 with mock.patch("mlflow.server.auth.auth_config") as mock_config: 3119 mock_config.admin_username = "admin" 3120 yield mock_config 3121 3122 3123 @pytest.fixture 3124 def enable_auth_cache(): 3125 # The credential cache is disabled by default; cache-behavior tests must opt in. 3126 cache = TTLCache(maxsize=10000, ttl=60) 3127 with mock.patch("mlflow.server.auth._USER_AUTH_CACHE", cache): 3128 yield cache 3129 3130 3131 def _make_request(path, authorization=None): 3132 request = mock.Mock() 3133 request.url.path = path 3134 request.headers = {} 3135 if authorization: 3136 request.headers["Authorization"] = authorization 3137 return request 3138 3139 3140 # -- Basic auth with internal token (trusted internal requests) -- 3141 3142 3143 def test_basic_auth_with_internal_token_returns_user( 3144 mock_auth_store, mock_auth_config, monkeypatch 3145 ): 3146 monkeypatch.setenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, "internal-secret") 3147 credentials = base64.b64encode(b"alice:internal-secret").decode("ascii") 3148 request = _make_request("/gateway/mlflow/v1/chat", f"Basic {credentials}") 3149 3150 user = _authenticate_fastapi_request(request) 3151 3152 assert user.username == "alice" 3153 mock_auth_store.get_user.assert_called_once_with("alice") 3154 mock_auth_store.authenticate_user.assert_not_called() 3155 3156 3157 def test_basic_auth_with_internal_token_deleted_user_returns_none( 3158 mock_auth_store, mock_auth_config, monkeypatch 3159 ): 3160 monkeypatch.setenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, "internal-secret") 3161 mock_auth_store.get_user.side_effect = MlflowException("User not found") 3162 credentials = base64.b64encode(b"deleted_user:internal-secret").decode("ascii") 3163 request = _make_request("/gateway/mlflow/v1/chat", f"Basic {credentials}") 3164 3165 user = _authenticate_fastapi_request(request) 3166 3167 assert user is None 3168 3169 3170 def test_basic_auth_with_wrong_password_falls_through_to_authenticate( 3171 mock_auth_store, mock_auth_config, monkeypatch 3172 ): 3173 monkeypatch.setenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, "internal-secret") 3174 credentials = base64.b64encode(b"alice:wrong-password").decode("ascii") 3175 request = _make_request("/gateway/mlflow/v1/chat", f"Basic {credentials}") 3176 3177 user = _authenticate_fastapi_request(request) 3178 3179 assert user.username == "alice" 3180 mock_auth_store.authenticate_user.assert_called_once_with("alice", "wrong-password") 3181 3182 3183 def test_basic_auth_internal_token_rejected_on_non_gateway_route( 3184 mock_auth_store, mock_auth_config, monkeypatch 3185 ): 3186 monkeypatch.setenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, "internal-secret") 3187 credentials = base64.b64encode(b"alice:internal-secret").decode("ascii") 3188 request = _make_request("/api/3.0/mlflow/experiments/list", f"Basic {credentials}") 3189 3190 _authenticate_fastapi_request(request) 3191 3192 # Internal token should NOT be accepted on non-gateway routes — falls through 3193 # to store.authenticate_user instead 3194 mock_auth_store.authenticate_user.assert_called_once_with("alice", "internal-secret") 3195 3196 3197 def test_basic_auth_no_internal_token_uses_normal_auth( 3198 mock_auth_store, mock_auth_config, monkeypatch 3199 ): 3200 monkeypatch.delenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, raising=False) 3201 credentials = base64.b64encode(b"alice:password123").decode("ascii") 3202 request = _make_request("/gateway/mlflow/v1/chat", f"Basic {credentials}") 3203 3204 user = _authenticate_fastapi_request(request) 3205 3206 assert user.username == "alice" 3207 mock_auth_store.authenticate_user.assert_called_once_with("alice", "password123") 3208 3209 3210 # -- Standard Basic auth -- 3211 3212 3213 def test_fastapi_valid_basic_auth(mock_auth_store, mock_auth_config, monkeypatch): 3214 monkeypatch.delenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, raising=False) 3215 credentials = base64.b64encode(b"alice:password123").decode("ascii") 3216 request = _make_request("/api/3.0/mlflow/experiments/list", f"Basic {credentials}") 3217 3218 user = _authenticate_fastapi_request(request) 3219 3220 assert user.username == "alice" 3221 mock_auth_store.authenticate_user.assert_called_once_with("alice", "password123") 3222 3223 3224 def test_fastapi_invalid_basic_auth(mock_auth_store, mock_auth_config, monkeypatch): 3225 monkeypatch.delenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, raising=False) 3226 mock_auth_store.authenticate_user.return_value = False 3227 credentials = base64.b64encode(b"alice:wrong").decode("ascii") 3228 request = _make_request("/api/3.0/mlflow/experiments/list", f"Basic {credentials}") 3229 3230 user = _authenticate_fastapi_request(request) 3231 3232 assert user is None 3233 3234 3235 # -- Non-Basic auth schemes -- 3236 3237 3238 def test_bearer_returns_none(mock_auth_store, mock_auth_config, monkeypatch): 3239 monkeypatch.setenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, "abc123") 3240 request = _make_request("/gateway/mlflow/v1/chat", "Bearer abc123") 3241 3242 user = _authenticate_fastapi_request(request) 3243 3244 assert user is None 3245 mock_auth_store.get_user.assert_not_called() 3246 3247 3248 # -- No auth header -- 3249 3250 3251 def test_fastapi_no_authorization_header(mock_auth_store, mock_auth_config): 3252 request = _make_request("/api/3.0/mlflow/experiments/list") 3253 3254 user = _authenticate_fastapi_request(request) 3255 3256 assert user is None 3257 3258 3259 def test_fastapi_malformed_authorization_header(mock_auth_store, mock_auth_config): 3260 request = _make_request("/api/3.0/mlflow/experiments/list", "garbage") 3261 3262 user = _authenticate_fastapi_request(request) 3263 3264 assert user is None 3265 3266 3267 # -- Basic auth credential cache -- 3268 3269 3270 def test_basic_auth_caches_successful_credentials( 3271 enable_auth_cache, mock_auth_store, mock_auth_config, monkeypatch 3272 ): 3273 monkeypatch.delenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, raising=False) 3274 credentials = base64.b64encode(b"alice:password123").decode("ascii") 3275 request = _make_request("/api/3.0/mlflow/experiments/list", f"Basic {credentials}") 3276 3277 user_a = _authenticate_fastapi_request(request) 3278 user_b = _authenticate_fastapi_request(request) 3279 3280 assert user_a.username == "alice" 3281 assert user_b.username == "alice" 3282 # Both PBKDF2 check and user fetch should run exactly once across the two requests. 3283 mock_auth_store.authenticate_user.assert_called_once_with("alice", "password123") 3284 mock_auth_store.get_user.assert_called_once_with("alice") 3285 3286 3287 def test_basic_auth_cache_does_not_store_failed_credentials( 3288 enable_auth_cache, mock_auth_store, mock_auth_config, monkeypatch 3289 ): 3290 monkeypatch.delenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, raising=False) 3291 mock_auth_store.authenticate_user.return_value = False 3292 credentials = base64.b64encode(b"alice:wrong").decode("ascii") 3293 request = _make_request("/api/3.0/mlflow/experiments/list", f"Basic {credentials}") 3294 3295 assert _authenticate_fastapi_request(request) is None 3296 assert _authenticate_fastapi_request(request) is None 3297 assert mock_auth_store.authenticate_user.call_count == 2 3298 3299 3300 def test_basic_auth_cache_keyed_by_username_and_password( 3301 enable_auth_cache, mock_auth_store, mock_auth_config, monkeypatch 3302 ): 3303 monkeypatch.delenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, raising=False) 3304 alice = base64.b64encode(b"alice:password123").decode("ascii") 3305 bob = base64.b64encode(b"bob:password123").decode("ascii") 3306 alice_wrong = base64.b64encode(b"alice:other-password").decode("ascii") 3307 3308 _authenticate_fastapi_request(_make_request("/x", f"Basic {alice}")) 3309 _authenticate_fastapi_request(_make_request("/x", f"Basic {bob}")) 3310 _authenticate_fastapi_request(_make_request("/x", f"Basic {alice_wrong}")) 3311 3312 assert mock_auth_store.authenticate_user.call_args_list == [ 3313 mock.call("alice", "password123"), 3314 mock.call("bob", "password123"), 3315 mock.call("alice", "other-password"), 3316 ] 3317 3318 3319 def test_basic_auth_returns_none_when_user_deleted_between_authenticate_and_get( 3320 enable_auth_cache, mock_auth_store, mock_auth_config, monkeypatch 3321 ): 3322 # TOCTOU: authenticate_user returned True but the user disappeared before get_user. 3323 monkeypatch.delenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, raising=False) 3324 mock_auth_store.get_user.side_effect = MlflowException("User not found") 3325 credentials = base64.b64encode(b"ghost:password123").decode("ascii") 3326 request = _make_request("/x", f"Basic {credentials}") 3327 3328 # Flask and FastAPI paths both must treat this as an auth failure, not surface 3329 # a 500 and, critically, must not cache the (ghost, password123) pair. 3330 assert _authenticate_fastapi_request(request) is None 3331 if auth_module._USER_AUTH_CACHE is not None: 3332 assert ( 3333 auth_module._auth_cache_key("ghost", "password123") not in auth_module._USER_AUTH_CACHE 3334 ) 3335 3336 3337 def test_flask_basic_auth_skips_get_user_when_cache_disabled( 3338 mock_auth_store, mock_auth_config, monkeypatch 3339 ): 3340 monkeypatch.delenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, raising=False) 3341 fake_flask_request = mock.Mock() 3342 fake_flask_request.authorization.username = "alice" 3343 fake_flask_request.authorization.password = "password123" 3344 3345 with ( 3346 mock.patch("mlflow.server.auth._USER_AUTH_CACHE", None), 3347 mock.patch("mlflow.server.auth.request", fake_flask_request), 3348 ): 3349 result = auth_module.authenticate_request_basic_auth() 3350 3351 assert result is fake_flask_request.authorization 3352 mock_auth_store.authenticate_user.assert_called_once_with("alice", "password123") 3353 # Cache disabled + Flask path only needs the yes/no answer → no user fetch. 3354 mock_auth_store.get_user.assert_not_called() 3355 3356 3357 def test_flask_basic_auth_shares_cache_with_fastapi_path( 3358 enable_auth_cache, mock_auth_store, mock_auth_config, monkeypatch 3359 ): 3360 monkeypatch.delenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, raising=False) 3361 # Prime the cache via the FastAPI path. 3362 credentials = base64.b64encode(b"alice:password123").decode("ascii") 3363 _authenticate_fastapi_request(_make_request("/x", f"Basic {credentials}")) 3364 mock_auth_store.authenticate_user.assert_called_once_with("alice", "password123") 3365 3366 # A subsequent Flask-side call for the same credentials must be served from 3367 # cache — no second PBKDF2 verification, no second user fetch. 3368 fake_flask_request = mock.Mock() 3369 fake_flask_request.authorization.username = "alice" 3370 fake_flask_request.authorization.password = "password123" 3371 with mock.patch("mlflow.server.auth.request", fake_flask_request): 3372 result = auth_module.authenticate_request_basic_auth() 3373 3374 assert result is fake_flask_request.authorization 3375 mock_auth_store.authenticate_user.assert_called_once_with("alice", "password123") 3376 3377 3378 def test_invalidate_user_auth_cache_drops_only_matching_username( 3379 enable_auth_cache, mock_auth_store, mock_auth_config, monkeypatch 3380 ): 3381 monkeypatch.delenv(_MLFLOW_INTERNAL_GATEWAY_AUTH_TOKEN.name, raising=False) 3382 alice = base64.b64encode(b"alice:password123").decode("ascii") 3383 alice_alt = base64.b64encode(b"alice:other-password").decode("ascii") 3384 bob = base64.b64encode(b"bob:password123").decode("ascii") 3385 3386 _authenticate_fastapi_request(_make_request("/x", f"Basic {alice}")) 3387 _authenticate_fastapi_request(_make_request("/x", f"Basic {alice_alt}")) 3388 _authenticate_fastapi_request(_make_request("/x", f"Basic {bob}")) 3389 assert mock_auth_store.authenticate_user.call_count == 3 3390 3391 auth_module._invalidate_user_auth_cache("alice") 3392 3393 # Alice's two cached credentials are re-checked; bob's cache entry stays hot. 3394 _authenticate_fastapi_request(_make_request("/x", f"Basic {alice}")) 3395 _authenticate_fastapi_request(_make_request("/x", f"Basic {alice_alt}")) 3396 _authenticate_fastapi_request(_make_request("/x", f"Basic {bob}")) 3397 assert mock_auth_store.authenticate_user.call_count == 5