mlflow--mlflow/mlflow/server/auth/entities.py

from mlflow.exceptions import MlflowException
from mlflow.server.auth.permissions import get_permission
from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME, resolve_entity_workspace_name


class User:
    def __init__(
        self,
        id_,
        username,
        password_hash,
        is_admin,
    ):
        self._id = id_
        self._username = username
        self._password_hash = password_hash
        self._is_admin = is_admin

    @property
    def id(self):
        return self._id

    @property
    def username(self):
        return self._username

    @property
    def password_hash(self):
        return self._password_hash

    @property
    def is_admin(self):
        return self._is_admin

    @is_admin.setter
    def is_admin(self, is_admin):
        self._is_admin = is_admin

    def to_json(self):
        return {
            "id": self.id,
            "username": self.username,
            "is_admin": self.is_admin,
        }

    @classmethod
    def from_json(cls, dictionary):
        return cls(
            id_=dictionary["id"],
            username=dictionary["username"],
            password_hash="REDACTED",
            is_admin=dictionary["is_admin"],
        )


class ExperimentPermission:
    def __init__(
        self,
        experiment_id,
        user_id,
        permission,
    ):
        self._experiment_id = experiment_id
        self._user_id = user_id
        self._permission = permission

    @property
    def experiment_id(self):
        return self._experiment_id

    @property
    def user_id(self):
        return self._user_id

    @property
    def permission(self):
        return self._permission

    @permission.setter
    def permission(self, permission):
        self._permission = permission

    def to_json(self):
        return {
            "experiment_id": self.experiment_id,
            "user_id": self.user_id,
            "permission": self.permission,
        }

    @classmethod
    def from_json(cls, dictionary):
        return cls(
            experiment_id=dictionary["experiment_id"],
            user_id=dictionary["user_id"],
            permission=dictionary["permission"],
        )


class RegisteredModelPermission:
    def __init__(
        self,
        name,
        user_id,
        permission,
        workspace=None,
    ):
        self._workspace = resolve_entity_workspace_name(workspace)
        self._name = name
        self._user_id = user_id
        self._permission = permission

    @property
    def workspace(self):
        return self._workspace

    @property
    def name(self):
        return self._name

    @property
    def user_id(self):
        return self._user_id

    @property
    def permission(self):
        return self._permission

    @permission.setter
    def permission(self, permission):
        self._permission = permission

    def to_json(self):
        return {
            "workspace": self.workspace,
            "name": self.name,
            "user_id": self.user_id,
            "permission": self.permission,
        }

    @classmethod
    def from_json(cls, dictionary):
        return cls(
            name=dictionary["name"],
            user_id=dictionary["user_id"],
            permission=dictionary["permission"],
            workspace=dictionary.get("workspace"),
        )


class ScorerPermission:
    def __init__(
        self,
        experiment_id,
        scorer_name,
        user_id,
        permission,
    ):
        self._experiment_id = experiment_id
        self._scorer_name = scorer_name
        self._user_id = user_id
        self._permission = permission

    @property
    def experiment_id(self):
        return self._experiment_id

    @property
    def scorer_name(self):
        return self._scorer_name

    @property
    def user_id(self):
        return self._user_id

    @property
    def permission(self):
        return self._permission

    @permission.setter
    def permission(self, permission):
        self._permission = permission

    def to_json(self):
        return {
            "experiment_id": self.experiment_id,
            "scorer_name": self.scorer_name,
            "user_id": self.user_id,
            "permission": self.permission,
        }

    @classmethod
    def from_json(cls, dictionary):
        return cls(
            experiment_id=dictionary["experiment_id"],
            scorer_name=dictionary["scorer_name"],
            user_id=dictionary["user_id"],
            permission=dictionary["permission"],
        )


class GatewaySecretPermission:
    def __init__(self, secret_id, user_id, permission):
        self._secret_id = secret_id
        self._user_id = user_id
        self._permission = permission

    @property
    def secret_id(self):
        return self._secret_id

    @property
    def user_id(self):
        return self._user_id

    @property
    def permission(self):
        return self._permission

    @permission.setter
    def permission(self, permission):
        self._permission = permission

    def to_json(self):
        return {
            "secret_id": self.secret_id,
            "user_id": self.user_id,
            "permission": self.permission,
        }

    @classmethod
    def from_json(cls, dictionary):
        return cls(
            secret_id=dictionary["secret_id"],
            user_id=dictionary["user_id"],
            permission=dictionary["permission"],
        )


class GatewayEndpointPermission:
    def __init__(self, endpoint_id, user_id, permission):
        self._endpoint_id = endpoint_id
        self._user_id = user_id
        self._permission = permission

    @property
    def endpoint_id(self):
        return self._endpoint_id

    @property
    def user_id(self):
        return self._user_id

    @property
    def permission(self):
        return self._permission

    @permission.setter
    def permission(self, permission):
        self._permission = permission

    def to_json(self):
        return {
            "endpoint_id": self.endpoint_id,
            "user_id": self.user_id,
            "permission": self.permission,
        }

    @classmethod
    def from_json(cls, dictionary):
        return cls(
            endpoint_id=dictionary["endpoint_id"],
            user_id=dictionary["user_id"],
            permission=dictionary["permission"],
        )


