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

zenml/zen_server/rbac/utils.py

@@ -34,6 +34,8 @@ from zenml.models import (
     Page,
     UserResponse,
     UserScopedResponse,
+    WorkspaceScopedRequest,
+    WorkspaceScopedResponse,
 )
 from zenml.zen_server.auth import get_auth_context
 from zenml.zen_server.rbac.models import Action, Resource, ResourceType
@@ -55,24 +57,39 @@ def dehydrate_page(page: Page[AnyResponse]) -> Page[AnyResponse]:
     Returns:
         The page with (potentially) dehydrated items.
     """
+    new_items = dehydrate_response_model_batch(page.items)
+    return page.model_copy(update={"items": new_items})
+
+
+def dehydrate_response_model_batch(
+    batch: List[AnyResponse],
+) -> List[AnyResponse]:
+    """Dehydrate all items of a batch.
+
+    Args:
+        batch: The batch to dehydrate.
+
+    Returns:
+        The batch with (potentially) dehydrated items.
+    """
     if not server_config().rbac_enabled:
-        return page
+        return batch
 
     auth_context = get_auth_context()
     assert auth_context
 
-    resource_list = [get_subresources_for_model(item) for item in page.items]
+    resource_list = [get_subresources_for_model(item) for item in batch]
     resources = set.union(*resource_list) if resource_list else set()
     permissions = rbac().check_permissions(
         user=auth_context.user, resources=resources, action=Action.READ
     )
 
-    new_items = [
+    new_batch = [
         dehydrate_response_model(item, permissions=permissions)
-        for item in page.items
+        for item in batch
     ]
 
-    return page.model_copy(update={"items": new_items})
+    return new_batch
 
 
 def dehydrate_response_model(
@@ -161,7 +178,7 @@ def _dehydrate_value(
         return value
 
 
-def has_permissions_for_model(model: AnyResponse, action: Action) -> bool:
+def has_permissions_for_model(model: AnyModel, action: Action) -> bool:
     """If the active user has permissions to perform the action on the model.
 
     Args:
@@ -201,7 +218,7 @@ def get_permission_denied_model(model: AnyResponse) -> AnyResponse:
 
 
 def batch_verify_permissions_for_models(
-    models: Sequence[AnyResponse],
+    models: Sequence[AnyModel],
     action: Action,
 ) -> None:
     """Batch permission verification for models.
@@ -229,7 +246,7 @@ def batch_verify_permissions_for_models(
     batch_verify_permissions(resources=resources, action=action)
 
 
-def verify_permission_for_model(model: AnyResponse, action: Action) -> None:
+def verify_permission_for_model(model: AnyModel, action: Action) -> None:
     """Verifies if a user has permission to perform an action on a model.
 
     Args:
@@ -283,6 +300,7 @@ def verify_permission(
     resource_type: str,
     action: Action,
     resource_id: Optional[UUID] = None,
+    workspace_id: Optional[UUID] = None,
 ) -> None:
     """Verifies if a user has permission to perform an action on a resource.
 
@@ -291,20 +309,27 @@ def verify_permission(
             action on.
         action: The action the user wants to perform.
         resource_id: ID of the resource the user wants to perform the action on.
+        workspace_id: ID of the workspace the user wants to perform the action
+            on. Only used for workspace scoped resources.
     """
-    resource = Resource(type=resource_type, id=resource_id)
+    resource = Resource(
+        type=resource_type, id=resource_id, workspace_id=workspace_id
+    )
     batch_verify_permissions(resources={resource}, action=action)
 
 
 def get_allowed_resource_ids(
     resource_type: str,
     action: Action = Action.READ,
+    workspace_id: Optional[UUID] = None,
 ) -> Optional[Set[UUID]]:
     """Get all resource IDs of a resource type that a user can access.
 
     Args:
         resource_type: The resource type.
         action: The action the user wants to perform on the resource.
+        workspace_id: Optional workspace ID to filter the resources by.
+            Required for workspace scoped resources.
 
     Returns:
         A list of resource IDs or `None` if the user has full access to the
@@ -321,7 +346,7 @@ def get_allowed_resource_ids(
         allowed_ids,
     ) = rbac().list_allowed_resource_ids(
         user=auth_context.user,
-        resource=Resource(type=resource_type),
+        resource=Resource(type=resource_type, workspace_id=workspace_id),
         action=action,
     )
 
@@ -331,7 +356,7 @@ def get_allowed_resource_ids(
     return {UUID(id) for id in allowed_ids}
 
 
-def get_resource_for_model(model: AnyResponse) -> Optional[Resource]:
+def get_resource_for_model(model: AnyModel) -> Optional[Resource]:
     """Get the resource associated with a model object.
 
     Args:
