1  import json
   2  from contextlib import contextmanager
   3  from types import SimpleNamespace
   4  from unittest.mock import MagicMock, Mock
   5  
   6  import pytest
   7  from flask import Response, request
   8  
   9  from mlflow.environment_variables import MLFLOW_ENABLE_WORKSPACES
  10  from mlflow.exceptions import MlflowException
  11  from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST
  12  from mlflow.server import auth as auth_module
  13  from mlflow.server.auth.permissions import EDIT, MANAGE, NO_PERMISSIONS, READ
  14  from mlflow.server.auth.routes import (
  15      CREATE_PROMPTLAB_RUN,
  16      GET_ARTIFACT,
  17      GET_METRIC_HISTORY_BULK,
  18      GET_METRIC_HISTORY_BULK_INTERVAL,
  19      GET_MODEL_VERSION_ARTIFACT,
  20      GET_TRACE_ARTIFACT,
  21      SEARCH_DATASETS,
  22      UPLOAD_ARTIFACT,
  23  )
  24  from mlflow.server.auth.sqlalchemy_store import SqlAlchemyStore
  25  from mlflow.utils import workspace_context
  26  from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME
  27  
  28  from tests.helper_functions import random_str
  29  
  30  
  31  def test_cleanup_workspace_permissions_handler(monkeypatch):
  32      mock_delete_workspace_perms = Mock()
  33      mock_delete_roles = Mock()
  34  
  35      monkeypatch.setattr(
  36          auth_module.store,
  37          "delete_workspace_permissions_for_workspace",
  38          mock_delete_workspace_perms,
  39          raising=True,
  40      )
  41      monkeypatch.setattr(
  42          auth_module.store,
  43          "delete_roles_for_workspace",
  44          mock_delete_roles,
  45          raising=True,
  46      )
  47  
  48      workspace_name = f"team-{random_str(10)}"
  49      with auth_module.app.test_request_context(
  50          f"/api/3.0/mlflow/workspaces/{workspace_name}", method="DELETE"
  51      ):
  52          request.view_args = {"workspace_name": workspace_name}
  53          response = Response(status=204)
  54          auth_module._after_request(response)
  55  
  56      mock_delete_workspace_perms.assert_called_once_with(workspace_name)
  57      mock_delete_roles.assert_called_once_with(workspace_name)
  58  
  59  
  60  def _create_workspace_response(workspace_name: str) -> Response:
  61      payload = {"workspace": {"name": workspace_name}}
  62      return Response(json.dumps(payload), status=201, content_type="application/json")
  63  
  64  
  65  def test_seed_default_workspace_roles_happy_path(monkeypatch):
  66      monkeypatch.setenv("MLFLOW_RBAC_SEED_DEFAULT_ROLES", "true")
  67      workspace_name = f"team-{random_str(10)}"
  68  
  69      created_roles: list[dict[str, object]] = []
  70      added_perms: list[dict[str, object]] = []
  71  
  72      def fake_create_role(name, workspace, description=None):
  73          role_id = len(created_roles) + 1
  74          created_roles.append({
  75              "id": role_id,
  76              "name": name,
  77              "workspace": workspace,
  78              "description": description,
  79          })
  80          return SimpleNamespace(id=role_id, name=name, workspace=workspace)
  81  
  82      def fake_add_role_permission(role_id, resource_type, resource_pattern, permission):
  83          added_perms.append({
  84              "role_id": role_id,
  85              "resource_type": resource_type,
  86              "resource_pattern": resource_pattern,
  87              "permission": permission,
  88          })
  89          return SimpleNamespace(id=role_id)
  90  
  91      monkeypatch.setattr(auth_module.store, "create_role", fake_create_role, raising=True)
  92      monkeypatch.setattr(
  93          auth_module.store, "add_role_permission", fake_add_role_permission, raising=True
  94      )
  95  
  96      with auth_module.app.test_request_context("/api/3.0/mlflow/workspaces", method="POST"):
  97          auth_module._seed_default_workspace_roles(_create_workspace_response(workspace_name))
  98  
  99      names = [r["name"] for r in created_roles]
 100      assert names == ["workspace-admin", "editor", "viewer"]
 101      assert all(r["workspace"] == workspace_name for r in created_roles)
 102  
 103      # All three roles use resource_type='workspace' (the only supported workspace-wide
 104      # resource_type in VALID_RESOURCE_TYPES). The permission level differentiates them.
 105      assert [(p["resource_type"], p["permission"]) for p in added_perms] == [
 106          ("workspace", MANAGE.name),
 107          ("workspace", EDIT.name),
 108          ("workspace", READ.name),
 109      ]
 110      assert all(p["resource_pattern"] == "*" for p in added_perms)
 111  
 112  
 113  def test_seed_default_workspace_roles_disabled_skips_seeding(monkeypatch):
 114      # With seeding off, no roles are created. ``CreateWorkspace`` is gated to
 115      # super-admins so the creator already bypasses RBAC — there is nothing to
 116      # fall back to.
 117      monkeypatch.setenv("MLFLOW_RBAC_SEED_DEFAULT_ROLES", "false")
 118      workspace_name = f"team-{random_str(10)}"
 119  
 120      mock_create_role = Mock()
 121      mock_add_role_permission = Mock()
 122      mock_assign_role_to_user = Mock()
 123      mock_set_workspace_permission = Mock()
 124  
 125      monkeypatch.setattr(auth_module.store, "create_role", mock_create_role, raising=True)
 126      monkeypatch.setattr(
 127          auth_module.store, "add_role_permission", mock_add_role_permission, raising=True
 128      )
 129      monkeypatch.setattr(
 130          auth_module.store, "assign_role_to_user", mock_assign_role_to_user, raising=True
 131      )
 132      monkeypatch.setattr(
 133          auth_module.store,
 134          "set_workspace_permission",
 135          mock_set_workspace_permission,
 136          raising=True,
 137      )
 138  
 139      with auth_module.app.test_request_context("/api/3.0/mlflow/workspaces", method="POST"):
 140          auth_module._seed_default_workspace_roles(_create_workspace_response(workspace_name))
 141  
 142      mock_create_role.assert_not_called()
 143      mock_add_role_permission.assert_not_called()
 144      mock_assign_role_to_user.assert_not_called()
 145      mock_set_workspace_permission.assert_not_called()
 146  
 147  
 148  def test_seed_default_workspace_roles_admin_creation_fails_still_seeds_others(monkeypatch):
 149      # Best-effort seeding: a failure on one role doesn't block the rest.
 150      monkeypatch.setenv("MLFLOW_RBAC_SEED_DEFAULT_ROLES", "true")
 151      workspace_name = f"team-{random_str(10)}"
 152  
 153      def fake_create_role(name, workspace, description=None):
 154          if name == "workspace-admin":
 155              raise MlflowException("simulated admin role failure")
 156          return SimpleNamespace(id=10, name=name, workspace=workspace)
 157  
 158      mock_add_role_permission = Mock()
 159  
 160      monkeypatch.setattr(auth_module.store, "create_role", fake_create_role, raising=True)
 161      monkeypatch.setattr(
 162          auth_module.store, "add_role_permission", mock_add_role_permission, raising=True
 163      )
 164  
 165      with auth_module.app.test_request_context("/api/3.0/mlflow/workspaces", method="POST"):
 166          auth_module._seed_default_workspace_roles(_create_workspace_response(workspace_name))
 167  
 168      # editor and viewer still got created (best-effort seeding).
 169      assert mock_add_role_permission.call_count == 2
 170  
 171  
 172  def test_seed_default_workspace_roles_permission_add_fails_rolls_back_role(monkeypatch):
 173      # create_role succeeds but add_role_permission raises — the orphan role must be
 174      # deleted so the workspace doesn't end up with a named role that grants nothing.
 175      monkeypatch.setenv("MLFLOW_RBAC_SEED_DEFAULT_ROLES", "true")
 176      workspace_name = f"team-{random_str(10)}"
 177  
 178      def fake_create_role(name, workspace, description=None):
 179          return SimpleNamespace(
 180              id={"workspace-admin": 1, "editor": 2, "viewer": 3}[name],
 181              name=name,
 182              workspace=workspace,
 183          )
 184  
 185      def fake_add_role_permission(role_id, resource_type, resource_pattern, permission):
 186          if role_id == 1:  # workspace-admin
 187              raise MlflowException("simulated add_role_permission failure")
 188          return SimpleNamespace(id=role_id)
 189  
 190      mock_delete_role = Mock()
 191  
 192      monkeypatch.setattr(auth_module.store, "create_role", fake_create_role, raising=True)
 193      monkeypatch.setattr(
 194          auth_module.store, "add_role_permission", fake_add_role_permission, raising=True
 195      )
 196      monkeypatch.setattr(auth_module.store, "delete_role", mock_delete_role, raising=True)
 197  
 198      with auth_module.app.test_request_context("/api/3.0/mlflow/workspaces", method="POST"):
 199          auth_module._seed_default_workspace_roles(_create_workspace_response(workspace_name))
 200  
 201      # Orphan workspace-admin role (id=1) was rolled back.
 202      mock_delete_role.assert_called_once_with(1)
 203  
 204  
 205  class _TrackingStore:
 206      def __init__(
 207          self,
 208          experiment_workspaces: dict[str, str],
 209          run_experiments: dict[str, str],
 210          trace_experiments: dict[str, str],
 211          experiment_names: dict[str, str] | None = None,
 212          logged_model_experiments: dict[str, str] | None = None,
 213          gateway_secret_workspaces: dict[str, str] | None = None,
 214          gateway_endpoint_workspaces: dict[str, str] | None = None,
 215          gateway_model_def_workspaces: dict[str, str] | None = None,
 216          engine=None,
 217          ManagedSessionMaker=None,
 218      ):
 219          self._experiment_workspaces = experiment_workspaces
 220          self._run_experiments = run_experiments
 221          self._trace_experiments = trace_experiments
 222          self._experiment_names = experiment_names or {}
 223          self._logged_model_experiments = logged_model_experiments or {}
 224          self._gateway_secret_workspaces = gateway_secret_workspaces or {}
 225          self._gateway_endpoint_workspaces = gateway_endpoint_workspaces or {}
 226          self._gateway_model_def_workspaces = gateway_model_def_workspaces or {}
 227          self.engine = engine
 228          self.ManagedSessionMaker = ManagedSessionMaker
 229  
 230      def get_experiment(self, experiment_id: str):
 231          return SimpleNamespace(workspace=self._experiment_workspaces[experiment_id])
 232  
 233      def get_experiment_by_name(self, experiment_name: str):
 234          experiment_id = self._experiment_names.get(experiment_name)
 235          if experiment_id is None:
 236              return None
 237          return SimpleNamespace(
 238              experiment_id=experiment_id,
 239              workspace=self._experiment_workspaces[experiment_id],
 240          )
 241  
 242      def get_run(self, run_id: str):
 243          return SimpleNamespace(info=SimpleNamespace(experiment_id=self._run_experiments[run_id]))
 244  
 245      def get_trace_info(self, request_id: str):
 246          return SimpleNamespace(experiment_id=self._trace_experiments[request_id])
 247  
 248      def get_logged_model(self, model_id: str):
 249          experiment_id = self._logged_model_experiments[model_id]
 250          return SimpleNamespace(experiment_id=experiment_id)
 251  
 252      def get_secret_info(self, secret_id: str | None = None, secret_name: str | None = None):
 253          if secret_id:
 254              if secret_id not in self._gateway_secret_workspaces:
 255                  raise MlflowException(
 256                      f"GatewaySecret not found ({secret_id})",
 257                      error_code=RESOURCE_DOES_NOT_EXIST,
 258                  )
 259              # Add workspace attribute so _get_resource_workspace can extract it
 260              return SimpleNamespace(
 261                  secret_id=secret_id, workspace=self._gateway_secret_workspaces[secret_id]
 262              )
 263          raise ValueError("Must provide secret_id or secret_name")
 264  
 265      def get_gateway_endpoint(self, endpoint_id: str | None = None, name: str | None = None):
 266          # For test simplicity we treat ``name`` as a synonym for ``endpoint_id``
 267          # (our fixture data uses the same string for both). This mirrors how the
 268          # real store resolves a name → id lookup before returning the endpoint.
 269          if lookup_id := (endpoint_id or name):
 270              if lookup_id not in self._gateway_endpoint_workspaces:
 271                  raise MlflowException(
 272                      f"GatewayEndpoint not found ({lookup_id})",
 273                      error_code=RESOURCE_DOES_NOT_EXIST,
 274                  )
 275              # Add workspace attribute so _get_resource_workspace can extract it
 276              return SimpleNamespace(
 277                  endpoint_id=lookup_id, workspace=self._gateway_endpoint_workspaces[lookup_id]
 278              )
 279          raise ValueError("Must provide endpoint_id or name")
 280  
 281      def get_gateway_model_definition(
 282          self, model_definition_id: str | None = None, name: str | None = None
 283      ):
 284          if model_definition_id:
 285              if model_definition_id not in self._gateway_model_def_workspaces:
 286                  raise MlflowException(
 287                      f"GatewayModelDefinition not found ({model_definition_id})",
 288                      error_code=RESOURCE_DOES_NOT_EXIST,
 289                  )
 290              # Add workspace attribute so _get_resource_workspace can extract it
 291              return SimpleNamespace(
 292                  model_definition_id=model_definition_id,
 293                  workspace=self._gateway_model_def_workspaces[model_definition_id],
 294              )
 295          raise ValueError("Must provide model_definition_id or name")
 296  
 297      def _create_mock_session(self):
 298          """Create a mock session that can query gateway SQL models."""
 299          mock_session = MagicMock()
 300  
 301          def _filter_by_secret_id(secret_id):
 302              if secret_id in self._gateway_secret_workspaces:
 303                  mock_result = MagicMock()
 304                  mock_result.first.return_value = SimpleNamespace(
 305                      workspace=self._gateway_secret_workspaces[secret_id]
 306                  )
 307                  return mock_result
 308              mock_result = MagicMock()
 309              mock_result.first.return_value = None
 310              return mock_result
 311  
 312          def _filter_by_endpoint_id(endpoint_id):
 313              if endpoint_id in self._gateway_endpoint_workspaces:
 314                  mock_result = MagicMock()
 315                  mock_result.first.return_value = SimpleNamespace(
 316                      workspace=self._gateway_endpoint_workspaces[endpoint_id]
 317                  )
 318                  return mock_result
 319              mock_result = MagicMock()
 320              mock_result.first.return_value = None
 321              return mock_result
 322  
 323          def _filter_by_model_def_id(model_definition_id):
 324              if model_definition_id in self._gateway_model_def_workspaces:
 325                  mock_result = MagicMock()
 326                  mock_result.first.return_value = SimpleNamespace(
 327                      workspace=self._gateway_model_def_workspaces[model_definition_id]
 328                  )
 329                  return mock_result
 330              mock_result = MagicMock()
 331              mock_result.first.return_value = None
 332              return mock_result
 333  
 334          def _query(model_class):
 335              mock_query_result = MagicMock()
 336              # Mock the filter method to return different results based on the filter
 337  
 338              def _mock_filter(*args, **kwargs):
 339                  if "secret_id" in kwargs:
 340                      return _filter_by_secret_id(kwargs["secret_id"])
 341                  elif "endpoint_id" in kwargs:
 342                      return _filter_by_endpoint_id(kwargs["endpoint_id"])
 343                  elif "model_definition_id" in kwargs:
 344                      return _filter_by_model_def_id(kwargs["model_definition_id"])
 345                  return mock_query_result
 346  
 347              mock_query_result.filter = _mock_filter
 348              return mock_query_result
 349  
 350          mock_session.query = _query
 351          return mock_session
 352  
 353      def _create_mock_session_maker(self):
 354          """Create a mock ManagedSessionMaker context manager."""
 355  
 356          @contextmanager
 357          def _mock_session_maker():
 358              yield self._create_mock_session()
 359  
 360          return _mock_session_maker
 361  
 362  
 363  class _RegistryStore:
 364      def __init__(self, model_workspaces: dict[str, str]):
 365          self._model_workspaces = model_workspaces
 366  
 367      def get_registered_model(self, name: str):
 368          return SimpleNamespace(workspace=self._model_workspaces[name])
 369  
 370  
 371  @pytest.fixture
 372  def workspace_permission_setup(tmp_path, monkeypatch):
 373      monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true")
 374      monkeypatch.setattr(
 375          auth_module,
 376          "auth_config",
 377          auth_module.auth_config._replace(default_permission=NO_PERMISSIONS.name),
 378      )
 379  
 380      db_uri = f"sqlite:///{tmp_path / 'auth-store.db'}"
 381      auth_store = SqlAlchemyStore()
 382      auth_store.init_db(db_uri)
 383      monkeypatch.setattr(auth_module, "store", auth_store, raising=False)
 384  
 385      username = "alice"
 386      auth_store.create_user(username, "supersecurepassword", is_admin=False)
 387  
 388      tracking_store = _TrackingStore(
 389          experiment_workspaces={"exp-1": "team-a", "exp-2": "team-a", "1": "team-a"},
 390          run_experiments={"run-1": "exp-1", "run-2": "exp-2"},
 391          trace_experiments={"trace-1": "exp-1"},
 392          experiment_names={"Primary Experiment": "exp-1"},
 393          logged_model_experiments={"model-1": "exp-1"},
 394          gateway_secret_workspaces={"secret-1": "team-a", "secret-2": "team-a"},
 395          gateway_endpoint_workspaces={"endpoint-1": "team-a", "endpoint-2": "team-a"},
 396          gateway_model_def_workspaces={"model-def-1": "team-a", "model-def-2": "team-a"},
 397          engine=MagicMock(),  # Mock engine for SQL model queries
 398      )
 399      # Set ManagedSessionMaker after creating the store
 400      tracking_store.ManagedSessionMaker = tracking_store._create_mock_session_maker()
 401      monkeypatch.setattr(auth_module, "_get_tracking_store", lambda: tracking_store)
 402  
 403      registry_store = _RegistryStore({"model-xyz": "team-a"})
 404      monkeypatch.setattr(auth_module, "_get_model_registry_store", lambda: registry_store)
 405  
 406      monkeypatch.setattr(
 407          auth_module,
 408          "authenticate_request",
 409          lambda: SimpleNamespace(username=username),
 410      )
 411  
 412      auth_store.set_workspace_permission("team-a", username, MANAGE.name)
 413  
 414      with workspace_context.WorkspaceContext("team-a"):
 415          yield {"store": auth_store, "username": username}
 416      auth_store.engine.dispose()
 417  
 418  
 419  def _set_workspace_permission(store: SqlAlchemyStore, username: str, permission: str):
 420      store.set_workspace_permission("team-a", username, permission)
 421  
 422  
 423  def test_workspace_permission_grants_default_access(monkeypatch):
 424      monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true")
 425  
 426      default_permission = MANAGE.name
 427      monkeypatch.setattr(
 428          auth_module,
 429          "auth_config",
 430          auth_module.auth_config._replace(
 431              default_permission=default_permission,
 432              grant_default_workspace_access=True,
 433          ),
 434          raising=False,
 435      )
 436  
 437      class DummyStore:
 438          def get_workspace_permission(self, workspace_name, username):
 439              return None
 440  
 441          def list_accessible_workspace_names(self, username):
 442              return []
 443  
 444      dummy_store = DummyStore()
 445      monkeypatch.setattr(auth_module, "store", dummy_store, raising=False)
 446  
 447      default_workspace = DEFAULT_WORKSPACE_NAME
 448      monkeypatch.setattr(auth_module, "_get_workspace_store", lambda: None, raising=False)
 449      monkeypatch.setattr(
 450          auth_module,
 451          "get_default_workspace_optional",
 452          lambda *args, **kwargs: (SimpleNamespace(name=default_workspace), True),
 453          raising=False,
 454      )
 455  
 456      auth = SimpleNamespace(username="alice")
 457      permission = auth_module._workspace_permission(auth.username, default_workspace)
 458      assert permission is not None
 459      assert permission.can_manage
 460  
 461      with workspace_context.WorkspaceContext(default_workspace):
 462          monkeypatch.setattr(auth_module, "authenticate_request", lambda: auth)
 463          assert auth_module.validate_can_create_experiment()
 464  
 465  
 466  def test_filter_list_workspaces_includes_default_when_autogrant(monkeypatch):
 467      monkeypatch.setattr(auth_module, "sender_is_admin", lambda: False)
 468      auth = SimpleNamespace(username="alice")
 469      monkeypatch.setattr(auth_module, "authenticate_request", lambda: auth)
 470      monkeypatch.setattr(
 471          auth_module,
 472          "auth_config",
 473          auth_module.auth_config._replace(
 474              grant_default_workspace_access=True,
 475              default_permission=READ.name,
 476          ),
 477          raising=False,
 478      )
 479  
 480      default_workspace = "team-default"
 481      monkeypatch.setattr(auth_module, "_get_workspace_store", lambda: None, raising=False)
 482      monkeypatch.setattr(
 483          auth_module,
 484          "get_default_workspace_optional",
 485          lambda *args, **kwargs: (SimpleNamespace(name=default_workspace), True),
 486          raising=False,
 487      )
 488  
 489      class DummyStore:
 490          def list_accessible_workspace_names(self, username):
 491              return []
 492  
 493      monkeypatch.setattr(auth_module, "store", DummyStore(), raising=False)
 494  
 495      response = Response(
 496          json.dumps({
 497              "workspaces": [
 498                  {"name": default_workspace},
 499                  {"name": "other-workspace"},
 500              ]
 501          }),
 502          mimetype="application/json",
 503      )
 504  
 505      auth_module.filter_list_workspaces(response)
 506      payload = json.loads(response.get_data(as_text=True))
 507      assert payload["workspaces"] == [{"name": default_workspace}]
 508  
 509  
 510  def test_filter_list_workspaces_filters_to_allowed(monkeypatch):
 511      monkeypatch.setattr(auth_module, "sender_is_admin", lambda: False)
 512      auth = SimpleNamespace(username="alice")
 513      monkeypatch.setattr(auth_module, "authenticate_request", lambda: auth)
 514      monkeypatch.setattr(
 515          auth_module,
 516          "auth_config",
 517          auth_module.auth_config._replace(
 518              grant_default_workspace_access=False,
 519          ),
 520          raising=False,
 521      )
 522  
 523      class DummyStore:
 524          def list_accessible_workspace_names(self, username):
 525              return ["team-a"]
 526  
 527      monkeypatch.setattr(auth_module, "store", DummyStore(), raising=False)
 528  
 529      response = Response(
 530          json.dumps({"workspaces": [{"name": "team-a"}, {"name": "team-b"}]}),
 531          mimetype="application/json",
 532      )
 533  
 534      auth_module.filter_list_workspaces(response)
 535      payload = json.loads(response.get_data(as_text=True))
 536      assert [ws["name"] for ws in payload["workspaces"]] == ["team-a"]
 537  
 538  
 539  def test_list_workspaces_filters_to_role_assigned_workspaces(tmp_path, monkeypatch):
 540      # End-to-end guard for the list_accessible_workspace_names fix: alice has NO
 541      # legacy workspace_permissions rows — her only workspace membership is via a
 542      # role assignment in ws-alpha. The ListWorkspaces filter must treat that role
 543      # assignment as workspace visibility and surface ws-alpha but not ws-beta.
 544      # Before the fix, the legacy-only query returned an empty set and alice saw
 545      # no workspaces in the UI.
 546      monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true")
 547      monkeypatch.setattr(auth_module, "sender_is_admin", lambda: False)
 548      monkeypatch.setattr(
 549          auth_module,
 550          "auth_config",
 551          auth_module.auth_config._replace(grant_default_workspace_access=False),
 552          raising=False,
 553      )
 554  
 555      db_uri = f"sqlite:///{tmp_path / 'auth-store.db'}"
 556      auth_store = SqlAlchemyStore()
 557      auth_store.init_db(db_uri)
 558      monkeypatch.setattr(auth_module, "store", auth_store, raising=False)
 559  
 560      alice = auth_store.create_user("alice", "supersecurepassword", is_admin=False)
 561      role = auth_store.create_role(name="viewer", workspace="ws-alpha")
 562      auth_store.add_role_permission(role.id, "experiment", "*", READ.name)
 563      auth_store.assign_role_to_user(alice.id, role.id)
 564  
 565      monkeypatch.setattr(
 566          auth_module, "authenticate_request", lambda: SimpleNamespace(username="alice")
 567      )
 568  
 569      response = Response(
 570          json.dumps({"workspaces": [{"name": "ws-alpha"}, {"name": "ws-beta"}]}),
 571          mimetype="application/json",
 572      )
 573  
 574      auth_module.filter_list_workspaces(response)
 575      payload = json.loads(response.get_data(as_text=True))
 576      assert [ws["name"] for ws in payload["workspaces"]] == ["ws-alpha"]
 577  
 578      auth_store.engine.dispose()
 579  
 580  
 581  def test_validate_can_view_workspace_allows_default_autogrant(monkeypatch):
 582      monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true")
 583      monkeypatch.setattr(auth_module, "sender_is_admin", lambda: False)
 584      auth = SimpleNamespace(username="alice")
 585      monkeypatch.setattr(auth_module, "authenticate_request", lambda: auth)
 586      monkeypatch.setattr(
 587          auth_module,
 588          "auth_config",
 589          auth_module.auth_config._replace(
 590              grant_default_workspace_access=True,
 591              default_permission=READ.name,
 592          ),
 593          raising=False,
 594      )
 595  
 596      default_workspace = "team-default"
 597      monkeypatch.setattr(auth_module, "_get_workspace_store", lambda: None, raising=False)
 598      monkeypatch.setattr(
 599          auth_module,
 600          "get_default_workspace_optional",
 601          lambda *args, **kwargs: (SimpleNamespace(name=default_workspace), True),
 602          raising=False,
 603      )
 604  
 605      class DummyStore:
 606          def list_accessible_workspace_names(self, username):
 607              return []
 608  
 609      monkeypatch.setattr(auth_module, "store", DummyStore(), raising=False)
 610  
 611      with auth_module.app.test_request_context(
 612          f"/api/3.0/mlflow/workspaces/{default_workspace}", method="GET"
 613      ):
 614          request.view_args = {"workspace_name": default_workspace}
 615          assert auth_module.validate_can_view_workspace()
 616  
 617      with auth_module.app.test_request_context(
 618          "/api/3.0/mlflow/workspaces/other-team", method="GET"
 619      ):
 620          request.view_args = {"workspace_name": "other-team"}
 621          assert not auth_module.validate_can_view_workspace()
 622  
 623  
 624  def test_experiment_validators_allow_manage_permission(workspace_permission_setup):
 625      store = workspace_permission_setup["store"]
 626      username = workspace_permission_setup["username"]
 627      _set_workspace_permission(store, username, MANAGE.name)
 628  
 629      with auth_module.app.test_request_context(
 630          "/api/2.0/mlflow/experiments/get", method="GET", query_string={"experiment_id": "exp-1"}
 631      ):
 632          assert auth_module.validate_can_read_experiment()
 633          assert auth_module.validate_can_update_experiment()
 634          assert auth_module.validate_can_delete_experiment()
 635          assert auth_module.validate_can_manage_experiment()
 636  
 637      with auth_module.app.test_request_context(
 638          "/api/2.0/mlflow/experiments/get-by-name",
 639          method="GET",
 640          query_string={"experiment_name": "Primary Experiment"},
 641      ):
 642          assert auth_module.validate_can_read_experiment_by_name()
 643  
 644      with workspace_context.WorkspaceContext("team-a"):
 645          assert auth_module.validate_can_create_experiment()
 646  
 647  
 648  def test_experiment_validators_read_permission_blocks_writes(workspace_permission_setup):
 649      store = workspace_permission_setup["store"]
 650      username = workspace_permission_setup["username"]
 651      _set_workspace_permission(store, username, READ.name)
 652  
 653      with auth_module.app.test_request_context(
 654          "/api/2.0/mlflow/experiments/get", method="GET", query_string={"experiment_id": "exp-1"}
 655      ):
 656          assert auth_module.validate_can_read_experiment()
 657          assert not auth_module.validate_can_update_experiment()
 658          assert not auth_module.validate_can_delete_experiment()
 659          assert not auth_module.validate_can_manage_experiment()
 660  
 661      with auth_module.app.test_request_context(
 662          "/api/2.0/mlflow/experiments/get-by-name",
 663          method="GET",
 664          query_string={"experiment_name": "Primary Experiment"},
 665      ):
 666          assert auth_module.validate_can_read_experiment_by_name()
 667  
 668      with workspace_context.WorkspaceContext("team-a"):
 669          assert not auth_module.validate_can_create_experiment()
 670  
 671  
 672  def test_experiment_artifact_proxy_validators_respect_permissions(workspace_permission_setup):
 673      store = workspace_permission_setup["store"]
 674      username = workspace_permission_setup["username"]
 675      _set_workspace_permission(store, username, MANAGE.name)
 676  
 677      with auth_module.app.test_request_context(
 678          "/ajax-api/2.0/mlflow-artifacts/artifacts/1/path",
 679          method="GET",
 680      ):
 681          request.view_args = {"artifact_path": "1/path"}
 682          assert auth_module.validate_can_read_experiment_artifact_proxy()
 683          assert auth_module.validate_can_update_experiment_artifact_proxy()
 684          assert auth_module.validate_can_delete_experiment_artifact_proxy()
 685  
 686      _set_workspace_permission(store, username, READ.name)
 687  
 688      with auth_module.app.test_request_context(
 689          "/ajax-api/2.0/mlflow-artifacts/artifacts/1/path",
 690          method="GET",
 691      ):
 692          request.view_args = {"artifact_path": "1/path"}
 693          assert auth_module.validate_can_read_experiment_artifact_proxy()
 694          assert not auth_module.validate_can_update_experiment_artifact_proxy()
 695          assert not auth_module.validate_can_delete_experiment_artifact_proxy()
 696  
 697  
 698  def test_experiment_artifact_proxy_without_experiment_id_uses_workspace_permissions(
 699      workspace_permission_setup,
 700  ):
 701      store = workspace_permission_setup["store"]
 702      username = workspace_permission_setup["username"]
 703      _set_workspace_permission(store, username, READ.name)
 704  
 705      with auth_module.app.test_request_context(
 706          "/ajax-api/2.0/mlflow-artifacts/artifacts/uploads/path",
 707          method="GET",
 708      ):
 709          request.view_args = {"artifact_path": "uploads/path"}
 710          assert auth_module.validate_can_read_experiment_artifact_proxy()
 711          assert not auth_module.validate_can_update_experiment_artifact_proxy()
 712  
 713  
 714  def test_experiment_artifact_proxy_without_experiment_id_denied_without_workspace_permission(
 715      workspace_permission_setup,
 716  ):
 717      store = workspace_permission_setup["store"]
 718      username = workspace_permission_setup["username"]
 719      _set_workspace_permission(store, username, NO_PERMISSIONS.name)
 720  
 721      with auth_module.app.test_request_context(
 722          "/ajax-api/2.0/mlflow-artifacts/artifacts/uploads/path",
 723          method="GET",
 724      ):
 725          request.view_args = {"artifact_path": "uploads/path"}
 726          assert not auth_module.validate_can_read_experiment_artifact_proxy()
 727  
 728  
 729  def test_filter_experiment_ids_respects_workspace_permissions(
 730      workspace_permission_setup, monkeypatch
 731  ):
 732      store = workspace_permission_setup["store"]
 733      username = workspace_permission_setup["username"]
 734      monkeypatch.setattr(auth_module, "sender_is_admin", lambda: False)
 735  
 736      experiment_ids = ["exp-1", "exp-2"]
 737      assert auth_module.filter_experiment_ids(experiment_ids) == experiment_ids
 738  
 739      _set_workspace_permission(store, username, NO_PERMISSIONS.name)
 740      assert auth_module.filter_experiment_ids(experiment_ids) == []
 741  
 742  
 743  def test_filter_experiment_ids_role_wildcard_grant(workspace_permission_setup, monkeypatch):
 744      # Role granting experiment(*) in the active workspace should include all experiments.
 745      store = workspace_permission_setup["store"]
 746      username = workspace_permission_setup["username"]
 747      user_id = store.get_user(username).id
 748      monkeypatch.setattr(auth_module, "sender_is_admin", lambda: False)
 749      # Start from NO_PERMISSIONS: workspace fallback would exclude everything.
 750      _set_workspace_permission(store, username, NO_PERMISSIONS.name)
 751  
 752      role = store.create_role(name="exp-reader", workspace="team-a")
 753      store.add_role_permission(role.id, "experiment", "*", "READ")
 754      store.assign_role_to_user(user_id, role.id)
 755  
 756      token = workspace_context.set_server_request_workspace("team-a")
 757      try:
 758          assert auth_module.filter_experiment_ids(["exp-1", "exp-2"]) == ["exp-1", "exp-2"]
 759      finally:
 760          workspace_context._WORKSPACE.reset(token)
 761  
 762  
 763  def test_filter_experiment_ids_role_specific_grant(workspace_permission_setup, monkeypatch):
 764      # Role granting a specific experiment id should include that id only (plus direct grants).
 765      store = workspace_permission_setup["store"]
 766      username = workspace_permission_setup["username"]
 767      user_id = store.get_user(username).id
 768      monkeypatch.setattr(auth_module, "sender_is_admin", lambda: False)
 769      _set_workspace_permission(store, username, NO_PERMISSIONS.name)
 770  
 771      role = store.create_role(name="exp-1-reader", workspace="team-a")
 772      store.add_role_permission(role.id, "experiment", "exp-1", "READ")
 773      store.assign_role_to_user(user_id, role.id)
 774  
 775      token = workspace_context.set_server_request_workspace("team-a")
 776      try:
 777          # Only exp-1 (via role); exp-2 is filtered out.
 778          assert auth_module.filter_experiment_ids(["exp-1", "exp-2"]) == ["exp-1"]
 779      finally:
 780          workspace_context._WORKSPACE.reset(token)
 781  
 782  
 783  def test_filter_experiment_ids_workspace_scope_role(workspace_permission_setup, monkeypatch):
 784      # Role with (resource_type='workspace', '*', READ) should grant access to all experiments.
 785      store = workspace_permission_setup["store"]
 786      username = workspace_permission_setup["username"]
 787      user_id = store.get_user(username).id
 788      monkeypatch.setattr(auth_module, "sender_is_admin", lambda: False)
 789      _set_workspace_permission(store, username, NO_PERMISSIONS.name)
 790  
 791      role = store.create_role(name="ws-reader", workspace="team-a")
 792      store.add_role_permission(role.id, "workspace", "*", "READ")
 793      store.assign_role_to_user(user_id, role.id)
 794  
 795      token = workspace_context.set_server_request_workspace("team-a")
 796      try:
 797          assert auth_module.filter_experiment_ids(["exp-1", "exp-2"]) == ["exp-1", "exp-2"]
 798      finally:
 799          workspace_context._WORKSPACE.reset(token)
 800  
 801  
 802  def test_run_validators_allow_manage_permission(workspace_permission_setup):
 803      store = workspace_permission_setup["store"]
 804      username = workspace_permission_setup["username"]
 805      _set_workspace_permission(store, username, MANAGE.name)
 806  
 807      with auth_module.app.test_request_context(
 808          "/api/2.0/mlflow/runs/get", method="GET", query_string={"run_id": "run-1"}
 809      ):
 810          assert auth_module.validate_can_read_run()
 811          assert auth_module.validate_can_update_run()
 812          assert auth_module.validate_can_delete_run()
 813          assert auth_module.validate_can_manage_run()
 814  
 815  
 816  def test_run_validators_read_permission_blocks_writes(workspace_permission_setup):
 817      store = workspace_permission_setup["store"]
 818      username = workspace_permission_setup["username"]
 819      _set_workspace_permission(store, username, READ.name)
 820  
 821      with auth_module.app.test_request_context(
 822          "/api/2.0/mlflow/runs/get", method="GET", query_string={"run_id": "run-1"}
 823      ):
 824          assert auth_module.validate_can_read_run()
 825          assert not auth_module.validate_can_update_run()
 826          assert not auth_module.validate_can_delete_run()
 827          assert not auth_module.validate_can_manage_run()
 828  
 829  
 830  def test_logged_model_validators_respect_permissions(workspace_permission_setup):
 831      store = workspace_permission_setup["store"]
 832      username = workspace_permission_setup["username"]
 833  
 834      _set_workspace_permission(store, username, MANAGE.name)
 835      with auth_module.app.test_request_context(
 836          "/api/2.0/mlflow/logged-models/get",
 837          method="GET",
 838          query_string={"model_id": "model-1"},
 839      ):
 840          assert auth_module.validate_can_read_logged_model()
 841          assert auth_module.validate_can_update_logged_model()
 842          assert auth_module.validate_can_delete_logged_model()
 843          assert auth_module.validate_can_manage_logged_model()
 844  
 845      _set_workspace_permission(store, username, READ.name)
 846      with auth_module.app.test_request_context(
 847          "/api/2.0/mlflow/logged-models/get",
 848          method="GET",
 849          query_string={"model_id": "model-1"},
 850      ):
 851          assert auth_module.validate_can_read_logged_model()
 852          assert not auth_module.validate_can_update_logged_model()
 853          assert not auth_module.validate_can_delete_logged_model()
 854          assert not auth_module.validate_can_manage_logged_model()
 855  
 856  
 857  def test_scorer_validators_use_workspace_permissions(workspace_permission_setup):
 858      store = workspace_permission_setup["store"]
 859      username = workspace_permission_setup["username"]
 860      _set_workspace_permission(store, username, MANAGE.name)
 861  
 862      with auth_module.app.test_request_context(
 863          "/api/3.0/mlflow/scorers/get",
 864          method="GET",
 865          query_string={"experiment_id": "exp-1", "name": "score-1"},
 866      ):
 867          assert auth_module.validate_can_read_scorer()
 868          assert auth_module.validate_can_update_scorer()
 869          assert auth_module.validate_can_delete_scorer()
 870          assert auth_module.validate_can_manage_scorer()
 871  
 872      with auth_module.app.test_request_context(
 873          "/api/3.0/mlflow/scorers/permissions/create",
 874          method="POST",
 875          json={
 876              "experiment_id": "exp-1",
 877              "scorer_name": "score-1",
 878              "username": "bob",
 879              "permission": "READ",
 880          },
 881      ):
 882          assert auth_module.validate_can_manage_scorer_permission()
 883  
 884  
 885  def test_scorer_validators_read_permission_blocks_writes(workspace_permission_setup):
 886      store = workspace_permission_setup["store"]
 887      username = workspace_permission_setup["username"]
 888      _set_workspace_permission(store, username, READ.name)
 889  
 890      with auth_module.app.test_request_context(
 891          "/api/3.0/mlflow/scorers/get",
 892          method="GET",
 893          query_string={"experiment_id": "exp-1", "name": "score-1"},
 894      ):
 895          assert auth_module.validate_can_read_scorer()
 896          assert not auth_module.validate_can_update_scorer()
 897          assert not auth_module.validate_can_delete_scorer()
 898          assert not auth_module.validate_can_manage_scorer()
 899  
 900      with auth_module.app.test_request_context(
 901          "/api/3.0/mlflow/scorers/permissions/create",
 902          method="POST",
 903          json={
 904              "experiment_id": "exp-1",
 905              "scorer_name": "score-1",
 906              "username": "bob",
 907              "permission": "READ",
 908          },
 909      ):
 910          assert not auth_module.validate_can_manage_scorer_permission()
 911  
 912  
 913  def test_registered_model_validators_require_manage_for_writes(workspace_permission_setup):
 914      store = workspace_permission_setup["store"]
 915      username = workspace_permission_setup["username"]
 916  
 917      with workspace_context.WorkspaceContext("team-a"):
 918          _set_workspace_permission(store, username, MANAGE.name)
 919          with auth_module.app.test_request_context(
 920              "/api/2.0/mlflow/registered-models/get",
 921              method="GET",
 922              query_string={"name": "model-xyz"},
 923          ):
 924              assert auth_module.validate_can_read_registered_model()
 925              assert auth_module.validate_can_update_registered_model()
 926              assert auth_module.validate_can_delete_registered_model()
 927              assert auth_module.validate_can_manage_registered_model()
 928          perm = auth_module._workspace_permission(
 929              auth_module.authenticate_request().username, "team-a"
 930          )
 931          assert perm is not None
 932          assert perm.can_manage
 933          assert workspace_context.get_request_workspace() == "team-a"
 934          assert auth_module.validate_can_create_registered_model()
 935  
 936          _set_workspace_permission(store, username, READ.name)
 937          with auth_module.app.test_request_context(
 938              "/api/2.0/mlflow/registered-models/get",
 939              method="GET",
 940              query_string={"name": "model-xyz"},
 941          ):
 942              assert auth_module.validate_can_read_registered_model()
 943              assert not auth_module.validate_can_update_registered_model()
 944              assert not auth_module.validate_can_delete_registered_model()
 945              assert not auth_module.validate_can_manage_registered_model()
 946          assert not auth_module.validate_can_create_registered_model()
 947  
 948  
 949  def test_validate_can_view_workspace_requires_access(workspace_permission_setup):
 950      store = workspace_permission_setup["store"]
 951      username = workspace_permission_setup["username"]
 952  
 953      with auth_module.app.test_request_context(
 954          "/api/3.0/mlflow/workspaces/team-a",
 955          method="GET",
 956      ):
 957          request.view_args = {"workspace_name": "team-a"}
 958          assert auth_module.validate_can_view_workspace()
 959  
 960      store.delete_workspace_permission("team-a", username)
 961  
 962      with auth_module.app.test_request_context(
 963          "/api/3.0/mlflow/workspaces/team-a",
 964          method="GET",
 965      ):
 966          request.view_args = {"workspace_name": "team-a"}
 967          assert not auth_module.validate_can_view_workspace()
 968  
 969  
 970  def test_run_artifact_validators_use_workspace_permissions(workspace_permission_setup):
 971      with auth_module.app.test_request_context(
 972          GET_ARTIFACT,
 973          method="GET",
 974          query_string={"run_id": "run-1"},
 975      ):
 976          assert auth_module.validate_can_read_run_artifact()
 977  
 978      with auth_module.app.test_request_context(
 979          UPLOAD_ARTIFACT,
 980          method="POST",
 981          query_string={"run_id": "run-1"},
 982      ):
 983          assert auth_module.validate_can_update_run_artifact()
 984  
 985  
 986  def test_model_version_artifact_validator_uses_workspace_permissions(workspace_permission_setup):
 987      with auth_module.app.test_request_context(
 988          GET_MODEL_VERSION_ARTIFACT,
 989          method="GET",
 990          query_string={"name": "model-xyz"},
 991      ):
 992          assert auth_module.validate_can_read_model_version_artifact()
 993  
 994  
 995  def test_metric_history_bulk_validator_uses_workspace_permissions(workspace_permission_setup):
 996      with auth_module.app.test_request_context(
 997          GET_METRIC_HISTORY_BULK,
 998          method="GET",
 999          query_string=[("run_id", "run-1"), ("run_id", "run-2")],
1000      ):
1001          assert auth_module.validate_can_read_metric_history_bulk()
1002  
1003  
1004  def test_metric_history_bulk_interval_validator_uses_workspace_permissions(
1005      workspace_permission_setup,
1006  ):
1007      with auth_module.app.test_request_context(
1008          GET_METRIC_HISTORY_BULK_INTERVAL,
1009          method="GET",
1010          query_string=[
1011              ("run_ids", "run-1"),
1012              ("run_ids", "run-2"),
1013              ("metric_key", "loss"),
1014          ],
1015      ):
1016          assert auth_module.validate_can_read_metric_history_bulk_interval()
1017  
1018  
1019  def test_search_datasets_validator_uses_workspace_permissions(workspace_permission_setup):
1020      with auth_module.app.test_request_context(
1021          SEARCH_DATASETS,
1022          method="POST",
1023          json={"experiment_ids": ["exp-1", "exp-2"]},
1024      ):
1025          assert auth_module.validate_can_search_datasets()
1026  
1027  
1028  def test_create_promptlab_run_validator_uses_workspace_permissions(workspace_permission_setup):
1029      with auth_module.app.test_request_context(
1030          CREATE_PROMPTLAB_RUN,
1031          method="POST",
1032          json={"experiment_id": "exp-2"},
1033      ):
1034          assert auth_module.validate_can_create_promptlab_run()
1035  
1036  
1037  def test_trace_artifact_validator_uses_workspace_permissions(workspace_permission_setup):
1038      with auth_module.app.test_request_context(
1039          GET_TRACE_ARTIFACT,
1040          method="GET",
1041          query_string={"request_id": "trace-1"},
1042      ):
1043          assert auth_module.validate_can_read_trace_artifact()
1044  
1045  
1046  def test_experiment_artifact_proxy_without_workspaces_falls_back_to_default(monkeypatch):
1047      monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "false")
1048      monkeypatch.setattr(
1049          auth_module,
1050          "auth_config",
1051          auth_module.auth_config._replace(default_permission=READ.name),
1052          raising=False,
1053      )
1054      monkeypatch.setattr(
1055          auth_module,
1056          "authenticate_request",
1057          lambda: SimpleNamespace(username="carol"),
1058      )
1059  
1060      with auth_module.app.test_request_context(
1061          "/ajax-api/2.0/mlflow-artifacts/artifacts/uploads/path",
1062          method="GET",
1063      ):
1064          request.view_args = {"artifact_path": "uploads/path"}
1065          assert auth_module.validate_can_read_experiment_artifact_proxy()
1066  
1067  
1068  def test_run_artifact_validators_denied_without_workspace_permission(workspace_permission_setup):
1069      store = workspace_permission_setup["store"]
1070      username = workspace_permission_setup["username"]
1071      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1072  
1073      with auth_module.app.test_request_context(
1074          GET_ARTIFACT,
1075          method="GET",
1076          query_string={"run_id": "run-1"},
1077      ):
1078          assert not auth_module.validate_can_read_run_artifact()
1079  
1080      with auth_module.app.test_request_context(
1081          UPLOAD_ARTIFACT,
1082          method="POST",
1083          query_string={"run_id": "run-1"},
1084      ):
1085          assert not auth_module.validate_can_update_run_artifact()
1086  
1087  
1088  def test_model_version_artifact_validator_denied_without_workspace_permission(
1089      workspace_permission_setup,
1090  ):
1091      store = workspace_permission_setup["store"]
1092      username = workspace_permission_setup["username"]
1093      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1094  
1095      with auth_module.app.test_request_context(
1096          GET_MODEL_VERSION_ARTIFACT,
1097          method="GET",
1098          query_string={"name": "model-xyz"},
1099      ):
1100          assert not auth_module.validate_can_read_model_version_artifact()
1101  
1102  
1103  def test_metric_history_bulk_validator_denied_without_workspace_permission(
1104      workspace_permission_setup,
1105  ):
1106      store = workspace_permission_setup["store"]
1107      username = workspace_permission_setup["username"]
1108      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1109  
1110      with auth_module.app.test_request_context(
1111          GET_METRIC_HISTORY_BULK,
1112          method="GET",
1113          query_string=[("run_id", "run-1"), ("run_id", "run-2")],
1114      ):
1115          assert not auth_module.validate_can_read_metric_history_bulk()
1116  
1117  
1118  def test_metric_history_bulk_interval_validator_denied_without_workspace_permission(
1119      workspace_permission_setup,
1120  ):
1121      store = workspace_permission_setup["store"]
1122      username = workspace_permission_setup["username"]
1123      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1124  
1125      with auth_module.app.test_request_context(
1126          GET_METRIC_HISTORY_BULK_INTERVAL,
1127          method="GET",
1128          query_string=[
1129              ("run_ids", "run-1"),
1130              ("run_ids", "run-2"),
1131              ("metric_key", "loss"),
1132          ],
1133      ):
1134          assert not auth_module.validate_can_read_metric_history_bulk_interval()
1135  
1136  
1137  def test_search_datasets_validator_denied_without_workspace_permission(
1138      workspace_permission_setup,
1139  ):
1140      store = workspace_permission_setup["store"]
1141      username = workspace_permission_setup["username"]
1142      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1143  
1144      with auth_module.app.test_request_context(
1145          SEARCH_DATASETS,
1146          method="POST",
1147          json={"experiment_ids": ["exp-1", "exp-2"]},
1148      ):
1149          assert not auth_module.validate_can_search_datasets()
1150  
1151  
1152  def test_create_promptlab_run_validator_denied_without_workspace_permission(
1153      workspace_permission_setup,
1154  ):
1155      store = workspace_permission_setup["store"]
1156      username = workspace_permission_setup["username"]
1157      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1158  
1159      with auth_module.app.test_request_context(
1160          CREATE_PROMPTLAB_RUN,
1161          method="POST",
1162          json={"experiment_id": "exp-2"},
1163      ):
1164          assert not auth_module.validate_can_create_promptlab_run()
1165  
1166  
1167  def test_trace_artifact_validator_denied_without_workspace_permission(
1168      workspace_permission_setup,
1169  ):
1170      store = workspace_permission_setup["store"]
1171      username = workspace_permission_setup["username"]
1172      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1173  
1174      with auth_module.app.test_request_context(
1175          GET_TRACE_ARTIFACT,
1176          method="GET",
1177          query_string={"request_id": "trace-1"},
1178      ):
1179          assert not auth_module.validate_can_read_trace_artifact()
1180  
1181  
1182  def test_cross_workspace_access_denied(workspace_permission_setup, monkeypatch):
1183      tracking_store = _TrackingStore(
1184          experiment_workspaces={"exp-other-ws": "team-b"},
1185          run_experiments={"run-other-ws": "exp-other-ws"},
1186          trace_experiments={},
1187      )
1188      monkeypatch.setattr(auth_module, "_get_tracking_store", lambda: tracking_store)
1189  
1190      with auth_module.app.test_request_context(
1191          "/api/2.0/mlflow/experiments/get",
1192          method="GET",
1193          query_string={"experiment_id": "exp-other-ws"},
1194      ):
1195          assert not auth_module.validate_can_read_experiment()
1196          assert not auth_module.validate_can_update_experiment()
1197          assert not auth_module.validate_can_delete_experiment()
1198  
1199      with auth_module.app.test_request_context(
1200          "/api/2.0/mlflow/runs/get",
1201          method="GET",
1202          query_string={"run_id": "run-other-ws"},
1203      ):
1204          assert not auth_module.validate_can_read_run()
1205          assert not auth_module.validate_can_update_run()
1206  
1207  
1208  def test_cross_workspace_registered_model_access_denied(workspace_permission_setup, monkeypatch):
1209      registry_store = _RegistryStore({"model-other-ws": "team-b"})
1210      monkeypatch.setattr(auth_module, "_get_model_registry_store", lambda: registry_store)
1211  
1212      with auth_module.app.test_request_context(
1213          "/api/2.0/mlflow/registered-models/get",
1214          method="GET",
1215          query_string={"name": "model-other-ws"},
1216      ):
1217          assert not auth_module.validate_can_read_registered_model()
1218          assert not auth_module.validate_can_update_registered_model()
1219          assert not auth_module.validate_can_delete_registered_model()
1220  
1221  
1222  def test_explicit_experiment_permission_overrides_workspace(
1223      workspace_permission_setup,
1224  ):
1225      store = workspace_permission_setup["store"]
1226      username = workspace_permission_setup["username"]
1227  
1228      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1229      store.create_experiment_permission("exp-1", username, READ.name)
1230  
1231      with auth_module.app.test_request_context(
1232          "/api/2.0/mlflow/experiments/get",
1233          method="GET",
1234          query_string={"experiment_id": "exp-1"},
1235      ):
1236          assert auth_module.validate_can_read_experiment()
1237          assert not auth_module.validate_can_update_experiment()
1238  
1239      with auth_module.app.test_request_context(
1240          "/api/2.0/mlflow/experiments/get",
1241          method="GET",
1242          query_string={"experiment_id": "exp-2"},
1243      ):
1244          assert not auth_module.validate_can_read_experiment()
1245  
1246  
1247  def test_cross_workspace_gateway_secret_access_denied(workspace_permission_setup, monkeypatch):
1248      tracking_store = _TrackingStore(
1249          experiment_workspaces={"exp-1": "team-a"},
1250          run_experiments={},
1251          trace_experiments={},
1252          gateway_secret_workspaces={"secret-other-ws": "team-b"},
1253          engine=MagicMock(),
1254      )
1255      tracking_store.ManagedSessionMaker = tracking_store._create_mock_session_maker()
1256      monkeypatch.setattr(auth_module, "_get_tracking_store", lambda: tracking_store)
1257  
1258      with auth_module.app.test_request_context(
1259          "/api/3.0/mlflow/gateway/secrets/get",
1260          method="GET",
1261          query_string={"secret_id": "secret-other-ws"},
1262      ):
1263          assert not auth_module.validate_can_read_gateway_secret()
1264          assert not auth_module.validate_can_update_gateway_secret()
1265          assert not auth_module.validate_can_delete_gateway_secret()
1266  
1267  
1268  def test_cross_workspace_gateway_endpoint_access_denied(workspace_permission_setup, monkeypatch):
1269      tracking_store = _TrackingStore(
1270          experiment_workspaces={"exp-1": "team-a"},
1271          run_experiments={},
1272          trace_experiments={},
1273          gateway_endpoint_workspaces={"endpoint-other-ws": "team-b"},
1274          engine=MagicMock(),
1275      )
1276      tracking_store.ManagedSessionMaker = tracking_store._create_mock_session_maker()
1277      monkeypatch.setattr(auth_module, "_get_tracking_store", lambda: tracking_store)
1278  
1279      with auth_module.app.test_request_context(
1280          "/api/3.0/mlflow/gateway/endpoints/get",
1281          method="GET",
1282          query_string={"endpoint_id": "endpoint-other-ws"},
1283      ):
1284          assert not auth_module.validate_can_read_gateway_endpoint()
1285          assert not auth_module.validate_can_update_gateway_endpoint()
1286          assert not auth_module.validate_can_delete_gateway_endpoint()
1287  
1288  
1289  def test_cross_workspace_gateway_model_definition_access_denied(
1290      workspace_permission_setup, monkeypatch
1291  ):
1292      tracking_store = _TrackingStore(
1293          experiment_workspaces={"exp-1": "team-a"},
1294          run_experiments={},
1295          trace_experiments={},
1296          gateway_model_def_workspaces={"model-def-other-ws": "team-b"},
1297          engine=MagicMock(),
1298      )
1299      tracking_store.ManagedSessionMaker = tracking_store._create_mock_session_maker()
1300      monkeypatch.setattr(auth_module, "_get_tracking_store", lambda: tracking_store)
1301  
1302      with auth_module.app.test_request_context(
1303          "/api/3.0/mlflow/gateway/model-definitions/get",
1304          method="GET",
1305          query_string={"model_definition_id": "model-def-other-ws"},
1306      ):
1307          assert not auth_module.validate_can_read_gateway_model_definition()
1308          assert not auth_module.validate_can_update_gateway_model_definition()
1309          assert not auth_module.validate_can_delete_gateway_model_definition()
1310  
1311  
1312  def test_workspace_permission_required_for_gateway_creation(workspace_permission_setup):
1313      store = workspace_permission_setup["store"]
1314      username = workspace_permission_setup["username"]
1315  
1316      # Remove workspace permission
1317      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1318  
1319      with auth_module.app.test_request_context(
1320          "/api/3.0/mlflow/gateway/endpoints/create",
1321          method="POST",
1322          json={"name": "test-endpoint", "model_configs": []},
1323      ):
1324          assert not auth_module.validate_can_create_gateway_endpoint()
1325  
1326      with auth_module.app.test_request_context(
1327          "/api/3.0/mlflow/gateway/model-definitions/create",
1328          method="POST",
1329          json={
1330              "name": "test-model",
1331              "secret_id": "secret-1",
1332              "provider": "openai",
1333              "model_name": "gpt-4",
1334          },
1335      ):
1336          assert not auth_module.validate_can_create_gateway_model_definition()
1337  
1338      # Restore workspace permission
1339      store.set_workspace_permission("team-a", username, MANAGE.name)
1340  
1341      with auth_module.app.test_request_context(
1342          "/api/3.0/mlflow/gateway/endpoints/create",
1343          method="POST",
1344          json={"name": "test-endpoint", "model_configs": []},
1345      ):
1346          assert auth_module.validate_can_create_gateway_endpoint()
1347  
1348  
1349  def test_prompt_optimization_job_validators_use_workspace_permissions(
1350      workspace_permission_setup, monkeypatch
1351  ):
1352      store = workspace_permission_setup["store"]
1353      username = workspace_permission_setup["username"]
1354  
1355      # Mock get_job to return a job associated with exp-1 (in team-a)
1356      mock_job = SimpleNamespace(params='{"experiment_id": "exp-1"}')
1357      monkeypatch.setattr(auth_module, "get_job", lambda job_id: mock_job)
1358  
1359      _set_workspace_permission(store, username, MANAGE.name)
1360  
1361      with auth_module.app.test_request_context(
1362          "/api/3.0/mlflow/prompt-optimization/jobs/get",
1363          method="GET",
1364          query_string={"job_id": "job-1"},
1365      ):
1366          assert auth_module.validate_can_read_prompt_optimization_job()
1367          assert auth_module.validate_can_update_prompt_optimization_job()
1368          assert auth_module.validate_can_delete_prompt_optimization_job()
1369  
1370  
1371  def test_prompt_optimization_job_validators_read_permission_blocks_writes(
1372      workspace_permission_setup, monkeypatch
1373  ):
1374      store = workspace_permission_setup["store"]
1375      username = workspace_permission_setup["username"]
1376  
1377      # Mock get_job to return a job associated with exp-1 (in team-a)
1378      mock_job = SimpleNamespace(params='{"experiment_id": "exp-1"}')
1379      monkeypatch.setattr(auth_module, "get_job", lambda job_id: mock_job)
1380  
1381      _set_workspace_permission(store, username, READ.name)
1382  
1383      with auth_module.app.test_request_context(
1384          "/api/3.0/mlflow/prompt-optimization/jobs/get",
1385          method="GET",
1386          query_string={"job_id": "job-1"},
1387      ):
1388          assert auth_module.validate_can_read_prompt_optimization_job()
1389          assert not auth_module.validate_can_update_prompt_optimization_job()
1390          assert not auth_module.validate_can_delete_prompt_optimization_job()
1391  
1392  
1393  def test_prompt_optimization_job_validators_denied_without_workspace_permission(
1394      workspace_permission_setup, monkeypatch
1395  ):
1396      store = workspace_permission_setup["store"]
1397      username = workspace_permission_setup["username"]
1398  
1399      # Mock get_job to return a job associated with exp-1 (in team-a)
1400      mock_job = SimpleNamespace(params='{"experiment_id": "exp-1"}')
1401      monkeypatch.setattr(auth_module, "get_job", lambda job_id: mock_job)
1402  
1403      _set_workspace_permission(store, username, NO_PERMISSIONS.name)
1404  
1405      with auth_module.app.test_request_context(
1406          "/api/3.0/mlflow/prompt-optimization/jobs/get",
1407          method="GET",
1408          query_string={"job_id": "job-1"},
1409      ):
1410          assert not auth_module.validate_can_read_prompt_optimization_job()
1411          assert not auth_module.validate_can_update_prompt_optimization_job()
1412          assert not auth_module.validate_can_delete_prompt_optimization_job()
1413  
1414  
1415  def test_graphql_permission_functions_use_workspace_permissions(workspace_permission_setup):
1416      store = workspace_permission_setup["store"]
1417      username = workspace_permission_setup["username"]
1418  
1419      _set_workspace_permission(store, username, MANAGE.name)
1420  
1421      # Test experiment permission
1422      assert auth_module._graphql_can_read_experiment("exp-1", username)
1423  
1424      # Test run permission (inherits from experiment)
1425      assert auth_module._graphql_can_read_run("run-1", username)
1426  
1427      # Test registered model permission
1428      assert auth_module._graphql_can_read_model("model-xyz", username)
1429  
1430  
1431  def test_graphql_permission_functions_denied_without_workspace_permission(
1432      workspace_permission_setup,
1433  ):
1434      store = workspace_permission_setup["store"]
1435      username = workspace_permission_setup["username"]
1436  
1437      _set_workspace_permission(store, username, NO_PERMISSIONS.name)
1438  
1439      # Test experiment permission denied
1440      assert not auth_module._graphql_can_read_experiment("exp-1", username)
1441  
1442      # Test run permission denied (inherits from experiment)
1443      assert not auth_module._graphql_can_read_run("run-1", username)
1444  
1445      # Test registered model permission denied
1446      assert not auth_module._graphql_can_read_model("model-xyz", username)
1447  
1448  
1449  def test_cross_workspace_graphql_access_denied(workspace_permission_setup, monkeypatch):
1450      # User has MANAGE in team-a but tries to access resources in team-b
1451      tracking_store = _TrackingStore(
1452          experiment_workspaces={"exp-other-ws": "team-b"},
1453          run_experiments={"run-other-ws": "exp-other-ws"},
1454          trace_experiments={},
1455      )
1456      monkeypatch.setattr(auth_module, "_get_tracking_store", lambda: tracking_store)
1457  
1458      registry_store = _RegistryStore({"model-other-ws": "team-b"})
1459      monkeypatch.setattr(auth_module, "_get_model_registry_store", lambda: registry_store)
1460  
1461      username = workspace_permission_setup["username"]
1462  
1463      # Should be denied access to resources in team-b
1464      assert not auth_module._graphql_can_read_experiment("exp-other-ws", username)
1465      assert not auth_module._graphql_can_read_run("run-other-ws", username)
1466      assert not auth_module._graphql_can_read_model("model-other-ws", username)
1467  
1468  
1469  # =============================================================================
1470  # Role-based permission coverage for gateway resources
1471  # =============================================================================
1472  #
1473  # The fixture grants workspace MANAGE by default. These tests first strip that
1474  # grant (set to NO_PERMISSIONS) so the only path to a positive permission is
1475  # the role assignment being exercised. That isolates the role-based resolver
1476  # from the legacy workspace_permissions fallback.
1477  
1478  
1479  def _assign_role_with_permission(
1480      store: SqlAlchemyStore, username: str, workspace: str, resource_type: str, permission: str
1481  ) -> None:
1482      """Create a role in ``workspace`` with a wildcard grant of ``permission`` on
1483      ``resource_type``, and assign ``username`` to it.
1484  
1485      Using ``random_str`` keeps the role names unique so multiple calls within a
1486      single test don't collide on the (workspace, name) unique constraint.
1487      """
1488      role = store.create_role(name=random_str(), workspace=workspace)
1489      store.add_role_permission(role.id, resource_type, "*", permission)
1490      user = store.get_user(username)
1491      store.assign_role_to_user(user.id, role.id)
1492  
1493  
1494  # ---- Gateway endpoint: role-based permission levels ----
1495  
1496  
1497  @pytest.mark.parametrize(
1498      ("granted", "expected_read", "expected_delete", "expected_manage"),
1499      [
1500          ("READ", True, False, False),
1501          ("USE", True, False, False),
1502          ("EDIT", True, False, False),
1503          ("MANAGE", True, True, True),
1504      ],
1505  )
1506  def test_role_grant_on_gateway_endpoint_gates_validator_capabilities(
1507      workspace_permission_setup, granted, expected_read, expected_delete, expected_manage
1508  ):
1509      """A role grant at permission level ``granted`` exposes exactly the
1510      capabilities that level implies on the endpoint validators — no more, no
1511      less. Catches regressions where a validator starts accepting a weaker
1512      permission than it should (or refuses a stronger one).
1513      """
1514      store = workspace_permission_setup["store"]
1515      username = workspace_permission_setup["username"]
1516  
1517      # Strip the default workspace MANAGE so the only positive grant is the role.
1518      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1519  
1520      _assign_role_with_permission(store, username, "team-a", "gateway_endpoint", granted)
1521  
1522      with auth_module.app.test_request_context(
1523          "/api/3.0/mlflow/gateway/endpoints/get",
1524          method="GET",
1525          query_string={"endpoint_id": "endpoint-1"},
1526      ):
1527          assert auth_module.validate_can_read_gateway_endpoint() is expected_read
1528          assert auth_module.validate_can_delete_gateway_endpoint() is expected_delete
1529          assert auth_module.validate_can_manage_gateway_endpoint() is expected_manage
1530  
1531  
1532  def test_role_grant_read_on_gateway_endpoint_does_not_permit_use(
1533      workspace_permission_setup,
1534  ):
1535      """Regression guard specific to the bug class the user called out:
1536      a user with only READ on a gateway endpoint should not be able to *invoke*
1537      it. USE is a stricter capability than READ and has its own validator.
1538      """
1539      store = workspace_permission_setup["store"]
1540      username = workspace_permission_setup["username"]
1541      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1542  
1543      _assign_role_with_permission(store, username, "team-a", "gateway_endpoint", "READ")
1544  
1545      # _validate_gateway_use_permission looks up the endpoint by name, resolves
1546      # the endpoint id, then checks ``can_use`` via the permission resolver.
1547      with auth_module.app.test_request_context("/"):
1548          assert auth_module._validate_gateway_use_permission("endpoint-1", username) is False
1549  
1550  
1551  def test_role_grant_use_on_gateway_endpoint_permits_use(workspace_permission_setup):
1552      store = workspace_permission_setup["store"]
1553      username = workspace_permission_setup["username"]
1554      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1555  
1556      _assign_role_with_permission(store, username, "team-a", "gateway_endpoint", "USE")
1557  
1558      with auth_module.app.test_request_context("/"):
1559          assert auth_module._validate_gateway_use_permission("endpoint-1", username) is True
1560  
1561  
1562  @pytest.mark.parametrize(
1563      ("granted", "expected_can_use"),
1564      [
1565          ("READ", False),  # READ does not imply USE.
1566          ("USE", True),
1567          ("EDIT", True),  # EDIT implies USE.
1568          ("MANAGE", True),  # MANAGE implies USE.
1569      ],
1570  )
1571  def test_role_grant_permission_level_determines_use_capability(
1572      workspace_permission_setup, granted, expected_can_use
1573  ):
1574      """Parametrized matrix for the USE capability specifically. READ should NOT
1575      let the user invoke; every stronger permission should.
1576      """
1577      store = workspace_permission_setup["store"]
1578      username = workspace_permission_setup["username"]
1579      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1580  
1581      _assign_role_with_permission(store, username, "team-a", "gateway_endpoint", granted)
1582  
1583      with auth_module.app.test_request_context("/"):
1584          assert (
1585              auth_module._validate_gateway_use_permission("endpoint-1", username) is expected_can_use
1586          )
1587  
1588  
1589  # ---- Workspace-wide role grants on gateway resources ----
1590  
1591  
1592  @pytest.mark.parametrize("granted", ["READ", "USE", "EDIT", "MANAGE"])
1593  def test_role_workspace_wide_grant_applies_to_gateway_endpoints(
1594      workspace_permission_setup, granted
1595  ):
1596      """``(workspace, *, X)`` grants apply to every resource type in the
1597      workspace — including gateway endpoints. Confirms the workspace-wide
1598      short-circuit isn't accidentally gated behind resource_type=='experiment'
1599      or similar, which would silently lock workspace admins out of gateway
1600      resources.
1601      """
1602      store = workspace_permission_setup["store"]
1603      username = workspace_permission_setup["username"]
1604      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1605  
1606      _assign_role_with_permission(store, username, "team-a", "workspace", granted)
1607  
1608      with auth_module.app.test_request_context(
1609          "/api/3.0/mlflow/gateway/endpoints/get",
1610          method="GET",
1611          query_string={"endpoint_id": "endpoint-1"},
1612      ):
1613          # All four levels grant READ.
1614          assert auth_module.validate_can_read_gateway_endpoint() is True
1615  
1616          # Only MANAGE grants can_delete / can_manage.
1617          assert auth_module.validate_can_delete_gateway_endpoint() is (granted == "MANAGE")
1618          assert auth_module.validate_can_manage_gateway_endpoint() is (granted == "MANAGE")
1619  
1620  
1621  def test_role_workspace_wide_read_does_not_imply_use_on_gateway_endpoint(
1622      workspace_permission_setup,
1623  ):
1624      """``(workspace, *, READ)`` grants READ on every resource but not USE.
1625      Users with a workspace-wide viewer role shouldn't be able to invoke
1626      gateway endpoints.
1627      """
1628      store = workspace_permission_setup["store"]
1629      username = workspace_permission_setup["username"]
1630      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1631  
1632      _assign_role_with_permission(store, username, "team-a", "workspace", "READ")
1633  
1634      with auth_module.app.test_request_context("/"):
1635          assert auth_module._validate_gateway_use_permission("endpoint-1", username) is False
1636  
1637  
1638  @pytest.mark.parametrize("granted", ["USE", "EDIT", "MANAGE"])
1639  def test_role_workspace_wide_non_read_grants_imply_use_on_gateway_endpoint(
1640      workspace_permission_setup, granted
1641  ):
1642      """``(workspace, *, {USE,EDIT,MANAGE})`` all imply USE → invocation
1643      allowed.
1644      """
1645      store = workspace_permission_setup["store"]
1646      username = workspace_permission_setup["username"]
1647      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1648  
1649      _assign_role_with_permission(store, username, "team-a", "workspace", granted)
1650  
1651      with auth_module.app.test_request_context("/"):
1652          assert auth_module._validate_gateway_use_permission("endpoint-1", username) is True
1653  
1654  
1655  # ---- Gateway secret and model definition parity ----
1656  
1657  
1658  @pytest.mark.parametrize(
1659      ("granted", "expected_read", "expected_delete"),
1660      [
1661          ("READ", True, False),
1662          ("EDIT", True, False),
1663          ("MANAGE", True, True),
1664      ],
1665  )
1666  def test_role_grant_on_gateway_secret_gates_validator(
1667      workspace_permission_setup, granted, expected_read, expected_delete
1668  ):
1669      store = workspace_permission_setup["store"]
1670      username = workspace_permission_setup["username"]
1671      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1672  
1673      _assign_role_with_permission(store, username, "team-a", "gateway_secret", granted)
1674  
1675      with auth_module.app.test_request_context(
1676          "/api/3.0/mlflow/gateway/secrets/get",
1677          method="GET",
1678          query_string={"secret_id": "secret-1"},
1679      ):
1680          assert auth_module.validate_can_read_gateway_secret() is expected_read
1681          assert auth_module.validate_can_delete_gateway_secret() is expected_delete
1682  
1683  
1684  @pytest.mark.parametrize(
1685      ("granted", "expected_read", "expected_delete"),
1686      [
1687          ("READ", True, False),
1688          ("EDIT", True, False),
1689          ("MANAGE", True, True),
1690      ],
1691  )
1692  def test_role_grant_on_gateway_model_definition_gates_validator(
1693      workspace_permission_setup, granted, expected_read, expected_delete
1694  ):
1695      store = workspace_permission_setup["store"]
1696      username = workspace_permission_setup["username"]
1697      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1698  
1699      _assign_role_with_permission(store, username, "team-a", "gateway_model_definition", granted)
1700  
1701      with auth_module.app.test_request_context(
1702          "/api/3.0/mlflow/gateway/model-definitions/get",
1703          method="GET",
1704          query_string={"model_definition_id": "model-def-1"},
1705      ):
1706          assert auth_module.validate_can_read_gateway_model_definition() is expected_read
1707          assert auth_module.validate_can_delete_gateway_model_definition() is expected_delete
1708  
1709  
1710  # ---- Cross-workspace isolation for role-based gateway grants ----
1711  
1712  
1713  def test_role_in_other_workspace_does_not_grant_gateway_endpoint_access(
1714      workspace_permission_setup,
1715  ):
1716      """A role in team-b with MANAGE on gateway_endpoints must not grant any
1717      access when resolving an endpoint that belongs to team-a. The resolver
1718      scopes role permissions to the role's workspace.
1719      """
1720      store = workspace_permission_setup["store"]
1721      username = workspace_permission_setup["username"]
1722      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1723  
1724      # Role with MANAGE in team-b — should NOT apply to team-a endpoints.
1725      _assign_role_with_permission(store, username, "team-b", "gateway_endpoint", "MANAGE")
1726  
1727      with auth_module.app.test_request_context(
1728          "/api/3.0/mlflow/gateway/endpoints/get",
1729          method="GET",
1730          query_string={"endpoint_id": "endpoint-1"},  # endpoint-1 is in team-a.
1731      ):
1732          assert auth_module.validate_can_read_gateway_endpoint() is False
1733          assert auth_module.validate_can_manage_gateway_endpoint() is False
1734  
1735  
1736  def test_role_in_other_workspace_does_not_grant_gateway_use(workspace_permission_setup):
1737      store = workspace_permission_setup["store"]
1738      username = workspace_permission_setup["username"]
1739      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1740  
1741      _assign_role_with_permission(store, username, "team-b", "gateway_endpoint", "USE")
1742  
1743      with auth_module.app.test_request_context("/"):
1744          # endpoint-1 is in team-a; role grant is in team-b.
1745          assert auth_module._validate_gateway_use_permission("endpoint-1", username) is False
1746  
1747  
1748  # ---- Multi-role union: best grant wins ----
1749  
1750  
1751  def test_role_union_best_permission_wins_for_gateway_endpoint(workspace_permission_setup):
1752      """Two roles: one grants READ, the other grants MANAGE. Validator should
1753      reflect the max (MANAGE).
1754      """
1755      store = workspace_permission_setup["store"]
1756      username = workspace_permission_setup["username"]
1757      store.set_workspace_permission("team-a", username, NO_PERMISSIONS.name)
1758  
1759      _assign_role_with_permission(store, username, "team-a", "gateway_endpoint", "READ")
1760      _assign_role_with_permission(store, username, "team-a", "gateway_endpoint", "MANAGE")
1761  
1762      with auth_module.app.test_request_context(
1763          "/api/3.0/mlflow/gateway/endpoints/get",
1764          method="GET",
1765          query_string={"endpoint_id": "endpoint-1"},
1766      ):
1767          assert auth_module.validate_can_manage_gateway_endpoint() is True
1768  
1769  
1770  # =============================================================================
1771  # Authorization for role management endpoints (Batch 5)
1772  # =============================================================================
1773  #
1774  # Four validators guard the role endpoints:
1775  #   - validate_can_manage_roles: create/update/delete role, add/remove/update
1776  #     role_permission, assign/unassign role. Super admin OR workspace admin
1777  #     in the resolved workspace.
1778  #   - validate_can_view_roles: get_role, list_role_permissions. Super admin
1779  #     OR any role assignment in the resolved workspace.
1780  #   - validate_can_list_roles: list_roles. Super admin unconditionally; for
1781  #     non-admins the request must scope to a workspace where the caller holds
1782  #     at least one role.
1783  #   - validate_can_view_user_roles: list_user_roles. Super admin, the target
1784  #     themselves, or a workspace admin over any workspace the target is in.
1785  #
1786  # _get_role_workspace_from_request resolves the workspace from role_id,
1787  # role_permission_id, or a literal ``workspace`` param. These tests exercise
1788  # all three shapes and every actor x endpoint combination.
1789  
1790  
1791  @pytest.fixture
1792  def role_auth_setup(tmp_path, monkeypatch):
1793      monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true")
1794      monkeypatch.setattr(
1795          auth_module,
1796          "auth_config",
1797          auth_module.auth_config._replace(default_permission=NO_PERMISSIONS.name),
1798      )
1799  
1800      db_uri = f"sqlite:///{tmp_path / 'auth-store.db'}"
1801      auth_store = SqlAlchemyStore()
1802      auth_store.init_db(db_uri)
1803      monkeypatch.setattr(auth_module, "store", auth_store, raising=False)
1804  
1805      auth_store.create_user("super_admin", "supersecurepassword", is_admin=True)
1806      for name in ("ws_admin_foo", "ws_admin_bar", "ws_member_foo", "outsider"):
1807          auth_store.create_user(name, "supersecurepassword", is_admin=False)
1808  
1809      admin_role_foo = auth_store.create_role(name="admin-foo", workspace="foo")
1810      auth_store.add_role_permission(admin_role_foo.id, "workspace", "*", MANAGE.name)
1811      auth_store.assign_role_to_user(auth_store.get_user("ws_admin_foo").id, admin_role_foo.id)
1812  
1813      admin_role_bar = auth_store.create_role(name="admin-bar", workspace="bar")
1814      auth_store.add_role_permission(admin_role_bar.id, "workspace", "*", MANAGE.name)
1815      auth_store.assign_role_to_user(auth_store.get_user("ws_admin_bar").id, admin_role_bar.id)
1816  
1817      member_role_foo = auth_store.create_role(name="member-foo", workspace="foo")
1818      auth_store.add_role_permission(member_role_foo.id, "experiment", "*", READ.name)
1819      auth_store.assign_role_to_user(auth_store.get_user("ws_member_foo").id, member_role_foo.id)
1820  
1821      role_foo = auth_store.create_role(name="target-foo", workspace="foo")
1822      role_bar = auth_store.create_role(name="target-bar", workspace="bar")
1823      rp_foo = auth_store.add_role_permission(role_foo.id, "experiment", "*", READ.name)
1824      rp_bar = auth_store.add_role_permission(role_bar.id, "experiment", "*", READ.name)
1825  
1826      def login_as(username: str) -> None:
1827          monkeypatch.setattr(
1828              auth_module,
1829              "authenticate_request",
1830              lambda: SimpleNamespace(username=username),
1831          )
1832  
1833      yield {
1834          "store": auth_store,
1835          "login_as": login_as,
1836          "role_foo_id": role_foo.id,
1837          "role_bar_id": role_bar.id,
1838          "role_permission_foo_id": rp_foo.id,
1839          "role_permission_bar_id": rp_bar.id,
1840      }
1841      auth_store.engine.dispose()
1842  
1843  
1844  def _request_context_for_shape(shape, role_auth_setup, workspace):
1845      match shape:
1846          case "role_id":
1847              role_id = (
1848                  role_auth_setup["role_foo_id"]
1849                  if workspace == "foo"
1850                  else role_auth_setup["role_bar_id"]
1851              )
1852              return auth_module.app.test_request_context(
1853                  "/api/3.0/mlflow/roles/get",
1854                  method="GET",
1855                  query_string={"role_id": str(role_id)},
1856              )
1857          case "role_permission_id":
1858              rp_id = (
1859                  role_auth_setup["role_permission_foo_id"]
1860                  if workspace == "foo"
1861                  else role_auth_setup["role_permission_bar_id"]
1862              )
1863              return auth_module.app.test_request_context(
1864                  "/api/3.0/mlflow/roles/permissions/update",
1865                  method="PATCH",
1866                  json={"role_permission_id": rp_id, "permission": READ.name},
1867              )
1868          case "workspace":
1869              return auth_module.app.test_request_context(
1870                  "/api/3.0/mlflow/roles/create",
1871                  method="POST",
1872                  json={"name": "new-role", "workspace": workspace},
1873              )
1874          case _:
1875              raise ValueError(f"Unknown shape: {shape}")
1876  
1877  
1878  # Authorization matrices are exercised with a single request shape (role_id);
1879  # shape-resolution itself is covered independently below so we don't multiply
1880  # every actor-case by three shape-cases.
1881  
1882  
1883  @pytest.mark.parametrize(
1884      ("actor", "workspace", "expected"),
1885      [
1886          # Super admin short-circuits regardless of workspace — one case suffices.
1887          ("super_admin", "foo", True),
1888          # Outsider has no role anywhere — one case suffices.
1889          ("outsider", "foo", False),
1890          # Workspace admins manage only their own workspace.
1891          ("ws_admin_foo", "foo", True),
1892          ("ws_admin_foo", "bar", False),
1893          ("ws_admin_bar", "foo", False),
1894          ("ws_admin_bar", "bar", True),
1895          # Plain role membership is not enough to manage — needs workspace MANAGE.
1896          ("ws_member_foo", "foo", False),
1897          ("ws_member_foo", "bar", False),
1898      ],
1899  )
1900  def test_validate_can_manage_roles_authorization(role_auth_setup, actor, workspace, expected):
1901      role_auth_setup["login_as"](actor)
1902      with _request_context_for_shape("role_id", role_auth_setup, workspace):
1903          assert auth_module.validate_can_manage_roles() is expected
1904  
1905  
1906  @pytest.mark.parametrize("workspace", ["foo", "bar"])
1907  @pytest.mark.parametrize("shape", ["role_id", "role_permission_id", "workspace"])
1908  def test_manage_roles_resolves_workspace_from_each_shape(role_auth_setup, shape, workspace):
1909      # Sanity check that _get_role_workspace_from_request correctly dispatches
1910      # on every request shape. Use ws_admin_foo — their answer differs by
1911      # workspace, so an incorrectly resolved (or swapped) workspace flips the
1912      # result and the test fails.
1913      role_auth_setup["login_as"]("ws_admin_foo")
1914      expected = workspace == "foo"
1915      with _request_context_for_shape(shape, role_auth_setup, workspace):
1916          assert auth_module.validate_can_manage_roles() is expected
1917  
1918  
1919  @pytest.mark.parametrize(
1920      ("actor", "workspace", "expected"),
1921      [
1922          ("super_admin", "foo", True),
1923          ("outsider", "foo", False),
1924          ("ws_admin_foo", "foo", True),
1925          ("ws_admin_foo", "bar", False),
1926          ("ws_admin_bar", "foo", False),
1927          ("ws_admin_bar", "bar", True),
1928          # Unlike manage, a plain workspace member can view roles.
1929          ("ws_member_foo", "foo", True),
1930          ("ws_member_foo", "bar", False),
1931      ],
1932  )
1933  def test_validate_can_view_roles_authorization(role_auth_setup, actor, workspace, expected):
1934      role_auth_setup["login_as"](actor)
1935      with _request_context_for_shape("role_id", role_auth_setup, workspace):
1936          assert auth_module.validate_can_view_roles() is expected
1937  
1938  
1939  @pytest.mark.parametrize(
1940      ("actor", "expected"),
1941      [
1942          ("super_admin", True),
1943          # Any non-admin is denied regardless of their workspace memberships —
1944          # one representative non-admin is enough.
1945          ("ws_admin_foo", False),
1946      ],
1947  )
1948  def test_validate_can_list_roles_unscoped_is_super_admin_only(role_auth_setup, actor, expected):
1949      # No workspace param: only super admins may list every role in the system.
1950      role_auth_setup["login_as"](actor)
1951      with auth_module.app.test_request_context("/api/3.0/mlflow/roles/list", method="GET"):
1952          assert auth_module.validate_can_list_roles() is expected
1953  
1954  
1955  @pytest.mark.parametrize(
1956      ("actor", "workspace", "expected"),
1957      [
1958          ("super_admin", "foo", True),
1959          ("outsider", "foo", False),
1960          ("ws_admin_foo", "foo", True),
1961          ("ws_admin_foo", "bar", False),
1962          ("ws_admin_bar", "foo", False),
1963          ("ws_admin_bar", "bar", True),
1964          ("ws_member_foo", "foo", True),
1965          ("ws_member_foo", "bar", False),
1966      ],
1967  )
1968  def test_validate_can_list_roles_workspace_scoped(role_auth_setup, actor, workspace, expected):
1969      role_auth_setup["login_as"](actor)
1970      with auth_module.app.test_request_context(
1971          "/api/3.0/mlflow/roles/list", method="GET", query_string={"workspace": workspace}
1972      ):
1973          assert auth_module.validate_can_list_roles() is expected
1974  
1975  
1976  def test_validate_can_list_roles_blank_workspace_denied_for_non_admin(role_auth_setup):
1977      # Blank workspace param hits a *different* branch from the missing-param
1978      # case: validate_can_list_roles checks ``workspace.strip()`` and denies
1979      # rather than raising, unlike _get_role_workspace_from_request which would
1980      # raise on blank workspace. Kept as a guard for that specific branch.
1981      role_auth_setup["login_as"]("ws_admin_foo")
1982      with auth_module.app.test_request_context(
1983          "/api/3.0/mlflow/roles/list",
1984          method="GET",
1985          query_string={"workspace": "   "},
1986      ):
1987          assert auth_module.validate_can_list_roles() is False
1988  
1989  
1990  def test_validate_can_view_user_roles_self_always_allowed(role_auth_setup):
1991      # A user can always read their own role list, even one with no roles.
1992      # Using ``outsider`` (zero roles) exercises the self-short-circuit without
1993      # any membership helping.
1994      role_auth_setup["login_as"]("outsider")
1995      with auth_module.app.test_request_context(
1996          "/api/3.0/mlflow/users/roles/list",
1997          method="GET",
1998          query_string={"username": "outsider"},
1999      ):
2000          assert auth_module.validate_can_view_user_roles() is True
2001  
2002  
2003  @pytest.mark.parametrize(
2004      ("requester", "target", "expected"),
2005      [
2006          ("super_admin", "ws_member_foo", True),
2007          ("ws_admin_foo", "ws_member_foo", True),
2008          ("ws_admin_bar", "ws_member_foo", False),
2009          ("ws_member_foo", "ws_admin_foo", False),
2010          ("outsider", "ws_member_foo", False),
2011      ],
2012  )
2013  def test_validate_can_view_user_roles_cross_user(role_auth_setup, requester, target, expected):
2014      role_auth_setup["login_as"](requester)
2015      with auth_module.app.test_request_context(
2016          "/api/3.0/mlflow/users/roles/list",
2017          method="GET",
2018          query_string={"username": target},
2019      ):
2020          assert auth_module.validate_can_view_user_roles() is expected
2021  
2022  
2023  def test_validate_can_view_user_roles_nonexistent_target_denied_for_non_admin(
2024      role_auth_setup,
2025  ):
2026      # Non-existent target: return False rather than leaking existence via the
2027      # RESOURCE_DOES_NOT_EXIST the handler would raise downstream.
2028      role_auth_setup["login_as"]("ws_admin_foo")
2029      with auth_module.app.test_request_context(
2030          "/api/3.0/mlflow/users/roles/list",
2031          method="GET",
2032          query_string={"username": "ghost"},
2033      ):
2034          assert auth_module.validate_can_view_user_roles() is False
2035  
2036  
2037  def test_validate_can_view_user_roles_nonexistent_target_allowed_for_super_admin(
2038      role_auth_setup,
2039  ):
2040      # Super admin short-circuits before the target lookup — they're authorized
2041      # regardless of whether the target exists (the handler then 404s cleanly).
2042      role_auth_setup["login_as"]("super_admin")
2043      with auth_module.app.test_request_context(
2044          "/api/3.0/mlflow/users/roles/list",
2045          method="GET",
2046          query_string={"username": "ghost"},
2047      ):
2048          assert auth_module.validate_can_view_user_roles() is True
2049  
2050  
2051  @pytest.mark.parametrize("shape", ["role_id", "role_permission_id"])
2052  def test_validate_can_manage_roles_nonexistent_resource_denied(role_auth_setup, shape):
2053      # A non-admin pointing at a role/role_permission that doesn't exist fails
2054      # closed: _get_role_workspace_from_request returns None and the validator
2055      # treats that as unauthorized rather than leaking existence.
2056      role_auth_setup["login_as"]("ws_admin_foo")
2057      bogus_id = 999_999
2058      if shape == "role_id":
2059          ctx = auth_module.app.test_request_context(
2060              "/api/3.0/mlflow/roles/get",
2061              method="GET",
2062              query_string={"role_id": str(bogus_id)},
2063          )
2064      else:
2065          ctx = auth_module.app.test_request_context(
2066              "/api/3.0/mlflow/roles/permissions/update",
2067              method="PATCH",
2068              json={"role_permission_id": bogus_id, "permission": READ.name},
2069          )
2070      with ctx:
2071          assert auth_module.validate_can_manage_roles() is False
2072  
2073  
2074  def test_validate_can_manage_roles_nonexistent_role_id_bypassed_by_super_admin(
2075      role_auth_setup,
2076  ):
2077      # Super admins skip the workspace resolution entirely — an unresolvable
2078      # role_id still produces True at the validator layer.
2079      role_auth_setup["login_as"]("super_admin")
2080      with auth_module.app.test_request_context(
2081          "/api/3.0/mlflow/roles/get",
2082          method="GET",
2083          query_string={"role_id": "999999"},
2084      ):
2085          assert auth_module.validate_can_manage_roles() is True
2086  
2087  
2088  def test_validate_can_manage_roles_missing_workspace_params_raises(role_auth_setup):
2089      # No role_id / role_permission_id / workspace in the request body: the
2090      # resolver raises INVALID_PARAMETER_VALUE — callers that hit this path have
2091      # a client bug, and we surface it instead of silently denying.
2092      role_auth_setup["login_as"]("ws_admin_foo")
2093      with auth_module.app.test_request_context(
2094          "/api/3.0/mlflow/roles/create", method="POST", json={}
2095      ):
2096          with pytest.raises(MlflowException, match="must include one of"):
2097              auth_module.validate_can_manage_roles()
2098  
2099  
2100  def test_validate_can_manage_roles_blank_workspace_raises(role_auth_setup):
2101      role_auth_setup["login_as"]("ws_admin_foo")
2102      with auth_module.app.test_request_context(
2103          "/api/3.0/mlflow/roles/create",
2104          method="POST",
2105          json={"name": "new-role", "workspace": "   "},
2106      ):
2107          with pytest.raises(MlflowException, match="non-empty string"):
2108              auth_module.validate_can_manage_roles()
2109  
2110  
2111  def test_validate_can_manage_roles_propagates_param_coercion_errors(role_auth_setup):
2112      # Integration check: a non-integer role_id in the request surfaces the
2113      # coercion error through the validator chain rather than silently denying.
2114      role_auth_setup["login_as"]("ws_admin_foo")
2115      with auth_module.app.test_request_context(
2116          "/api/3.0/mlflow/roles/get",
2117          method="GET",
2118          query_string={"role_id": "not-an-int"},
2119      ):
2120          with pytest.raises(MlflowException, match="must be an integer"):
2121              auth_module.validate_can_manage_roles()