class GatewayModelDefinitionPermission:
    def __init__(self, model_definition_id, user_id, permission):
        self._model_definition_id = model_definition_id
        self._user_id = user_id
        self._permission = permission

    @property
    def model_definition_id(self):
        return self._model_definition_id

    @property
    def user_id(self):
        return self._user_id

    @property
    def permission(self):
        return self._permission

    @permission.setter
    def permission(self, permission):
        self._permission = permission

    def to_json(self):
        return {
            "model_definition_id": self.model_definition_id,
            "user_id": self.user_id,
            "permission": self.permission,
        }

    @classmethod
    def from_json(cls, dictionary):
        return cls(
            model_definition_id=dictionary["model_definition_id"],
            user_id=dictionary["user_id"],
            permission=dictionary["permission"],
        )


class Role:
    def __init__(
        self,
        id_,
        name,
        workspace,
        description=None,
        permissions=None,
    ):
        self._id = id_
        self._name = name
        self._workspace = workspace
        self._description = description
        self._permissions = permissions or []

    @property
    def id(self):
        return self._id

    @property
    def name(self):
        return self._name

    @name.setter
    def name(self, name):
        self._name = name

    @property
    def workspace(self):
        return self._workspace

    @property
    def description(self):
        return self._description

    @description.setter
    def description(self, description):
        self._description = description

    @property
    def permissions(self):
        return self._permissions

    def to_json(self):
        return {
            "id": self.id,
            "name": self.name,
            "workspace": self.workspace,
            "description": self.description,
            "permissions": [p.to_json() for p in self.permissions],
        }

    @classmethod
    def from_json(cls, dictionary):
        return cls(
            id_=dictionary["id"],
            name=dictionary["name"],
            workspace=dictionary.get("workspace", DEFAULT_WORKSPACE_NAME),
            description=dictionary.get("description"),
            permissions=[RolePermission.from_json(p) for p in dictionary.get("permissions", [])],
        )


class RolePermission:
    def __init__(self, id_, role_id, resource_type, resource_pattern, permission):
        self._id = id_
        self._role_id = role_id
        self._resource_type = resource_type
        self._resource_pattern = resource_pattern
        self._permission = permission

    @property
    def id(self):
        return self._id

    @property
    def role_id(self):
        return self._role_id

    @property
    def resource_type(self):
        return self._resource_type

    @property
    def resource_pattern(self):
        return self._resource_pattern

    @property
    def permission(self):
        return self._permission

    @permission.setter
    def permission(self, permission):
        self._permission = permission

    def to_json(self):
        return {
            "id": self.id,
            "role_id": self.role_id,
            "resource_type": self.resource_type,
            "resource_pattern": self.resource_pattern,
            "permission": self.permission,
        }

    @classmethod
    def from_json(cls, dictionary):
        return cls(
            id_=dictionary["id"],
            role_id=dictionary["role_id"],
            resource_type=dictionary["resource_type"],
            resource_pattern=dictionary["resource_pattern"],
            permission=dictionary["permission"],
        )


class UserRoleAssignment:
    def __init__(self, id_, user_id, role_id):
        self._id = id_
        self._user_id = user_id
        self._role_id = role_id

    @property
    def id(self):
        return self._id

    @property
    def user_id(self):
        return self._user_id

    @property
    def role_id(self):
        return self._role_id

    def to_json(self):
        return {
            "id": self.id,
            "user_id": self.user_id,
            "role_id": self.role_id,
        }

    @classmethod
    def from_json(cls, dictionary):
        return cls(
            id_=dictionary["id"],
            user_id=dictionary["user_id"],
            role_id=dictionary["role_id"],
        )


class GetUserPermissionResult:
    """Response shape for ``GET /mlflow/users/permissions/get``. ``allowed``
    mirrors ``Permission.can_use`` (regular access tier); ``permission`` is the
    resolved effective permission name.
    """

    def __init__(self, allowed: bool, permission: str):
        self._allowed = allowed
        self._permission = permission

    @property
    def allowed(self) -> bool:
        return self._allowed

    @property
    def permission(self) -> str:
        return self._permission

    def to_json(self):
        return {"allowed": self.allowed, "permission": self.permission}

    @classmethod
    def from_json(cls, dictionary):
        return cls(allowed=dictionary["allowed"], permission=dictionary["permission"])


class WorkspacePermission:
    def __init__(self, workspace, user_id, permission):
        if workspace is None or user_id is None or permission is None:
            raise MlflowException.invalid_parameter_value(
                "workspace, user_id, and permission are required."
            )
        self._workspace = workspace
        self._user_id = user_id
        self._permission = permission

    @property
    def workspace(self):
        return self._workspace

    @property
    def user_id(self):
        return self._user_id

    @property
    def permission(self):
        return self._permission

    @property
    def can_use(self):
        return get_permission(self.permission).can_use

    def to_json(self):
        return {
            "workspace": self.workspace,
            "user_id": self.user_id,
            "permission": self.permission,
            "can_use": self.can_use,
        }

    @classmethod
    def from_json(cls, dictionary):
        required_fields = ["workspace", "user_id", "permission"]
        if missing := [field for field in required_fields if field not in dictionary]:
            raise MlflowException.invalid_parameter_value(
                f"Missing required fields: {', '.join(missing)}"
            )
        return cls(
            workspace=dictionary["workspace"],
            user_id=dictionary["user_id"],
            permission=dictionary["permission"],
        )