@@ -346,12 +371,26 @@ def get_resource_for_model(model: AnyResponse) -> Optional[Resource]:
         # This model is not tied to any RBAC resource type
         return None
 
-    return Resource(type=resource_type, id=model.id)
+    workspace_id: Optional[UUID] = None
+    if isinstance(model, WorkspaceScopedResponse):
+        # A workspace scoped response is always scoped to a specific workspace
+        workspace_id = model.workspace.id
+    elif isinstance(model, WorkspaceScopedRequest):
+        # A workspace scoped request is always scoped to a specific workspace
+        workspace_id = model.workspace
+
+    resource_id: Optional[UUID] = None
+    if isinstance(model, BaseIdentifiedResponse):
+        resource_id = model.id
+
+    return Resource(
+        type=resource_type, id=resource_id, workspace_id=workspace_id
+    )
 
 
 def get_surrogate_permission_model_for_model(
-    model: AnyResponse, action: str
-) -> BaseIdentifiedResponse[Any, Any, Any]:
+    model: BaseModel, action: str
+) -> BaseModel:
     """Get a surrogate permission model for a model.
 
     In some cases a different model instead of the original model is used to
@@ -379,7 +418,7 @@ def get_surrogate_permission_model_for_model(
 
 
 def get_resource_type_for_model(
-    model: AnyResponse,
+    model: AnyModel,
 ) -> Optional[ResourceType]:
     """Get the resource type associated with a model object.
 
