/ tests / server / auth / test_sqlalchemy_store.py
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)