test_sqlalchemy_store.py
1 import pytest 2 3 from mlflow.exceptions import MlflowException 4 from mlflow.protos.databricks_pb2 import ( 5 INVALID_PARAMETER_VALUE, 6 RESOURCE_ALREADY_EXISTS, 7 RESOURCE_DOES_NOT_EXIST, 8 ErrorCode, 9 ) 10 from mlflow.server.auth.entities import ( 11 ExperimentPermission, 12 GatewayEndpointPermission, 13 GatewayModelDefinitionPermission, 14 GatewaySecretPermission, 15 RegisteredModelPermission, 16 ScorerPermission, 17 User, 18 ) 19 from mlflow.server.auth.permissions import ALL_PERMISSIONS, EDIT, MANAGE, READ 20 from mlflow.server.auth.sqlalchemy_store import SqlAlchemyStore 21 from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME 22 23 from tests.helper_functions import random_str 24 25 pytestmark = pytest.mark.notrackingurimock 26 27 28 @pytest.fixture 29 def store(tmp_sqlite_uri): 30 store = SqlAlchemyStore() 31 store.init_db(tmp_sqlite_uri) 32 return store 33 34 35 def _user_maker(store, username, password, is_admin=False): 36 return store.create_user(username, password, is_admin) 37 38 39 def _ep_maker(store, experiment_id, username, permission): 40 return store.create_experiment_permission(experiment_id, username, permission) 41 42 43 def _rmp_maker(store, name, username, permission): 44 return store.create_registered_model_permission(name, username, permission) 45 46 47 def _sp_maker(store, experiment_id, scorer_name, username, permission): 48 return store.create_scorer_permission(experiment_id, scorer_name, username, permission) 49 50 51 def _gsp_maker(store, secret_id, username, permission): 52 return store.create_gateway_secret_permission(secret_id, username, permission) 53 54 55 def _gep_maker(store, endpoint_id, username, permission): 56 return store.create_gateway_endpoint_permission(endpoint_id, username, permission) 57 58 59 def _gmdp_maker(store, model_definition_id, username, permission): 60 return store.create_gateway_model_definition_permission( 61 model_definition_id, username, permission 62 ) 63 64 65 def test_create_user(store): 66 username1 = random_str() 67 password1 = random_str() 68 user1 = _user_maker(store, username1, password1) 69 assert user1.username == username1 70 assert user1.password_hash != password1 71 assert user1.is_admin is False 72 73 # error on duplicate 74 with pytest.raises( 75 MlflowException, match=rf"User \(username={username1}\) already exists" 76 ) as exception_context: 77 _user_maker(store, username1, password1) 78 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS) 79 80 # slightly different name is ok 81 username2 = username1 + "_2" 82 password2 = password1 + "_2" 83 user2 = _user_maker(store, username2, password2, is_admin=True) 84 assert user2.username == username2 85 assert user2.password_hash != password2 86 assert user2.is_admin is True 87 88 # invalid username will fail 89 with pytest.raises(MlflowException, match=r"Username cannot be empty") as exception_context: 90 _user_maker(store, None, None) 91 assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) 92 with pytest.raises(MlflowException, match=r"Username cannot be empty") as exception_context: 93 _user_maker(store, "", "") 94 assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) 95 96 97 def test_has_user(store): 98 username1 = random_str() 99 password1 = random_str() 100 _user_maker(store, username1, password1) 101 assert store.has_user(username=username1) is True 102 103 # error on non-existent user 104 username2 = random_str() 105 assert store.has_user(username=username2) is False 106 107 108 def test_get_user(store): 109 username1 = random_str() 110 password1 = random_str() 111 _user_maker(store, username1, password1) 112 user1 = store.get_user(username=username1) 113 assert isinstance(user1, User) 114 assert user1.username == username1 115 116 # error on non-existent user 117 username2 = random_str() 118 with pytest.raises( 119 MlflowException, match=rf"User with username={username2} not found" 120 ) as exception_context: 121 store.get_user(username=username2) 122 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 123 124 125 def test_list_user(store): 126 username1 = "1" + random_str() 127 password1 = "1" + random_str() 128 _user_maker(store, username1, password1) 129 130 username2 = "2" + random_str() 131 password2 = "2" + random_str() 132 _user_maker(store, username2, password2) 133 134 username3 = "3" + random_str() 135 password3 = "3" + random_str() 136 _user_maker(store, username3, password3) 137 138 users = store.list_users() 139 users.sort(key=lambda u: u.username) 140 141 assert len(users) == 3 142 assert isinstance(users[0], User) 143 assert users[0].username == username1 144 assert users[1].username == username2 145 assert users[2].username == username3 146 147 148 def test_authenticate_user(store): 149 username1 = random_str() 150 password1 = random_str() 151 _user_maker(store, username1, password1) 152 assert store.authenticate_user(username1, password1) 153 assert not store.authenticate_user(username1, random_str()) 154 # non existent user 155 assert not store.authenticate_user(random_str(), random_str()) 156 157 158 def test_update_user(store): 159 username1 = random_str() 160 password1 = random_str() 161 _user_maker(store, username1, password1) 162 password2 = random_str() 163 store.update_user(username1, password=password2) 164 assert not store.authenticate_user(username1, password1) 165 assert store.authenticate_user(username1, password2) 166 167 store.update_user(username1, is_admin=True) 168 assert store.get_user(username1).is_admin 169 store.update_user(username1, is_admin=False) 170 assert not store.get_user(username1).is_admin 171 172 173 def test_delete_user(store): 174 username1 = random_str() 175 password1 = random_str() 176 _user_maker(store, username1, password1) 177 store.delete_user(username1) 178 179 with pytest.raises( 180 MlflowException, 181 match=rf"User with username={username1} not found", 182 ) as exception_context: 183 store.get_user(username1) 184 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 185 186 187 def test_create_experiment_permission(store): 188 username1 = random_str() 189 password1 = random_str() 190 user1 = _user_maker(store, username1, password1) 191 192 experiment_id1 = random_str() 193 user_id1 = user1.id 194 permission1 = READ.name 195 ep1 = _ep_maker(store, experiment_id1, username1, permission1) 196 assert ep1.experiment_id == experiment_id1 197 assert ep1.user_id == user_id1 198 assert ep1.permission == permission1 199 200 # error on duplicate 201 with pytest.raises( 202 MlflowException, 203 match=rf"Experiment permission \(experiment_id={experiment_id1}, " 204 rf"username={username1}\) already exists", 205 ) as exception_context: 206 _ep_maker(store, experiment_id1, username1, permission1) 207 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS) 208 209 # slightly different name is ok 210 experiment_id2 = random_str() 211 ep2 = _ep_maker(store, experiment_id2, username1, permission1) 212 assert ep2.experiment_id == experiment_id2 213 assert ep2.user_id == user_id1 214 assert ep2.permission == permission1 215 216 # all permissions are ok 217 for perm in ALL_PERMISSIONS: 218 experiment_id3 = random_str() 219 ep3 = _ep_maker(store, experiment_id3, username1, perm) 220 assert ep3.experiment_id == experiment_id3 221 assert ep3.user_id == user_id1 222 assert ep3.permission == perm 223 224 # invalid permission will fail 225 experiment_id4 = random_str() 226 with pytest.raises(MlflowException, match=r"Invalid permission") as exception_context: 227 _ep_maker(store, experiment_id4, username1, "some_invalid_permission_string") 228 assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) 229 230 231 def test_get_experiment_permission(store): 232 username1 = random_str() 233 password1 = random_str() 234 user1 = _user_maker(store, username1, password1) 235 236 experiment_id1 = random_str() 237 user_id1 = user1.id 238 permission1 = READ.name 239 _ep_maker(store, experiment_id1, username1, permission1) 240 ep1 = store.get_experiment_permission(experiment_id1, username1) 241 assert isinstance(ep1, ExperimentPermission) 242 assert ep1.experiment_id == experiment_id1 243 assert ep1.user_id == user_id1 244 assert ep1.permission == permission1 245 246 # error on non-existent row 247 experiment_id2 = random_str() 248 with pytest.raises( 249 MlflowException, 250 match=rf"Experiment permission with experiment_id={experiment_id2} " 251 rf"and username={username1} not found", 252 ) as exception_context: 253 store.get_experiment_permission(experiment_id2, username1) 254 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 255 256 257 def test_list_experiment_permission(store): 258 username1 = random_str() 259 password1 = random_str() 260 _user_maker(store, username1, password1) 261 262 experiment_id1 = "1" + random_str() 263 _ep_maker(store, experiment_id1, username1, READ.name) 264 265 experiment_id2 = "2" + random_str() 266 _ep_maker(store, experiment_id2, username1, READ.name) 267 268 experiment_id3 = "3" + random_str() 269 _ep_maker(store, experiment_id3, username1, READ.name) 270 271 eps = store.list_experiment_permissions(username1) 272 eps.sort(key=lambda ep: ep.experiment_id) 273 274 assert len(eps) == 3 275 assert isinstance(eps[0], ExperimentPermission) 276 assert eps[0].experiment_id == experiment_id1 277 assert eps[1].experiment_id == experiment_id2 278 assert eps[2].experiment_id == experiment_id3 279 280 281 def test_update_experiment_permission(store): 282 username1 = random_str() 283 password1 = random_str() 284 _user_maker(store, username1, password1) 285 286 experiment_id1 = random_str() 287 permission1 = READ.name 288 _ep_maker(store, experiment_id1, username1, permission1) 289 290 permission2 = EDIT.name 291 store.update_experiment_permission(experiment_id1, username1, permission2) 292 ep1 = store.get_experiment_permission(experiment_id1, username1) 293 assert ep1.permission == permission2 294 295 # invalid permission will fail 296 with pytest.raises(MlflowException, match=r"Invalid permission") as exception_context: 297 store.update_experiment_permission( 298 experiment_id1, username1, "some_invalid_permission_string" 299 ) 300 assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) 301 302 303 def test_delete_experiment_permission(store): 304 username1 = random_str() 305 password1 = random_str() 306 _user_maker(store, username1, password1) 307 308 experiment_id1 = random_str() 309 permission1 = READ.name 310 _ep_maker(store, experiment_id1, username1, permission1) 311 312 store.delete_experiment_permission(experiment_id1, username1) 313 with pytest.raises( 314 MlflowException, 315 match=rf"Experiment permission with experiment_id={experiment_id1} " 316 rf"and username={username1} not found", 317 ) as exception_context: 318 store.get_experiment_permission(experiment_id1, username1) 319 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 320 321 322 def test_create_registered_model_permission(store): 323 username1 = random_str() 324 password1 = random_str() 325 user1 = _user_maker(store, username1, password1) 326 327 name1 = random_str() 328 user_id1 = user1.id 329 permission1 = READ.name 330 rmp1 = _rmp_maker(store, name1, username1, permission1) 331 assert rmp1.name == name1 332 assert rmp1.user_id == user_id1 333 assert rmp1.permission == permission1 334 assert rmp1.workspace == DEFAULT_WORKSPACE_NAME 335 336 # error on duplicate 337 duplicate_permission_pattern = ( 338 rf"(?s)Registered model permission " 339 rf"with workspace={DEFAULT_WORKSPACE_NAME}, name={name1} " 340 rf"and username={username1} already exists" 341 ) 342 with pytest.raises( 343 MlflowException, 344 match=duplicate_permission_pattern, 345 ) as exception_context: 346 _rmp_maker(store, name1, username1, permission1) 347 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS) 348 349 # slightly different name is ok 350 name2 = random_str() 351 rmp2 = _rmp_maker(store, name2, username1, permission1) 352 assert rmp2.name == name2 353 assert rmp2.user_id == user_id1 354 assert rmp2.permission == permission1 355 assert rmp2.workspace == DEFAULT_WORKSPACE_NAME 356 357 # all permissions are ok 358 for perm in ALL_PERMISSIONS: 359 name3 = random_str() 360 rmp3 = _rmp_maker(store, name3, username1, perm) 361 assert rmp3.name == name3 362 assert rmp3.user_id == user_id1 363 assert rmp3.permission == perm 364 assert rmp3.workspace == DEFAULT_WORKSPACE_NAME 365 366 # invalid permission will fail 367 name4 = random_str() 368 with pytest.raises(MlflowException, match=r"Invalid permission") as exception_context: 369 _rmp_maker(store, name4, username1, "some_invalid_permission_string") 370 assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) 371 372 373 def test_get_registered_model_permission(store): 374 username1 = random_str() 375 password1 = random_str() 376 user1 = _user_maker(store, username1, password1) 377 378 name1 = random_str() 379 user_id1 = user1.id 380 permission1 = READ.name 381 _rmp_maker(store, name1, username1, permission1) 382 rmp1 = store.get_registered_model_permission(name1, username1) 383 assert isinstance(rmp1, RegisteredModelPermission) 384 assert rmp1.name == name1 385 assert rmp1.user_id == user_id1 386 assert rmp1.permission == permission1 387 assert rmp1.workspace == DEFAULT_WORKSPACE_NAME 388 389 # error on non-existent row 390 name2 = random_str() 391 missing_permission_message = ( 392 "Registered model permission with " 393 f"workspace={DEFAULT_WORKSPACE_NAME}, name={name2} " 394 f"and username={username1} not found" 395 ) 396 with pytest.raises( 397 MlflowException, 398 match=missing_permission_message, 399 ) as exception_context: 400 store.get_registered_model_permission(name2, username1) 401 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 402 403 404 def test_list_registered_model_permission(store): 405 username1 = random_str() 406 password1 = random_str() 407 _user_maker(store, username1, password1) 408 409 name1 = "1" + random_str() 410 _rmp_maker(store, name1, username1, READ.name) 411 412 name2 = "2" + random_str() 413 _rmp_maker(store, name2, username1, READ.name) 414 415 name3 = "3" + random_str() 416 _rmp_maker(store, name3, username1, READ.name) 417 418 rmps = store.list_registered_model_permissions(username1) 419 rmps.sort(key=lambda rmp: rmp.name) 420 421 assert len(rmps) == 3 422 assert isinstance(rmps[0], RegisteredModelPermission) 423 assert rmps[0].name == name1 424 assert rmps[0].workspace == DEFAULT_WORKSPACE_NAME 425 assert rmps[1].name == name2 426 assert rmps[1].workspace == DEFAULT_WORKSPACE_NAME 427 assert rmps[2].name == name3 428 assert rmps[2].workspace == DEFAULT_WORKSPACE_NAME 429 430 431 def test_update_registered_model_permission(store): 432 username1 = random_str() 433 password1 = random_str() 434 _user_maker(store, username1, password1) 435 436 name1 = random_str() 437 permission1 = READ.name 438 _rmp_maker(store, name1, username1, permission1) 439 440 permission2 = EDIT.name 441 store.update_registered_model_permission(name1, username1, permission2) 442 rmp1 = store.get_registered_model_permission(name1, username1) 443 assert rmp1.permission == permission2 444 assert rmp1.workspace == DEFAULT_WORKSPACE_NAME 445 446 # invalid permission will fail 447 with pytest.raises(MlflowException, match=r"Invalid permission") as exception_context: 448 store.update_registered_model_permission(name1, username1, "some_invalid_permission_string") 449 assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) 450 451 452 def test_delete_registered_model_permission(store): 453 username1 = random_str() 454 password1 = random_str() 455 _user_maker(store, username1, password1) 456 457 name1 = random_str() 458 permission1 = READ.name 459 _rmp_maker(store, name1, username1, permission1) 460 461 store.delete_registered_model_permission(name1, username1) 462 missing_permission_message = ( 463 "Registered model permission with " 464 f"workspace={DEFAULT_WORKSPACE_NAME}, name={name1} " 465 f"and username={username1} not found" 466 ) 467 with pytest.raises( 468 MlflowException, 469 match=missing_permission_message, 470 ) as exception_context: 471 store.get_registered_model_permission(name1, username1) 472 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 473 474 475 def test_rename_registered_model_permission(store): 476 # create 2 users and create 2 permission for the model registry with the same name 477 model_name = random_str() 478 username1 = random_str() 479 password1 = random_str() 480 _user_maker(store, username1, password1) 481 _rmp_maker(store, model_name, username1, MANAGE.name) 482 483 username2 = random_str() 484 password2 = random_str() 485 _user_maker(store, username2, password2) 486 _rmp_maker(store, model_name, username2, READ.name) 487 488 new_name = random_str() 489 490 store.rename_registered_model_permissions(model_name, new_name) 491 492 # get permission by model registry new name and all user must have the same new name 493 perm_user_1 = store.get_registered_model_permission(new_name, username1) 494 perm_user_2 = store.get_registered_model_permission(new_name, username2) 495 assert isinstance(perm_user_1, RegisteredModelPermission) 496 assert isinstance(perm_user_2, RegisteredModelPermission) 497 assert perm_user_1.name == new_name 498 assert perm_user_2.name == new_name 499 500 assert perm_user_1.permission == MANAGE.name 501 assert perm_user_1.workspace == DEFAULT_WORKSPACE_NAME 502 assert perm_user_2.permission == READ.name 503 assert perm_user_2.workspace == DEFAULT_WORKSPACE_NAME 504 505 506 def test_create_scorer_permission(store): 507 username1 = random_str() 508 password1 = random_str() 509 user1 = _user_maker(store, username1, password1) 510 511 experiment_id1 = random_str() 512 scorer_name1 = random_str() 513 user_id1 = user1.id 514 permission1 = READ.name 515 sp1 = _sp_maker(store, experiment_id1, scorer_name1, username1, permission1) 516 assert sp1.experiment_id == experiment_id1 517 assert sp1.scorer_name == scorer_name1 518 assert sp1.user_id == user_id1 519 assert sp1.permission == permission1 520 521 with pytest.raises( 522 MlflowException, 523 match=rf"Scorer permission \(experiment_id={experiment_id1}, scorer_name={scorer_name1}, " 524 rf"username={username1}\) already exists", 525 ) as exception_context: 526 _sp_maker(store, experiment_id1, scorer_name1, username1, permission1) 527 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS) 528 529 experiment_id2 = random_str() 530 sp2 = _sp_maker(store, experiment_id2, scorer_name1, username1, permission1) 531 assert sp2.experiment_id == experiment_id2 532 assert sp2.scorer_name == scorer_name1 533 assert sp2.user_id == user_id1 534 assert sp2.permission == permission1 535 536 for perm in ALL_PERMISSIONS: 537 experiment_id3 = random_str() 538 scorer_name3 = random_str() 539 sp3 = _sp_maker(store, experiment_id3, scorer_name3, username1, perm) 540 assert sp3.experiment_id == experiment_id3 541 assert sp3.scorer_name == scorer_name3 542 assert sp3.user_id == user_id1 543 assert sp3.permission == perm 544 545 experiment_id4 = random_str() 546 scorer_name4 = random_str() 547 with pytest.raises(MlflowException, match=r"Invalid permission") as exception_context: 548 _sp_maker(store, experiment_id4, scorer_name4, username1, "some_invalid_permission_string") 549 assert exception_context.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) 550 551 552 def test_get_scorer_permission(store): 553 username1 = random_str() 554 password1 = random_str() 555 user1 = _user_maker(store, username1, password1) 556 557 experiment_id1 = random_str() 558 scorer_name1 = random_str() 559 user_id1 = user1.id 560 permission1 = READ.name 561 _sp_maker(store, experiment_id1, scorer_name1, username1, permission1) 562 sp1 = store.get_scorer_permission(experiment_id1, scorer_name1, username1) 563 assert sp1.experiment_id == experiment_id1 564 assert sp1.scorer_name == scorer_name1 565 assert sp1.user_id == user_id1 566 assert sp1.permission == permission1 567 568 experiment_id2 = random_str() 569 with pytest.raises( 570 MlflowException, 571 match=rf"Scorer permission with experiment_id={experiment_id2}, " 572 rf"scorer_name={scorer_name1}, and username={username1} not found", 573 ) as exception_context: 574 store.get_scorer_permission(experiment_id2, scorer_name1, username1) 575 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 576 577 578 def test_list_scorer_permission(store): 579 username1 = random_str() 580 password1 = random_str() 581 _user_maker(store, username1, password1) 582 583 experiment_id1 = random_str() 584 scorer_name1 = random_str() 585 permission1 = READ.name 586 _sp_maker(store, experiment_id1, scorer_name1, username1, permission1) 587 588 experiment_id2 = random_str() 589 scorer_name2 = random_str() 590 permission2 = EDIT.name 591 _sp_maker(store, experiment_id2, scorer_name2, username1, permission2) 592 593 sps = store.list_scorer_permissions(username1) 594 assert len(sps) == 2 595 assert isinstance(sps[0], ScorerPermission) 596 assert isinstance(sps[1], ScorerPermission) 597 598 599 def test_update_scorer_permission(store): 600 username1 = random_str() 601 password1 = random_str() 602 user1 = _user_maker(store, username1, password1) 603 604 experiment_id1 = random_str() 605 scorer_name1 = random_str() 606 user_id1 = user1.id 607 permission1 = READ.name 608 _sp_maker(store, experiment_id1, scorer_name1, username1, permission1) 609 610 permission2 = MANAGE.name 611 sp2 = store.update_scorer_permission(experiment_id1, scorer_name1, username1, permission2) 612 assert sp2.experiment_id == experiment_id1 613 assert sp2.scorer_name == scorer_name1 614 assert sp2.user_id == user_id1 615 assert sp2.permission == permission2 616 617 618 def test_delete_scorer_permission(store): 619 username1 = random_str() 620 password1 = random_str() 621 _user_maker(store, username1, password1) 622 623 experiment_id1 = random_str() 624 scorer_name1 = random_str() 625 permission1 = READ.name 626 _sp_maker(store, experiment_id1, scorer_name1, username1, permission1) 627 628 store.delete_scorer_permission(experiment_id1, scorer_name1, username1) 629 630 with pytest.raises( 631 MlflowException, 632 match=rf"Scorer permission with experiment_id={experiment_id1}, " 633 rf"scorer_name={scorer_name1}, and username={username1} not found", 634 ) as exception_context: 635 store.get_scorer_permission(experiment_id1, scorer_name1, username1) 636 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 637 638 639 def test_delete_scorer_permissions_for_scorer(store): 640 username1 = random_str() 641 password1 = random_str() 642 _user_maker(store, username1, password1) 643 644 username2 = random_str() 645 password2 = random_str() 646 _user_maker(store, username2, password2) 647 648 experiment_id1 = random_str() 649 scorer_name1 = random_str() 650 _sp_maker(store, experiment_id1, scorer_name1, username1, MANAGE.name) 651 _sp_maker(store, experiment_id1, scorer_name1, username2, READ.name) 652 653 store.delete_scorer_permissions_for_scorer(experiment_id1, scorer_name1) 654 655 with pytest.raises(MlflowException, match=r"not found") as exception_context: 656 store.get_scorer_permission(experiment_id1, scorer_name1, username1) 657 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 658 659 with pytest.raises(MlflowException, match=r"not found") as exception_context: 660 store.get_scorer_permission(experiment_id1, scorer_name1, username2) 661 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 662 663 664 # Gateway Secret Permission Tests 665 666 667 def test_create_gateway_secret_permission(store): 668 username1 = random_str() 669 password1 = random_str() 670 user1 = _user_maker(store, username1, password1) 671 672 secret_id1 = random_str() 673 user_id1 = user1.id 674 permission1 = READ.name 675 gsp1 = _gsp_maker(store, secret_id1, username1, permission1) 676 assert gsp1.secret_id == secret_id1 677 assert gsp1.user_id == user_id1 678 assert gsp1.permission == permission1 679 680 # error on duplicate 681 with pytest.raises( 682 MlflowException, 683 match=rf"\(secret_id={secret_id1}, username={username1}\) already exists", 684 ) as exception_context: 685 _gsp_maker(store, secret_id1, username1, permission1) 686 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS) 687 688 # slightly different secret_id is ok 689 secret_id2 = random_str() 690 gsp2 = _gsp_maker(store, secret_id2, username1, permission1) 691 assert gsp2.secret_id == secret_id2 692 assert gsp2.user_id == user_id1 693 assert gsp2.permission == permission1 694 695 # all permissions are ok 696 for perm in ALL_PERMISSIONS: 697 secret_id = random_str() 698 gsp = _gsp_maker(store, secret_id, username1, perm) 699 assert gsp.permission == perm 700 701 702 def test_get_gateway_secret_permission(store): 703 username1 = random_str() 704 password1 = random_str() 705 _user_maker(store, username1, password1) 706 707 secret_id1 = random_str() 708 permission1 = READ.name 709 gsp1 = _gsp_maker(store, secret_id1, username1, permission1) 710 gsp2 = store.get_gateway_secret_permission(secret_id1, username1) 711 assert isinstance(gsp2, GatewaySecretPermission) 712 assert gsp2.secret_id == gsp1.secret_id 713 assert gsp2.user_id == gsp1.user_id 714 assert gsp2.permission == gsp1.permission 715 716 # error on non-existent permission 717 with pytest.raises(MlflowException, match=r"not found") as exception_context: 718 store.get_gateway_secret_permission(secret_id1, "random") 719 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 720 721 722 def test_update_gateway_secret_permission(store): 723 username1 = random_str() 724 password1 = random_str() 725 _user_maker(store, username1, password1) 726 727 secret_id1 = random_str() 728 permission1 = READ.name 729 _gsp_maker(store, secret_id1, username1, permission1) 730 store.update_gateway_secret_permission(secret_id1, username1, MANAGE.name) 731 gsp = store.get_gateway_secret_permission(secret_id1, username1) 732 assert gsp.permission == MANAGE.name 733 734 # error on non-existent permission 735 with pytest.raises( 736 MlflowException, 737 match=r"not found", 738 ) as exception_context: 739 store.update_gateway_secret_permission(secret_id1, "random", MANAGE.name) 740 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 741 742 743 def test_delete_gateway_secret_permission(store): 744 username1 = random_str() 745 password1 = random_str() 746 _user_maker(store, username1, password1) 747 748 secret_id1 = random_str() 749 permission1 = READ.name 750 _gsp_maker(store, secret_id1, username1, permission1) 751 store.delete_gateway_secret_permission(secret_id1, username1) 752 753 # error on non-existent permission 754 with pytest.raises( 755 MlflowException, 756 match=f"secret_id={secret_id1} and username={username1} not found", 757 ) as exception_context: 758 store.get_gateway_secret_permission(secret_id1, username1) 759 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 760 761 762 def test_delete_gateway_secret_permissions_for_secret(store): 763 username1 = random_str() 764 password1 = random_str() 765 _user_maker(store, username1, password1) 766 767 username2 = random_str() 768 password2 = random_str() 769 _user_maker(store, username2, password2) 770 771 secret_id1 = random_str() 772 _gsp_maker(store, secret_id1, username1, MANAGE.name) 773 _gsp_maker(store, secret_id1, username2, READ.name) 774 775 store.delete_gateway_secret_permissions_for_secret(secret_id1) 776 777 with pytest.raises(MlflowException, match=r"not found") as exception_context: 778 store.get_gateway_secret_permission(secret_id1, username1) 779 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 780 781 with pytest.raises(MlflowException, match=r"not found") as exception_context: 782 store.get_gateway_secret_permission(secret_id1, username2) 783 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 784 785 786 # Gateway Endpoint Permission Tests 787 788 789 def test_create_gateway_endpoint_permission(store): 790 username1 = random_str() 791 password1 = random_str() 792 user1 = _user_maker(store, username1, password1) 793 794 endpoint_id1 = random_str() 795 user_id1 = user1.id 796 permission1 = READ.name 797 gep1 = _gep_maker(store, endpoint_id1, username1, permission1) 798 assert gep1.endpoint_id == endpoint_id1 799 assert gep1.user_id == user_id1 800 assert gep1.permission == permission1 801 802 # error on duplicate 803 with pytest.raises( 804 MlflowException, 805 match=rf"\(endpoint_id={endpoint_id1}, username={username1}\) already exists", 806 ) as exception_context: 807 _gep_maker(store, endpoint_id1, username1, permission1) 808 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS) 809 810 # slightly different endpoint_id is ok 811 endpoint_id2 = random_str() 812 gep2 = _gep_maker(store, endpoint_id2, username1, permission1) 813 assert gep2.endpoint_id == endpoint_id2 814 assert gep2.user_id == user_id1 815 assert gep2.permission == permission1 816 817 # all permissions are ok 818 for perm in ALL_PERMISSIONS: 819 endpoint_id = random_str() 820 gep = _gep_maker(store, endpoint_id, username1, perm) 821 assert gep.permission == perm 822 823 824 def test_get_gateway_endpoint_permission(store): 825 username1 = random_str() 826 password1 = random_str() 827 _user_maker(store, username1, password1) 828 829 endpoint_id1 = random_str() 830 permission1 = READ.name 831 gep1 = _gep_maker(store, endpoint_id1, username1, permission1) 832 gep2 = store.get_gateway_endpoint_permission(endpoint_id1, username1) 833 assert isinstance(gep2, GatewayEndpointPermission) 834 assert gep2.endpoint_id == gep1.endpoint_id 835 assert gep2.user_id == gep1.user_id 836 assert gep2.permission == gep1.permission 837 838 # error on non-existent permission 839 with pytest.raises( 840 MlflowException, 841 match=r"not found", 842 ) as exception_context: 843 store.get_gateway_endpoint_permission(endpoint_id1, "random") 844 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 845 846 847 def test_update_gateway_endpoint_permission(store): 848 username1 = random_str() 849 password1 = random_str() 850 _user_maker(store, username1, password1) 851 852 endpoint_id1 = random_str() 853 permission1 = READ.name 854 _gep_maker(store, endpoint_id1, username1, permission1) 855 store.update_gateway_endpoint_permission(endpoint_id1, username1, MANAGE.name) 856 gep = store.get_gateway_endpoint_permission(endpoint_id1, username1) 857 assert gep.permission == MANAGE.name 858 859 # error on non-existent permission 860 with pytest.raises( 861 MlflowException, 862 match=r"not found", 863 ) as exception_context: 864 store.update_gateway_endpoint_permission(endpoint_id1, "random", MANAGE.name) 865 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 866 867 868 def test_delete_gateway_endpoint_permission(store): 869 username1 = random_str() 870 password1 = random_str() 871 _user_maker(store, username1, password1) 872 873 endpoint_id1 = random_str() 874 permission1 = READ.name 875 _gep_maker(store, endpoint_id1, username1, permission1) 876 store.delete_gateway_endpoint_permission(endpoint_id1, username1) 877 878 # error on non-existent permission 879 with pytest.raises( 880 MlflowException, 881 match=f"endpoint_id={endpoint_id1} and username={username1} not found", 882 ) as exception_context: 883 store.get_gateway_endpoint_permission(endpoint_id1, username1) 884 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 885 886 887 def test_delete_gateway_endpoint_permissions_for_endpoint(store): 888 username1 = random_str() 889 password1 = random_str() 890 _user_maker(store, username1, password1) 891 892 username2 = random_str() 893 password2 = random_str() 894 _user_maker(store, username2, password2) 895 896 endpoint_id1 = random_str() 897 _gep_maker(store, endpoint_id1, username1, MANAGE.name) 898 _gep_maker(store, endpoint_id1, username2, READ.name) 899 900 store.delete_gateway_endpoint_permissions_for_endpoint(endpoint_id1) 901 902 with pytest.raises(MlflowException, match=r"not found") as exception_context: 903 store.get_gateway_endpoint_permission(endpoint_id1, username1) 904 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 905 906 with pytest.raises(MlflowException, match=r"not found") as exception_context: 907 store.get_gateway_endpoint_permission(endpoint_id1, username2) 908 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 909 910 911 # Gateway Model Definition Permission Tests 912 913 914 def test_create_gateway_model_definition_permission(store): 915 username1 = random_str() 916 password1 = random_str() 917 user1 = _user_maker(store, username1, password1) 918 919 model_definition_id1 = random_str() 920 user_id1 = user1.id 921 permission1 = READ.name 922 gmdp1 = _gmdp_maker(store, model_definition_id1, username1, permission1) 923 assert gmdp1.model_definition_id == model_definition_id1 924 assert gmdp1.user_id == user_id1 925 assert gmdp1.permission == permission1 926 927 # error on duplicate 928 with pytest.raises( 929 MlflowException, 930 match="already exists", 931 ) as exception_context: 932 _gmdp_maker(store, model_definition_id1, username1, permission1) 933 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS) 934 935 # slightly different model_definition_id is ok 936 model_definition_id2 = random_str() 937 gmdp2 = _gmdp_maker(store, model_definition_id2, username1, permission1) 938 assert gmdp2.model_definition_id == model_definition_id2 939 assert gmdp2.user_id == user_id1 940 assert gmdp2.permission == permission1 941 942 # all permissions are ok 943 for perm in ALL_PERMISSIONS: 944 model_definition_id = random_str() 945 gmdp = _gmdp_maker(store, model_definition_id, username1, perm) 946 assert gmdp.permission == perm 947 948 949 def test_get_gateway_model_definition_permission(store): 950 username1 = random_str() 951 password1 = random_str() 952 _user_maker(store, username1, password1) 953 954 model_definition_id1 = random_str() 955 permission1 = READ.name 956 gmdp1 = _gmdp_maker(store, model_definition_id1, username1, permission1) 957 gmdp2 = store.get_gateway_model_definition_permission(model_definition_id1, username1) 958 assert isinstance(gmdp2, GatewayModelDefinitionPermission) 959 assert gmdp2.model_definition_id == gmdp1.model_definition_id 960 assert gmdp2.user_id == gmdp1.user_id 961 assert gmdp2.permission == gmdp1.permission 962 963 # error on non-existent permission 964 with pytest.raises( 965 MlflowException, 966 match=r"not found", 967 ) as exception_context: 968 store.get_gateway_model_definition_permission(model_definition_id1, "random") 969 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 970 971 972 def test_update_gateway_model_definition_permission(store): 973 username1 = random_str() 974 password1 = random_str() 975 _user_maker(store, username1, password1) 976 977 model_definition_id1 = random_str() 978 permission1 = READ.name 979 _gmdp_maker(store, model_definition_id1, username1, permission1) 980 store.update_gateway_model_definition_permission(model_definition_id1, username1, MANAGE.name) 981 gmdp = store.get_gateway_model_definition_permission(model_definition_id1, username1) 982 assert gmdp.permission == MANAGE.name 983 984 # error on non-existent permission 985 with pytest.raises( 986 MlflowException, 987 match=r"not found", 988 ) as exception_context: 989 store.update_gateway_model_definition_permission( 990 model_definition_id1, "random", MANAGE.name 991 ) 992 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 993 994 995 def test_delete_gateway_model_definition_permission(store): 996 username1 = random_str() 997 password1 = random_str() 998 _user_maker(store, username1, password1) 999 1000 model_definition_id1 = random_str() 1001 permission1 = READ.name 1002 _gmdp_maker(store, model_definition_id1, username1, permission1) 1003 store.delete_gateway_model_definition_permission(model_definition_id1, username1) 1004 1005 # error on non-existent permission 1006 with pytest.raises( 1007 MlflowException, 1008 match=f"model_definition_id={model_definition_id1} and username={username1} not found", 1009 ) as exception_context: 1010 store.get_gateway_model_definition_permission(model_definition_id1, username1) 1011 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 1012 1013 1014 def test_delete_gateway_model_definition_permissions_for_model_definition(store): 1015 username1 = random_str() 1016 password1 = random_str() 1017 _user_maker(store, username1, password1) 1018 1019 username2 = random_str() 1020 password2 = random_str() 1021 _user_maker(store, username2, password2) 1022 1023 model_definition_id1 = random_str() 1024 _gmdp_maker(store, model_definition_id1, username1, MANAGE.name) 1025 _gmdp_maker(store, model_definition_id1, username2, READ.name) 1026 1027 store.delete_gateway_model_definition_permissions_for_model_definition(model_definition_id1) 1028 1029 with pytest.raises(MlflowException, match=r"not found") as exception_context: 1030 store.get_gateway_model_definition_permission(model_definition_id1, username1) 1031 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST) 1032 1033 with pytest.raises(MlflowException, match=r"not found") as exception_context: 1034 store.get_gateway_model_definition_permission(model_definition_id1, username2) 1035 assert exception_context.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)