@@ -391,64 +430,113 @@ def get_resource_type_for_model(
         is not associated with any resource type.
     """
     from zenml.models import (
+        ActionRequest,
         ActionResponse,
+        ArtifactRequest,
         ArtifactResponse,
+        ArtifactVersionRequest,
         ArtifactVersionResponse,
+        CodeRepositoryRequest,
         CodeRepositoryResponse,
+        ComponentRequest,
         ComponentResponse,
+        EventSourceRequest,
         EventSourceResponse,
+        FlavorRequest,
         FlavorResponse,
+        ModelRequest,
         ModelResponse,
+        ModelVersionRequest,
         ModelVersionResponse,
+        PipelineBuildRequest,
         PipelineBuildResponse,
+        PipelineDeploymentRequest,
         PipelineDeploymentResponse,
+        PipelineRequest,
         PipelineResponse,
+        PipelineRunRequest,
         PipelineRunResponse,
+        RunMetadataRequest,
+        RunTemplateRequest,
         RunTemplateResponse,
+        SecretRequest,
         SecretResponse,
+        ServiceAccountRequest,
         ServiceAccountResponse,
+        ServiceConnectorRequest,
         ServiceConnectorResponse,
+        ServiceRequest,
         ServiceResponse,
+        StackRequest,
         StackResponse,
+        TagRequest,
         TagResponse,
+        TriggerExecutionRequest,
         TriggerExecutionResponse,
+        TriggerRequest,
         TriggerResponse,
+        WorkspaceRequest,
+        WorkspaceResponse,
     )
 
     mapping: Dict[
         Any,
         ResourceType,
     ] = {
+        ActionRequest: ResourceType.ACTION,
         ActionResponse: ResourceType.ACTION,
+        ArtifactRequest: ResourceType.ARTIFACT,
+        ArtifactResponse: ResourceType.ARTIFACT,
+        ArtifactVersionRequest: ResourceType.ARTIFACT_VERSION,
+        ArtifactVersionResponse: ResourceType.ARTIFACT_VERSION,
+        CodeRepositoryRequest: ResourceType.CODE_REPOSITORY,
+        CodeRepositoryResponse: ResourceType.CODE_REPOSITORY,
+        ComponentRequest: ResourceType.STACK_COMPONENT,
+        ComponentResponse: ResourceType.STACK_COMPONENT,
+        EventSourceRequest: ResourceType.EVENT_SOURCE,
         EventSourceResponse: ResourceType.EVENT_SOURCE,
+        FlavorRequest: ResourceType.FLAVOR,
         FlavorResponse: ResourceType.FLAVOR,
-        ServiceConnectorResponse: ResourceType.SERVICE_CONNECTOR,
-        ComponentResponse: ResourceType.STACK_COMPONENT,
-        StackResponse: ResourceType.STACK,
-        PipelineResponse: ResourceType.PIPELINE,
-        CodeRepositoryResponse: ResourceType.CODE_REPOSITORY,
-        SecretResponse: ResourceType.SECRET,
+        ModelRequest: ResourceType.MODEL,
         ModelResponse: ResourceType.MODEL,
+        ModelVersionRequest: ResourceType.MODEL_VERSION,
         ModelVersionResponse: ResourceType.MODEL_VERSION,
-        ArtifactResponse: ResourceType.ARTIFACT,
-        ArtifactVersionResponse: ResourceType.ARTIFACT_VERSION,
-        # WorkspaceResponse: ResourceType.WORKSPACE,
-        # UserResponse: ResourceType.USER,
-        PipelineDeploymentResponse: ResourceType.PIPELINE_DEPLOYMENT,
+        PipelineBuildRequest: ResourceType.PIPELINE_BUILD,
         PipelineBuildResponse: ResourceType.PIPELINE_BUILD,
+        PipelineDeploymentRequest: ResourceType.PIPELINE_DEPLOYMENT,
+        PipelineDeploymentResponse: ResourceType.PIPELINE_DEPLOYMENT,
+        PipelineRequest: ResourceType.PIPELINE,
+        PipelineResponse: ResourceType.PIPELINE,
+        PipelineRunRequest: ResourceType.PIPELINE_RUN,
         PipelineRunResponse: ResourceType.PIPELINE_RUN,
+        RunMetadataRequest: ResourceType.RUN_METADATA,
+        RunTemplateRequest: ResourceType.RUN_TEMPLATE,
         RunTemplateResponse: ResourceType.RUN_TEMPLATE,
+        SecretRequest: ResourceType.SECRET,
+        SecretResponse: ResourceType.SECRET,
+        ServiceAccountRequest: ResourceType.SERVICE_ACCOUNT,
+        ServiceAccountResponse: ResourceType.SERVICE_ACCOUNT,
+        ServiceConnectorRequest: ResourceType.SERVICE_CONNECTOR,
+        ServiceConnectorResponse: ResourceType.SERVICE_CONNECTOR,
+        ServiceRequest: ResourceType.SERVICE,
+        ServiceResponse: ResourceType.SERVICE,
+        StackRequest: ResourceType.STACK,
+        StackResponse: ResourceType.STACK,
+        TagRequest: ResourceType.TAG,
         TagResponse: ResourceType.TAG,
+        TriggerRequest: ResourceType.TRIGGER,
         TriggerResponse: ResourceType.TRIGGER,
+        TriggerExecutionRequest: ResourceType.TRIGGER_EXECUTION,
         TriggerExecutionResponse: ResourceType.TRIGGER_EXECUTION,
-        ServiceAccountResponse: ResourceType.SERVICE_ACCOUNT,
-        ServiceResponse: ResourceType.SERVICE,
+        WorkspaceResponse: ResourceType.WORKSPACE,
+        WorkspaceRequest: ResourceType.WORKSPACE,
+        # UserResponse: ResourceType.USER,
     }
 
     return mapping.get(type(model))
 
 
-def is_owned_by_authenticated_user(model: AnyResponse) -> bool:
+def is_owned_by_authenticated_user(model: AnyModel) -> bool:
     """Returns whether the currently authenticated user owns the model.
 
     Args:
@@ -618,3 +706,42 @@ def update_resource_membership(
     rbac().update_resource_membership(
         user=user, resource=resource, actions=actions
     )
+
+
+def delete_model_resource(model: AnyModel) -> None:
+    """Delete resource membership information for a model.
+
+    Args:
+        model: The model for which to delete the resource membership information.
+    """
+    delete_model_resources(models=[model])
+
+
+def delete_model_resources(models: List[AnyModel]) -> None:
+    """Delete resource membership information for a list of models.
+
+    Args:
+        models: The models for which to delete the resource membership information.
+    """
+    if not server_config().rbac_enabled:
+        return
+
+    resources = set()
+    for model in models:
+        if resource := get_resource_for_model(model):
+            resources.add(resource)
+
+    delete_resources(resources=list(resources))
+
+
+def delete_resources(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.
+    """
+    if not server_config().rbac_enabled:
+        return
+
+    rbac().delete_resources(resources=resources)

zenml/zen_server/rbac/zenml_cloud_rbac.py

@@ -13,10 +13,10 @@
 #  permissions and limitations under the License.
 """Cloud RBAC implementation."""
 
-from typing import TYPE_CHECKING, Dict, List, Set, Tuple
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
 
 from zenml.zen_server.cloud_utils import cloud_connection
-from zenml.zen_server.rbac.models import Action, Resource
+from zenml.zen_server.rbac.models import Action, Resource, ResourceType
 from zenml.zen_server.rbac.rbac_interface import RBACInterface
 from zenml.zen_server.utils import server_config
 
@@ -27,6 +27,7 @@ if TYPE_CHECKING:
 PERMISSIONS_ENDPOINT = "/rbac/check_permissions"
 ALLOWED_RESOURCE_IDS_ENDPOINT = "/rbac/allowed_resource_ids"
 RESOURCE_MEMBERSHIP_ENDPOINT = "/rbac/resource_members"
+RESOURCES_ENDPOINT = "/rbac/resources"
 
 SERVER_SCOPE_IDENTIFIER = "server"
 
@@ -42,12 +43,7 @@ def _convert_to_cloud_resource(resource: Resource) -> str:
     Returns:
         The converted resource.
     """
-    resource_string = f"{SERVER_ID}@{SERVER_SCOPE_IDENTIFIER}:{resource.type}"
-
-    if resource.id:
-        resource_string += f"/{resource.id}"
-
-    return resource_string
+    return f"{SERVER_ID}@{SERVER_SCOPE_IDENTIFIER}:{resource}"
 
 
 def _convert_from_cloud_resource(cloud_resource: str) -> Resource:
@@ -62,16 +58,38 @@ def _convert_from_cloud_resource(cloud_resource: str) -> Resource:
     Returns:
         The converted resource.
     """
-    scope, resource_type_and_id = cloud_resource.rsplit(":", maxsplit=1)
+    scope, workspace_resource_type_and_id = cloud_resource.split(
+        ":", maxsplit=1
+    )
 
     if scope != f"{SERVER_ID}@{SERVER_SCOPE_IDENTIFIER}":
         raise ValueError("Invalid scope for server resource.")
 
+    workspace_id: Optional[str] = None
+    if ":" in workspace_resource_type_and_id:
+        (
+            workspace_id,
+            resource_type_and_id,
+        ) = workspace_resource_type_and_id.split(":", maxsplit=1)
+    else:
+        workspace_id = None
+        resource_type_and_id = workspace_resource_type_and_id
+
+    resource_id: Optional[str] = None
     if "/" in resource_type_and_id:
         resource_type, resource_id = resource_type_and_id.split("/")
-        return Resource(type=resource_type, id=resource_id)
     else:
-        return Resource(type=resource_type_and_id)
+        resource_type = resource_type_and_id
+
+    if resource_type == ResourceType.WORKSPACE and workspace_id is not None:
+        # 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 = None
+
+    return Resource(
+        type=resource_type, id=resource_id, workspace_id=workspace_id
+    )
 
 
 class ZenMLCloudRBAC(RBACInterface):
@@ -184,3 +202,17 @@ class ZenMLCloudRBAC(RBACInterface):
             "actions": [str(action) for action in actions],
         }
         self._connection.post(endpoint=RESOURCE_MEMBERSHIP_ENDPOINT, data=data)
+
+    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.
+        """
+        params = {
+            "resources": [
+                _convert_to_cloud_resource(resource) for resource in resources
+            ],
+        }
+        self._connection.delete(endpoint=RESOURCES_ENDPOINT, params=params)

zenml/zen_server/routers/actions_endpoints.py

@@ -40,6 +40,7 @@ from zenml.zen_server.rbac.endpoint_utils import (
 from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import (
     dehydrate_response_model,
+    delete_model_resource,
     verify_permission_for_model,
 )
 from zenml.zen_server.utils import (
@@ -58,7 +59,6 @@ router = APIRouter(
 
 @router.get(
     "",
-    response_model=Page[ActionResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -130,7 +130,6 @@ def list_actions(
 
 @router.get(
     "/{action_id}",
-    response_model=ActionResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -178,7 +177,6 @@ def get_action(
 
 @router.post(
     "",
-    response_model=ActionResponse,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -219,14 +217,12 @@ def create_action(
 
     return verify_permissions_and_create_entity(
         request_model=action,
-        resource_type=ResourceType.ACTION,
         create_method=action_handler.create_action,
     )
 
 
 @router.put(
     "/{action_id}",
-    response_model=ActionResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -322,3 +318,5 @@ def delete_action(
         action=action,
         force=force,
     )
+
+    delete_model_resource(action)

zenml/zen_server/routers/artifact_endpoint.py

@@ -50,7 +50,6 @@ artifact_router = APIRouter(
 
 @artifact_router.get(
     "",
-    response_model=Page[ArtifactResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -82,7 +81,6 @@ def list_artifacts(
 
 @artifact_router.post(
     "",
-    response_model=ArtifactResponse,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -100,14 +98,12 @@ def create_artifact(
     """
     return verify_permissions_and_create_entity(
         request_model=artifact,
-        resource_type=ResourceType.ARTIFACT,
         create_method=zen_store().create_artifact,
     )
 
 
 @artifact_router.get(
     "/{artifact_id}",
-    response_model=ArtifactResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -135,7 +131,6 @@ def get_artifact(
 
 @artifact_router.put(
     "/{artifact_id}",
-    response_model=ArtifactResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions

zenml/zen_server/routers/artifact_version_endpoints.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for artifact versions."""
 
-from typing import List
+from typing import List, Union
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security
@@ -46,6 +46,7 @@ from zenml.zen_server.rbac.utils import (
 from zenml.zen_server.utils import (
     handle_exceptions,
     make_dependable,
+    set_filter_workspace_scope,
     zen_store,
 )
 
@@ -58,7 +59,6 @@ artifact_version_router = APIRouter(
 
 @artifact_version_router.get(
     "",
-    response_model=Page[ArtifactVersionResponse],
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -81,8 +81,14 @@ def list_artifact_versions(
     Returns:
         The artifact versions according to query filters.
     """
+    # A workspace scoped request must always be scoped to a specific
+    # workspace. This is required for the RBAC check to work.
+    set_filter_workspace_scope(artifact_version_filter_model)
+    assert isinstance(artifact_version_filter_model.workspace, UUID)
+
     allowed_artifact_ids = get_allowed_resource_ids(
-        resource_type=ResourceType.ARTIFACT
+        resource_type=ResourceType.ARTIFACT,
+        workspace_id=artifact_version_filter_model.workspace,
     )
     artifact_version_filter_model.configure_rbac(
         authenticated_user_id=auth_context.user.id,
@@ -97,7 +103,6 @@ def list_artifact_versions(
 
 @artifact_version_router.post(
     "",
-    response_model=ArtifactVersionResponse,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -115,7 +120,6 @@ def create_artifact_version(
     """
     return verify_permissions_and_create_entity(
         request_model=artifact_version,
-        resource_type=ResourceType.ARTIFACT_VERSION,
         create_method=zen_store().create_artifact_version,
     )
 
@@ -139,14 +143,12 @@ def batch_create_artifact_version(
     """
     return verify_permissions_and_batch_create_entity(
         batch=artifact_versions,
-        resource_type=ResourceType.ARTIFACT_VERSION,
         create_method=zen_store().batch_create_artifact_versions,
     )
 
 
 @artifact_version_router.get(
     "/{artifact_version_id}",
-    response_model=ArtifactVersionResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -174,7 +176,6 @@ def get_artifact_version(
 
 @artifact_version_router.put(
     "/{artifact_version_id}",
-    response_model=ArtifactVersionResponse,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions
@@ -227,24 +228,29 @@ def delete_artifact_version(
 )
 @handle_exceptions
 def prune_artifact_versions(
+    workspace_name_or_id: Union[str, UUID],
     only_versions: bool = True,
     _: AuthContext = Security(authorize),
 ) -> None:
     """Prunes unused artifact versions and their artifacts.
 
     Args:
+        workspace_name_or_id: The workspace name or ID to prune artifact
+            versions for.
         only_versions: Only delete artifact versions, keeping artifacts
     """
+    workspace_id = zen_store().get_workspace(workspace_name_or_id).id
+
     verify_permissions_and_prune_entities(
         resource_type=ResourceType.ARTIFACT_VERSION,
         prune_method=zen_store().prune_artifact_versions,
         only_versions=only_versions,
+        workspace_id=workspace_id,
     )
 
 
 @artifact_version_router.get(
     "/{artifact_version_id}" + VISUALIZE,
-    response_model=LoadedVisualization,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
 @handle_exceptions