zenml-nightly 0.75.0.dev20250312__py3-none-any.whl → 0.75.0.dev20250313__py3-none-any.whl

zenml/zen_server/rbac/endpoint_utils.py

@@ -13,21 +13,17 @@
 #  permissions and limitations under the License.
 """High-level helper functions to write endpoints with RBAC."""
 
-from typing import Any, Callable, List, TypeVar, Union
+from typing import Any, Callable, List, Optional, Tuple, TypeVar, Union
 from uuid import UUID
 
-from pydantic import BaseModel
-
-from zenml.constants import (
-    REQUIRES_CUSTOM_RESOURCE_REPORTING,
-)
-from zenml.exceptions import IllegalOperationError
 from zenml.models import (
     BaseFilter,
     BaseIdentifiedResponse,
     BaseRequest,
+    BaseUpdate,
     Page,
     UserScopedRequest,
+    WorkspaceScopedFilter,
 )
 from zenml.zen_server.auth import get_auth_context
 from zenml.zen_server.feature_gate.endpoint_utils import (
@@ -36,36 +32,42 @@ from zenml.zen_server.feature_gate.endpoint_utils import (
 )
 from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import (
+    batch_verify_permissions_for_models,
     dehydrate_page,
     dehydrate_response_model,
+    dehydrate_response_model_batch,
+    delete_model_resource,
     get_allowed_resource_ids,
+    get_resource_type_for_model,
     verify_permission,
     verify_permission_for_model,
 )
-from zenml.zen_server.utils import server_config
+from zenml.zen_server.utils import server_config, set_filter_workspace_scope
 
 AnyRequest = TypeVar("AnyRequest", bound=BaseRequest)
 AnyResponse = TypeVar("AnyResponse", bound=BaseIdentifiedResponse)  # type: ignore[type-arg]
+AnyOtherResponse = TypeVar("AnyOtherResponse", bound=BaseIdentifiedResponse)  # type: ignore[type-arg]
 AnyFilter = TypeVar("AnyFilter", bound=BaseFilter)
-AnyUpdate = TypeVar("AnyUpdate", bound=BaseModel)
+AnyUpdate = TypeVar("AnyUpdate", bound=BaseUpdate)
 UUIDOrStr = TypeVar("UUIDOrStr", UUID, Union[UUID, str])
 
 
 def verify_permissions_and_create_entity(
     request_model: AnyRequest,
-    resource_type: ResourceType,
     create_method: Callable[[AnyRequest], AnyResponse],
+    surrogate_models: Optional[List[AnyOtherResponse]] = None,
+    skip_entitlements: bool = False,
 ) -> AnyResponse:
     """Verify permissions and create the entity if authorized.
 
     Args:
         request_model: The entity request model.
-        resource_type: The resource type of the entity to create.
         create_method: The method to create the entity.
-
-    Raises:
-        IllegalOperationError: If the request model has a different owner then
-            the currently authenticated user.
+        surrogate_models: Optional list of surrogate models to verify
+            UPDATE permissions for instead of verifying CREATE permissions for
+            the request model.
+        skip_entitlements: Whether to skip the entitlement check and usage
+            increment.
 
     Returns:
         A model of the created entity.
@@ -74,43 +76,47 @@ def verify_permissions_and_create_entity(
         auth_context = get_auth_context()
         assert auth_context
 
-        if request_model.user != auth_context.user.id:
-            raise IllegalOperationError(
-                f"Not allowed to create resource '{resource_type}' for a "
-                "different user."
-            )
-    verify_permission(resource_type=resource_type, action=Action.CREATE)
+        # Ignore the user field set in the request model, if any, and set it to
+        # the current user's ID instead. This is just a precaution, given that
+        # the SQLZenStore also does this same validation on all request models.
+        request_model.user = auth_context.user.id
 
-    needs_usage_increment = (
-        resource_type in server_config().reportable_resources
-        and resource_type not in REQUIRES_CUSTOM_RESOURCE_REPORTING
-    )
-    if needs_usage_increment:
-        check_entitlement(resource_type)
+    if surrogate_models:
+        batch_verify_permissions_for_models(
+            models=surrogate_models, action=Action.UPDATE
+        )
+    else:
+        verify_permission_for_model(model=request_model, action=Action.CREATE)
+
+    resource_type = get_resource_type_for_model(request_model)
+
+    if resource_type:
+        needs_usage_increment = (
+            not skip_entitlements
+            and resource_type in server_config().reportable_resources
+        )
+        if needs_usage_increment:
+            check_entitlement(resource_type)
 
     created = create_method(request_model)
 
-    if needs_usage_increment:
+    if resource_type and needs_usage_increment:
         report_usage(resource_type, resource_id=created.id)
 
-    return created
+    return dehydrate_response_model(created)
 
 
 def verify_permissions_and_batch_create_entity(
     batch: List[AnyRequest],
-    resource_type: ResourceType,
     create_method: Callable[[List[AnyRequest]], List[AnyResponse]],
 ) -> List[AnyResponse]:
     """Verify permissions and create a batch of entities if authorized.
 
     Args:
         batch: The batch to create.
-        resource_type: The resource type of the entities to create.
         create_method: The method to create the entities.
 
     Raises:
-        IllegalOperationError: If the request model has a different owner then
-            the currently authenticated user.
         RuntimeError: If the resource type is usage-tracked.
 
     Returns:
@@ -119,23 +125,73 @@ def verify_permissions_and_batch_create_entity(
     auth_context = get_auth_context()
     assert auth_context
 
+    resource_types = set()
     for request_model in batch:
+        resource_type = get_resource_type_for_model(request_model)
+        if resource_type:
+            resource_types.add(resource_type)
+
         if isinstance(request_model, UserScopedRequest):
-            if request_model.user != auth_context.user.id:
-                raise IllegalOperationError(
-                    f"Not allowed to create resource '{resource_type}' for a "
-                    "different user."
-                )
+            # Ignore the user field set in the request model, if any, and set it
+            # to the current user's ID instead. This is just a precaution, given
+            # that the SQLZenStore also does this same validation on all request
+            # models.
+            request_model.user = auth_context.user.id
 
-    verify_permission(resource_type=resource_type, action=Action.CREATE)
+    batch_verify_permissions_for_models(models=batch, action=Action.CREATE)
 
-    if resource_type in server_config().reportable_resources:
+    if resource_types & set(server_config().reportable_resources):
         raise RuntimeError(
-            "Batch requests are currently not possible with usage-tracked features."
+            "Batch requests are currently not possible with usage-tracked "
+            "features."
         )
 
     created = create_method(batch)
-    return created
+    return dehydrate_response_model_batch(created)
+
+
+def verify_permissions_and_get_or_create_entity(
+    request_model: AnyRequest,
+    get_or_create_method: Callable[
+        [AnyRequest, Optional[Callable[[], None]]], Tuple[AnyResponse, bool]
+    ],
+) -> Tuple[AnyResponse, bool]:
+    """Verify permissions and create the entity if authorized.
+
+    Args:
+        request_model: The entity request model.
+        get_or_create_method: The method to get or create the entity.
+
+    Returns:
+        The entity and a boolean indicating whether the entity was created.
+    """
+    if isinstance(request_model, UserScopedRequest):
+        auth_context = get_auth_context()
+        assert auth_context
+
+        # Ignore the user field set in the request model, if any, and set it to
+        # the current user's ID instead. This is just a precaution, given that
+        # the SQLZenStore also does this same validation on all request models.
+        request_model.user = auth_context.user.id
+
+    resource_type = get_resource_type_for_model(request_model)
+    needs_usage_increment = (
+        resource_type and resource_type in server_config().reportable_resources
+    )
+
+    def _pre_creation_hook() -> None:
+        verify_permission_for_model(model=request_model, action=Action.CREATE)
+        if resource_type and needs_usage_increment:
+            check_entitlement(resource_type=resource_type)
+
+    model, created = get_or_create_method(request_model, _pre_creation_hook)
+
+    if not created:
+        verify_permission_for_model(model=model, action=Action.READ)
+    elif resource_type and needs_usage_increment:
+        report_usage(resource_type, resource_id=model.id)
+
+    return dehydrate_response_model(model), created
 
 
 def verify_permissions_and_get_entity(
@@ -174,11 +230,30 @@ def verify_permissions_and_list_entities(
 
     Returns:
         A page of entity models.
+
+    Raises:
+        ValueError: If the filter's workspace scope is not set or is not a UUID.
     """
     auth_context = get_auth_context()
     assert auth_context
 
-    allowed_ids = get_allowed_resource_ids(resource_type=resource_type)
+    workspace_id: Optional[UUID] = None
+    if isinstance(filter_model, WorkspaceScopedFilter):
+        # A workspace scoped filter must always be scoped to a specific
+        # workspace. This is required for the RBAC check to work.
+        set_filter_workspace_scope(filter_model)
+        if not filter_model.workspace or not isinstance(
+            filter_model.workspace, UUID
+        ):
+            raise ValueError(
+                "Workspace scope must be a UUID, got "
+                f"{type(filter_model.workspace)}."
+            )
+        workspace_id = filter_model.workspace
+
+    allowed_ids = get_allowed_resource_ids(
+        resource_type=resource_type, workspace_id=workspace_id
+    )
     filter_model.configure_rbac(
         authenticated_user_id=auth_context.user.id, id=allowed_ids
     )
@@ -191,6 +266,7 @@ def verify_permissions_and_update_entity(
     update_model: AnyUpdate,
     get_method: Callable[[UUIDOrStr, bool], AnyResponse],
     update_method: Callable[[UUIDOrStr, AnyUpdate], AnyResponse],
+    **update_method_kwargs: Any,
 ) -> AnyResponse:
     """Verify permissions and update an entity.
 
@@ -199,6 +275,7 @@ def verify_permissions_and_update_entity(
         update_model: The entity update model.
         get_method: The method to fetch the entity.
         update_method: The method to update the entity.
+        update_method_kwargs: Keyword arguments to pass to the update method.
 
     Returns:
         A model of the updated entity.
@@ -206,7 +283,9 @@ def verify_permissions_and_update_entity(
     # We don't need the hydrated version here
     model = get_method(id, False)
     verify_permission_for_model(model, action=Action.UPDATE)
-    updated_model = update_method(model.id, update_model)
+    updated_model = update_method(
+        model.id, update_model, **update_method_kwargs
+    )
     return dehydrate_response_model(updated_model)
 
 
@@ -214,6 +293,7 @@ def verify_permissions_and_delete_entity(
     id: UUIDOrStr,
     get_method: Callable[[UUIDOrStr, bool], AnyResponse],
     delete_method: Callable[[UUIDOrStr], None],
+    **delete_method_kwargs: Any,
 ) -> AnyResponse:
     """Verify permissions and delete an entity.
 
@@ -221,6 +301,7 @@ def verify_permissions_and_delete_entity(
         id: The ID of the entity to delete.
         get_method: The method to fetch the entity.
         delete_method: The method to delete the entity.
+        delete_method_kwargs: Keyword arguments to pass to the delete method.
 
     Returns:
         The deleted entity.
@@ -228,7 +309,8 @@ def verify_permissions_and_delete_entity(
     # We don't need the hydrated version here
     model = get_method(id, False)
     verify_permission_for_model(model, action=Action.DELETE)
-    delete_method(model.id)
+    delete_method(model.id, **delete_method_kwargs)
+    delete_model_resource(model)
 
     return model
 
@@ -236,6 +318,7 @@ def verify_permissions_and_delete_entity(
 def verify_permissions_and_prune_entities(
     resource_type: ResourceType,
     prune_method: Callable[..., None],
+    workspace_id: Optional[UUID] = None,
     **kwargs: Any,
 ) -> None:
     """Verify permissions and prune entities of certain type.
@@ -243,7 +326,12 @@ def verify_permissions_and_prune_entities(
     Args:
         resource_type: The resource type of the entities to prune.
         prune_method: The method to prune the entities.
+        workspace_id: The workspace ID to prune the entities for.
         kwargs: Keyword arguments to pass to the prune method.
     """
-    verify_permission(resource_type=resource_type, action=Action.PRUNE)
+    verify_permission(
+        resource_type=resource_type,
+        action=Action.PRUNE,
+        workspace_id=workspace_id,
+    )
     prune_method(**kwargs)

zenml/zen_server/rbac/models.py

@@ -16,7 +16,11 @@
 from typing import Optional
 from uuid import UUID
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    model_validator,
+)
 
 from zenml.utils.enum_utils import StrEnum
 
@@ -69,9 +73,28 @@ class ResourceType(StrEnum):
     TAG = "tag"
     TRIGGER = "trigger"
     TRIGGER_EXECUTION = "trigger_execution"
+    WORKSPACE = "workspace"
     # Deactivated for now
     # USER = "user"
-    # WORKSPACE = "workspace"
+
+    def is_workspace_scoped(self) -> bool:
+        """Check if a resource type is workspace scoped.
+
+        Returns:
+            Whether the resource type is workspace scoped.
+        """
+        return self not in [
+            self.FLAVOR,
+            self.SECRET,
+            self.SERVICE_CONNECTOR,
+            self.STACK,
+            self.STACK_COMPONENT,
+            self.TAG,
+            self.SERVICE_ACCOUNT,
+            self.WORKSPACE,
+            # Deactivated for now
+            # self.USER,
+        ]
 
 
 class Resource(BaseModel):
@@ -79,6 +102,7 @@ class Resource(BaseModel):
 
     type: str
     id: Optional[UUID] = None
+    workspace_id: Optional[UUID] = None
 
     def __str__(self) -> str:
         """Convert to a string.
@@ -86,10 +110,48 @@ class Resource(BaseModel):
         Returns:
             Resource string representation.
         """
-        representation = self.type
+        workspace_id = self.workspace_id
+        if self.type == ResourceType.WORKSPACE and self.id:
+            # TODO: For now, we duplicate the workspace ID in the string
+            # representation when describing a workspace instance, because
+            # this is what is expected by the RBAC implementation.
+            workspace_id = self.id
+
+        if workspace_id:
+            representation = f"{workspace_id}:"
+        else:
+            representation = ""
+        representation += self.type
         if self.id:
             representation += f"/{self.id}"
 
         return representation
 
+    @model_validator(mode="after")
+    def validate_workspace_id(self) -> "Resource":
+        """Validate that workspace_id is set in combination with workspace-scoped resource types.
+
+        Raises:
+            ValueError: If workspace_id is not set for a workspace-scoped
+                resource or set for an unscoped resource.
+
+        Returns:
+            The validated resource.
+        """
+        resource_type = ResourceType(self.type)
+
+        if resource_type.is_workspace_scoped() and not self.workspace_id:
+            raise ValueError(
+                "workspace_id must be set for workspace-scoped resource type "
+                f"'{self.type}'"
+            )
+
+        if not resource_type.is_workspace_scoped() and self.workspace_id:
+            raise ValueError(
+                "workspace_id must not be set for global resource type "
+                f"'{self.type}'"
+            )
+
+        return self
+
     model_config = ConfigDict(frozen=True)

zenml/zen_server/rbac/rbac_interface.py

@@ -73,3 +73,12 @@ class RBACInterface(ABC):
             actions: The actions that the user should be able to perform on the
                 resource.
         """
+
+    @abstractmethod
+    def delete_resources(self, resources: List[Resource]) -> None:
+        """Delete resource membership information for a list of resources.
+
+        Args:
+            resources: The resources for which to delete the resource membership
+                information.
+        """

zenml/zen_server/rbac/rbac_sql_zen_store.py

@@ -63,7 +63,9 @@ class RBACSqlZenStore(SqlZenStore):
 
         try:
             verify_permission(
-                resource_type=ResourceType.MODEL, action=Action.CREATE
+                resource_type=ResourceType.MODEL,
+                action=Action.CREATE,
+                workspace_id=model_request.workspace,
             )
             check_entitlement(resource_type=ResourceType.MODEL)
         except Exception as e:
@@ -76,7 +78,10 @@ class RBACSqlZenStore(SqlZenStore):
             )
         else:
             try:
-                model_response = self.get_model(model_request.name)
+                model_response = self.get_model_by_name_or_id(
+                    model_name_or_id=model_request.name,
+                    workspace=model_request.workspace,
+                )
                 created = False
             except KeyError:
                 # The model does not exist. We now raise the error that
@@ -145,17 +150,20 @@ class RBACSqlZenStore(SqlZenStore):
 
         try:
             verify_permission(
-                resource_type=ResourceType.MODEL_VERSION, action=Action.CREATE
+                resource_type=ResourceType.MODEL_VERSION,
+                action=Action.CREATE,
+                workspace_id=model_version_request.workspace,
             )
         except Exception as e:
             allow_creation = False
             error = e
 
         if allow_creation:
-            created, model_version_response = (
-                super()._get_or_create_model_version(
-                    model_version_request, producer_run_id=producer_run_id
-                )
+            (
+                created,
+                model_version_response,
+            ) = super()._get_or_create_model_version(
+                model_version_request, producer_run_id=producer_run_id
             )
         else:
             try: